diff --git a/src/agent/src/rpc.rs b/src/agent/src/rpc.rs index 87cdc4f84b..c4d164cd36 100644 --- a/src/agent/src/rpc.rs +++ b/src/agent/src/rpc.rs @@ -368,7 +368,6 @@ impl AgentService { let eid = req.exec_id.clone(); let s = self.sandbox.clone(); let mut sandbox = s.lock().await; - let mut init = false; info!( sl!(), @@ -377,11 +376,7 @@ impl AgentService { "exec-id" => eid.clone(), ); - if eid.is_empty() { - init = true; - } - - let p = find_process(&mut sandbox, cid.as_str(), eid.as_str(), init)?; + let p = sandbox.find_container_process(cid.as_str(), eid.as_str())?; let mut signal = Signal::try_from(req.signal as i32).map_err(|e| { anyhow!(e).context(format!( @@ -424,7 +419,7 @@ impl AgentService { let exit_rx = { let mut sandbox = s.lock().await; - let p = find_process(&mut sandbox, cid.as_str(), eid.as_str(), false)?; + let p = sandbox.find_container_process(cid.as_str(), eid.as_str())?; p.exit_watchers.push(exit_send); pid = p.pid; @@ -482,7 +477,7 @@ impl AgentService { let writer = { let s = self.sandbox.clone(); let mut sandbox = s.lock().await; - let p = find_process(&mut sandbox, cid.as_str(), eid.as_str(), false)?; + let p = sandbox.find_container_process(cid.as_str(), eid.as_str())?; // use ptmx io if p.term_master.is_some() { @@ -515,7 +510,7 @@ impl AgentService { let s = self.sandbox.clone(); let mut sandbox = s.lock().await; - let p = find_process(&mut sandbox, cid.as_str(), eid.as_str(), false)?; + let p = sandbox.find_container_process(cid.as_str(), eid.as_str())?; if p.term_master.is_some() { term_exit_notifier = p.term_exit_notifier.clone(); @@ -783,12 +778,14 @@ impl protocols::agent_ttrpc::AgentService for AgentService { let s = Arc::clone(&self.sandbox); let mut sandbox = s.lock().await; - let p = find_process(&mut sandbox, cid.as_str(), eid.as_str(), false).map_err(|e| { - ttrpc_error( - ttrpc::Code::INVALID_ARGUMENT, - format!("invalid argument: {:?}", e), - ) - })?; + let p = sandbox + .find_container_process(cid.as_str(), eid.as_str()) + .map_err(|e| { + ttrpc_error( + ttrpc::Code::INVALID_ARGUMENT, + format!("invalid argument: {:?}", e), + ) + })?; p.close_stdin(); @@ -807,12 +804,14 @@ impl protocols::agent_ttrpc::AgentService for AgentService { let eid = req.exec_id.clone(); let s = Arc::clone(&self.sandbox); let mut sandbox = s.lock().await; - let p = find_process(&mut sandbox, cid.as_str(), eid.as_str(), false).map_err(|e| { - ttrpc_error( - ttrpc::Code::UNAVAILABLE, - format!("invalid argument: {:?}", e), - ) - })?; + let p = sandbox + .find_container_process(cid.as_str(), eid.as_str()) + .map_err(|e| { + ttrpc_error( + ttrpc::Code::UNAVAILABLE, + format!("invalid argument: {:?}", e), + ) + })?; if let Some(fd) = p.term_master { unsafe { @@ -1366,26 +1365,6 @@ async fn read_stream(reader: Arc>>, l: usize) -> Resu Ok(content) } -fn find_process<'a>( - sandbox: &'a mut Sandbox, - cid: &'a str, - eid: &'a str, - init: bool, -) -> Result<&'a mut Process> { - let ctr = sandbox - .get_container(cid) - .ok_or_else(|| anyhow!("Invalid container id"))?; - - if init || eid.is_empty() { - return ctr - .processes - .get_mut(&ctr.init_process_pid) - .ok_or_else(|| anyhow!("cannot find init process!")); - } - - ctr.get_process(eid).map_err(|_| anyhow!("Invalid exec id")) -} - pub fn start(s: Arc>, server_address: &str) -> Result { let agent_service = Box::new(AgentService { sandbox: s }) as Box; diff --git a/src/agent/src/sandbox.rs b/src/agent/src/sandbox.rs index 5264640147..e13d10e349 100644 --- a/src/agent/src/sandbox.rs +++ b/src/agent/src/sandbox.rs @@ -226,6 +226,21 @@ impl Sandbox { None } + pub fn find_container_process(&mut self, cid: &str, eid: &str) -> Result<&mut Process> { + let ctr = self + .get_container(cid) + .ok_or_else(|| anyhow!("Invalid container id"))?; + + if eid.is_empty() { + return ctr + .processes + .get_mut(&ctr.init_process_pid) + .ok_or_else(|| anyhow!("cannot find init process!")); + } + + ctr.get_process(eid).map_err(|_| anyhow!("Invalid exec id")) + } + #[instrument] pub async fn destroy(&mut self) -> Result<()> { for ctr in self.containers.values_mut() { @@ -454,6 +469,7 @@ mod tests { use nix::mount::MsFlags; use oci::{Linux, Root, Spec}; use rustjail::container::LinuxContainer; + use rustjail::process::Process; use rustjail::specconv::CreateOpts; use slog::Logger; use std::fs::{self, File}; @@ -785,4 +801,50 @@ mod tests { let ret = s.destroy().await; assert!(ret.is_ok()); } + + #[tokio::test] + async fn test_find_container_process() { + skip_if_not_root!(); + let logger = slog::Logger::root(slog::Discard, o!()); + let mut s = Sandbox::new(&logger).unwrap(); + let cid = "container-123"; + + let mut linux_container = create_linuxcontainer(); + 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).unwrap(), + ); + // add exec process + linux_container.processes.insert( + 123, + Process::new(&logger, &oci::Process::default(), "exec-123", false, 1).unwrap(), + ); + + s.add_container(linux_container); + + // empty exec-id will return init process + let p = s.find_container_process(cid, ""); + assert!(p.is_ok(), "Expecting Ok, Got {:?}", p); + let p = p.unwrap(); + assert_eq!("1", p.exec_id, "exec_id should be 1"); + assert!(p.init, "init flag should be true"); + + // get exist exec-id will return the exec process + let p = s.find_container_process(cid, "exec-123"); + assert!(p.is_ok(), "Expecting Ok, Got {:?}", p); + let p = p.unwrap(); + assert_eq!("exec-123", p.exec_id, "exec_id should be exec-123"); + assert!(!p.init, "init flag should be false"); + + // get not exist exec-id will return error + let p = s.find_container_process(cid, "exec-456"); + assert!(p.is_err(), "Expecting Error, Got {:?}", p); + + // container does not exist + let p = s.find_container_process("not-exist-cid", ""); + assert!(p.is_err(), "Expecting Error, Got {:?}", p); + } }