do_memory_mcp/sandbox/
isolation.rs1#![allow(unsafe_code)]
18
19use anyhow::Result;
20use std::process::Command;
21#[cfg(unix)]
22use std::process::Stdio;
23#[cfg(unix)]
24use tracing::debug;
25#[cfg(not(unix))]
26use tracing::warn;
27
28#[derive(Debug, Clone)]
30pub struct IsolationConfig {
31 pub drop_to_uid: Option<u32>,
33 pub drop_to_gid: Option<u32>,
35 pub max_memory_bytes: Option<usize>,
37 pub max_cpu_seconds: Option<u64>,
39 pub max_processes: Option<usize>,
41}
42
43impl Default for IsolationConfig {
44 fn default() -> Self {
45 Self {
46 drop_to_uid: None,
47 drop_to_gid: None,
48 max_memory_bytes: Some(128 * 1024 * 1024), max_cpu_seconds: Some(5), max_processes: Some(1), }
52 }
53}
54
55pub fn apply_isolation(
57 #[cfg_attr(not(unix), allow(unused_mut))] mut cmd: Command,
58 config: &IsolationConfig,
59) -> Result<Command> {
60 #[cfg(unix)]
62 {
63 use std::os::unix::process::CommandExt;
64
65 let mut ulimit_args = Vec::new();
67
68 if let Some(max_mem) = config.max_memory_bytes {
70 let max_mem_kb = max_mem / 1024;
71 ulimit_args.push(format!("-v {}", max_mem_kb));
72 }
73
74 if let Some(max_cpu) = config.max_cpu_seconds {
76 ulimit_args.push(format!("-t {}", max_cpu));
77 }
78
79 if let Some(max_proc) = config.max_processes {
81 ulimit_args.push(format!("-u {}", max_proc));
82 }
83
84 ulimit_args.push("-f 0".to_string()); ulimit_args.push("-c 0".to_string()); debug!("Applying ulimit restrictions: {:?}", ulimit_args);
91
92 if !ulimit_args.is_empty() {
94 let program = cmd.get_program().to_string_lossy().to_string();
96 let args: Vec<String> = cmd
97 .get_args()
98 .map(|s| s.to_string_lossy().to_string())
99 .collect();
100
101 let mut wrapped = Command::new("sh");
103 wrapped.arg("-c");
104
105 let ulimit_cmd = ulimit_args.join("; ulimit ");
107 let exec_cmd = format!(
108 "{} {}",
109 program,
110 args.iter()
111 .map(|a| shell_escape(a))
112 .collect::<Vec<_>>()
113 .join(" ")
114 );
115 let full_cmd = format!("ulimit {}; {}", ulimit_cmd, exec_cmd);
116
117 wrapped.arg(full_cmd);
118
119 wrapped.stdin(Stdio::null());
121 wrapped.stdout(Stdio::piped());
122 wrapped.stderr(Stdio::piped());
123
124 cmd = wrapped;
125 }
126
127 if let Some(uid) = config.drop_to_uid {
129 debug!("Dropping privileges to UID: {}", uid);
130
131 let _gid = config.drop_to_gid;
133
134 unsafe {
141 cmd.pre_exec(move || {
142 #[cfg(target_os = "linux")]
144 {
145 use libc::{setgid, setuid};
146
147 if let Some(gid_val) = _gid {
149 if setgid(gid_val) != 0 {
150 return Err(std::io::Error::last_os_error());
151 }
152 }
153
154 if setuid(uid) != 0 {
156 return Err(std::io::Error::last_os_error());
157 }
158 }
159
160 Ok(())
161 });
162 }
163 }
164 }
165
166 #[cfg(not(unix))]
168 {
169 warn!("Process isolation not fully supported on this platform");
170 let _ = config; }
172
173 Ok(cmd)
174}
175
176#[cfg(unix)]
178fn shell_escape(arg: &str) -> String {
179 format!("'{}'", arg.replace('\'', "'\\''"))
181}
182
183pub fn is_running_as_root() -> bool {
185 #[cfg(unix)]
186 {
187 unsafe { libc::geteuid() == 0 }
191 }
192
193 #[cfg(not(unix))]
194 {
195 false
197 }
198}
199
200pub fn current_uid() -> Option<u32> {
202 #[cfg(unix)]
203 {
204 Some(unsafe { libc::getuid() })
207 }
208
209 #[cfg(not(unix))]
210 {
211 None
212 }
213}
214
215pub fn current_gid() -> Option<u32> {
217 #[cfg(unix)]
218 {
219 Some(unsafe { libc::getgid() })
222 }
223
224 #[cfg(not(unix))]
225 {
226 None
227 }
228}
229
230pub fn recommend_safe_uid() -> Option<(u32, u32)> {
232 if is_running_as_root() {
233 Some((65534, 65534))
235 } else {
236 None
238 }
239}
240
241#[cfg(test)]
242mod tests {
243 use super::*;
244
245 #[test]
246 fn test_isolation_config_default() {
247 let config = IsolationConfig::default();
248 assert!(config.max_memory_bytes.is_some());
249 assert!(config.max_cpu_seconds.is_some());
250 assert_eq!(config.max_processes, Some(1));
251 }
252
253 #[test]
254 #[cfg(unix)]
255 fn test_shell_escape() {
256 assert_eq!(shell_escape("simple"), "'simple'");
257 assert_eq!(shell_escape("with spaces"), "'with spaces'");
258 assert_eq!(shell_escape("with'quote"), "'with'\\''quote'");
259 assert_eq!(
260 shell_escape("complex'test'string"),
261 "'complex'\\''test'\\''string'"
262 );
263 }
264
265 #[test]
266 fn test_current_uid_gid() {
267 #[cfg(unix)]
268 {
269 assert!(current_uid().is_some());
271 assert!(current_gid().is_some());
272 }
273
274 #[cfg(not(unix))]
275 {
276 assert!(current_uid().is_none());
278 assert!(current_gid().is_none());
279 }
280 }
281
282 #[test]
283 fn test_recommend_safe_uid() {
284 let recommendation = recommend_safe_uid();
285 if is_running_as_root() {
286 assert!(recommendation.is_some());
287 assert_eq!(recommendation.unwrap(), (65534, 65534));
288 } else {
289 assert!(recommendation.is_none());
290 }
291 }
292
293 #[test]
294 fn test_apply_isolation_basic() {
295 let config = IsolationConfig::default();
296 let cmd = Command::new("echo");
297
298 let result = apply_isolation(cmd, &config);
299 assert!(result.is_ok());
300 }
301}