Disallow or quote all specials in bat args

library/std/src/sys/pal/windows/args.rs

@@ -7,7 +7,7 @@
 mod tests;
 
 use super::os::current_exe;
-use crate::ffi::OsString;
+use crate::ffi::{OsStr, OsString};
 use crate::fmt;
 use crate::io;
 use crate::num::NonZero;
@@ -17,6 +17,7 @@
 use crate::sys::process::ensure_no_nuls;
 use crate::sys::{c, to_u16s};
 use crate::sys_common::wstr::WStrUnits;
+use crate::sys_common::AsInner;
 use crate::vec;
 
 use crate::iter;
@@ -262,16 +263,92 @@ pub(crate) fn append_arg(cmd: &mut Vec<u16>, arg: &Arg, force_quotes: bool) -> i
     Ok(())
 }
 
+fn append_bat_arg(cmd: &mut Vec<u16>, arg: &OsStr, mut quote: bool) -> io::Result<()> {
+    ensure_no_nuls(arg)?;
+    // 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.
+    //
+    // We also need to quote the argument if it ends with `\` to guard against
+    // bat usage such as `"%~2"` (i.e. force quote arguments) otherwise a
+    // trailing slash will escape the closing quote.
+    if arg.is_empty() || arg.as_encoded_bytes().last() == Some(&b'\\') {
+        quote = true;
+    }
+    for cp in arg.as_inner().inner.code_points() {
+        if let Some(cp) = cp.to_char() {
+            // Rather than trying to find every ascii symbol that must be quoted,
+            // we assume that all ascii symbols must be quoted unless they're known to be good.
+            // We also quote Unicode control blocks for good measure.
+            // Note an unquoted `\` is fine so long as the argument isn't otherwise quoted.
+            static UNQUOTED: &str = r"#$*+-./:?@\_";
+            let ascii_needs_quotes =
+                cp.is_ascii() && !(cp.is_ascii_alphanumeric() || UNQUOTED.contains(cp));
+            if ascii_needs_quotes || cp.is_control() {
+                quote = true;
+            }
+        }
+    }
+
+    if quote {
+        cmd.push('"' as u16);
+    }
+    // Loop through the string, escaping `\` only if followed by `"`.
+    // And escaping `"` by doubling them.
+    let mut backslashes: usize = 0;
+    for x in arg.encode_wide() {
+        if x == '\\' as u16 {
+            backslashes += 1;
+        } else {
+            if x == '"' as u16 {
+                // Add n backslashes to total 2n before internal `"`.
+                cmd.extend((0..backslashes).map(|_| '\\' as u16));
+                // Appending an additional double-quote acts as an escape.
+                cmd.push(b'"' as u16)
+            } else if x == '%' as u16 || x == '\r' as u16 {
+                // yt-dlp hack: replaces `%` with `%%cd:~,%` to stop %VAR% being expanded as an environment variable.
+                //
+                // # Explanation
+                //
+                // cmd supports extracting a substring from a variable using the following syntax:
+                //     %variable:~start_index,end_index%
+                //
+                // In the above command `cd` is used as the variable and the start_index and end_index are left blank.
+                // `cd` is a built-in variable that dynamically expands to the current directory so it's always available.
+                // Explicitly omitting both the start and end index creates a zero-length substring.
+                //
+                // Therefore it all resolves to nothing. However, by doing this no-op we distract cmd.exe
+                // from potentially expanding %variables% in the argument.
+                cmd.extend_from_slice(&[
+                    '%' as u16, '%' as u16, 'c' as u16, 'd' as u16, ':' as u16, '~' as u16,
+                    ',' 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(())
+}
+
 pub(crate) fn make_bat_command_line(
     script: &[u16],
     args: &[Arg],
     force_quotes: bool,
 ) -> io::Result<Vec<u16>> {
+    const INVALID_ARGUMENT_ERROR: io::Error =
+        io::const_io_error!(io::ErrorKind::InvalidInput, r#"batch file arguments are invalid"#);
     // Set the start of the command line to `cmd.exe /c "`
     // It is necessary to surround the command in an extra pair of quotes,
     // hence the trailing quote here. It will be closed after all arguments
     // have been added.
-    let mut cmd: Vec<u16> = "cmd.exe /d /c \"".encode_utf16().collect();
+    // Using /e:ON enables "command extensions" which is essential for the `%` hack to work.
+    let mut cmd: Vec<u16> = "cmd.exe /e:ON /v:OFF /d /c \"".encode_utf16().collect();
 
     // Push the script name surrounded by its quote pair.
     cmd.push(b'"' as u16);
@@ -291,19 +368,23 @@ pub(crate) fn make_bat_command_line(
     // reconstructed by the batch script by default.
     for arg in args {
         cmd.push(' ' as u16);
-        // Make sure to always quote special command prompt characters, including:
-        // * Characters `cmd /?` says require quotes.
-        // * `%` for environment variables, as in `%TMP%`.
-        // * `|<>` pipe/redirect characters.
-        const SPECIAL: &[u8] = b"\t &()[]{}^=;!'+,`~%|<>";
-        let force_quotes = match arg {
-            Arg::Regular(arg) if !force_quotes => {
-                arg.as_encoded_bytes().iter().any(|c| SPECIAL.contains(c))
+        match arg {
+            Arg::Regular(arg_os) => {
+                let arg_bytes = arg_os.as_encoded_bytes();
+                // Disallow \r and \n as they may truncate the arguments.
+                const DISALLOWED: &[u8] = b"\r\n";
+                if arg_bytes.iter().any(|c| DISALLOWED.contains(c)) {
+                    return Err(INVALID_ARGUMENT_ERROR);
                 }
-            _ => force_quotes,
-        };
+                append_bat_arg(&mut cmd, arg_os, force_quotes)?;
+            }
+            _ => {
+                // Raw arguments are passed on as-is.
+                // It's the user's responsibility to properly handle arguments in this case.
                 append_arg(&mut cmd, arg, force_quotes)?;
             }
+        };
+    }
 
     // Close the quote we left opened earlier.
     cmd.push(b'"' as u16);

tests/ui/std/windows-bat-args.rs

@@ -0,0 +1,90 @@
+// only-windows
+// run-pass
+// run-flags:--parent-process
+
+use std::env;
+use std::io::ErrorKind::{self, InvalidInput};
+use std::path::{Path, PathBuf};
+use std::process::Command;
+
+fn main() {
+    if env::args().nth(1).as_deref() == Some("--parent-process") {
+        parent();
+    } else {
+        child();
+    }
+}
+
+fn child() {
+    if env::args().len() == 1 {
+        panic!("something went wrong :/");
+    }
+    for arg in env::args().skip(1) {
+        print!("{arg}\0");
+    }
+}
+
+fn parent() {
+    let mut bat = PathBuf::from(file!());
+    bat.set_file_name("windows-bat-args1.bat");
+    let bat1 = String::from(bat.to_str().unwrap());
+    bat.set_file_name("windows-bat-args2.bat");
+    let bat2 = String::from(bat.to_str().unwrap());
+    bat.set_file_name("windows-bat-args3.bat");
+    let bat3 = String::from(bat.to_str().unwrap());
+    let bat = [bat1.as_str(), bat2.as_str(), bat3.as_str()];
+
+    check_args(&bat, &["a", "b"]).unwrap();
+    check_args(&bat, &["c is for cat", "d is for dog"]).unwrap();
+    check_args(&bat, &["\"", " \""]).unwrap();
+    check_args(&bat, &["\\", "\\"]).unwrap();
+    check_args(&bat, &[">file.txt"]).unwrap();
+    check_args(&bat, &["whoami.exe"]).unwrap();
+    check_args(&bat, &["&a.exe"]).unwrap();
+    check_args(&bat, &["&echo hello "]).unwrap();
+    check_args(&bat, &["&echo hello", "&whoami", ">file.txt"]).unwrap();
+    check_args(&bat, &["!TMP!"]).unwrap();
+    check_args(&bat, &["key=value"]).unwrap();
+    check_args(&bat, &["\"key=value\""]).unwrap();
+    check_args(&bat, &["key = value"]).unwrap();
+    check_args(&bat, &["key=[\"value\"]"]).unwrap();
+    check_args(&bat, &["", "a=b"]).unwrap();
+    check_args(&bat, &["key=\"foo bar\""]).unwrap();
+    check_args(&bat, &["key=[\"my_value]"]).unwrap();
+    check_args(&bat, &["key=[\"my_value\",\"other-value\"]"]).unwrap();
+    check_args(&bat, &["key\\=value"]).unwrap();
+    check_args(&bat, &["key=\"&whoami\""]).unwrap();
+    check_args(&bat, &["key=\"value\"=5"]).unwrap();
+    check_args(&bat, &["key=[\">file.txt\"]"]).unwrap();
+    assert_eq!(check_args(&bat, &["\n"]), Err(InvalidInput));
+    assert_eq!(check_args(&bat, &["\r"]), Err(InvalidInput));
+    check_args(&bat, &["%hello"]).unwrap();
+    check_args(&bat, &["%PATH%"]).unwrap();
+    check_args(&bat, &["%%cd:~,%"]).unwrap();
+    check_args(&bat, &["%PATH%PATH%"]).unwrap();
+    check_args(&bat, &["\">file.txt"]).unwrap();
+    check_args(&bat, &["abc\"&echo hello"]).unwrap();
+    check_args(&bat, &["123\">file.txt"]).unwrap();
+    check_args(&bat, &["\"&echo hello&whoami.exe"]).unwrap();
+    check_args(&bat, &[r#"hello^"world"#, "hello &echo oh no >file.txt"]).unwrap();
+}
+
+// Check if the arguments roundtrip through a bat file and back into a Rust process.
+// Our Rust process outptuts the arguments as null terminated strings.
+#[track_caller]
+fn check_args(bats: &[&str], args: &[&str]) -> Result<(), ErrorKind> {
+    for bat in bats {
+        let output = Command::new(&bat).args(args).output().map_err(|e| e.kind())?;
+        assert!(output.status.success());
+        let child_args = String::from_utf8(output.stdout).unwrap();
+        let mut child_args: Vec<&str> =
+            child_args.strip_suffix('\0').unwrap().split('\0').collect();
+        // args3.bat can append spurious empty arguments, so trim them here.
+        child_args.truncate(
+            child_args.iter().rposition(|s| !s.is_empty()).unwrap_or(child_args.len() - 1) + 1,
+        );
+        assert_eq!(&child_args, &args);
+        assert!(!Path::new("file.txt").exists());
+    }
+    Ok(())
+}

tests/ui/std/windows-bat-args1.bat

@@ -0,0 +1 @@
+@a.exe %*

tests/ui/std/windows-bat-args2.bat

@@ -0,0 +1 @@
+@a.exe %1 %2 %3 %4 %5 %6 %7 %8 %9

tests/ui/std/windows-bat-args3.bat

@@ -0,0 +1 @@
+@a.exe "%~1" "%~2" "%~3" "%~4" "%~5" "%~6" "%~7" "%~8" "%~9"