diff --git a/src/agent/rustjail/src/container.rs b/src/agent/rustjail/src/container.rs index bc4dd09073..291f19af1f 100644 --- a/src/agent/rustjail/src/container.rs +++ b/src/agent/rustjail/src/container.rs @@ -32,6 +32,7 @@ use crate::cgroups::{DevicesCgroupInfo, Manager}; use crate::console; use crate::log_child; use crate::process::Process; +use crate::process::ProcessOperations; #[cfg(feature = "seccomp")] use crate::seccomp; use crate::selinux; @@ -261,7 +262,7 @@ pub struct LinuxContainer { pub init_process_start_time: u64, pub uid_map_path: String, pub gid_map_path: String, - pub processes: HashMap, + pub processes: HashMap, pub status: ContainerStatus, pub created: SystemTime, pub logger: Logger, @@ -933,17 +934,13 @@ impl BaseContainer for LinuxContainer { } fn processes(&self) -> Result> { - Ok(self.processes.keys().cloned().collect()) + Ok(self.processes.values().map(|p| p.pid).collect()) } fn get_process(&mut self, eid: &str) -> Result<&mut Process> { - for (_, v) in self.processes.iter_mut() { - if eid == v.exec_id.as_str() { - return Ok(v); - } - } - - Err(anyhow!("invalid eid {}", eid)) + self.processes + .get_mut(eid) + .ok_or_else(|| anyhow!("invalid eid {}", eid)) } fn stats(&self) -> Result { @@ -967,6 +964,12 @@ impl BaseContainer for LinuxContainer { async fn start(&mut self, mut p: Process) -> Result<()> { let logger = self.logger.new(o!("eid" => p.exec_id.clone())); + + // Check if exec_id is already in use to prevent collisions + if self.processes.contains_key(p.exec_id.as_str()) { + return Err(anyhow!("exec_id '{}' already exists", p.exec_id)); + } + let tty = p.tty; let fifo_file = format!("{}/{}", &self.root, EXEC_FIFO_FILENAME); info!(logger, "enter container.start!"); @@ -1235,7 +1238,7 @@ impl BaseContainer for LinuxContainer { let spec = self.config.spec.as_mut().unwrap(); update_namespaces(&self.logger, spec, p.pid)?; } - self.processes.insert(p.pid, p); + self.processes.insert(p.exec_id.clone(), p); info!(logger, "wait on child log handler"); let _ = log_handler @@ -1261,13 +1264,13 @@ impl BaseContainer for LinuxContainer { let spec = self.config.spec.as_ref().unwrap(); let st = self.oci_state()?; - for pid in self.processes.keys() { - match signal::kill(Pid::from_raw(*pid), Some(Signal::SIGKILL)) { + for process in self.processes.values() { + match signal::kill(process.pid(), Some(Signal::SIGKILL)) { Err(Errno::ESRCH) => { info!( self.logger, "kill encounters ESRCH, pid: {}, container: {}", - pid, + process.pid(), self.id.clone() ); continue; @@ -2084,10 +2087,11 @@ mod tests { #[tokio::test] async fn test_linuxcontainer_get_process() { let _ = new_linux_container_and_then(|mut c: LinuxContainer| { - c.processes.insert( - 1, - Process::new(&sl(), &oci::Process::default(), "123", true, 1, None).unwrap(), - ); + let process = + Process::new(&sl(), &oci::Process::default(), "123", true, 1, None).unwrap(); + let exec_id = process.exec_id.clone(); + c.processes.insert(exec_id, process); + let p = c.get_process("123"); assert!(p.is_ok(), "Expecting Ok, Got {:?}", p); Ok(()) diff --git a/src/agent/src/rpc.rs b/src/agent/src/rpc.rs index ba41b3da4e..92d5a316ce 100644 --- a/src/agent/src/rpc.rs +++ b/src/agent/src/rpc.rs @@ -554,7 +554,7 @@ impl AgentService { req: protocols::agent::WaitProcessRequest, ) -> Result { let cid = req.container_id; - let eid = req.exec_id; + let mut eid = req.exec_id; let mut resp = WaitProcessResponse::new(); info!( @@ -587,7 +587,7 @@ impl AgentService { .get_container(&cid) .ok_or_else(|| anyhow!("Invalid container id"))?; - let p = match ctr.processes.get_mut(&pid) { + let p = match ctr.processes.values_mut().find(|p| p.pid == pid) { Some(p) => p, None => { // Lost race, pick up exit code from channel @@ -600,6 +600,8 @@ impl AgentService { } }; + eid = p.exec_id.clone(); + // need to close all fd // ignore errors for some fd might be closed by stream p.cleanup_process_stream(); @@ -611,7 +613,7 @@ impl AgentService { let _ = s.send(p.exit_code).await; } - ctr.processes.remove(&pid); + ctr.processes.remove(&eid); Ok(resp) } @@ -2670,7 +2672,7 @@ mod tests { } linux_container .processes - .insert(exec_process_id, exec_process); + .insert(exec_process.exec_id.clone(), exec_process); sandbox.add_container(linux_container); } diff --git a/src/agent/src/sandbox.rs b/src/agent/src/sandbox.rs index b8ba9e45ef..b168fbedf9 100644 --- a/src/agent/src/sandbox.rs +++ b/src/agent/src/sandbox.rs @@ -272,8 +272,10 @@ impl Sandbox { pub fn find_process(&mut self, pid: pid_t) -> Option<&mut Process> { for (_, c) in self.containers.iter_mut() { - if let Some(p) = c.processes.get_mut(&pid) { - return Some(p); + for p in c.processes.values_mut() { + if p.pid == pid { + return Some(p); + } } } @@ -286,9 +288,11 @@ impl Sandbox { .ok_or_else(|| anyhow!(ERR_INVALID_CONTAINER_ID))?; if eid.is_empty() { + let init_pid = ctr.init_process_pid; return ctr .processes - .get_mut(&ctr.init_process_pid) + .values_mut() + .find(|p| p.pid == init_pid) .ok_or_else(|| anyhow!("cannot find init process!")); } @@ -1014,23 +1018,26 @@ mod tests { linux_container.init_process_pid = 1; linux_container.id = cid.to_string(); // add init process - linux_container.processes.insert( - 1, - Process::new(&logger, &oci::Process::default(), "1", true, 1, None).unwrap(), - ); + let mut init_process = + Process::new(&logger, &oci::Process::default(), "1", true, 1, None).unwrap(); + init_process.pid = 1; + linux_container + .processes + .insert("1".to_string(), init_process); // add exec process - linux_container.processes.insert( - 123, - Process::new( - &logger, - &oci::Process::default(), - "exec-123", - false, - 1, - None, - ) - .unwrap(), - ); + let mut exec_process = Process::new( + &logger, + &oci::Process::default(), + "exec-123", + false, + 1, + None, + ) + .unwrap(); + exec_process.pid = 123; + linux_container + .processes + .insert("exec-123".to_string(), exec_process); s.add_container(linux_container); @@ -1081,8 +1088,8 @@ mod tests { .unwrap(); // processes interally only have pids when manually set test_process.pid = test_pid; - - linux_container.processes.insert(test_pid, test_process); + let test_exec_id = test_process.exec_id.clone(); + linux_container.processes.insert(test_exec_id, test_process); s.add_container(linux_container);