Mercurial > hg
view rust/chg/src/message.rs @ 44685:90e05b304902
rust-chg: silence warning about dated coding style
Differential Revision: https://phab.mercurial-scm.org/D8384
author | Yuya Nishihara <yuya@tcha.org> |
---|---|
date | Tue, 31 Mar 2020 23:25:26 +0900 |
parents | 82adc720c0a3 |
children | 61fda2dbc522 |
line wrap: on
line source
// Copyright 2018 Yuya Nishihara <yuya@tcha.org> // // 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::{BufMut, Bytes, BytesMut}; use std::error; use std::ffi::{OsStr, OsString}; use std::io; use std::os::unix::ffi::OsStrExt; use std::path::PathBuf; 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) ))), } } /// Client-side instruction requested by the server. #[derive(Clone, Debug, Eq, PartialEq)] pub enum Instruction { Exit(i32), Reconnect, Redirect(PathBuf), Unlink(PathBuf), } /// Parses validation result into instructions. pub fn parse_instructions(data: Bytes) -> io::Result<Vec<Instruction>> { let mut instructions = Vec::new(); for l in data.split(|&c| c == b'\0') { if l.is_empty() { continue; } let mut s = l.splitn(2, |&c| c == b' '); let inst = match (s.next().unwrap(), s.next()) { (b"exit", Some(arg)) => decode_latin1(arg) .parse() .map(Instruction::Exit) .map_err(|_| new_parse_error(format!("invalid exit code: {:?}", arg)))?, (b"reconnect", None) => Instruction::Reconnect, (b"redirect", Some(arg)) => { Instruction::Redirect(OsStr::from_bytes(arg).to_owned().into()) } (b"unlink", Some(arg)) => Instruction::Unlink(OsStr::from_bytes(arg).to_owned().into()), _ => { return Err(new_parse_error(format!("unknown command: {:?}", l))); } }; instructions.push(inst); } Ok(instructions) } // 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<dyn 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 parse_instructions_good() { let src = [ b"exit 123".as_ref(), b"reconnect".as_ref(), b"redirect /whatever".as_ref(), b"unlink /someother".as_ref(), ] .join(&0); let insts = vec![ Instruction::Exit(123), Instruction::Reconnect, Instruction::Redirect(path_buf_from(b"/whatever")), Instruction::Unlink(path_buf_from(b"/someother")), ]; assert_eq!(parse_instructions(Bytes::from(src)).unwrap(), insts); } #[test] fn parse_instructions_empty() { assert_eq!(parse_instructions(Bytes::new()).unwrap(), vec![]); assert_eq!( parse_instructions(Bytes::from_static(b"\0")).unwrap(), vec![] ); } #[test] fn parse_instructions_malformed_exit_code() { assert!(parse_instructions(Bytes::from_static(b"exit foo")).is_err()); } #[test] fn parse_instructions_missing_argument() { assert!(parse_instructions(Bytes::from_static(b"exit")).is_err()); assert!(parse_instructions(Bytes::from_static(b"redirect")).is_err()); assert!(parse_instructions(Bytes::from_static(b"unlink")).is_err()); } #[test] fn parse_instructions_unknown_command() { assert!(parse_instructions(Bytes::from_static(b"quit 0")).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)) } fn path_buf_from(s: &[u8]) -> PathBuf { os_string_from(s).into() } }