rust-chg: add helper to pack environment variables

On my machine, "printenv | wc -c" says 3422. That's the only reason why
the initial buffer capacity is set to 4kB.

Differential Revision: https://phab.mercurial-scm.org/D8363

rust/chg/src/message.rs

@@ -1,133 +1,232 b''
 // Copyright 2018 Yuya Nishihara <[e-mail]>
 //
 // This software may be used and distributed according to the terms of the
 // GNU General Public License version 2 or any later version.
 
 //! Utility for parsing and building command-server messages.
 
-use bytes::Bytes;
+use bytes::{BufMut, Bytes, BytesMut};
 use std::error;
 use std::ffi::{OsStr, OsString};
 use std::io;
 use std::os::unix::ffi::OsStrExt;
 
 pub use tokio_hglib::message::*; // re-exports
 
 /// Shell command type requested by the server.
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 pub enum CommandType {
     /// Pager should be spawned.
     Pager,
     /// Shell command should be executed to send back the result code.
     System,
 }
 
 /// Shell command requested by the server.
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub struct CommandSpec {
     pub command: OsString,
     pub current_dir: OsString,
     pub envs: Vec<(OsString, OsString)>,
 }
 
 /// Parses "S" channel request into command type and spec.
 pub fn parse_command_spec(data: Bytes) -> io::Result<(CommandType, CommandSpec)> {
     let mut split = data.split(|&c| c == b'\0');
     let ctype = parse_command_type(split.next().ok_or(new_parse_error("missing type"))?)?;
     let command = split.next().ok_or(new_parse_error("missing command"))?;
     let current_dir = split.next().ok_or(new_parse_error("missing current dir"))?;
 
     let mut envs = Vec::new();
     for l in split {
         let mut s = l.splitn(2, |&c| c == b'=');
         let k = s.next().unwrap();
         let v = s.next().ok_or(new_parse_error("malformed env"))?;
         envs.push((
             OsStr::from_bytes(k).to_owned(),
             OsStr::from_bytes(v).to_owned(),
         ));
     }
 
     let spec = CommandSpec {
         command: OsStr::from_bytes(command).to_owned(),
         current_dir: OsStr::from_bytes(current_dir).to_owned(),
         envs: envs,
     };
     Ok((ctype, spec))
 }
 
 fn parse_command_type(value: &[u8]) -> io::Result<CommandType> {
     match value {
         b"pager" => Ok(CommandType::Pager),
         b"system" => Ok(CommandType::System),
         _ => Err(new_parse_error(format!(
             "unknown command type: {}",
             decode_latin1(value)
         ))),
     }
 }
 
+// allocate large buffer as environment variables can be quite long
+const INITIAL_PACKED_ENV_VARS_CAPACITY: usize = 4096;
+
+/// Packs environment variables of platform encoding into bytes.
+///
+/// # Panics
+///
+/// Panics if key or value contains `\0` character, or key contains '='
+/// character.
+pub fn pack_env_vars_os<I, P>(vars: I) -> Bytes
+where
+    I: IntoIterator<Item = (P, P)>,
+    P: AsRef<OsStr>,
+{
+    let mut vars_iter = vars.into_iter();
+    if let Some((k, v)) = vars_iter.next() {
+        let mut dst = BytesMut::with_capacity(INITIAL_PACKED_ENV_VARS_CAPACITY);
+        pack_env_into(&mut dst, k.as_ref(), v.as_ref());
+        for (k, v) in vars_iter {
+            dst.reserve(1);
+            dst.put_u8(b'\0');
+            pack_env_into(&mut dst, k.as_ref(), v.as_ref());
+        }
+        dst.freeze()
+    } else {
+        Bytes::new()
+    }
+}
+
+fn pack_env_into(dst: &mut BytesMut, k: &OsStr, v: &OsStr) {
+    assert!(!k.as_bytes().contains(&0), "key shouldn't contain NUL");
+    assert!(!k.as_bytes().contains(&b'='), "key shouldn't contain '='");
+    assert!(!v.as_bytes().contains(&0), "value shouldn't contain NUL");
+    dst.reserve(k.as_bytes().len() + 1 + v.as_bytes().len());
+    dst.put_slice(k.as_bytes());
+    dst.put_u8(b'=');
+    dst.put_slice(v.as_bytes());
+}
+
 fn decode_latin1<S>(s: S) -> String
 where
     S: AsRef<[u8]>,
 {
     s.as_ref().iter().map(|&c| c as char).collect()
 }
 
 fn new_parse_error<E>(error: E) -> io::Error
 where
     E: Into<Box<error::Error + Send + Sync>>,
 {
     io::Error::new(io::ErrorKind::InvalidData, error)
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
     use std::os::unix::ffi::OsStringExt;
+    use std::panic;
 
     #[test]
     fn parse_command_spec_good() {
         let src = [
             b"pager".as_ref(),
             b"less -FRX".as_ref(),
             b"/tmp".as_ref(),
             b"LANG=C".as_ref(),
             b"HGPLAIN=".as_ref(),
         ]
         .join(&0);
         let spec = CommandSpec {
             command: os_string_from(b"less -FRX"),
             current_dir: os_string_from(b"/tmp"),
             envs: vec![
                 (os_string_from(b"LANG"), os_string_from(b"C")),
                 (os_string_from(b"HGPLAIN"), os_string_from(b"")),
             ],
         };
         assert_eq!(
             parse_command_spec(Bytes::from(src)).unwrap(),
             (CommandType::Pager, spec)
         );
     }
 
     #[test]
     fn parse_command_spec_too_short() {
         assert!(parse_command_spec(Bytes::from_static(b"")).is_err());
         assert!(parse_command_spec(Bytes::from_static(b"pager")).is_err());
         assert!(parse_command_spec(Bytes::from_static(b"pager\0less")).is_err());
     }
 
     #[test]
     fn parse_command_spec_malformed_env() {
         assert!(parse_command_spec(Bytes::from_static(b"pager\0less\0/tmp\0HOME")).is_err());
     }
 
     #[test]
     fn parse_command_spec_unknown_type() {
         assert!(parse_command_spec(Bytes::from_static(b"paper\0less")).is_err());
     }
 
+    #[test]
+    fn pack_env_vars_os_good() {
+        assert_eq!(
+            pack_env_vars_os(vec![] as Vec<(OsString, OsString)>),
+            Bytes::new()
+        );
+        assert_eq!(
+            pack_env_vars_os(vec![os_string_pair_from(b"FOO", b"bar")]),
+            Bytes::from_static(b"FOO=bar")
+        );
+        assert_eq!(
+            pack_env_vars_os(vec![
+                os_string_pair_from(b"FOO", b""),
+                os_string_pair_from(b"BAR", b"baz")
+            ]),
+            Bytes::from_static(b"FOO=\0BAR=baz")
+        );
+    }
+
+    #[test]
+    fn pack_env_vars_os_large_key() {
+        let mut buf = vec![b'A'; INITIAL_PACKED_ENV_VARS_CAPACITY];
+        let envs = vec![os_string_pair_from(&buf, b"")];
+        buf.push(b'=');
+        assert_eq!(pack_env_vars_os(envs), Bytes::from(buf));
+    }
+
+    #[test]
+    fn pack_env_vars_os_large_value() {
+        let mut buf = vec![b'A', b'='];
+        buf.resize(INITIAL_PACKED_ENV_VARS_CAPACITY + 1, b'a');
+        let envs = vec![os_string_pair_from(&buf[..1], &buf[2..])];
+        assert_eq!(pack_env_vars_os(envs), Bytes::from(buf));
+    }
+
+    #[test]
+    fn pack_env_vars_os_nul_eq() {
+        assert!(panic::catch_unwind(|| {
+            pack_env_vars_os(vec![os_string_pair_from(b"\0", b"")])
+        })
+        .is_err());
+        assert!(panic::catch_unwind(|| {
+            pack_env_vars_os(vec![os_string_pair_from(b"FOO", b"\0bar")])
+        })
+        .is_err());
+        assert!(panic::catch_unwind(|| {
+            pack_env_vars_os(vec![os_string_pair_from(b"FO=", b"bar")])
+        })
+        .is_err());
+        assert_eq!(
+            pack_env_vars_os(vec![os_string_pair_from(b"FOO", b"=ba")]),
+            Bytes::from_static(b"FOO==ba")
+        );
+    }
+
     fn os_string_from(s: &[u8]) -> OsString {
         OsString::from_vec(s.to_vec())
     }
+
+    fn os_string_pair_from(k: &[u8], v: &[u8]) -> (OsString, OsString) {
+        (os_string_from(k), os_string_from(v))
+    }
 }