migrate to rust

src/cmd/init/sshd.rs

@@ -0,0 +1,110 @@
+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 crate::bootstrap::config::{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;
+        });
+    }
+}