use std::cmp::Ordering;
use std::collections::btree_map::{BTreeMap, Entry};
use std::ffi::{OsStr, OsString};
use std::io;
use std::iter;
use std::os::windows::ffi::OsStrExt;
use std::path::{Path, PathBuf};
use windows_sys::Win32::Foundation::TRUE;
use windows_sys::Win32::Globalization::{
CompareStringOrdinal, CSTR_EQUAL, CSTR_GREATER_THAN, CSTR_LESS_THAN,
};
const QUOTE: u16 = b'"' as u16;
const BACKSLASH: u16 = b'\\' as u16;
const SPACE: u16 = b' ' as u16;
#[derive(Debug)]
enum Arg {
Regular(OsString),
Raw(OsString),
}
#[derive(Debug)]
enum EnvOp {
Set(OsString, OsString),
Remove(OsString),
}
#[derive(Debug)]
struct EnvKey {
wide: Vec<u16>,
len: i32,
}
impl EnvKey {
fn new(key: &OsStr) -> io::Result<Self> {
let wide: Vec<u16> = key.encode_wide().collect();
let len = match i32::try_from(wide.len()) {
Ok(len) => len,
Err(err) => return Err(io::Error::new(io::ErrorKind::InvalidInput, err)),
};
Ok(Self { wide, len })
}
}
impl Ord for EnvKey {
fn cmp(&self, other: &Self) -> Ordering {
match unsafe {
CompareStringOrdinal(
self.wide.as_ptr(),
self.len,
other.wide.as_ptr(),
other.len,
TRUE,
)
} {
CSTR_LESS_THAN => Ordering::Less,
CSTR_EQUAL => Ordering::Equal,
CSTR_GREATER_THAN => Ordering::Greater,
_ => self.wide.cmp(&other.wide),
}
}
}
impl PartialOrd for EnvKey {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl PartialEq for EnvKey {
fn eq(&self, other: &Self) -> bool {
self.cmp(other) == Ordering::Equal
}
}
impl Eq for EnvKey {}
#[derive(Debug)]
pub(super) struct Command {
program: OsString,
args: Vec<Arg>,
env_clear: bool,
env_ops: Vec<EnvOp>,
cwd: Option<PathBuf>,
#[cfg(all(test, any(feature = "blocking", feature = "tokio")))]
kill_on_drop: bool,
}
impl Command {
pub(super) fn new(program: impl AsRef<OsStr>) -> Self {
Self {
program: program.as_ref().to_os_string(),
args: Vec::new(),
env_clear: false,
env_ops: Vec::new(),
cwd: None,
#[cfg(all(test, any(feature = "blocking", feature = "tokio")))]
kill_on_drop: false,
}
}
pub(super) fn arg(&mut self, arg: impl AsRef<OsStr>) -> &mut Self {
self.args.push(Arg::Regular(arg.as_ref().to_os_string()));
self
}
pub(super) fn args<I, S>(&mut self, args: I) -> &mut Self
where
I: IntoIterator<Item = S>,
S: AsRef<OsStr>,
{
for arg in args {
self.arg(arg);
}
self
}
pub(super) fn raw_arg(&mut self, text: impl AsRef<OsStr>) -> &mut Self {
self.args.push(Arg::Raw(text.as_ref().to_os_string()));
self
}
pub(super) fn env(&mut self, key: impl AsRef<OsStr>, value: impl AsRef<OsStr>) -> &mut Self {
self.env_ops.push(EnvOp::Set(
key.as_ref().to_os_string(),
value.as_ref().to_os_string(),
));
self
}
pub(super) fn envs<I, K, V>(&mut self, vars: I) -> &mut Self
where
I: IntoIterator<Item = (K, V)>,
K: AsRef<OsStr>,
V: AsRef<OsStr>,
{
for (key, value) in vars {
self.env(key, value);
}
self
}
pub(super) fn env_remove(&mut self, key: impl AsRef<OsStr>) -> &mut Self {
self.env_ops
.push(EnvOp::Remove(key.as_ref().to_os_string()));
self
}
pub(super) fn env_clear(&mut self) -> &mut Self {
self.env_clear = true;
self.env_ops.clear();
self
}
pub(super) fn current_dir(&mut self, dir: impl AsRef<Path>) -> &mut Self {
self.cwd = Some(dir.as_ref().to_path_buf());
self
}
#[cfg(all(test, any(feature = "blocking", feature = "tokio")))]
pub(super) fn kill_on_drop(&mut self, kill: bool) -> &mut Self {
self.kill_on_drop = kill;
self
}
#[cfg(any(feature = "blocking", feature = "tokio"))]
pub(super) fn get_program(&self) -> &OsStr {
&self.program
}
pub(super) fn get_current_dir(&self) -> Option<&Path> {
self.cwd.as_deref()
}
#[cfg(all(test, any(feature = "blocking", feature = "tokio")))]
pub(super) const fn get_kill_on_drop(&self) -> bool {
self.kill_on_drop
}
pub(super) fn build_command_line(&self) -> io::Result<Vec<u16>> {
ensure_no_nuls(&self.program)?;
if self.program.as_encoded_bytes().contains(&b'"') {
return Err(invalid_input(
"program name must not contain a double quote",
));
}
let mut cmd: Vec<u16> = Vec::new();
cmd.push(QUOTE);
cmd.extend(self.program.encode_wide());
cmd.push(QUOTE);
for arg in &self.args {
cmd.push(SPACE);
match arg {
Arg::Regular(arg) => append_regular_arg(&mut cmd, arg)?,
Arg::Raw(text) => {
ensure_no_nuls(text)?;
cmd.extend(text.encode_wide());
},
}
}
cmd.push(0);
Ok(cmd)
}
pub(super) fn build_environment_block(&self) -> io::Result<Option<Vec<u16>>> {
if !self.env_clear && self.env_ops.is_empty() {
return Ok(None);
}
let mut map: BTreeMap<EnvKey, (OsString, OsString)> = BTreeMap::new();
if !self.env_clear {
for (key, value) in std::env::vars_os() {
map.insert(EnvKey::new(&key)?, (key, value));
}
}
for op in &self.env_ops {
match op {
EnvOp::Set(key, value) => {
validate_env_key(key)?;
ensure_no_nuls(value)?;
match map.entry(EnvKey::new(key)?) {
Entry::Occupied(mut entry) => entry.get_mut().1.clone_from(value),
Entry::Vacant(entry) => {
entry.insert((key.clone(), value.clone()));
},
}
},
EnvOp::Remove(key) => {
validate_env_key(key)?;
map.remove(&EnvKey::new(key)?);
},
}
}
let mut block: Vec<u16> = Vec::new();
for (key, value) in map.values() {
block.extend(key.encode_wide());
block.push(u16::from(b'='));
block.extend(value.encode_wide());
block.push(0);
}
if block.is_empty() {
block.push(0);
}
block.push(0);
Ok(Some(block))
}
}
pub(super) fn to_wide_nul(s: &OsStr) -> io::Result<Vec<u16>> {
ensure_no_nuls(s)?;
let mut wide: Vec<u16> = s.encode_wide().collect();
wide.push(0);
Ok(wide)
}
fn invalid_input(msg: &'static str) -> io::Error {
io::Error::new(io::ErrorKind::InvalidInput, msg)
}
fn ensure_no_nuls(s: &OsStr) -> io::Result<()> {
if s.encode_wide().any(|unit| unit == 0) {
Err(invalid_input("nul character found in provided data"))
} else {
Ok(())
}
}
fn validate_env_key(key: &OsStr) -> io::Result<()> {
if key.is_empty() {
return Err(invalid_input("environment variable name must not be empty"));
}
ensure_no_nuls(key)?;
if key.as_encoded_bytes().contains(&b'=') {
return Err(invalid_input(
"environment variable name must not contain `=`",
));
}
Ok(())
}
fn append_regular_arg(cmd: &mut Vec<u16>, arg: &OsStr) -> io::Result<()> {
ensure_no_nuls(arg)?;
let bytes = arg.as_encoded_bytes();
let quote = bytes.is_empty()
|| bytes
.iter()
.any(|&b| b == b' ' || b == b'\t' || b == b'"' || b == b'|');
if quote {
cmd.push(QUOTE);
}
let mut backslashes: usize = 0;
for unit in arg.encode_wide() {
if unit == BACKSLASH {
backslashes += 1;
} else {
if unit == QUOTE {
cmd.extend(iter::repeat(BACKSLASH).take(backslashes + 1));
}
backslashes = 0;
}
cmd.push(unit);
}
if quote {
cmd.extend(iter::repeat(BACKSLASH).take(backslashes));
cmd.push(QUOTE);
}
Ok(())
}
#[cfg(test)]
#[path = "command_tests.rs"]
mod tests;