krataguest/
exec.rs

1use std::{collections::HashMap, process::Stdio};
2
3use anyhow::{anyhow, Result};
4use krata::idm::{
5    client::IdmClientStreamResponseHandle,
6    internal::{
7        exec_stream_request_update::Update, request::Request as RequestType,
8        ExecStreamResponseUpdate,
9    },
10    internal::{response::Response as ResponseType, Request, Response},
11};
12use tokio::{
13    io::{AsyncReadExt, AsyncWriteExt},
14    join,
15    process::Command,
16};
17
18pub struct GuestExecTask {
19    pub handle: IdmClientStreamResponseHandle<Request>,
20}
21
22impl GuestExecTask {
23    pub async fn run(&self) -> Result<()> {
24        let mut receiver = self.handle.take().await?;
25
26        let Some(ref request) = self.handle.initial.request else {
27            return Err(anyhow!("request was empty"));
28        };
29
30        let RequestType::ExecStream(update) = request else {
31            return Err(anyhow!("request was not an exec update"));
32        };
33
34        let Some(Update::Start(ref start)) = update.update else {
35            return Err(anyhow!("first request did not contain a start update"));
36        };
37
38        let mut cmd = start.command.clone();
39        if cmd.is_empty() {
40            return Err(anyhow!("command line was empty"));
41        }
42        let exe = cmd.remove(0);
43        let mut env = HashMap::new();
44        for entry in &start.environment {
45            env.insert(entry.key.clone(), entry.value.clone());
46        }
47
48        if !env.contains_key("PATH") {
49            env.insert(
50                "PATH".to_string(),
51                "/bin:/usr/bin:/usr/local/bin".to_string(),
52            );
53        }
54
55        let dir = if start.working_directory.is_empty() {
56            "/".to_string()
57        } else {
58            start.working_directory.clone()
59        };
60
61        let mut child = Command::new(exe)
62            .args(cmd)
63            .envs(env)
64            .current_dir(dir)
65            .stdin(Stdio::piped())
66            .stdout(Stdio::piped())
67            .stderr(Stdio::piped())
68            .kill_on_drop(true)
69            .spawn()
70            .map_err(|error| anyhow!("failed to spawn: {}", error))?;
71
72        let mut stdin = child
73            .stdin
74            .take()
75            .ok_or_else(|| anyhow!("stdin was missing"))?;
76        let mut stdout = child
77            .stdout
78            .take()
79            .ok_or_else(|| anyhow!("stdout was missing"))?;
80        let mut stderr = child
81            .stderr
82            .take()
83            .ok_or_else(|| anyhow!("stderr was missing"))?;
84
85        let stdout_handle = self.handle.clone();
86        let stdout_task = tokio::task::spawn(async move {
87            let mut stdout_buffer = vec![0u8; 8 * 1024];
88            loop {
89                let Ok(size) = stdout.read(&mut stdout_buffer).await else {
90                    break;
91                };
92                if size > 0 {
93                    let response = Response {
94                        response: Some(ResponseType::ExecStream(ExecStreamResponseUpdate {
95                            exited: false,
96                            exit_code: 0,
97                            error: String::new(),
98                            stdout: stdout_buffer[0..size].to_vec(),
99                            stderr: vec![],
100                        })),
101                    };
102                    let _ = stdout_handle.respond(response).await;
103                } else {
104                    break;
105                }
106            }
107        });
108
109        let stderr_handle = self.handle.clone();
110        let stderr_task = tokio::task::spawn(async move {
111            let mut stderr_buffer = vec![0u8; 8 * 1024];
112            loop {
113                let Ok(size) = stderr.read(&mut stderr_buffer).await else {
114                    break;
115                };
116                if size > 0 {
117                    let response = Response {
118                        response: Some(ResponseType::ExecStream(ExecStreamResponseUpdate {
119                            exited: false,
120                            exit_code: 0,
121                            error: String::new(),
122                            stdout: vec![],
123                            stderr: stderr_buffer[0..size].to_vec(),
124                        })),
125                    };
126                    let _ = stderr_handle.respond(response).await;
127                } else {
128                    break;
129                }
130            }
131        });
132
133        let stdin_task = tokio::task::spawn(async move {
134            loop {
135                let Some(request) = receiver.recv().await else {
136                    break;
137                };
138
139                let Some(RequestType::ExecStream(update)) = request.request else {
140                    continue;
141                };
142
143                let Some(Update::Stdin(update)) = update.update else {
144                    continue;
145                };
146
147                if stdin.write_all(&update.data).await.is_err() {
148                    break;
149                }
150            }
151        });
152
153        let exit = child.wait().await?;
154        let code = exit.code().unwrap_or(-1);
155
156        let _ = join!(stdout_task, stderr_task);
157        stdin_task.abort();
158
159        let response = Response {
160            response: Some(ResponseType::ExecStream(ExecStreamResponseUpdate {
161                exited: true,
162                exit_code: code,
163                error: String::new(),
164                stdout: vec![],
165                stderr: vec![],
166            })),
167        };
168        self.handle.respond(response).await?;
169
170        Ok(())
171    }
172}
krataguest/exec.rs

krataguest/
exec.rs
