river/src/process.rs

365 lines
12 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

use super::error::{errno_str, Error, Result};
use libc;
use std::env;
use std::ffi::CString;
use std::fs;
use std::future::Future;
use std::io;
use std::mem;
use std::path::PathBuf;
use std::pin::Pin;
use std::ptr;
use std::sync::mpsc;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use std::thread;
use std::time::SystemTime;
#[derive(Clone)]
pub struct Process {
pub pid: i32,
pub workdir: PathBuf,
pub time_limit: i32,
pub memory_limit: i32,
pub stdin_file: Option<PathBuf>,
pub cmd: String,
tx: Arc<Mutex<mpsc::Sender<ProcessStatus>>>,
rx: Arc<Mutex<mpsc::Receiver<ProcessStatus>>>,
}
impl Process {
pub fn new() -> Process {
let (tx, rx) = mpsc::channel();
Process {
pid: -1,
time_limit: -1,
memory_limit: -1,
stdin_file: None,
cmd: "".to_string(),
workdir: PathBuf::from(""),
tx: Arc::new(Mutex::new(tx)),
rx: Arc::new(Mutex::new(rx)),
}
}
pub fn set_pid(&mut self, pid: i32) {
self.pid = pid;
}
}
impl Future for Process {
type Output = Result<ProcessStatus>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<ProcessStatus>> {
let process = Pin::into_inner(self);
// 如果 pid == -1则说明子进程还没有运行开始进程
if process.pid == -1 {
let (tx, rx) = mpsc::channel();
let process_clone = process.clone();
let waker = cx.waker().clone();
thread::spawn(move || {
let pid;
unsafe {
pid = libc::fork();
}
if pid == 0 {
run(process_clone);
} else if pid > 0 {
tx.send(pid).unwrap();
let status = wait(pid, process_clone.workdir);
let status_tx = process_clone.tx.lock().unwrap();
status_tx.send(status).unwrap();
waker.wake();
} else {
panic!("How dare you!");
}
});
// 等待子线程启动子进程并返回 pid
let pid = match rx.recv() {
Ok(val) => val,
Err(_) => return Poll::Ready(Err(Error::ChannelRecvError)),
};
process.set_pid(pid);
return Poll::Pending;
} else {
// 再次进入 poll说明子进程已经结束通知了 wake
// 此时 channel 应该是有数据的
let status = match process.rx.lock() {
Ok(rx) => match rx.recv() {
Ok(val) => val,
Err(_) => return Poll::Ready(Err(Error::ChannelRecvError)),
},
Err(_) => return Poll::Ready(Err(Error::ChannelRecvError)),
};
return Poll::Ready(Ok(status));
}
}
}
pub struct ExecArgs {
pub pathname: *const libc::c_char,
pub argv: *const *const libc::c_char,
pub envp: *const *const libc::c_char,
}
impl ExecArgs {
fn build(cmd: &String) -> Result<ExecArgs> {
let cmds: Vec<&str> = cmd.split_whitespace().collect();
let pathname = cmds[0].clone();
let pathname_str = match CString::new(pathname) {
Ok(value) => value,
Err(err) => return Err(Error::StringToCStringError(err)),
};
let pathname = pathname_str.as_ptr();
let mut argv_vec: Vec<*const libc::c_char> = vec![];
for item in cmds.iter() {
let cstr = match CString::new(item.clone()) {
Ok(value) => value,
Err(err) => return Err(Error::StringToCStringError(err)),
};
let cptr = cstr.as_ptr();
// 需要使用 mem::forget 来标记
// 否则在此次循环结束后cstr 就会被回收,后续 exec 函数无法通过指针获取到字符串内容
mem::forget(cstr);
argv_vec.push(cptr);
}
// argv 与 envp 的参数需要使用 NULL 来标记结束
argv_vec.push(ptr::null());
let argv: *const *const libc::c_char = argv_vec.as_ptr() as *const *const libc::c_char;
// env 环境变量传递当前进程环境变量
let mut envp_vec: Vec<*const libc::c_char> = vec![];
for (key, value) in env::vars_os() {
let mut key = match key.to_str() {
Some(val) => val.to_string(),
None => return Err(Error::OsStringToStringError(key)),
};
let value = match value.to_str() {
Some(val) => val.to_string(),
None => return Err(Error::OsStringToStringError(value)),
};
key.push_str("=");
key.push_str(&value);
let cstr = match CString::new(key) {
Ok(value) => value,
Err(err) => return Err(Error::StringToCStringError(err)),
};
let cptr = cstr.as_ptr();
// 需要使用 mem::forget 来标记
// 否则在此次循环结束后cstr 就会被回收,后续 exec 函数无法通过指针获取到字符串内容
mem::forget(cstr);
envp_vec.push(cptr);
}
envp_vec.push(ptr::null());
let envp = envp_vec.as_ptr() as *const *const libc::c_char;
mem::forget(pathname_str);
mem::forget(argv_vec);
mem::forget(envp_vec);
Ok(ExecArgs {
pathname,
argv,
envp,
})
}
}
impl Drop for ExecArgs {
fn drop(&mut self) {
// TODO: 将不安全的指针类型转换回内置类型,以便由 Rust 自动回收资源
// TODO: 优先级较低,因为目前只在子进程里进行这个操作,且操作后会很快 exec操作系统会回收这些内存
println!("Dropping!");
}
}
fn run(process: Process) {
// 子进程里崩溃也无法返回,崩溃就直接崩溃了
let exec_args = ExecArgs::build(&process.cmd).unwrap();
// 修改工作目录
env::set_current_dir(process.workdir).unwrap();
let mut rl = libc::rlimit {
rlim_cur: 0,
rlim_max: 0,
};
// 实际运行时间限制设置为 CPU 时间 + 2 * 2尽量在防止恶意代码占用评测资源的情况下给正常用户的代码最宽松的环境
let rt = libc::itimerval {
it_interval: libc::timeval {
tv_sec: 0,
tv_usec: 0,
},
it_value: libc::timeval {
tv_sec: i64::from(process.time_limit / 1000 + 2) * 2,
tv_usec: 0,
},
};
unsafe {
// 重定向文件描述符
if let Some(file) = process.stdin_file {
let filename = file.to_str().unwrap();
dup(&filename, libc::STDIN_FILENO, libc::O_RDONLY, 0o644)
}
dup(
"stdout.txt",
libc::STDOUT_FILENO,
libc::O_CREAT | libc::O_RDWR,
0o644,
);
dup(
"stderr.txt",
libc::STDERR_FILENO,
libc::O_CREAT | libc::O_RDWR,
0o644,
);
// 墙上时钟限制
if libc::setitimer(libc::ITIMER_REAL, &rt, ptr::null_mut()) == -1 {
eprintln!("setitimer failure!");
eprintln!("{:?}", io::Error::last_os_error().raw_os_error());
panic!("How dare you!");
}
// CPU 时间限制,粒度为 S
rl.rlim_cur = (process.time_limit / 1000 + 1) as u64;
rl.rlim_max = rl.rlim_cur + 1;
if libc::setrlimit(libc::RLIMIT_CPU, &rl) != 0 {
eprintln!("setrlimit failure!");
eprintln!("{:?}", io::Error::last_os_error().raw_os_error());
panic!("How dare you!");
}
// 设置内存限制
rl.rlim_cur = (process.memory_limit * 1024) as u64;
rl.rlim_max = rl.rlim_cur + 1024;
if libc::setrlimit(libc::RLIMIT_DATA, &rl) != 0 {
eprintln!("setrlimit failure!");
eprintln!("{:?}", io::Error::last_os_error().raw_os_error());
panic!("How dare you!");
}
if libc::setrlimit(libc::RLIMIT_AS, &rl) != 0 {
eprintln!("setrlimit failure!");
eprintln!("{:?}", io::Error::last_os_error().raw_os_error());
panic!("How dare you!");
}
libc::execve(exec_args.pathname, exec_args.argv, exec_args.envp);
}
panic!("How dare you!");
}
unsafe fn dup(filename: &str, to: libc::c_int, flag: libc::c_int, mode: libc::c_int) {
let filename_str = CString::new(filename).unwrap();
let filename = filename_str.as_ptr();
let fd = libc::open(filename, flag, mode);
if fd < 0 {
let err = io::Error::last_os_error().raw_os_error();
eprintln!("open failure!");
eprintln!("{:?}", io::Error::last_os_error().raw_os_error());
panic!(errno_str(err));
}
if libc::dup2(fd, to) < 0 {
let err = io::Error::last_os_error().raw_os_error();
eprintln!("dup2 failure!");
eprintln!("{:?}", io::Error::last_os_error().raw_os_error());
panic!(errno_str(err));
}
}
fn wait(pid: i32, workdir: PathBuf) -> ProcessStatus {
let start = SystemTime::now();
let mut status = 0;
let mut rusage = libc::rusage {
ru_utime: libc::timeval {
tv_sec: 0 as libc::time_t,
tv_usec: 0 as libc::suseconds_t,
},
ru_stime: libc::timeval {
tv_sec: 0 as libc::time_t,
tv_usec: 0 as libc::suseconds_t,
},
ru_maxrss: 0 as libc::c_long,
ru_ixrss: 0 as libc::c_long,
ru_idrss: 0 as libc::c_long,
ru_isrss: 0 as libc::c_long,
ru_minflt: 0 as libc::c_long,
ru_majflt: 0 as libc::c_long,
ru_nswap: 0 as libc::c_long,
ru_inblock: 0 as libc::c_long,
ru_oublock: 0 as libc::c_long,
ru_msgsnd: 0 as libc::c_long,
ru_msgrcv: 0 as libc::c_long,
ru_nsignals: 0 as libc::c_long,
ru_nvcsw: 0 as libc::c_long,
ru_nivcsw: 0 as libc::c_long,
};
unsafe {
let val = libc::wait4(pid, &mut status, 0, &mut rusage);
if val < 0 {
eprintln!("{:?}", io::Error::last_os_error().raw_os_error());
panic!("How dare you!");
}
}
let mut exit_code = 0;
let exited = unsafe { libc::WIFEXITED(status) };
if exited {
exit_code = unsafe { libc::WEXITSTATUS(status) };
}
let signal = unsafe {
if libc::WIFSIGNALED(status) {
libc::WTERMSIG(status)
} else if libc::WIFSTOPPED(status) {
libc::WSTOPSIG(status)
} else {
0
}
};
// TODO: 添加 CGroup 的量度
let time_used = rusage.ru_utime.tv_sec * 1000
+ i64::from(rusage.ru_utime.tv_usec) / 1000
+ rusage.ru_stime.tv_sec * 1000
+ i64::from(rusage.ru_stime.tv_usec) / 1000;
let memory_used = rusage.ru_maxrss;
let real_time_used = match start.elapsed() {
Ok(elapsed) => elapsed.as_millis(),
// 这种地方如果出错了,确实没有办法解决
// 只能崩溃再见了
// How dare you!
Err(_) => panic!("How dare you!"),
};
let stdout = match fs::read_to_string(workdir.join("stdout.txt")) {
Ok(val) => val,
Err(_) => panic!("How dare you!"),
};
if let Err(_) = fs::remove_file(workdir.join("stdout.txt")) {
panic!("How dare you!");
}
let stderr = match fs::read_to_string(workdir.join("stderr.txt")) {
Ok(val) => val,
Err(_) => panic!("How dare you!"),
};
if let Err(_) = fs::remove_file(workdir.join("stderr.txt")) {
panic!("How dare you!");
}
return ProcessStatus {
rusage: rusage,
exit_code: exit_code,
status: status,
signal: signal,
time_used: time_used,
memory_used: memory_used,
real_time_used: real_time_used,
stdout: stdout,
stderr: stderr,
};
}
#[derive(Clone)]
pub struct ProcessStatus {
pub rusage: libc::rusage,
pub exit_code: i32,
pub status: i32,
pub signal: i32,
pub time_used: i64,
pub memory_used: i64,
pub real_time_used: u128,
pub stdout: String,
pub stderr: String,
}