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
Branch
default
Lint
No Linters Available
Unit
No Unit Test Coverage

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 20965

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))
}
} }