initrd/src/cmd/init/sshd.rs

use log::{info, warn};
use std::fs;
use std::io::Write;
use std::os::unix::fs::PermissionsExt;
use std::process::Stdio;
use tokio::net;
use tokio::process::Command;

use super::retry_or_ignore;
use dkl::bootstrap::{Config, SSHServer};

pub async fn start(cfg: &Config) {
    retry_or_ignore(async || {
        info!("ssh: writing authorized keys");

        let ssh_dir = "/root/.ssh";
        let authorized_keys = format!("{ssh_dir}/authorized_keys");

        fs::create_dir_all(ssh_dir)?;
        fs::set_permissions(ssh_dir, fs::Permissions::from_mode(0o700))?;

        let mut ak = Vec::new();

        for auth in &cfg.auths {
            let Some(ref key) = auth.ssh_key else {
                continue;
            };
            writeln!(ak, "{key} {}", auth.name)?;
        }

        fs::write(authorized_keys, ak)?;
        Ok(())
    })
    .await;

    let cfg = cfg.ssh.clone();
    retry_or_ignore(async move || {
        // don't pre-start sshd as it should rarely be useful at this stage, use inetd-style.
        let listen_addr = cfg.listen.clone();
        info!("ssh: starting listener on {listen_addr}");

        let listener = net::TcpListener::bind(listen_addr).await?;

        tokio::spawn(handle_ssh_connections(listener, cfg.clone()));

        Ok(())
    })
    .await;
}

async fn handle_ssh_connections(listener: net::TcpListener, cfg: SSHServer) {
    let mut sshd_args = Vec::new();

    sshd_args.extend(["-i", "-E", "/var/log/sshd.log"]);

    let mut options = Vec::new();

    if let Some(ref user_ca) = cfg.user_ca {
        options.push(format!("TrustedUserCAKeys={user_ca}"));
    }

    for opt in &options {
        sshd_args.extend(["-o", &opt]);
    }

    let mut keygen_done = false;

    loop {
        let (stream, remote) = match listener.accept().await {
            Ok(v) => v,
            Err(e) => {
                warn!("ssh: listener stopped: {e}");
                return;
            }
        };

        if !keygen_done {
            // make sure we have ssh host keys even if not provided
            if (Command::new("ssh-keygen").arg("-A").status().await)
                .inspect_err(|e| warn!("ssh-keygen failed: {e}"))
                .is_ok_and(|s| s.success())
            {
                keygen_done = true
            }
        }

        use std::os::unix::io::{AsRawFd, FromRawFd};
        let fd = stream.as_raw_fd();

        let mut cmd = Command::new("/usr/sbin/sshd");
        cmd.args(&sshd_args);

        cmd.stdin(unsafe { Stdio::from_raw_fd(fd) });
        cmd.stdout(unsafe { Stdio::from_raw_fd(fd) });
        cmd.stderr(Stdio::null());

        let Ok(mut child) =
            (cmd.spawn()).inspect_err(|e| warn!("ssh: failed to start server: {e}"))
        else {
            continue;
        };

        let pid = child.id().unwrap();
        info!("ssh: new connection from {remote}, sshd PID {pid}");

        tokio::spawn(async move {
            let _ = child.wait().await;
        });
    }
}