This commit is contained in:
MeiK2333 2021-01-11 17:50:57 +08:00
parent 6e6b50a745
commit 5371ff4ed5
3 changed files with 103 additions and 70 deletions

View File

@ -10,8 +10,9 @@ use std::result;
#[allow(dead_code)] #[allow(dead_code)]
#[derive(Debug)] #[derive(Debug)]
pub enum Error { pub enum Error {
IOError(io::Error), IOError(io::Error),
StringToCStringError(NulError), StringToCStringError(NulError),
ParseIntError(std::num::ParseIntError),
} }
pub type Result<T> = result::Result<T, Error>; pub type Result<T> = result::Result<T, Error>;
@ -19,30 +20,30 @@ pub type Result<T> = result::Result<T, Error>;
// 创建一个简单的包装 // 创建一个简单的包装
#[macro_export] #[macro_export]
macro_rules! try_io { macro_rules! try_io {
($expression:expr) => { ($expression:expr) => {
match $expression { match $expression {
Ok(val) => val, Ok(val) => val,
Err(e) => return Err(Error::IOError(e)), Err(e) => return Err(Error::IOError(e)),
};
}; };
};
} }
impl fmt::Display for Error { impl fmt::Display for Error {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match *self { match *self {
Error::IOError(ref e) => write!(f, "IOError: {}", errno_str(e.raw_os_error())), Error::IOError(ref e) => write!(f, "IOError: {}", errno_str(e.raw_os_error())),
_ => write!(f, "{:?}", self), _ => write!(f, "{:?}", self),
}
} }
}
} }
pub fn errno_str(errno: Option<i32>) -> String { pub fn errno_str(errno: Option<i32>) -> String {
match errno { match errno {
Some(no) => { Some(no) => {
let stre = unsafe { strerror(no) }; let stre = unsafe { strerror(no) };
let c_str: &CStr = unsafe { CStr::from_ptr(stre) }; let c_str: &CStr = unsafe { CStr::from_ptr(stre) };
c_str.to_str().unwrap().to_string() c_str.to_str().unwrap().to_string()
}
_ => "Unknown Error!".to_string(),
} }
_ => "Unknown Error!".to_string(),
}
} }

View File

@ -19,7 +19,7 @@ mod process;
mod seccomp; mod seccomp;
pub mod river { pub mod river {
tonic::include_proto!("river"); tonic::include_proto!("river");
} }
#[derive(Debug, Default)] #[derive(Debug, Default)]
@ -27,43 +27,43 @@ pub struct RiverService {}
#[tonic::async_trait] #[tonic::async_trait]
impl River for RiverService { impl River for RiverService {
type JudgeStream = type JudgeStream =
Pin<Box<dyn Stream<Item = Result<JudgeResponse, Status>> + Send + Sync + 'static>>; Pin<Box<dyn Stream<Item = Result<JudgeResponse, Status>> + Send + Sync + 'static>>;
async fn judge( async fn judge(
&self, &self,
request: Request<tonic::Streaming<JudgeRequest>>, request: Request<tonic::Streaming<JudgeRequest>>,
) -> Result<Response<Self::JudgeStream>, Status> { ) -> Result<Response<Self::JudgeStream>, Status> {
let mut _stream = request.into_inner(); let mut _stream = request.into_inner();
let output = async_stream::try_stream! { let output = async_stream::try_stream! {
yield JudgeResponse { yield JudgeResponse {
state: Some(river::judge_response::State::Status(river::JudgeStatus::Pending as i32)) state: Some(river::judge_response::State::Status(river::JudgeStatus::Pending as i32))
};
}; };
};
Ok(Response::new(Box::pin(output) as Self::JudgeStream)) Ok(Response::new(Box::pin(output) as Self::JudgeStream))
} }
} }
#[tokio::main] #[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> { async fn main() -> Result<(), Box<dyn std::error::Error>> {
let env = Env::default() let env = Env::default()
.filter_or("LOG_LEVEL", "debug,h2=info,hyper=info") .filter_or("LOG_LEVEL", "debug,h2=info,hyper=info")
.write_style_or("LOG_STYLE", "always"); .write_style_or("LOG_STYLE", "always");
env_logger::init_from_env(env); env_logger::init_from_env(env);
let addr = "0.0.0.0:4003".parse()?; let addr = "0.0.0.0:4003".parse()?;
let river = RiverService::default(); let river = RiverService::default();
info!("listen on: {}", addr); info!("listen on: {}", addr);
Server::builder() Server::builder()
.concurrency_limit_per_connection(5) .concurrency_limit_per_connection(5)
.add_service(RiverServer::new(river)) .add_service(RiverServer::new(river))
.serve(addr) .serve(addr)
.await?; .await?;
Ok(()) Ok(())
} }

View File

@ -1,6 +1,6 @@
use crate::cgroup::CGroupSet; use crate::cgroup::CGroupSet;
use crate::config::{STDERR_FILENAME, STDIN_FILENAME, STDOUT_FILENAME}; use crate::config::{STDERR_FILENAME, STDIN_FILENAME, STDOUT_FILENAME};
use crate::error::{errno_str, Result}; use crate::error::{errno_str, Error, Result};
use crate::exec_args::ExecArgs; use crate::exec_args::ExecArgs;
use crate::seccomp; use crate::seccomp;
use libc; use libc;
@ -38,7 +38,7 @@ macro_rules! c_str_ptr {
} }
#[derive(Clone)] #[derive(Clone)]
struct Process { pub struct Process {
pub cmd: String, pub cmd: String,
pub time_limit: i32, pub time_limit: i32,
pub memory_limit: i32, pub memory_limit: i32,
@ -46,14 +46,15 @@ struct Process {
} }
#[derive(Clone)] #[derive(Clone)]
struct RunnerStatus { pub struct RunnerStatus {
pub rusage: libc::rusage, pub rusage: libc::rusage,
pub exit_code: i32, pub exit_code: i32,
pub status: i32, pub status: i32,
pub signal: i32, pub signal: i32,
pub time_used: i64, pub time_used: i64,
pub memory_used: i64,
pub real_time_used: u128, pub real_time_used: u128,
pub memory_used: i64,
pub cgroup_memory_used: i64,
pub errmsg: String, pub errmsg: String,
} }
@ -122,7 +123,6 @@ impl Drop for Runner {
} }
} }
// TODO: unwrap -> Result<Error>
impl Future for Runner { impl Future for Runner {
type Output = Result<RunnerStatus>; type Output = Result<RunnerStatus>;
@ -150,14 +150,10 @@ impl Future for Runner {
// 设置 cgroup 限制 // 设置 cgroup 限制
// 此处为父进程做的策略,所以没有与子进程的安全策略放一块 // 此处为父进程做的策略,所以没有与子进程的安全策略放一块
runner.cgroup_set.apply(pid).unwrap(); runner.cgroup_set.apply(pid).unwrap();
runner runner.cgroup_set.memory.set(
.cgroup_set "memory.limit_in_bytes",
.memory &format!("{}", runner.process.memory_limit as i64 * 1024),
.set( )?;
"memory.limit_in_bytes",
&format!("{}", runner.process.memory_limit as i64 * 1024),
)
.unwrap();
let tx = runner.tx.clone(); let tx = runner.tx.clone();
// 因为 wait 会阻塞等待结果,因此此处使用 thread::spawn 来 wait以防止主流程被阻塞 // 因为 wait 会阻塞等待结果,因此此处使用 thread::spawn 来 wait以防止主流程被阻塞
@ -166,22 +162,25 @@ impl Future for Runner {
// 子流程运行结束后通知主流程 // 子流程运行结束后通知主流程
let status_tx = tx.lock().unwrap(); let status_tx = tx.lock().unwrap();
status_tx.send(status).unwrap(); status_tx.send(status).unwrap();
debug!("waker");
waker.wake(); waker.wake();
}); });
return Poll::Pending; return Poll::Pending;
} else { } else {
let status = runner.rx.lock().unwrap().recv().unwrap(); let mut status = runner.rx.lock().unwrap().recv().unwrap();
// TODO: 将 CGroup 的结果合并到 status 中 let mem_used = match runner
let mem_used = runner
.cgroup_set .cgroup_set
.memory .memory
.get("memory.max_usage_in_bytes") .get("memory.max_usage_in_bytes")?
.unwrap()
.trim() .trim()
.parse::<i32>() .parse::<i64>()
.unwrap(); {
debug!("cgroup memory used: {} KiB", mem_used / 1024); Ok(val) => val,
Err(e) => return Poll::Ready(Err(Error::ParseIntError(e))),
};
status.cgroup_memory_used = mem_used / 1024;
debug!("cpu time used: {} ms", status.time_used);
debug!("real time used: {} ms", status.real_time_used);
debug!("cgroup memory used: {} KiB", status.cgroup_memory_used);
debug!("rusage memory used: {} KiB", status.memory_used); debug!("rusage memory used: {} KiB", status.memory_used);
return Poll::Ready(Ok(status)); return Poll::Ready(Ok(status));
} }
@ -223,6 +222,7 @@ impl Runner {
RunnerStatus { RunnerStatus {
errmsg: String::from(""), errmsg: String::from(""),
memory_used: memory_used, memory_used: memory_used,
cgroup_memory_used: -1,
time_used: time_used, time_used: time_used,
real_time_used: real_time_used, real_time_used: real_time_used,
exit_code: exit_code, exit_code: exit_code,
@ -239,7 +239,7 @@ extern "C" fn runit(process: *mut libc::c_void) -> i32 {
let exec_args = ExecArgs::build(&process.cmd).unwrap(); let exec_args = ExecArgs::build(&process.cmd).unwrap();
unsafe { unsafe {
security(&process); security(&process);
// fd_dup(); fd_dup();
syscall_or_panic!(libc::execve( syscall_or_panic!(libc::execve(
exec_args.pathname, exec_args.pathname,
exec_args.argv, exec_args.argv,
@ -453,6 +453,7 @@ fn judge_system_error(errmsg: String) -> RunnerStatus {
signal: -1, signal: -1,
time_used: -1, time_used: -1,
memory_used: -1, memory_used: -1,
cgroup_memory_used: -1,
real_time_used: 0, real_time_used: 0,
errmsg: errmsg, errmsg: errmsg,
} }
@ -463,7 +464,7 @@ mod tests {
use super::*; use super::*;
use tempfile::tempdir_in; use tempfile::tempdir_in;
#[tokio::test] #[tokio::test]
async fn test1() { async fn test_echo() {
let pwd = tempdir_in("/tmp").unwrap(); let pwd = tempdir_in("/tmp").unwrap();
println!("{:?}", pwd); println!("{:?}", pwd);
let process = Process::new( let process = Process::new(
@ -475,4 +476,35 @@ mod tests {
let result = Runner::from(process).unwrap().await.unwrap(); let result = Runner::from(process).unwrap().await.unwrap();
assert_eq!(result.exit_code, 0); assert_eq!(result.exit_code, 0);
} }
#[tokio::test]
async fn test_sleep() {
let pwd = tempdir_in("/tmp").unwrap();
println!("{:?}", pwd);
let process = Process::new(
String::from("/bin/sleep 1"),
2000,
65535,
pwd.path().to_path_buf(),
);
let result = Runner::from(process).unwrap().await.unwrap();
assert!(result.real_time_used > 1000);
assert!(result.real_time_used < 1500);
}
#[tokio::test]
async fn test_output() {
let pwd = tempdir_in("/tmp").unwrap();
println!("{:?}", pwd);
let process = Process::new(
String::from("/bin/echo Hello World!"),
1000,
65535,
pwd.path().to_path_buf(),
);
let runner = Runner::from(process).unwrap();
let _ = runner.await.unwrap();
let out = std::fs::read_to_string(pwd.path().join(STDOUT_FILENAME)).unwrap();
assert_eq!(out, "Hello World!\n");
}
} }