This is an archive of the discontinued Mercurial Phabricator instance.

Differential D8363

rust-chg: add helper to pack environment variables
ClosedPublic

Authored by yuja on Apr 2 2020, 7:12 AM.

Details

Reviewers
Alphare
Group Reviewers
hg-reviewers
Commits
rHG8a7beeea655f: rust-chg: add helper to pack environment variables
Summary

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

Diff Detail

Repository
rHG Mercurial
Lint
Automatic diff as part of commit; lint not applicable.
Unit
Automatic diff as part of commit; unit tests not applicable.

Event Timeline

yuja created this revision.Apr 2 2020, 7:12 AM
Alphare accepted this revision.Apr 3 2020, 8:24 AM
yuja added a commit: rHG8a7beeea655f: rust-chg: add helper to pack environment variables.Apr 3 2020, 3:17 PM
This revision was not accepted when it landed; it landed in state Needs Review.
Closed by commit rHG8a7beeea655f: rust-chg: add helper to pack environment variables (authored by yuja). · Explain Why
This revision was automatically updated to reflect the committed changes.

Revision Contents
Changeset List

PathPackages
Mrust/chg/src/message.rs (101 lines)

Diff 20992

rust/chg/src/message.rs

// Copyright 2018 Yuya Nishihara <yuya@tcha.org> // Copyright 2018 Yuya Nishihara <yuya@tcha.org>
// //
// This software may be used and distributed according to the terms of the // This software may be used and distributed according to the terms of the
// GNU General Public License version 2 or any later version. // GNU General Public License version 2 or any later version.
//! Utility for parsing and building command-server messages. //! Utility for parsing and building command-server messages.
use bytes::Bytes; use bytes::{BufMut, Bytes, BytesMut};
use std::error; use std::error;
use std::ffi::{OsStr, OsString}; use std::ffi::{OsStr, OsString};
use std::io; use std::io;
use std::os::unix::ffi::OsStrExt; use std::os::unix::ffi::OsStrExt;
pub use tokio_hglib::message::*; // re-exports pub use tokio_hglib::message::*; // re-exports
/// Shell command type requested by the server. /// Shell command type requested by the server.
b"system" => Ok(CommandType::System), b"system" => Ok(CommandType::System),
_ => Err(new_parse_error(format!( _ => Err(new_parse_error(format!(
"unknown command type: {}", "unknown command type: {}",
decode_latin1(value) 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 fn decode_latin1<S>(s: S) -> String
where where
S: AsRef<[u8]>, S: AsRef<[u8]>,
{ {
s.as_ref().iter().map(|&c| c as char).collect() s.as_ref().iter().map(|&c| c as char).collect()
} }
fn new_parse_error<E>(error: E) -> io::Error fn new_parse_error<E>(error: E) -> io::Error
where where
E: Into<Box<error::Error + Send + Sync>>, E: Into<Box<error::Error + Send + Sync>>,
{ {
io::Error::new(io::ErrorKind::InvalidData, error) io::Error::new(io::ErrorKind::InvalidData, error)
} }
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use std::os::unix::ffi::OsStringExt; use std::os::unix::ffi::OsStringExt;
use std::panic;
#[test] #[test]
fn parse_command_spec_good() { fn parse_command_spec_good() {
let src = [ let src = [
b"pager".as_ref(), b"pager".as_ref(),
b"less -FRX".as_ref(), b"less -FRX".as_ref(),
b"/tmp".as_ref(), b"/tmp".as_ref(),
b"LANG=C".as_ref(), b"LANG=C".as_ref(),
assert!(parse_command_spec(Bytes::from_static(b"pager\0less\0/tmp\0HOME")).is_err()); assert!(parse_command_spec(Bytes::from_static(b"pager\0less\0/tmp\0HOME")).is_err());
} }
#[test] #[test]
fn parse_command_spec_unknown_type() { fn parse_command_spec_unknown_type() {
assert!(parse_command_spec(Bytes::from_static(b"paper\0less")).is_err()); 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 { fn os_string_from(s: &[u8]) -> OsString {
OsString::from_vec(s.to_vec()) 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))
}
} }