From 5c413d09ae992d0c34020612906ae83ba3184ecd Mon Sep 17 00:00:00 2001 From: MeiK Date: Mon, 11 Mar 2024 20:25:52 +0800 Subject: [PATCH] =?UTF-8?q?=E6=94=AF=E6=8C=81=E6=B5=81=E9=87=8D=E5=AE=9A?= =?UTF-8?q?=E5=90=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 1 + README.md | 10 ++- src/sys/windows/mod.rs | 120 ++++++++++++++++++++++++++------- src/sys/windows/utils.rs | 42 +++++++----- tests/windows/hello.ps1 | 1 - tests/windows/output.ps1 | 1 - tests/windows/output/output.rs | 3 + tests/windows/sleep.ps1 | 1 - tests/windows/sleep/sleep.rs | 7 ++ tests/windows/stderr/stderr.rs | 3 + tests/windows/stdin/stdin.rs | 7 ++ 11 files changed, 150 insertions(+), 46 deletions(-) delete mode 100644 tests/windows/hello.ps1 delete mode 100644 tests/windows/output.ps1 create mode 100644 tests/windows/output/output.rs delete mode 100644 tests/windows/sleep.ps1 create mode 100644 tests/windows/sleep/sleep.rs create mode 100644 tests/windows/stderr/stderr.rs create mode 100644 tests/windows/stdin/stdin.rs diff --git a/.gitignore b/.gitignore index e837a2f..38a694b 100644 --- a/.gitignore +++ b/.gitignore @@ -14,3 +14,4 @@ Cargo.lock *.pdb .idea .vscode +*.exe diff --git a/README.md b/README.md index fb17777..8af4b6c 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,7 @@ | 特性 | Linux | Windows | macOS | |-------------|-------|---------|-------| | 执行指定命令 | ~ | √ | ~ | -| 流重定向 | ~ | ~ | ~ | +| 流重定向 | ~ | √ | ~ | | 运行时间统计 | ~ | √ | ~ | | 运行 CPU 时间统计 | ~ | √ | ~ | | 运行内存统计 | ~ | √ | ~ | @@ -25,3 +25,11 @@ | 平滑退出 | ~ | ~ | ~ | **注意：** Windows 平台下运行 CPU 时间限制与运行内存限制不能保证精确，请不要以此为基准进行判断。 + +## 测试 + +```bash +cargo test -- --test-threads=1 +``` + +测试涉及文件操作，建议顺序执行测试用例（并发限制为 1） diff --git a/src/sys/windows/mod.rs b/src/sys/windows/mod.rs index 6d84231..702ccf4 100644 --- a/src/sys/windows/mod.rs +++ b/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 { + unsafe fn wait_it( + &mut self, + info: &STARTUPINFOA, + information: &PROCESS_INFORMATION, + ) -> Result { 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,34 +236,18 @@ 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); - } - /** * 执行不存在的可执行文件 */ @@ -252,7 +255,8 @@ mod tests { #[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) + } + } } diff --git a/src/sys/windows/utils.rs b/src/sys/windows/utils.rs index 72ba234..ce59759 100644 --- a/src/sys/windows/utils.rs +++ b/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 { + pub unsafe fn handle_from_file(string: &String, wr: char) -> Result { let mut string = string.clone(); - let sa = SECURITY_ATTRIBUTES { nLength: size_of::() 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)), + }; } } diff --git a/tests/windows/hello.ps1 b/tests/windows/hello.ps1 deleted file mode 100644 index 06afcd4..0000000 --- a/tests/windows/hello.ps1 +++ /dev/null @@ -1 +0,0 @@ -../../target/debug/river.exe -- C:/WINDOWS/System32/WindowsPowerShell/v1.0/powershell.exe -- echo hello \ No newline at end of file diff --git a/tests/windows/output.ps1 b/tests/windows/output.ps1 deleted file mode 100644 index fd487fb..0000000 --- a/tests/windows/output.ps1 +++ /dev/null @@ -1 +0,0 @@ -../../target/debug/river.exe -vvvv -o output.txt -- C:/WINDOWS/System32/WindowsPowerShell/v1.0/powershell.exe -- echo hello \ No newline at end of file diff --git a/tests/windows/output/output.rs b/tests/windows/output/output.rs new file mode 100644 index 0000000..47ad8c6 --- /dev/null +++ b/tests/windows/output/output.rs @@ -0,0 +1,3 @@ +fn main() { + println!("Hello World!"); +} diff --git a/tests/windows/sleep.ps1 b/tests/windows/sleep.ps1 deleted file mode 100644 index 836d4c5..0000000 --- a/tests/windows/sleep.ps1 +++ /dev/null @@ -1 +0,0 @@ -../../target/debug/river.exe -t 1000 -- C:/WINDOWS/System32/WindowsPowerShell/v1.0/powershell.exe -- sleep 3 \ No newline at end of file diff --git a/tests/windows/sleep/sleep.rs b/tests/windows/sleep/sleep.rs new file mode 100644 index 0000000..e6da1e8 --- /dev/null +++ b/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); +} diff --git a/tests/windows/stderr/stderr.rs b/tests/windows/stderr/stderr.rs new file mode 100644 index 0000000..724666e --- /dev/null +++ b/tests/windows/stderr/stderr.rs @@ -0,0 +1,3 @@ +fn main() { + eprintln!("Hello World!"); +} diff --git a/tests/windows/stdin/stdin.rs b/tests/windows/stdin/stdin.rs new file mode 100644 index 0000000..d3c4bb6 --- /dev/null +++ b/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); +}