krata/crates/zone/src/exec.rs

use std::{collections::HashMap, process::Stdio};

use crate::childwait::ChildWait;
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 libc::c_int;
use pty_process::{Pty, Size};
use tokio::process::Child;
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    join,
    process::Command,
    select,
};
use tokio_util::sync::CancellationToken;

pub struct ZoneExecTask {
    pub wait: ChildWait,
    pub handle: IdmClientStreamResponseHandle<Request>,
}

impl ZoneExecTask {
    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:/sbin:/usr/sbin".to_string(),
            );
        }

        let dir = if start.working_directory.is_empty() {
            "/".to_string()
        } else {
            start.working_directory.clone()
        };

        let mut wait_subscription = self.wait.subscribe().await?;

        let code: c_int;
        if start.tty {
            let pty = Pty::new().map_err(|error| anyhow!("unable to allocate pty: {}", error))?;
            let size = start
                .terminal_size
                .map(|x| Size::new(x.rows as u16, x.columns as u16))
                .unwrap_or_else(|| Size::new(24, 80));
            pty.resize(size)?;
            let pts = pty
                .pts()
                .map_err(|error| anyhow!("unable to allocate pts: {}", error))?;
            let child = std::panic::catch_unwind(move || {
                let pts = pts;
                pty_process::Command::new(exe)
                    .args(cmd)
                    .envs(env)
                    .current_dir(dir)
                    .spawn(&pts)
            })
            .map_err(|_| anyhow!("internal error"))
            .map_err(|error| anyhow!("failed to spawn: {}", error))??;
            let mut child = ChildDropGuard {
                inner: child,
                kill: true,
            };
            let pid = child
                .inner
                .id()
                .ok_or_else(|| anyhow!("pid is not provided"))?;
            let (mut read, mut write) = pty.into_split();
            let pty_read_handle = self.handle.clone();
            let pty_read_task = tokio::task::spawn(async move {
                let mut stdout_buffer = vec![0u8; 8 * 1024];
                loop {
                    let Ok(size) = read.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 _ = pty_read_handle.respond(response).await;
                    } else {
                        break;
                    }
                }
            });

            let cancel = CancellationToken::new();
            let stdin_cancel = cancel.clone();
            let stdin_task = tokio::task::spawn(async move {
                loop {
                    let Some(request) = receiver.recv().await else {
                        stdin_cancel.cancel();
                        break;
                    };

                    let Some(RequestType::ExecStream(update)) = request.request else {
                        continue;
                    };

                    match update.update {
                        Some(Update::Stdin(update)) => {
                            if !update.data.is_empty()
                                && write.write_all(&update.data).await.is_err()
                            {
                                break;
                            }

                            if update.closed {
                                break;
                            }
                        }
                        Some(Update::TerminalResize(size)) => {
                            let _ = write.resize(Size::new(size.rows as u16, size.columns as u16));
                        }
                        _ => {
                            continue;
                        }
                    }
                }
            });

            code = loop {
                select! {
                    result = wait_subscription.recv() => match result {
                        Ok(event) => {
                            if event.pid.as_raw() as u32 == pid {
                                child.kill = false;
                                break event.status;
                            }
                        }
                        _ => {
                            child.inner.start_kill()?;
                            child.kill = false;
                            break -1;
                        }
                    },
                    _ = cancel.cancelled() => {
                        child.inner.start_kill()?;
                        child.kill = false;
                        break -1;
                    }
                }
            };

            let _ = join!(pty_read_task);
            stdin_task.abort();
        } else {
            let mut child = std::panic::catch_unwind(|| {
                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(|_| anyhow!("internal error"))
            .map_err(|error| anyhow!("failed to spawn: {}", error))??;

            let pid = child.id().ok_or_else(|| anyhow!("pid is not provided"))?;
            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 cancel = CancellationToken::new();
            let stdin_cancel = cancel.clone();
            let stdin_task = tokio::task::spawn(async move {
                loop {
                    let Some(request) = receiver.recv().await else {
                        stdin_cancel.cancel();
                        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 data_task = tokio::task::spawn(async move {
                let _ = join!(stdout_task, stderr_task);
                stdin_task.abort();
            });

            code = loop {
                select! {
                    result = wait_subscription.recv() => match result {
                        Ok(event) => {
                            if event.pid.as_raw() as u32 == pid {
                                break event.status;
                            }
                        }
                        _ => {
                            child.start_kill()?;
                            break -1;
                        }
                    },
                    _ = cancel.cancelled() => {
                        child.start_kill()?;
                        break -1;
                    }
                }
            };
            data_task.await?;
        }
        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(())
    }
}

struct ChildDropGuard {
    pub inner: Child,
    pub kill: bool,
}

impl Drop for ChildDropGuard {
    fn drop(&mut self) {
        if self.kill {
            drop(self.inner.start_kill());
        }
    }
}