feat: implement guest exec (#107)

crates/guest/src/background.rs

@@ -1,16 +1,17 @@
 use crate::{
     childwait::{ChildEvent, ChildWait},
     death,
+    exec::GuestExecTask,
     metrics::MetricsCollector,
 };
 use anyhow::Result;
 use cgroups_rs::Cgroup;
 use krata::idm::{
-    client::IdmInternalClient,
+    client::{IdmClientStreamResponseHandle, IdmInternalClient},
     internal::{
         event::Event as EventType, request::Request as RequestType,
-        response::Response as ResponseType, Event, ExitEvent, MetricsResponse, PingResponse,
-        Request, Response,
+        response::Response as ResponseType, Event, ExecStreamResponseUpdate, ExitEvent,
+        MetricsResponse, PingResponse, Request, Response,
     },
 };
 use log::debug;
@@ -41,11 +42,11 @@ impl GuestBackground {
     pub async fn run(&mut self) -> Result<()> {
         let mut event_subscription = self.idm.subscribe().await?;
         let mut requests_subscription = self.idm.requests().await?;
+        let mut request_streams_subscription = self.idm.request_streams().await?;
         loop {
             select! {
                 x = event_subscription.recv() => match x {
                     Ok(_event) => {
-
                     },
 
                     Err(broadcast::error::RecvError::Closed) => {
@@ -73,6 +74,21 @@ impl GuestBackground {
                     }
                 },
 
+                x = request_streams_subscription.recv() => match x {
+                    Ok(handle) => {
+                        self.handle_idm_stream_request(handle).await?;
+                    },
+
+                    Err(broadcast::error::RecvError::Closed) => {
+                        debug!("idm packet channel closed");
+                        break;
+                    },
+
+                    _ => {
+                        continue;
+                    }
+                },
+
                 event = self.wait.recv() => match event {
                     Some(event) => self.child_event(event).await?,
                     None => {
@@ -107,7 +123,33 @@ impl GuestBackground {
                 self.idm.respond(id, response).await?;
             }
 
-            None => {}
+            _ => {}
+        }
+        Ok(())
+    }
+
+    async fn handle_idm_stream_request(
+        &mut self,
+        handle: IdmClientStreamResponseHandle<Request>,
+    ) -> Result<()> {
+        if let Some(RequestType::ExecStream(_)) = &handle.initial.request {
+            tokio::task::spawn(async move {
+                let exec = GuestExecTask { handle };
+                if let Err(error) = exec.run().await {
+                    let _ = exec
+                        .handle
+                        .respond(Response {
+                            response: Some(ResponseType::ExecStream(ExecStreamResponseUpdate {
+                                exited: true,
+                                error: error.to_string(),
+                                exit_code: -1,
+                                stdout: vec![],
+                                stderr: vec![],
+                            })),
+                        })
+                        .await;
+                }
+            });
         }
         Ok(())
     }

crates/guest/src/exec.rs

@@ -0,0 +1,172 @@
+use std::{collections::HashMap, process::Stdio};
+
+use anyhow::{anyhow, Result};
+use krata::idm::{
+    client::IdmClientStreamResponseHandle,
+    internal::{
+        exec_stream_request_update::Update, request::Request as RequestType,
+        ExecStreamResponseUpdate,
+    },
+    internal::{response::Response as ResponseType, Request, Response},
+};
+use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt},
+    join,
+    process::Command,
+};
+
+pub struct GuestExecTask {
+    pub handle: IdmClientStreamResponseHandle<Request>,
+}
+
+impl GuestExecTask {
+    pub async fn run(&self) -> Result<()> {
+        let mut receiver = self.handle.take().await?;
+
+        let Some(ref request) = self.handle.initial.request else {
+            return Err(anyhow!("request was empty"));
+        };
+
+        let RequestType::ExecStream(update) = request else {
+            return Err(anyhow!("request was not an exec update"));
+        };
+
+        let Some(Update::Start(ref start)) = update.update else {
+            return Err(anyhow!("first request did not contain a start update"));
+        };
+
+        let mut cmd = start.command.clone();
+        if cmd.is_empty() {
+            return Err(anyhow!("command line was empty"));
+        }
+        let exe = cmd.remove(0);
+        let mut env = HashMap::new();
+        for entry in &start.environment {
+            env.insert(entry.key.clone(), entry.value.clone());
+        }
+
+        if !env.contains_key("PATH") {
+            env.insert(
+                "PATH".to_string(),
+                "/bin:/usr/bin:/usr/local/bin".to_string(),
+            );
+        }
+
+        let dir = if start.working_directory.is_empty() {
+            "/".to_string()
+        } else {
+            start.working_directory.clone()
+        };
+
+        let mut child = Command::new(exe)
+            .args(cmd)
+            .envs(env)
+            .current_dir(dir)
+            .stdin(Stdio::piped())
+            .stdout(Stdio::piped())
+            .stderr(Stdio::piped())
+            .kill_on_drop(true)
+            .spawn()
+            .map_err(|error| anyhow!("failed to spawn: {}", error))?;
+
+        let mut stdin = child
+            .stdin
+            .take()
+            .ok_or_else(|| anyhow!("stdin was missing"))?;
+        let mut stdout = child
+            .stdout
+            .take()
+            .ok_or_else(|| anyhow!("stdout was missing"))?;
+        let mut stderr = child
+            .stderr
+            .take()
+            .ok_or_else(|| anyhow!("stderr was missing"))?;
+
+        let stdout_handle = self.handle.clone();
+        let stdout_task = tokio::task::spawn(async move {
+            let mut stdout_buffer = vec![0u8; 8 * 1024];
+            loop {
+                let Ok(size) = stdout.read(&mut stdout_buffer).await else {
+                    break;
+                };
+                if size > 0 {
+                    let response = Response {
+                        response: Some(ResponseType::ExecStream(ExecStreamResponseUpdate {
+                            exited: false,
+                            exit_code: 0,
+                            error: String::new(),
+                            stdout: stdout_buffer[0..size].to_vec(),
+                            stderr: vec![],
+                        })),
+                    };
+                    let _ = stdout_handle.respond(response).await;
+                } else {
+                    break;
+                }
+            }
+        });
+
+        let stderr_handle = self.handle.clone();
+        let stderr_task = tokio::task::spawn(async move {
+            let mut stderr_buffer = vec![0u8; 8 * 1024];
+            loop {
+                let Ok(size) = stderr.read(&mut stderr_buffer).await else {
+                    break;
+                };
+                if size > 0 {
+                    let response = Response {
+                        response: Some(ResponseType::ExecStream(ExecStreamResponseUpdate {
+                            exited: false,
+                            exit_code: 0,
+                            error: String::new(),
+                            stdout: vec![],
+                            stderr: stderr_buffer[0..size].to_vec(),
+                        })),
+                    };
+                    let _ = stderr_handle.respond(response).await;
+                } else {
+                    break;
+                }
+            }
+        });
+
+        let stdin_task = tokio::task::spawn(async move {
+            loop {
+                let Some(request) = receiver.recv().await else {
+                    break;
+                };
+
+                let Some(RequestType::ExecStream(update)) = request.request else {
+                    continue;
+                };
+
+                let Some(Update::Stdin(update)) = update.update else {
+                    continue;
+                };
+
+                if stdin.write_all(&update.data).await.is_err() {
+                    break;
+                }
+            }
+        });
+
+        let exit = child.wait().await?;
+        let code = exit.code().unwrap_or(-1);
+
+        let _ = join!(stdout_task, stderr_task);
+        stdin_task.abort();
+
+        let response = Response {
+            response: Some(ResponseType::ExecStream(ExecStreamResponseUpdate {
+                exited: true,
+                exit_code: code,
+                error: String::new(),
+                stdout: vec![],
+                stderr: vec![],
+            })),
+        };
+        self.handle.respond(response).await?;
+
+        Ok(())
+    }
+}

crates/guest/src/init.rs

@@ -479,7 +479,7 @@ impl GuestInit {
             env.insert("TERM".to_string(), "xterm".to_string());
         }
 
-        let path = GuestInit::resolve_executable(&env, path.into())?;
+        let path = resolve_executable(&env, path.into())?;
         let Some(file_name) = path.file_name() else {
             return Err(anyhow!("cannot get file name of command path"));
         };
@@ -537,27 +537,6 @@ impl GuestInit {
         map
     }
 
-    fn resolve_executable(env: &HashMap<String, String>, path: PathBuf) -> Result<PathBuf> {
-        if path.is_absolute() {
-            return Ok(path);
-        }
-
-        if path.is_file() {
-            return Ok(path.absolutize()?.to_path_buf());
-        }
-
-        if let Some(path_var) = env.get("PATH") {
-            for item in path_var.split(':') {
-                let mut exe_path: PathBuf = item.into();
-                exe_path.push(&path);
-                if exe_path.is_file() {
-                    return Ok(exe_path);
-                }
-            }
-        }
-        Ok(path)
-    }
-
     fn env_list(env: HashMap<String, String>) -> Vec<String> {
         env.iter()
             .map(|(key, value)| format!("{}={}", key, value))
@@ -613,3 +592,24 @@ impl GuestInit {
         Ok(())
     }
 }
+
+pub fn resolve_executable(env: &HashMap<String, String>, path: PathBuf) -> Result<PathBuf> {
+    if path.is_absolute() {
+        return Ok(path);
+    }
+
+    if path.is_file() {
+        return Ok(path.absolutize()?.to_path_buf());
+    }
+
+    if let Some(path_var) = env.get("PATH") {
+        for item in path_var.split(':') {
+            let mut exe_path: PathBuf = item.into();
+            exe_path.push(&path);
+            if exe_path.is_file() {
+                return Ok(exe_path);
+            }
+        }
+    }
+    Ok(path)
+}

crates/guest/src/lib.rs

@@ -6,6 +6,7 @@ use xenstore::{XsdClient, XsdInterface};
 
 pub mod background;
 pub mod childwait;
+pub mod exec;
 pub mod init;
 pub mod metrics;