Merge pull request #10 from algoux/dev

dev

.gitignore

@@ -14,3 +14,4 @@ Cargo.lock
 *.pdb
 .idea
 .vscode
+*.exe

README.md

@@ -7,7 +7,7 @@
 | 特性          | Linux | Windows | macOS |
 |-------------|-------|---------|-------|
 | 执行指定命令      | ~     | √       | ~     |
-| 流重定向        | ~     | ~       | ~     |
+| 流重定向        | ~     | √       | ~     |
 | 运行时间统计      | ~     | √       | ~     |
 | 运行 CPU 时间统计 | ~     | √       | ~     |
 | 运行内存统计      | ~     | √       | ~     |
@@ -25,3 +25,11 @@
 | 平滑退出        | ~     | ~       | ~     |
 
 **注意：** Windows 平台下运行 CPU 时间限制与运行内存限制不能保证精确，请不要以此为基准进行判断。
+
+## 测试
+
+```bash
+cargo test -- --test-threads=1
+```
+
+测试涉及文件操作，建议顺序执行测试用例（并发限制为 1）

src/sys/windows/mod.rs

@@ -2,7 +2,7 @@ use std::ffi::c_void;
 use std::mem::size_of;
 
 use windows::core::PSTR;
-use windows::Win32::Foundation::{FILETIME, WAIT_FAILED, WAIT_TIMEOUT};
+use windows::Win32::Foundation::{CloseHandle, FILETIME, WAIT_FAILED, WAIT_TIMEOUT};
 use windows::Win32::System::JobObjects::{
     AssignProcessToJobObject, CreateJobObjectA, JobObjectBasicLimitInformation,
     SetInformationJobObject, JOBOBJECT_BASIC_LIMIT_INFORMATION, JOB_OBJECT_LIMIT_PRIORITY_CLASS,
@@ -90,19 +90,27 @@ impl Sandbox {
     }
 
     unsafe fn redirect_fd(&mut self, info: &mut STARTUPINFOA) -> Result<()> {
+        // 重定向 stdin
+        if let Some(file) = &self.input {
+            info.hStdInput = handle_from_file(file, 'r')?;
+        }
         // 重定向 stdout
         if let Some(file) = &self.output {
-            info.hStdOutput = handle_from_file(file)?;
+            info.hStdOutput = handle_from_file(file, 'w')?;
         }
         // 重定向 stderr
         if let Some(file) = &self.error {
-            info.hStdError = handle_from_file(file)?;
+            info.hStdError = handle_from_file(file, 'w')?;
         }
 
         Ok(())
     }
 
-    unsafe fn wait_it(&mut self, information: &PROCESS_INFORMATION) -> Result<Status> {
+    unsafe fn wait_it(
+        &mut self,
+        info: &STARTUPINFOA,
+        information: &PROCESS_INFORMATION,
+    ) -> Result<Status> {
         let mut status: Status = Default::default();
         let timeout = if let Some(t) = self.time_limit {
             t
@@ -120,6 +128,17 @@ impl Sandbox {
             return Err(E(file!().to_string(), line!(), "WAIT_FAILED".to_string()));
         }
 
+        // 关闭文件流
+        if !info.hStdInput.is_invalid() {
+            winapi!(CloseHandle(info.hStdInput));
+        }
+        if !info.hStdOutput.is_invalid() {
+            winapi!(CloseHandle(info.hStdOutput));
+        }
+        if !info.hStdError.is_invalid() {
+            winapi!(CloseHandle(info.hStdError));
+        }
+
         let mut pmc: PROCESS_MEMORY_COUNTERS = Default::default();
 
         // 获取内存使用情况
@@ -217,42 +236,27 @@ impl SandboxImpl for Sandbox {
             ));
         }
 
-        self.wait_it(&information)
+        self.wait_it(&info, &information)
     }
 }
 
 #[cfg(test)]
 mod tests {
+    use std::fs;
+
     use crate::sys::windows::Sandbox;
     use crate::sys::SandboxImpl;
     use crate::Opts;
 
-    #[test]
-    fn hello() {
-        assert_eq!(1 + 1, 2);
-    }
-
-    /**
-     * 启动记事本
-     */
-    #[test]
-    fn notepad() {
-        let mut opts: Opts = Opts::default();
-        opts.command.push("C:/Windows/notepad.exe".to_string());
-        let status = unsafe { Sandbox::with_opts(opts).run().unwrap() };
-        assert_eq!(status.status, 0);
-        assert_eq!(status.exit_code, 0);
-        assert_eq!(status.signal, 0);
-    }
-
     /**
      * 执行不存在的可执行文件
      */
     #[test]
     #[should_panic]
     fn not_found() {
         let mut opts: Opts = Opts::default();
-        opts.command.push("Z:/not-found.exe".to_string());
+        opts.command
+            .push("./tests/windows/not-found.exe".to_string());
         unsafe {
             Sandbox::with_opts(opts).run().unwrap();
         }
@@ -264,10 +268,76 @@ mod tests {
     #[test]
     fn time_limit() {
         let mut opts: Opts = Opts::default();
-        opts.command.push("C:/Windows/notepad.exe".to_string());
+        opts.command
+            .push("./tests/windows/sleep/sleep.exe".to_string());
         opts.time_limit = Some(1000);
         let status = unsafe { Sandbox::with_opts(opts).run().unwrap() };
         assert!(status.time_used >= 1000);
         assert!(status.time_used < 2000);
     }
+
+    /**
+     * 测试 stdout
+     */
+    #[test]
+    fn output() {
+        let filename = "./output.txt";
+        let mut opts: Opts = Opts::default();
+        opts.command
+            .push("./tests/windows/output/output.exe".to_string());
+        opts.output = Option::from(filename.to_string());
+        unsafe { Sandbox::with_opts(opts).run().unwrap() };
+
+        if let Ok(content) = fs::read_to_string(filename) {
+            assert_eq!(content.trim(), "Hello World!");
+            fs::remove_file(filename).unwrap()
+        } else {
+            assert!(false)
+        }
+    }
+
+    /**
+     * 测试 stderr
+     */
+    #[test]
+    fn stderr() {
+        let filename = "./stderr.txt";
+        let mut opts: Opts = Opts::default();
+        opts.command
+            .push("./tests/windows/stderr/stderr.exe".to_string());
+        opts.error = Option::from(filename.to_string());
+        unsafe { Sandbox::with_opts(opts).run().unwrap() };
+
+        if let Ok(content) = fs::read_to_string(filename) {
+            assert_eq!("Hello World!", content.trim());
+            fs::remove_file(filename).unwrap()
+        } else {
+            assert!(false)
+        }
+    }
+
+    /**
+     * 测试 stdin
+     */
+    #[test]
+    fn stdin() {
+        let filename = "./stdin.txt";
+        let out_filename = "./stdin-stdout.txt";
+        let content = "Hello Stdin!";
+        fs::write(filename, content).unwrap();
+        let mut opts: Opts = Opts::default();
+        opts.command
+            .push("./tests/windows/stdin/stdin.exe".to_string());
+        opts.input = Option::from(filename.to_string());
+        opts.output = Option::from(out_filename.to_string());
+        unsafe { Sandbox::with_opts(opts).run().unwrap() };
+
+        if let Ok(c) = fs::read_to_string(out_filename) {
+            assert_eq!(c.trim(), content.trim());
+            fs::remove_file(filename).unwrap();
+            fs::remove_file(out_filename).unwrap()
+        } else {
+            assert!(false)
+        }
+    }
 }

src/sys/windows/utils.rs

@@ -1,12 +1,16 @@
 pub(crate) mod utils {
-    use crate::error::Error::WinError;
-    use crate::error::Result;
     use std::mem::size_of;
     use std::ptr;
+
     use windows::core::{PCSTR, PSTR};
-    use windows::Win32::Foundation::{GENERIC_WRITE, HANDLE, TRUE};
+    use windows::Win32::Foundation::{GENERIC_READ, GENERIC_WRITE, HANDLE, TRUE};
     use windows::Win32::Security::SECURITY_ATTRIBUTES;
-    use windows::Win32::Storage::FileSystem::{CreateFileA, CREATE_NEW, FILE_ATTRIBUTE_NORMAL};
+    use windows::Win32::Storage::FileSystem::{
+        CreateFileA, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, OPEN_EXISTING,
+    };
+
+    use crate::error::Error::WinError;
+    use crate::error::Result;
 
     pub unsafe fn string_to_pcstr(string: &mut String) -> PCSTR {
         string.push('\0');
@@ -18,30 +22,34 @@ pub(crate) mod utils {
         PSTR(string.as_mut_ptr())
     }
 
-    pub unsafe fn handle_from_file(string: &String) -> Result<HANDLE> {
+    pub unsafe fn handle_from_file(string: &String, wr: char) -> Result<HANDLE> {
         let mut string = string.clone();
-
         let sa = SECURITY_ATTRIBUTES {
             nLength: size_of::<SECURITY_ATTRIBUTES>() as u32,
             lpSecurityDescriptor: ptr::null_mut(),
-            bInheritHandle: TRUE,  // 指明这个 handle 需要被子进程继承
+            bInheritHandle: TRUE, // 指明这个 handle 需要被子进程继承
+        };
+        let mode = if wr == 'w' {
+            GENERIC_WRITE
+        } else {
+            GENERIC_READ
+        };
+        let exist = if wr == 'w' {
+            CREATE_ALWAYS
+        } else {
+            OPEN_EXISTING
         };
-
         return match CreateFileA(
             string_to_pcstr(&mut string),
-            GENERIC_WRITE.0,
+            mode.0,
             Default::default(),
             Some(&sa),
-            CREATE_NEW,
+            exist,
             FILE_ATTRIBUTE_NORMAL,
             HANDLE::default(),
         ) {
-            Ok(h_file) => {
-                Ok(h_file)
-            }
-            Err(e) => {
-                Err(WinError(String::from(file!()), line!(), e))
-            }
-        }
+            Ok(h_file) => Ok(h_file),
+            Err(e) => Err(WinError(String::from(string), line!(), e)),
+        };
     }
 }

tests/windows/output/output.rs

@@ -0,0 +1,3 @@
+fn main() {
+    println!("Hello World!");
+}

tests/windows/sleep/sleep.rs

@@ -0,0 +1,7 @@
+use std::thread;
+use std::time::Duration;
+
+fn main() {
+    let sleep_duration = Duration::from_secs(30);
+    thread::sleep(sleep_duration);
+}

tests/windows/stderr/stderr.rs

@@ -0,0 +1,3 @@
+fn main() {
+    eprintln!("Hello World!");
+}

tests/windows/stdin/stdin.rs

@@ -0,0 +1,7 @@
+use std::io;
+
+fn main() {
+    let mut input = String::new();
+    io::stdin().read_line(&mut input).unwrap();
+    println!("{}", input);
+}