-
Notifications
You must be signed in to change notification settings - Fork 168
/
limit.rs
73 lines (64 loc) · 1.86 KB
/
limit.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
use libc;
use std::{
os::unix::{io::RawFd, process::CommandExt},
process::Command,
};
pub trait SetLimit {
fn mem_limit(&mut self, size: u64) -> &mut Self;
fn setsid(&mut self) -> &mut Self;
fn pipe_stdin(&mut self, fd: RawFd, is_stdin: bool) -> &mut Self;
//fn dup2(&mut self, src: libc::c_int, dst: libc::c_int) -> &mut Self;
//fn close_fd(&mut self, fd: libc::c_int) -> &mut Self;
}
impl SetLimit for Command {
fn mem_limit(&mut self, size: u64) -> &mut Self {
if size == 0 {
return self;
}
let func = move || {
let size = size << 20;
let mem_limit: libc::rlim_t = size;
let r = libc::rlimit {
rlim_cur: mem_limit,
rlim_max: mem_limit,
};
let r0 = libc::rlimit {
rlim_cur: 0,
rlim_max: 0,
};
unsafe {
libc::setrlimit(libc::RLIMIT_AS, &r);
// libc::setrlimit(libc::RLIMIT_DATA, &r);
libc::setrlimit(libc::RLIMIT_CORE, &r0);
};
Ok(())
};
unsafe { self.pre_exec(func) }
}
fn setsid(&mut self) -> &mut Self {
let func = move || {
unsafe {
libc::setsid();
};
Ok(())
};
unsafe { self.pre_exec(func) }
}
fn pipe_stdin(&mut self, fd: RawFd, is_stdin: bool) -> &mut Self {
if is_stdin {
let func = move || {
let ret = unsafe { libc::dup2(fd, libc::STDIN_FILENO) };
if ret < 0 {
panic!("dup2() failded");
}
unsafe {
libc::close(fd);
}
Ok(())
};
unsafe { self.pre_exec(func) }
} else {
self
}
}
}