[WIP] size limit for command line

Artoria2e5 · Artoria2e5 · commit 191e31998b02 · 2020-07-29T02:15:57.000+08:00
diff --git a/library/std/src/sys/unix/process/process_common.rs b/library/std/src/sys/unix/process/process_common.rs
@@ -13,7 +13,7 @@ use crate::sys_common::process::CommandEnv;
 #[cfg(not(target_os = "fuchsia"))]
 use crate::sys::fs::OpenOptions;
 
-use libc::{c_char, c_int, gid_t, uid_t, EXIT_FAILURE, EXIT_SUCCESS};
+use libc::{c_char, c_int, gid_t, strlen, uid_t, EXIT_FAILURE, EXIT_SUCCESS};
 
 cfg_if::cfg_if! {
     if #[cfg(target_os = "fuchsia")] {
@@ -75,6 +75,7 @@ pub struct Command {
     args: Vec<CString>,
     argv: Argv,
     env: CommandEnv,
+    arg_max: Option<isize>,
 
     cwd: Option<CString>,
     uid: Option<uid_t>,
@@ -137,6 +138,7 @@ impl Command {
             args: vec![program.clone()],
             program,
             env: Default::default(),
+            arg_max: Default::default(),
             cwd: None,
             uid: None,
             gid: None,
@@ -202,6 +204,49 @@ impl Command {
         self.gid
     }
 
+    pub fn get_size(&mut self) -> io::Result<usize> {
+        use crate::mem;
+        let argv = self.argv.0;
+        let argv_size: usize = argv.iter().map(|x| unsafe { strlen(*x) + 1 }).sum::<usize>()
+            + (argv.len() + 1) * mem::size_of::<*const u8>();
+
+        // Envp size calculation is approximate.
+        let env = self.env.capture();
+        let env_size: usize = env
+            .iter()
+            .map(|(k, v)| unsafe {
+                os2c(k.as_ref(), &mut self.saw_nul).to_bytes().len()
+                    + os2c(v.as_ref(), &mut self.saw_nul).to_bytes().len()
+                    + 2
+            })
+            .sum::<usize>()
+            + (env.len() + 1) * mem::size_of::<*const u8>();
+
+        Ok(argv_size + env_size)
+    }
+
+    pub fn check_size(&mut self, refresh: bool) -> io::Result<bool> {
+        use crate::sys;
+        use core::convert::TryInto;
+        if refresh || self.arg_max.is_none() {
+            let (limit, errno) = unsafe {
+                let old_errno = sys::os::errno();
+                sys::os::set_errno(0);
+                let limit = libc::sysconf(libc::_SC_ARG_MAX);
+                let errno = sys::os::errno();
+                sys::os::set_errno(old_errno);
+                (limit, errno)
+            };
+
+            if errno != 0 {
+                return Err(io::Error::from_raw_os_error(errno));
+            } else {
+                self.arg_max = limit.try_into().ok();
+            }
+        }
+        Ok(self.arg_max.unwrap() < 0 || self.get_size()? < (self.arg_max.unwrap() as usize))
+    }
+
     pub fn get_closures(&mut self) -> &mut Vec<Box<dyn FnMut() -> io::Result<()> + Send + Sync>> {
         &mut self.closures
     }
diff --git a/library/std/src/sys/windows/process.rs b/library/std/src/sys/windows/process.rs
@@ -65,6 +65,8 @@ fn ensure_no_nuls<T: AsRef<OsStr>>(str: T) -> io::Result<T> {
     }
 }
 
+// 32768 minus NUL plus starting space in our implementation
+const CMDLINE_MAX: usize = 32768;
 pub struct Command {
     program: OsString,
     args: Vec<OsString>,
@@ -75,6 +77,8 @@ pub struct Command {
     stdin: Option<Stdio>,
     stdout: Option<Stdio>,
     stderr: Option<Stdio>,
+    cmdline: Vec<u16>,
+    cmdline_error: Option<io::Error>,
 }
 
 pub enum Stdio {
@@ -106,11 +110,32 @@ impl Command {
             stdin: None,
             stdout: None,
             stderr: None,
+            cmdline: Vec::new(),
+            cmdline_error: None,
         }
     }
 
+    pub fn maybe_arg(&mut self, arg: &OsStr) -> io::Result<()> {
+        self.args.push(arg.to_os_string());
+        self.cmdline.push(' ' as u16);
+        let result = append_arg(&mut cmd, arg, false);
+        if result.is_err() {
+            self.cmdline.truncate(self.cmdline.len() - 1);
+            result
+        } else if self.cmdline.size() >= CMDLINE_MAX {
+            // Roll back oversized
+            self.cmdline.truncate(self.cmdline.len() - 1 - result.unwrap());
+            Err(io::Error::new(ErrorKind::InvalidInput, "oversized cmdline"))
+        }
+        Ok()
+    }
     pub fn arg(&mut self, arg: &OsStr) {
-        self.args.push(arg.to_os_string())
+        if self.cmdline_error.is_none() {
+            let result = self.maybe_arg(self, arg);
+            if result.is_err() {
+                self.cmdline_error = Some(result.expect_err());
+            }
+        }
     }
     pub fn env_mut(&mut self) -> &mut CommandEnv {
         &mut self.env
@@ -136,34 +161,45 @@ impl Command {
         default: Stdio,
         needs_stdin: bool,
     ) -> io::Result<(Process, StdioPipes)> {
+        if self.cmdline_error.is_some() {
+            return self.cmdline_error.unwrap();
+        }
+
         let maybe_env = self.env.capture_if_changed();
         // To have the spawning semantics of unix/windows stay the same, we need
         // to read the *child's* PATH if one is provided. See #15149 for more
         // details.
-        let program = maybe_env.as_ref().and_then(|env| {
-            if let Some(v) = env.get(OsStr::new("PATH")) {
-                // Split the value and test each path to see if the
-                // program exists.
-                for path in split_paths(&v) {
-                    let path = path
-                        .join(self.program.to_str().unwrap())
-                        .with_extension(env::consts::EXE_EXTENSION);
-                    if fs::metadata(&path).is_ok() {
-                        return Some(path.into_os_string());
+        let program = maybe_env
+            .as_ref()
+            .and_then(|env| {
+                if let Some(v) = env.get(OsStr::new("PATH")) {
+                    // Split the value and test each path to see if the
+                    // program exists.
+                    for path in split_paths(&v) {
+                        let path = path
+                            .join(self.program.to_str().unwrap())
+                            .with_extension(env::consts::EXE_EXTENSION);
+                        if fs::metadata(&path).is_ok() {
+                            return Some(path.into_os_string());
+                        }
                     }
                 }
-            }
-            None
-        });
+                None
+            })
+            .as_ref()
+            .unwrap_or(&self.program);
+
+        // Prepare and terminate the application name and the cmdline
+        // XXX: this won't work for 16-bit, might be preferable to do a extend_from_slice
+        let program_str: Vec<u16> = Vec::new();
+        append_arg(&mut program_str, program, true)?;
+        program_str.push(0);
+        self.cmdline.push(0);
 
         let mut si = zeroed_startupinfo();
         si.cb = mem::size_of::<c::STARTUPINFO>() as c::DWORD;
         si.dwFlags = c::STARTF_USESTDHANDLES;
 
-        let program = program.as_ref().unwrap_or(&self.program);
-        let mut cmd_str = make_command_line(program, &self.args)?;
-        cmd_str.push(0); // add null terminator
-
         // stolen from the libuv code.
         let mut flags = self.flags | c::CREATE_UNICODE_ENVIRONMENT;
         if self.detach {
@@ -201,8 +237,8 @@ impl Command {
 
         unsafe {
             cvt(c::CreateProcessW(
-                ptr::null(),
-                cmd_str.as_mut_ptr(),
+                program_str.as_mut_ptr(),
+                self.cmdline.as_mut_ptr().offset(1), // Skip the starting space
                 ptr::null_mut(),
                 ptr::null_mut(),
                 c::TRUE,
@@ -221,6 +257,14 @@ impl Command {
 
         Ok((Process { handle: Handle::new(pi.hProcess) }, pipes))
     }
+
+    pub fn get_size(&mut self) -> io::Result<usize> {
+        let (_, cmd_str) = self.prepare_command_line()?;
+        Ok(cmd_str.len())
+    }
+    pub fn check_size(&mut self, _refresh: bool) -> io::Result<bool> {
+        Ok(self.get_size()? < 32767)
+    }
 }
 
 impl fmt::Debug for Command {
@@ -445,6 +489,44 @@ fn zeroed_process_information() -> c::PROCESS_INFORMATION {
     }
 }
 
+fn append_arg(cmd: &mut Vec<u16>, arg: &OsStr, force_quotes: bool) -> io::Result<usize> {
+    let mut addsize: usize = 0;
+    // 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 =
+        force_quotes || arg_bytes.iter().any(|c| *c == b' ' || *c == b'\t') || arg_bytes.is_empty();
+    if quote {
+        cmd.push('"' as u16);
+        addsize += 1;
+    }
+
+    let mut backslashes: usize = 0;
+    for x in arg.encode_wide() {
+        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));
+                addsize += backslashes + 1;
+            }
+            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);
+        addsize += backslashes + 1;
+    }
+    Ok(addsize)
+}
+
 // Produces a wide string *without terminating null*; returns an error if
 // `prog` or any of the `args` contain a nul.
 fn make_command_line(prog: &OsStr, args: &[OsString]) -> io::Result<Vec<u16>> {
@@ -459,41 +541,6 @@ fn make_command_line(prog: &OsStr, args: &[OsString]) -> io::Result<Vec<u16>> {
         append_arg(&mut cmd, arg, false)?;
     }
     return Ok(cmd);
-
-    fn append_arg(cmd: &mut Vec<u16>, arg: &OsStr, force_quotes: bool) -> 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 = force_quotes
-            || arg_bytes.iter().any(|c| *c == b' ' || *c == b'\t')
-            || arg_bytes.is_empty();
-        if quote {
-            cmd.push('"' as u16);
-        }
-
-        let mut backslashes: usize = 0;
-        for x in arg.encode_wide() {
-            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(())
-    }
 }
 
 fn make_envp(maybe_env: Option<BTreeMap<EnvKey, OsString>>) -> io::Result<(*mut c_void, Vec<u16>)> {
