Refactor: Move argument building into args

This commit is contained in:
Chris Denton 2022-03-23 04:18:47 +00:00
parent 3ea44938e2
commit d59cf5629e
No known key found for this signature in database
GPG Key ID: 713472F2F45627DE
2 changed files with 70 additions and 69 deletions

View File

@ -8,12 +8,14 @@ mod tests;
use crate::ffi::OsString;
use crate::fmt;
use crate::io;
use crate::marker::PhantomData;
use crate::num::NonZeroU16;
use crate::os::windows::prelude::*;
use crate::path::PathBuf;
use crate::ptr::NonNull;
use crate::sys::c;
use crate::sys::process::ensure_no_nuls;
use crate::sys::windows::os::current_exe;
use crate::vec;
@ -234,3 +236,66 @@ impl Iterator for WStrUnits<'_> {
}
}
}
#[derive(Debug)]
pub(crate) enum Arg {
/// Add quotes (if needed)
Regular(OsString),
/// Append raw string without quoting
Raw(OsString),
}
enum Quote {
// Every arg is quoted
Always,
// Whitespace and empty args are quoted
Auto,
// Arg appended without any changes (#29494)
Never,
}
pub(crate) fn append_arg(cmd: &mut Vec<u16>, arg: &Arg, force_quotes: bool) -> io::Result<()> {
let (arg, quote) = match arg {
Arg::Regular(arg) => (arg, if force_quotes { Quote::Always } else { Quote::Auto }),
Arg::Raw(arg) => (arg, Quote::Never),
};
// If an argument has 0 characters then we need to quote it to ensure
// that it actually gets passed through on the command line or otherwise
// it will be dropped entirely when parsed on the other end.
ensure_no_nuls(arg)?;
let arg_bytes = arg.bytes();
let (quote, escape) = match quote {
Quote::Always => (true, true),
Quote::Auto => {
(arg_bytes.iter().any(|c| *c == b' ' || *c == b'\t') || arg_bytes.is_empty(), true)
}
Quote::Never => (false, false),
};
if quote {
cmd.push('"' as u16);
}
let mut backslashes: usize = 0;
for x in arg.encode_wide() {
if escape {
if x == '\\' as u16 {
backslashes += 1;
} else {
if x == '"' as u16 {
// Add n+1 backslashes to total 2n+1 before internal '"'.
cmd.extend((0..=backslashes).map(|_| '\\' as u16));
}
backslashes = 0;
}
}
cmd.push(x);
}
if quote {
// Add n backslashes to total 2n before ending '"'.
cmd.extend((0..backslashes).map(|_| '\\' as u16));
cmd.push('"' as u16);
}
Ok(())
}

View File

@ -17,6 +17,7 @@ use crate::os::windows::ffi::{OsStrExt, OsStringExt};
use crate::os::windows::io::{AsRawHandle, FromRawHandle, IntoRawHandle};
use crate::path::{Path, PathBuf};
use crate::ptr;
use crate::sys::args::{self, Arg};
use crate::sys::c;
use crate::sys::c::NonZeroDWORD;
use crate::sys::cvt;
@ -27,7 +28,7 @@ use crate::sys::pipe::{self, AnonPipe};
use crate::sys::stdio;
use crate::sys_common::mutex::StaticMutex;
use crate::sys_common::process::{CommandEnv, CommandEnvs};
use crate::sys_common::{AsInner, IntoInner};
use crate::sys_common::IntoInner;
use libc::{c_void, EXIT_FAILURE, EXIT_SUCCESS};
@ -147,7 +148,7 @@ impl AsRef<OsStr> for EnvKey {
}
}
fn ensure_no_nuls<T: AsRef<OsStr>>(str: T) -> io::Result<T> {
pub(crate) fn ensure_no_nuls<T: AsRef<OsStr>>(str: T) -> io::Result<T> {
if str.as_ref().encode_wide().any(|b| b == 0) {
Err(io::const_io_error!(ErrorKind::InvalidInput, "nul byte found in provided data"))
} else {
@ -181,14 +182,6 @@ pub struct StdioPipes {
pub stderr: Option<AnonPipe>,
}
#[derive(Debug)]
enum Arg {
/// Add quotes (if needed)
Regular(OsString),
/// Append raw string without quoting
Raw(OsString),
}
impl Command {
pub fn new(program: &OsStr) -> Command {
Command {
@ -724,15 +717,6 @@ fn zeroed_process_information() -> c::PROCESS_INFORMATION {
}
}
enum Quote {
// Every arg is quoted
Always,
// Whitespace and empty args are quoted
Auto,
// Arg appended without any changes (#29494)
Never,
}
// Produces a wide string *without terminating null*; returns an error if
// `prog` or any of the `args` contain a nul.
fn make_command_line(
@ -763,57 +747,9 @@ fn make_command_line(
for arg in args {
cmd.push(' ' as u16);
let (arg, quote) = match arg {
Arg::Regular(arg) => (arg, if force_quotes { Quote::Always } else { Quote::Auto }),
Arg::Raw(arg) => (arg, Quote::Never),
};
append_arg(&mut cmd, arg, quote)?;
}
if is_batch_file {
cmd.push(b'"' as u16);
}
return Ok(cmd);
fn append_arg(cmd: &mut Vec<u16>, arg: &OsStr, quote: Quote) -> io::Result<()> {
// If an argument has 0 characters then we need to quote it to ensure
// that it actually gets passed through on the command line or otherwise
// it will be dropped entirely when parsed on the other end.
ensure_no_nuls(arg)?;
let arg_bytes = &arg.as_inner().inner.as_inner();
let (quote, escape) = match quote {
Quote::Always => (true, true),
Quote::Auto => {
(arg_bytes.iter().any(|c| *c == b' ' || *c == b'\t') || arg_bytes.is_empty(), true)
}
Quote::Never => (false, false),
};
if quote {
cmd.push('"' as u16);
}
let mut backslashes: usize = 0;
for x in arg.encode_wide() {
if escape {
if x == '\\' as u16 {
backslashes += 1;
} else {
if x == '"' as u16 {
// Add n+1 backslashes to total 2n+1 before internal '"'.
cmd.extend((0..=backslashes).map(|_| '\\' as u16));
}
backslashes = 0;
}
}
cmd.push(x);
}
if quote {
// Add n backslashes to total 2n before ending '"'.
cmd.extend((0..backslashes).map(|_| '\\' as u16));
cmd.push('"' as u16);
}
Ok(())
args::append_arg(&mut cmd, arg, force_quotes)?;
}
Ok(cmd)
}
fn make_envp(maybe_env: Option<BTreeMap<EnvKey, OsString>>) -> io::Result<(*mut c_void, Vec<u16>)> {