agent: refactor find_process function and add test cases

Delete redundant parameter init in find_process function and
add test case for it.

Fixes: #3117

Signed-off-by: bin <bin@hyper.sh>

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() {
+        let p = sandbox.find_container_process(cid.as_str(), eid.as_str())?;
-            init = true;
-        }
-
-        let p = find_process(&mut sandbox, cid.as_str(), eid.as_str(), init)?;
 
         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| {
+        let p = sandbox
-            ttrpc_error(
+            .find_container_process(cid.as_str(), eid.as_str())
-                ttrpc::Code::INVALID_ARGUMENT,
+            .map_err(|e| {
-                format!("invalid argument: {:?}", 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| {
+        let p = sandbox
-            ttrpc_error(
+            .find_container_process(cid.as_str(), eid.as_str())
-                ttrpc::Code::UNAVAILABLE,
+            .map_err(|e| {
-                format!("invalid argument: {:?}", 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<Mutex<ReadHalf<PipeStream>>>, 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<Mutex<Sandbox>>, server_address: &str) -> Result<TtrpcServer> {
     let agent_service = Box::new(AgentService { sandbox: s })
         as Box<dyn protocols::agent_ttrpc::AgentService + Send + Sync>;

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);
+    }
 }