mirror of
https://github.com/MeiK2333/river.git
synced 2025-11-04 14:49:40 +08:00
Update
This commit is contained in:
parent
6e6b50a745
commit
5371ff4ed5
39
src/error.rs
39
src/error.rs
@ -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(),
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
54
src/main.rs
54
src/main.rs
@ -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(())
|
||||||
}
|
}
|
||||||
|
|||||||
@ -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");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user