vmexec 0.7.1

use std::fmt::Debug;
use std::fs;
use std::io::{ErrorKind, Write};
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::{env, os::unix::fs::OpenOptionsExt, path::Path, process, sync::Arc, time::Duration};

use base64ct::LineEnding;
use eyre::{Context, Result, eyre};
use russh::keys::PrivateKey;
use russh::keys::ssh_key::private::Ed25519Keypair;
use russh::{ChannelMsg, Disconnect, keys::key::PrivateKeyWithHashAlg};
use serde::{Deserialize, Serialize};
use termion::raw::IntoRawMode;
use tokio::io::{AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::time::Instant;
use tokio_fd::AsyncFd;
use tokio_vsock::{VsockAddr, VsockStream};
use tracing::{debug, info, instrument, warn};

use crate::error::Error;
use crate::runner::CancellationTokens;
use crate::types::{EnvVar, Interactive, Io};
use crate::utils::safe_flush;

/// File with the SSH private key in the secrets dir
pub(crate) const SSH_PRIVKEY_FILENAME: &str = "id_ed25519";

/// Default timeout for SSH connection to VM
pub const DEFAULT_SSH_TIMEOUT: Duration = Duration::from_secs(20);

#[derive(Clone, Debug)]
pub(crate) struct PersistedSshKeypair {
    pub pubkey_str: String,
    pub privkey_str: String,
    pub privkey_path: PathBuf,
}

impl PersistedSshKeypair {
    /// Build a keypair from an OpenSSH-encoded private key
    ///
    /// The public key is derived from this.
    fn from_privkey(privkey_str: &str, privkey_path: &Path) -> Result<Self> {
        let privkey = PrivateKey::from_openssh(privkey_str)?;

        Ok(Self {
            pubkey_str: privkey.public_key().to_openssh()?,
            privkey_str: privkey_str.to_string(),
            privkey_path: privkey_path.to_path_buf(),
        })
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct SshLaunchOpts {
    #[serde(skip)]
    pub privkey: String,
    pub tty: bool,
    pub interactive: Interactive,
    pub timeout: Duration,
    pub env_vars: Vec<EnvVar>,
    pub workdir: Option<PathBuf>,
    pub args: Vec<String>,
    pub cid: u32,

    /// The command's stdout goes in here
    ///
    /// Note: All output is also written to the VM's session log.
    pub stdout: Io,

    /// The command's stdout goes in here
    ///
    /// Note: All output is also written to the VM's session log.
    pub stderr: Io,
}

/// End result of a command ran inside a VM
#[derive(Debug)]
pub struct CommandOutput {
    pub exit_code: u32,

    /// Captured stdout (empty unless with [`Io::Piped`])
    pub stdout: Vec<u8>,

    /// Captured stderr (empty unless with [`Io::Piped`])
    pub stderr: Vec<u8>,
}

/// Retrieve or create SSH keypair from `path` to be used with the virtual machine
#[instrument]
pub(crate) fn ensure_ssh_key(dir: &Path) -> Result<PersistedSshKeypair> {
    let privkey_path = dir.join(SSH_PRIVKEY_FILENAME);

    // A key that exists but doesn't parse is an error rather than a reason to replace it since any
    // VM still trusting the old key would no longer be able to authenticate.
    if let Ok(privkey_str) = fs::read_to_string(&privkey_path) {
        return PersistedSshKeypair::from_privkey(&privkey_str, &privkey_path)
            .wrap_err(format!("Couldn't read the SSH key at {privkey_path:?}"));
    }

    let privkey_str = PrivateKey::from(Ed25519Keypair::random(&mut rand::rng()))
        .to_openssh(LineEnding::default())?
        .to_string();

    // Write to a temp file and rename it into place so that a concurrent process sees either no key
    // or a complete one.
    let tmp_path = privkey_path.with_extension(format!("tmp.{}", process::id()));
    debug!("Writing SSH private key to {privkey_path:?}");
    fs::OpenOptions::new()
        .write(true)
        .create_new(true)
        .mode(0o600)
        .open(&tmp_path)
        .wrap_err(format!("Couldn't create {tmp_path:?}"))?
        .write_all(privkey_str.as_bytes())?;
    fs::rename(&tmp_path, &privkey_path)?;

    PersistedSshKeypair::from_privkey(&privkey_str, &privkey_path)
}

#[derive(Debug, Clone)]
struct SshClient {}

// More SSH event handlers can be defined in this trait
//
// In this example, we're only using Channel, so these aren't needed.
impl russh::client::Handler for SshClient {
    type Error = russh::Error;

    #[instrument]
    async fn check_server_key(
        &mut self,
        _server_public_key: &russh::keys::PublicKeyOrCertificate,
    ) -> Result<bool, Self::Error> {
        Ok(true)
    }
}

/// This struct is a convenience wrapper around a russh client that handles the input/output event
/// loop
pub(crate) struct Session {
    session: russh::client::Handle<SshClient>,
    tty_state: Pty,
}

/// State for handling a pseudo-tty allocated by SSH, if enabled.
enum Pty {
    Enabled {
        /// Last known size of the user's terminal on the host, outside the VM.
        host_terminal_size: (u16, u16),
    },
    Disabled,
}

impl Pty {
    fn is_enabled(&self) -> bool {
        match self {
            Pty::Enabled { .. } => true,
            Pty::Disabled => false,
        }
    }
}

/// Functionality for asynchronously reading from stdin, if enabled.
enum StdinReader {
    Enabled { fd: AsyncFd, buf: Vec<u8> },
    Closed,
    Disabled,
}

impl StdinReader {
    /// If reading from stdin is enabled, try to read into its buffer and return the bytes read along with a reference to the buffer.
    /// This allows us to conveniently do an optional read in tokio's select! macro.
    async fn maybe_read(&mut self) -> Option<(std::io::Result<usize>, &[u8])> {
        match self {
            StdinReader::Enabled { fd, buf } => Some((fd.read(buf).await, buf)),
            StdinReader::Disabled => None,
            StdinReader::Closed => None,
        }
    }
}

/// Whether a failed connection attempt is worth retrying
pub(crate) enum ConnectError {
    /// The VM likely isn't up yet so try again.
    Transient,

    /// Retrying won't help.
    Fatal(Error),
}

fn classify_vsock_error(e: &std::io::Error) -> ConnectError {
    match e.raw_os_error() {
        Some(n) if n == nix::errno::Errno::EAFNOSUPPORT as i32 => {
            ConnectError::Fatal(Error::VsockUnavailable)
        }
        Some(n) if n == nix::errno::Errno::ENODEV as i32 => ConnectError::Transient,
        // Since kernel 7.2 (commit bb26ed5f3a8b), connecting while the guest's vsock driver isn't
        // up yet returns EHOSTUNREACH immediately instead of blocking until ETIMEDOUT.
        // I think it should still be safe to always treat EHOSTUNREACH as transient because when
        // else would we hit this with a vsock?
        Some(n) if n == nix::errno::Errno::EHOSTUNREACH as i32 => ConnectError::Transient,
        _ => match e.kind() {
            ErrorKind::TimedOut | ErrorKind::ConnectionRefused | ErrorKind::ConnectionReset => {
                ConnectError::Transient
            }
            _ => ConnectError::Fatal(Error::Other(eyre!(
                "Unexpected error connecting to VM: {e}"
            ))),
        },
    }
}

fn classify_ssh_error(e: &russh::Error) -> ConnectError {
    match e {
        russh::Error::IO(e) => match e.kind() {
            ErrorKind::ConnectionRefused | ErrorKind::ConnectionReset => ConnectError::Transient,
            _ => ConnectError::Fatal(Error::Other(eyre!(
                "Unexpected error connecting to VM via SSH: {e}"
            ))),
        },
        russh::Error::Disconnect => ConnectError::Transient,
        e => ConnectError::Fatal(Error::Other(eyre!(
            "Unexpected error connecting to VM via SSH: {e}"
        ))),
    }
}

/// Write a chunk to the session log
///
/// On failure, warn and disable further logging.
async fn write_session_log(log: &mut Option<&mut tokio::fs::File>, data: &[u8]) {
    if let Some(f) = log.as_deref_mut() {
        let res = async {
            f.write_all(data).await?;
            // Flush every chunk so that the log can be followed live and nothing is lost on
            // abnormal exit.
            f.flush().await
        }
        .await;
        if let Err(e) = res {
            warn!("Failed to write to session log, disabling logging: {e}");
            *log = None;
        }
    }
}

impl Session {
    #[instrument(skip(privkey))]
    async fn connect(
        privkey: PrivateKey,
        cid: u32,
        port: u32,
        timeout: Duration,
        allocate_tty: bool,
    ) -> Result<Self, Error> {
        let config = russh::client::Config {
            keepalive_interval: Some(Duration::from_secs(5)),
            ..<_>::default()
        };

        let config = Arc::new(config);
        let sh = SshClient {};

        let vsock_addr = VsockAddr::new(cid, port);
        let now = Instant::now();
        debug!("Connecting to SSH via vsock");
        let mut session = loop {
            tokio::time::sleep(Duration::from_millis(100)).await;
            if now.elapsed() > timeout {
                return Err(Error::SshTimeout(timeout));
            }

            let stream = match VsockStream::connect(vsock_addr).await {
                Ok(stream) => stream,
                Err(ref e) => match classify_vsock_error(e) {
                    ConnectError::Transient => continue,
                    ConnectError::Fatal(e) => return Err(e),
                },
            };

            match russh::client::connect_stream(config.clone(), stream, sh.clone()).await {
                Ok(x) => break x,
                Err(ref e) => match classify_ssh_error(e) {
                    ConnectError::Transient => continue,
                    ConnectError::Fatal(e) => return Err(e),
                },
            }
        };
        debug!("Authenticating via SSH");

        // use publickey authentication
        let auth_res = session
            .authenticate_publickey("root", PrivateKeyWithHashAlg::new(Arc::new(privkey), None))
            .await
            .map_err(eyre::Report::from)?;

        if !auth_res.success() {
            return Err(Error::Other(eyre!(
                "Authentication (with publickey) failed, this can happen if you deleted the automatically generated SSH key in .local/state. \
                Try running `vmexec prune` to delete your local warmup images and then try running this command again."
            )));
        }

        let tty_state = if allocate_tty {
            Pty::Enabled {
                host_terminal_size: termion::terminal_size().wrap_err("Requested a TTY inside the VM, but vmexec doesn't seem to be running in a terminal")?,
            }
        } else {
            Pty::Disabled
        };

        Ok(Self { session, tty_state })
    }

    #[instrument(skip(self, session_log))]
    async fn call(
        &mut self,
        interactive: Interactive,
        env: Vec<EnvVar>,
        command: &str,
        mut session_log: Option<&mut tokio::fs::File>,
        stdout: Io,
        stderr: Io,
    ) -> Result<CommandOutput> {
        let mut channel = self.session.channel_open_session().await?;

        if let Pty::Enabled { host_terminal_size } = &self.tty_state {
            // Request an interactive PTY from the server
            channel
                .request_pty(
                    true,
                    &env::var("TERM").unwrap_or("xterm-256color".into()),
                    host_terminal_size.0 as u32,
                    host_terminal_size.1 as u32,
                    0,
                    0,
                    &[], // ideally you want to pass the actual terminal modes here
                )
                .await?;
        }

        for e in env {
            channel.set_env(true, e.key, e.value).await?;
        }

        //channel.request_shell(true).await?;
        channel.exec(true, command).await?;

        let code;
        let mut stdin_reader = match interactive {
            Interactive::Always => {
                let buf = vec![0; 1024];
                let fd = tokio_fd::AsyncFd::try_from(nix::libc::STDIN_FILENO).wrap_err("Requested stdin to be piped to the process inside the VM, but failed to open stdin.")?;
                StdinReader::Enabled { fd, buf }
            }
            Interactive::Never => StdinReader::Disabled,
            Interactive::Auto => {
                let fd = tokio_fd::AsyncFd::try_from(nix::libc::STDIN_FILENO);
                if let Ok(fd) = fd {
                    let buf = vec![0; 1024];
                    StdinReader::Enabled { fd, buf }
                } else {
                    StdinReader::Disabled
                }
            }
        };
        // Slightly nasty types but I didn't find a better way. This also requires explicit drop()s
        // at the end of the function because of the borrowed writers.
        let mut stdout_buf = Vec::new();
        let mut stderr_buf = Vec::new();
        let mut stdout: Box<dyn AsyncWrite + Send + Unpin + '_> = match stdout {
            Io::Inherit => Box::new(tokio::io::stdout()),
            Io::Piped => Box::new(&mut stdout_buf),
            Io::Null => Box::new(tokio::io::sink()),
        };
        let mut stderr: Box<dyn AsyncWrite + Send + Unpin + '_> = match stderr {
            Io::Inherit => Box::new(tokio::io::stderr()),
            Io::Piped => Box::new(&mut stderr_buf),
            Io::Null => Box::new(tokio::io::sink()),
        };

        // TODO maybe have entirely separate code paths for interactive vs non-interactive?
        // We don't need to handle terminal resizing and stdin at all if we have no tty and are not
        // interactive.

        loop {
            // Handle one of the possible events:
            tokio::select! {
                // Handle terminal resize
                _ = tokio::time::sleep(Duration::from_millis(500)), if self.tty_state.is_enabled() => {
                    if let Pty::Enabled{host_terminal_size} = &self.tty_state {
                        let new_terminal_size = termion::terminal_size()?;
                        if host_terminal_size != &new_terminal_size {
                            debug!("Terminal size change detected");
                            self.tty_state = Pty::Enabled { host_terminal_size: new_terminal_size };
                            channel.window_change(new_terminal_size.0 as u32, new_terminal_size.1 as u32, 0, 0).await?;
                        }
                    }
                },
                // There's terminal input available from the user
                Some((read_bytes, buf)) = stdin_reader.maybe_read() => {
                    match read_bytes {
                        Ok(0) => {
                            stdin_reader = StdinReader::Closed;
                            channel.eof().await?;
                        },
                        // Send it to the server
                        Ok(n) => channel.data(&buf[..n]).await?,
                        Err(e) => return Err(e.into()),
                    };
                },
                // There's an event available on the session channel
                Some(msg) = channel.wait() => {
                    match msg {
                        // Write data to the terminal
                        ChannelMsg::Data { ref data } => {
                            stdout.write_all(data).await?;
                            safe_flush(&mut stdout).await?;
                            write_session_log(&mut session_log, data).await;
                        }
                        ChannelMsg::ExtendedData { ref data, ext: 1 } => {
                            // ext == 1 means it's stderr content
                            // https://github.com/Eugeny/russh/discussions/258
                            stderr.write_all(data).await?;
                            safe_flush(&mut stderr).await?;
                            write_session_log(&mut session_log, data).await;
                        }
                        // The command has returned an exit code
                        ChannelMsg::ExitStatus { exit_status } => {
                            code = exit_status;
                            match stdin_reader {
                                StdinReader::Enabled { .. } => channel.eof().await?,
                                StdinReader::Closed => {},
                                StdinReader::Disabled => channel.eof().await?,
                            };
                            break;
                        }
                        _ => {}
                    }
                },
            }
        }
        // As pointed out above, we need to explicitly drop these.
        drop(stdout);
        drop(stderr);
        Ok(CommandOutput {
            exit_code: code,
            stdout: stdout_buf,
            stderr: stderr_buf,
        })
    }

    #[instrument(skip(self))]
    async fn close(&mut self) -> Result<()> {
        self.session
            .disconnect(Disconnect::ByApplication, "", "English")
            .await?;
        Ok(())
    }
}

pub(crate) async fn create_ssh_connection(
    ssh_launch_opts: &SshLaunchOpts,
) -> Result<Session, Error> {
    let privkey =
        PrivateKey::from_openssh(ssh_launch_opts.privkey.clone()).map_err(eyre::Report::from)?;

    // Session is a wrapper around a russh client.
    let ssh = Session::connect(
        privkey,
        ssh_launch_opts.cid,
        22,
        ssh_launch_opts.timeout,
        ssh_launch_opts.tty,
    )
    .await?;
    Ok(ssh)
}

/// Connect SSH and run a command that checks whether the system is ready for operation and then
/// shuts down.
///
/// The `qemu_should_exit` is used to tell the QEMU process to wait for process completion once
/// we've told the VM to shutdown. This is not the same as the cancellation token! We use the
/// cancellation token to actively cancel QEMU. In contrast, we use `qemu_should_exit` to signal
/// that QEMU is expected to exit on its own. Usually, we asssume QEMU exitting on its own is a
/// sign that something is wrong which is why we need a signal in cases where we expect it to exit.
#[instrument(skip(ssh_launch_opts))]
pub(crate) async fn connect_ssh_for_warmup(
    qemu_should_exit: Arc<AtomicBool>,
    ssh_launch_opts: SshLaunchOpts,
) -> Result<()> {
    let mut ssh = create_ssh_connection(&ssh_launch_opts).await?;
    info!("Connected");

    // First we'll wait until the system has fully booted up.
    let is_running_exitcode = ssh
        .call(
            Interactive::Never,
            vec![],
            "systemctl is-system-running --wait --quiet",
            None,
            Io::Inherit,
            Io::Inherit,
        )
        .await?;
    debug!(
        "systemctl is-system-running --wait exit code {}",
        is_running_exitcode.exit_code
    );

    // TODO: Here we'll add a stupid hack to deal with
    // https://github.com/linux-pam/linux-pam/issues/885 for the time being.
    ssh.call(
        Interactive::Never,
        vec![],
        "echo 127.0.0.1 unknown >> /etc/hosts",
        None,
        Io::Inherit,
        Io::Inherit,
    )
    .await?;

    // Allow the --env option to work by allowing SSH to accept all sent environment variables.
    ssh.call(
        Interactive::Never,
        vec![],
        "echo AcceptEnv * >> /etc/ssh/sshd_config",
        None,
        Io::Inherit,
        Io::Inherit,
    )
    .await?;

    // Then shut the system down.
    ssh.call(
        Interactive::Never,
        vec![],
        "systemctl poweroff",
        None,
        Io::Inherit,
        Io::Inherit,
    )
    .await?;
    debug!("Shutting down system");

    // Tell the QEMU handler it's now fine to wait for exit.
    qemu_should_exit.store(true, Ordering::SeqCst);

    // Ignore whatever error we might get from this as we want to close the connection at this
    // point anyway.
    let _ = ssh.close().await;
    Ok(())
}

/// Connect SSH and run a user-provided command.
///
/// If requested, this will be an interactive session.
///
/// The `cancellation_tokens` are used to cancel a running QEMU task in case there's a problem with
/// the SSH connection or upon command completion. The QEMU task can also use it to cancel the SSH
/// task.
#[instrument(skip(cancellation_tokens, ssh_launch_opts))]
pub(crate) async fn connect_ssh_for_command(
    cancellation_tokens: Option<CancellationTokens>,
    ssh_launch_opts: SshLaunchOpts,
    session_log_path: PathBuf,
) -> Result<Option<CommandOutput>> {
    let mut ssh = create_ssh_connection(&ssh_launch_opts)
        .await
        .inspect_err(|_| {
            if let Some(cancel_tokens) = cancellation_tokens.clone() {
                cancel_tokens.qemu.cancel();
            }
        })?;
    info!("Connected via SSH");

    // Open in append mode so that multiple commands on the same VM add to the same log.
    let mut session_log = Some(
        tokio::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(&session_log_path)
            .await
            .wrap_err(format!("Couldn't open session log at {session_log_path:?}"))?,
    );

    let output = {
        // We're using `termion` to put the terminal into raw mode, so that we can
        // display the output of interactive applications correctly.
        let _raw_term = if ssh_launch_opts.tty {
            Some(std::io::stdout().into_raw_mode()?)
        } else {
            None
        };

        let escaped_args = &ssh_launch_opts
            .args
            .into_iter()
            // arguments are escaped manually since the SSH protocol doesn't support quoting
            .map(|x| shell_escape::escape(x.into()))
            .collect::<Vec<_>>()
            .join(" ");

        // Handle workdir
        let escaped_args = if let Some(workdir) = ssh_launch_opts.workdir {
            let w = workdir.to_string_lossy();
            format!("cd {w} && {escaped_args}")
        } else {
            escaped_args.to_string()
        };

        if let Some(ref cancel_tokens) = cancellation_tokens {
            let ssh_output = tokio::select! {
                _ = cancel_tokens.ssh.cancelled() => {
                    debug!("SSH task was cancelled");
                    return Ok(None)
                }
                val = ssh.call(ssh_launch_opts.interactive, ssh_launch_opts.env_vars, &escaped_args, session_log.as_mut(), ssh_launch_opts.stdout, ssh_launch_opts.stderr) => {
                    val
                }
            };
            cancel_tokens.qemu.cancel();
            ssh_output?
        } else {
            let ssh_output = tokio::select! {
                val = ssh.call(ssh_launch_opts.interactive, ssh_launch_opts.env_vars, &escaped_args, session_log.as_mut(), ssh_launch_opts.stdout, ssh_launch_opts.stderr) => {
                    val
                }
            };
            ssh_output?
        }
    };

    info!("Exit code: {:?}", output.exit_code);
    if let Some(cancel_tokens) = cancellation_tokens {
        cancel_tokens.qemu.cancel();
    }
    ssh.close().await?;
    Ok(Some(output))
}