sail-rs 0.1.1

//! SSH enablement for sailboxes, shared by the SDK and the CLI.
//!
//! Resolve a local public key, reserve guest port 22 as TCP ingress, install the
//! key, and (re)start `sshd`. The host key is generated once with `ssh-keygen
//! -A`, so re-running is safe and a client's `known_hosts` stays valid. openssh
//! is baked into every base and built image, so setup only regenerates host keys
//! and hardens the config; it never installs the server.

use std::time::{Duration, Instant};

use crate::error::SailError;
use crate::exec::{ExecParams, ExecProcess, EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS};
use crate::Client;

const SSHD_KEY_DIR_SETUP: &str = "mkdir -p /root/.ssh";
const AUTHORIZED_KEYS_PATH: &str = "/root/.ssh/authorized_keys";
const SSHD_SETUP: &str = "ssh-keygen -A && \
chown root:root /root /root/.ssh /root/.ssh/authorized_keys && \
chmod 700 /root/.ssh && \
passwd -d root && \
sed -i 's/^#*PermitRootLogin.*/PermitRootLogin prohibit-password/' /etc/ssh/sshd_config && \
sed -i 's/^#*PasswordAuthentication.*/PasswordAuthentication no/' /etc/ssh/sshd_config && \
mkdir -p /run/sshd";
const SSHD_SETUP_TIMEOUT_SECONDS: u32 = 60;
const SSHD_START: &str = "/usr/sbin/sshd -D -e";
const SSH_PUBLIC_KEY_FILENAMES: &[&str] = &["id_ed25519.pub", "id_rsa.pub"];
const KEY_PREFIXES: &[&str] = &["ssh-", "ecdsa-", "sk-"];

/// The public TCP endpoint a sailbox's SSH listener is reachable at.
#[derive(Debug, Clone)]
pub struct SshEndpoint {
    /// Hostname to dial.
    pub host: String,
    /// Port to dial.
    pub port: u32,
}

/// Resolve an SSH public key from an explicit value or `~/.ssh`. A value with a
/// space and a known key prefix is treated as a literal key; any other non-empty
/// value is a path to a `.pub` file. With no value, `~/.ssh/id_ed25519.pub` then
/// `id_rsa.pub` are tried.
pub fn resolve_public_key(value: Option<&str>) -> Result<String, SailError> {
    if let Some(raw) = value {
        let text = raw.trim();
        if !text.is_empty() {
            if text.contains(' ') && KEY_PREFIXES.iter().any(|p| text.starts_with(p)) {
                return validate_key(text);
            }
            let path = crate::credentials::expand_user(text);
            let key = std::fs::read_to_string(&path).map_err(|_| SailError::InvalidArgument {
                message: format!(
                    "ssh public key {text:?} is neither a recognized public key nor an existing file path"
                ),
            })?;
            let key = key.trim();
            if key.is_empty() {
                return Err(SailError::InvalidArgument {
                    message: format!("ssh public key file {text:?} is empty"),
                });
            }
            return validate_key(key);
        }
    }
    let home = dirs::home_dir().unwrap_or_default();
    let mut searched = Vec::new();
    for name in SSH_PUBLIC_KEY_FILENAMES {
        let candidate = home.join(".ssh").join(name);
        searched.push(candidate.display().to_string());
        if let Ok(key) = std::fs::read_to_string(&candidate) {
            let key = key.trim();
            if !key.is_empty() {
                return validate_key(key);
            }
        }
    }
    Err(SailError::InvalidArgument {
        message: format!(
            "no SSH public key found at {}; provide one explicitly",
            searched.join(" or ")
        ),
    })
}

/// Verify a value is a well-formed SSH public key: `<type> <base64 blob>` whose
/// blob begins with a length-prefixed copy of the type. Catching a malformed key
/// here turns it into a clear error instead of a silent sshd auth failure once
/// the box is up.
fn validate_key(key: &str) -> Result<String, SailError> {
    use base64::prelude::{Engine as _, BASE64_STANDARD};

    let invalid = |message: &str| SailError::InvalidArgument {
        message: message.to_string(),
    };
    let mut parts = key.split_whitespace();
    let key_type = parts.next().unwrap_or_default();
    let blob_b64 = parts.next().unwrap_or_default();
    if !KEY_PREFIXES.iter().any(|p| key_type.starts_with(p)) {
        return Err(invalid(
            "value does not look like an SSH public key (ssh-ed25519/ssh-rsa/ecdsa-.../sk-...)",
        ));
    }
    let blob = BASE64_STANDARD
        .decode(blob_b64)
        .map_err(|_| invalid("SSH public key body is not valid base64"))?;
    let type_len = blob
        .get(..4)
        .map_or(0, |b| u32::from_be_bytes([b[0], b[1], b[2], b[3]]) as usize);
    if blob.get(4..4 + type_len) != Some(key_type.as_bytes()) {
        return Err(invalid("SSH public key type does not match its body"));
    }
    Ok(key.to_string())
}

/// Build the idempotent `authorized_keys` merge command for an uploaded temp key.
fn authorized_keys_merge_command(tmp_path: &str) -> String {
    let authorized = shell_words::quote(AUTHORIZED_KEYS_PATH);
    let tmp = shell_words::quote(tmp_path);
    format!(
        "set -e; touch {authorized}; \
if ! grep -qxF -f {tmp} {authorized}; then \
if [ -s {authorized} ] && [ \"$(tail -c 1 {authorized})\" != '' ]; then \
printf '\\n' >> {authorized}; \
fi; \
cat {tmp} >> {authorized}; \
fi; \
rm -f {tmp}; \
chmod 600 {authorized}"
    )
}

fn ssh_exec_params(
    exec_endpoint: &str,
    sailbox_id: &str,
    argv: Vec<String>,
    timeout_seconds: u32,
) -> ExecParams {
    ExecParams {
        sailbox_id: sailbox_id.to_string(),
        exec_endpoint: exec_endpoint.to_string(),
        argv,
        timeout_seconds,
        idempotency_key: String::new(),
        open_stdin: false,
        pty: false,
        term: String::new(),
        cols: 0,
        rows: 0,
        retry_timeout: EXEC_TRANSIENT_RETRY_TIMEOUT_SECONDS,
        extra_metadata: Vec::new(),
    }
}

impl Client {
    /// Make a running sailbox reachable over SSH, returning the endpoint when
    /// `wait` is set (else `None`). Resumes (wakes) the sailbox to reach it,
    /// reserves guest port 22 as TCP ingress if not already exposed, installs
    /// `public_key`, and (re)starts `sshd`. Idempotent and safe to re-run.
    pub async fn enable_ssh(
        &self,
        sailbox_id: &str,
        public_key: &str,
        wait: bool,
        timeout: Duration,
    ) -> Result<Option<SshEndpoint>, SailError> {
        let exec_endpoint = self.exec_endpoint(sailbox_id).await?;

        // Reuse an existing port-22 listener (preserving any caller allowlist);
        // expose it as TCP only if absent.
        match self.get_listener(sailbox_id, 22).await {
            Ok(_) => {}
            Err(SailError::NotFound { .. }) => {
                self.expose_listener(
                    sailbox_id,
                    22,
                    crate::sailbox::types::IngressProtocol::Tcp,
                    &[],
                )
                .await?;
            }
            Err(err) => return Err(err),
        }

        self.ssh_exec_check(
            &exec_endpoint,
            sailbox_id,
            SSHD_KEY_DIR_SETUP,
            30,
            "ssh key dir setup",
        )
        .await?;

        // Upload the key to a temp path, then atomically merge it into
        // authorized_keys without removing existing keys.
        let tmp_path = format!(
            "/root/.ssh/.sail-authorized-key-{}.tmp",
            uuid::Uuid::new_v4().simple()
        );
        let mut writer =
            self.worker()
                .write_file(&exec_endpoint, sailbox_id, &tmp_path, true, Some(0o600));
        writer
            .write_chunk(format!("{public_key}\n").into_bytes())
            .await?;
        writer.finish().await?;

        self.ssh_exec_check(
            &exec_endpoint,
            sailbox_id,
            &authorized_keys_merge_command(&tmp_path),
            30,
            "ssh key install",
        )
        .await?;
        self.ssh_exec_check(
            &exec_endpoint,
            sailbox_id,
            SSHD_SETUP,
            SSHD_SETUP_TIMEOUT_SECONDS,
            "sshd setup",
        )
        .await?;

        // Start sshd detached; the daemon outlives this exec.
        let start = format!("nohup {SSHD_START} </dev/null >/dev/null 2>&1 &");
        let proc = ExecProcess::start(
            self.worker(),
            ssh_exec_params(
                &exec_endpoint,
                sailbox_id,
                vec!["/bin/sh".to_string(), "-c".to_string(), start],
                30,
            ),
        )
        .await?;
        proc.wait().await?;

        if !wait {
            return Ok(None);
        }
        self.wait_for_ssh_listener(sailbox_id, timeout)
            .await
            .map(Some)
    }

    /// Run a single shell command in the guest and fail on a non-zero exit.
    async fn ssh_exec_check(
        &self,
        exec_endpoint: &str,
        sailbox_id: &str,
        command: &str,
        timeout_seconds: u32,
        label: &str,
    ) -> Result<(), SailError> {
        let proc = ExecProcess::start(
            self.worker(),
            ssh_exec_params(
                exec_endpoint,
                sailbox_id,
                vec!["/bin/sh".to_string(), "-c".to_string(), command.to_string()],
                timeout_seconds,
            ),
        )
        .await?;
        let result = proc.wait().await?;
        if result.return_code != 0 {
            let detail = if result.stderr.trim().is_empty() {
                result.stdout.trim()
            } else {
                result.stderr.trim()
            };
            return Err(SailError::Internal {
                message: format!("{label} failed (exit {}): {detail}", result.return_code),
            });
        }
        Ok(())
    }

    /// Poll the port-22 listener until the SSH endpoint actually accepts, up to
    /// `timeout`. The route flips ACTIVE off the sailbox's running status, not a
    /// guest-port probe, so it can report a host:port before the freshly-started
    /// `sshd` is listening. Connecting and reading the SSH banner is the signal
    /// that the printed command will work on first use.
    async fn wait_for_ssh_listener(
        &self,
        sailbox_id: &str,
        timeout: Duration,
    ) -> Result<SshEndpoint, SailError> {
        let deadline = Instant::now() + timeout;
        loop {
            if let Ok(listener) = self.get_listener(sailbox_id, 22).await {
                if !listener.public_host.is_empty()
                    && listener.public_port != 0
                    && listener.is_active()
                    && ssh_endpoint_accepts(&listener.public_host, listener.public_port).await
                {
                    return Ok(SshEndpoint {
                        host: listener.public_host,
                        port: listener.public_port,
                    });
                }
            }
            if Instant::now() >= deadline {
                return Err(SailError::Internal {
                    message: "timed out waiting for the SSH port to become reachable".to_string(),
                });
            }
            tokio::time::sleep(Duration::from_secs(1)).await;
        }
    }
}

/// Whether the public endpoint answers with an SSH identification banner. A
/// fresh sshd (or the relay before the guest dial succeeds) accepts the TCP
/// connection but sends nothing, so the banner is the signal that ssh will work.
async fn ssh_endpoint_accepts(host: &str, port: u32) -> bool {
    use tokio::io::AsyncReadExt;
    use tokio::net::TcpStream;

    let probe = Duration::from_secs(5);
    let addr = format!("{host}:{port}");
    let Ok(Ok(mut stream)) = tokio::time::timeout(probe, TcpStream::connect(&addr)).await else {
        return false;
    };
    // The banner ("SSH-2.0-...") can arrive split across reads, so accumulate up
    // to its 4-byte prefix before deciding rather than rejecting a partial read.
    let mut buf = [0u8; 4];
    let mut filled = 0;
    while filled < 4 {
        match tokio::time::timeout(probe, stream.read(&mut buf[filled..])).await {
            Ok(Ok(0)) | Err(_) => break,
            Ok(Ok(n)) => filled += n,
            Ok(Err(_)) => break,
        }
    }
    buf[..filled].starts_with(b"SSH-")
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn literal_key_is_accepted() {
        let key =
            "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIAABAgMEBQYHCAkKCwwNDg8QERITFBUWFxgZGhscHR4f comment";
        assert_eq!(resolve_public_key(Some(key)).unwrap(), key);
    }

    #[test]
    fn non_key_non_path_is_rejected() {
        assert!(matches!(
            resolve_public_key(Some("not-a-key")),
            Err(SailError::InvalidArgument { .. })
        ));
    }

    #[test]
    fn key_with_malformed_body_is_rejected() {
        assert!(matches!(
            resolve_public_key(Some("ssh-ed25519 AAAAC3Nz comment")),
            Err(SailError::InvalidArgument { .. })
        ));
    }

    #[test]
    fn merge_command_quotes_paths_and_chmods() {
        let cmd = authorized_keys_merge_command("/root/.ssh/.tmp");
        assert!(cmd.contains("/root/.ssh/authorized_keys"));
        assert!(cmd.contains("chmod 600"));
    }
}