feat: implement guest exec (#107)

crates/ctl/src/cli/exec.rs

@@ -0,0 +1,70 @@
+use std::collections::HashMap;
+
+use anyhow::Result;
+
+use clap::Parser;
+use krata::v1::{
+    common::{GuestTaskSpec, GuestTaskSpecEnvVar},
+    control::{control_service_client::ControlServiceClient, ExecGuestRequest},
+};
+
+use tonic::{transport::Channel, Request};
+
+use crate::console::StdioConsoleStream;
+
+use super::resolve_guest;
+
+#[derive(Parser)]
+#[command(about = "Execute a command inside the guest")]
+pub struct ExecCommand {
+    #[arg[short, long, help = "Environment variables"]]
+    env: Option<Vec<String>>,
+    #[arg(short = 'w', long, help = "Working directory")]
+    working_directory: Option<String>,
+    #[arg(help = "Guest to exec inside, either the name or the uuid")]
+    guest: String,
+    #[arg(
+        allow_hyphen_values = true,
+        trailing_var_arg = true,
+        help = "Command to run inside the guest"
+    )]
+    command: Vec<String>,
+}
+
+impl ExecCommand {
+    pub async fn run(self, mut client: ControlServiceClient<Channel>) -> Result<()> {
+        let guest_id: String = resolve_guest(&mut client, &self.guest).await?;
+        let initial = ExecGuestRequest {
+            guest_id,
+            task: Some(GuestTaskSpec {
+                environment: env_map(&self.env.unwrap_or_default())
+                    .iter()
+                    .map(|(key, value)| GuestTaskSpecEnvVar {
+                        key: key.clone(),
+                        value: value.clone(),
+                    })
+                    .collect(),
+                command: self.command,
+                working_directory: self.working_directory.unwrap_or_default(),
+            }),
+            data: vec![],
+        };
+
+        let stream = StdioConsoleStream::stdin_stream_exec(initial).await;
+
+        let response = client.exec_guest(Request::new(stream)).await?.into_inner();
+
+        let code = StdioConsoleStream::exec_output(response).await?;
+        std::process::exit(code);
+    }
+}
+
+fn env_map(env: &[String]) -> HashMap<String, String> {
+    let mut map = HashMap::<String, String>::new();
+    for item in env {
+        if let Some((key, value)) = item.split_once('=') {
+            map.insert(key.to_string(), value.to_string());
+        }
+    }
+    map
+}

crates/ctl/src/cli/idm_snoop.rs

@@ -106,13 +106,19 @@ pub fn convert_idm_snoop(reply: SnoopIdmReply) -> Option<IdmSnoopLine> {
                 .ok()
                 .and_then(|event| proto2dynamic(event).ok()),
 
-            IdmTransportPacketForm::Request => internal::Request::decode(&packet.data)
-                .ok()
-                .and_then(|event| proto2dynamic(event).ok()),
+            IdmTransportPacketForm::Request
+            | IdmTransportPacketForm::StreamRequest
+            | IdmTransportPacketForm::StreamRequestUpdate => {
+                internal::Request::decode(&packet.data)
+                    .ok()
+                    .and_then(|event| proto2dynamic(event).ok())
+            }
 
-            IdmTransportPacketForm::Response => internal::Response::decode(&packet.data)
-                .ok()
-                .and_then(|event| proto2dynamic(event).ok()),
+            IdmTransportPacketForm::Response | IdmTransportPacketForm::StreamResponseUpdate => {
+                internal::Response::decode(&packet.data)
+                    .ok()
+                    .and_then(|event| proto2dynamic(event).ok())
+            }
 
             _ => None,
         }
@@ -132,6 +138,11 @@ pub fn convert_idm_snoop(reply: SnoopIdmReply) -> Option<IdmSnoopLine> {
             IdmTransportPacketForm::Event => "event".to_string(),
             IdmTransportPacketForm::Request => "request".to_string(),
             IdmTransportPacketForm::Response => "response".to_string(),
+            IdmTransportPacketForm::StreamRequest => "stream-request".to_string(),
+            IdmTransportPacketForm::StreamRequestUpdate => "stream-request-update".to_string(),
+            IdmTransportPacketForm::StreamRequestClosed => "stream-request-closed".to_string(),
+            IdmTransportPacketForm::StreamResponseUpdate => "stream-response-update".to_string(),
+            IdmTransportPacketForm::StreamResponseClosed => "stream-response-closed".to_string(),
             _ => format!("unknown-{}", packet.form),
         },
         data: base64::prelude::BASE64_STANDARD.encode(&packet.data),

crates/ctl/src/cli/launch.rs

@@ -29,7 +29,7 @@ pub enum LaunchImageFormat {
 
 #[derive(Parser)]
 #[command(about = "Launch a new guest")]
-pub struct LauchCommand {
+pub struct LaunchCommand {
     #[arg(long, default_value = "squashfs", help = "Image format")]
     image_format: LaunchImageFormat,
     #[arg(long, help = "Overwrite image cache on pull")]
@@ -68,6 +68,8 @@ pub struct LauchCommand {
     kernel: Option<String>,
     #[arg(short = 'I', long, help = "OCI initrd image for guest to use")]
     initrd: Option<String>,
+    #[arg(short = 'w', long, help = "Working directory")]
+    working_directory: Option<String>,
     #[arg(help = "Container image for guest to use")]
     oci: String,
     #[arg(
@@ -78,7 +80,7 @@ pub struct LauchCommand {
     command: Vec<String>,
 }
 
-impl LauchCommand {
+impl LaunchCommand {
     pub async fn run(
         self,
         mut client: ControlServiceClient<Channel>,
@@ -130,6 +132,7 @@ impl LauchCommand {
                         })
                         .collect(),
                     command: self.command,
+                    working_directory: self.working_directory.unwrap_or_default(),
                 }),
                 annotations: vec![],
             }),

crates/ctl/src/cli/mod.rs

@@ -1,5 +1,6 @@
 pub mod attach;
 pub mod destroy;
+pub mod exec;
 pub mod identify_host;
 pub mod idm_snoop;
 pub mod launch;
@@ -21,10 +22,10 @@ use krata::{
 use tonic::{transport::Channel, Request};
 
 use self::{
-    attach::AttachCommand, destroy::DestroyCommand, identify_host::IdentifyHostCommand,
-    idm_snoop::IdmSnoopCommand, launch::LauchCommand, list::ListCommand, logs::LogsCommand,
-    metrics::MetricsCommand, pull::PullCommand, resolve::ResolveCommand, top::TopCommand,
-    watch::WatchCommand,
+    attach::AttachCommand, destroy::DestroyCommand, exec::ExecCommand,
+    identify_host::IdentifyHostCommand, idm_snoop::IdmSnoopCommand, launch::LaunchCommand,
+    list::ListCommand, logs::LogsCommand, metrics::MetricsCommand, pull::PullCommand,
+    resolve::ResolveCommand, top::TopCommand, watch::WatchCommand,
 };
 
 #[derive(Parser)]
@@ -47,7 +48,7 @@ pub struct ControlCommand {
 
 #[derive(Subcommand)]
 pub enum Commands {
-    Launch(LauchCommand),
+    Launch(LaunchCommand),
     Destroy(DestroyCommand),
     List(ListCommand),
     Attach(AttachCommand),
@@ -59,6 +60,7 @@ pub enum Commands {
     IdmSnoop(IdmSnoopCommand),
     Top(TopCommand),
     IdentifyHost(IdentifyHostCommand),
+    Exec(ExecCommand),
 }
 
 impl ControlCommand {
@@ -114,6 +116,10 @@ impl ControlCommand {
             Commands::IdentifyHost(identify) => {
                 identify.run(client).await?;
             }
+
+            Commands::Exec(exec) => {
+                exec.run(client).await?;
+            }
         }
         Ok(())
     }

crates/ctl/src/console.rs

@@ -1,4 +1,4 @@
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 use async_stream::stream;
 use crossterm::{
     terminal::{disable_raw_mode, enable_raw_mode, is_raw_mode_enabled},
@@ -8,12 +8,15 @@ use krata::{
     events::EventStream,
     v1::{
         common::GuestStatus,
-        control::{watch_events_reply::Event, ConsoleDataReply, ConsoleDataRequest},
+        control::{
+            watch_events_reply::Event, ConsoleDataReply, ConsoleDataRequest, ExecGuestReply,
+            ExecGuestRequest,
+        },
     },
 };
 use log::debug;
 use tokio::{
-    io::{stdin, stdout, AsyncReadExt, AsyncWriteExt},
+    io::{stderr, stdin, stdout, AsyncReadExt, AsyncWriteExt},
     task::JoinHandle,
 };
 use tokio_stream::{Stream, StreamExt};
@@ -45,6 +48,31 @@ impl StdioConsoleStream {
         }
     }
 
+    pub async fn stdin_stream_exec(
+        initial: ExecGuestRequest,
+    ) -> impl Stream<Item = ExecGuestRequest> {
+        let mut stdin = stdin();
+        stream! {
+            yield initial;
+
+            let mut buffer = vec![0u8; 60];
+            loop {
+                let size = match stdin.read(&mut buffer).await {
+                    Ok(size) => size,
+                    Err(error) => {
+                        debug!("failed to read stdin: {}", error);
+                        break;
+                    }
+                };
+                let data = buffer[0..size].to_vec();
+                if size == 1 && buffer[0] == 0x1d {
+                    break;
+                }
+                yield ExecGuestRequest { guest_id: String::default(), task: None, data };
+            }
+        }
+    }
+
     pub async fn stdout(mut stream: Streaming<ConsoleDataReply>) -> Result<()> {
         if stdin().is_tty() {
             enable_raw_mode()?;
@@ -62,6 +90,32 @@ impl StdioConsoleStream {
         Ok(())
     }
 
+    pub async fn exec_output(mut stream: Streaming<ExecGuestReply>) -> Result<i32> {
+        let mut stdout = stdout();
+        let mut stderr = stderr();
+        while let Some(reply) = stream.next().await {
+            let reply = reply?;
+            if !reply.stdout.is_empty() {
+                stdout.write_all(&reply.stdout).await?;
+                stdout.flush().await?;
+            }
+
+            if !reply.stderr.is_empty() {
+                stderr.write_all(&reply.stderr).await?;
+                stderr.flush().await?;
+            }
+
+            if reply.exited {
+                if reply.error.is_empty() {
+                    return Ok(reply.exit_code);
+                } else {
+                    return Err(anyhow!("exec failed: {}", reply.error));
+                }
+            }
+        }
+        Ok(-1)
+    }
+
     pub async fn guest_exit_hook(
         id: String,
         events: EventStream,