Wire native Dosh auth runtime

src/bin/dosh-server.rs

@@ -1,16 +1,21 @@
 use anyhow::{Context, Result, anyhow};
 use clap::{Parser, Subcommand};
 use dosh::auth::{
-    build_bootstrap, encode_bootstrap, load_or_create_server_secret, open_attach_ticket,
-    verify_bootstrap,
+    build_attach_ticket, build_bootstrap, encode_bootstrap, load_or_create_server_secret, now_secs,
+    open_attach_ticket, verify_bootstrap,
 };
 use dosh::config::{ServerConfig, load_server_config};
 use dosh::crypto;
-use dosh::native::{host_public_key, host_public_key_line, load_or_create_host_key};
+use dosh::native::{
+    NativeAuthOk, NativeServerHello, derive_native_session_key, generate_native_ephemeral,
+    host_public_key, host_public_key_line, load_or_create_host_key, sign_server_hello,
+    verify_native_user_auth_from_config,
+};
 use dosh::protocol::{
     self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input,
-    PacketKind, ReplayWindow, Resize, ResumeRequest, SERVER_TO_CLIENT, TicketAttachBody,
-    TicketAttachEnvelope, TicketAttachOkEnvelope,
+    NativeAuthOkBody, NativeClientHelloBody, NativeServerHelloBody, NativeUserAuthBody, PacketKind,
+    ReplayWindow, Resize, ResumeRequest, SERVER_TO_CLIENT, TicketAttachBody, TicketAttachEnvelope,
+    TicketAttachOkEnvelope,
 };
 use dosh::pty::{PtyHandle, PtyOutput, spawn_pty_session};
 use std::collections::{HashMap, VecDeque};
@@ -157,6 +162,7 @@ struct ServerState {
     secret: [u8; 32],
     pty_tx: mpsc::UnboundedSender<PtyOutput>,
     sessions: HashMap<String, Session>,
+    pending_native: HashMap<[u8; 16], PendingNativeAuth>,
 }
 
 struct Session {
@@ -190,6 +196,14 @@ struct PendingFrame {
     attempts: u8,
 }
 
+struct PendingNativeAuth {
+    client: dosh::native::NativeClientHello,
+    server: NativeServerHello,
+    session_key: [u8; 32],
+    peer: SocketAddr,
+    created_at: Instant,
+}
+
 impl ServerState {
     fn new(
         config: ServerConfig,
@@ -201,6 +215,7 @@ impl ServerState {
             secret,
             pty_tx,
             sessions: HashMap::new(),
+            pending_native: HashMap::new(),
         }
     }
 
@@ -237,6 +252,10 @@ async fn handle_packet(
 ) -> Result<()> {
     let packet = protocol::decode(raw)?;
     match packet.header.kind {
+        PacketKind::NativeClientHello => {
+            handle_native_client_hello(state, socket, peer, packet.body).await
+        }
+        PacketKind::NativeUserAuth => handle_native_user_auth(state, socket, peer, &packet).await,
         PacketKind::BootstrapAttachRequest => {
             handle_bootstrap_attach(state, socket, peer, packet.body).await
         }
@@ -258,6 +277,246 @@ async fn handle_packet(
     }
 }
 
+async fn handle_native_client_hello(
+    state: &Arc<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
+    peer: SocketAddr,
+    body: Vec<u8>,
+) -> Result<()> {
+    let req: NativeClientHelloBody = protocol::from_body(&body)?;
+    let built: Result<([u8; 16], NativeServerHello)> = {
+        let mut locked = state.lock().expect("server state poisoned");
+        if !locked.config.native_auth {
+            Err(anyhow!("native auth disabled"))
+        } else if !req
+            .hello
+            .supported_aead
+            .iter()
+            .any(|algorithm| algorithm == "chacha20poly1305")
+        {
+            Err(anyhow!("native auth requires chacha20poly1305"))
+        } else if !req
+            .hello
+            .supported_user_key_algorithms
+            .iter()
+            .any(|algorithm| algorithm == "ssh-ed25519")
+        {
+            Err(anyhow!("native auth requires ssh-ed25519"))
+        } else {
+            let current_user = std::env::var("USER").unwrap_or_else(|_| "unknown".to_string());
+            if req.hello.requested_user != current_user {
+                Err(anyhow!("native auth user mismatch"))
+            } else {
+                let host_signing = load_or_create_host_key(&locked.config)?;
+                let (server_secret, server_public) = generate_native_ephemeral();
+                let mut hello = NativeServerHello {
+                    protocol_version: dosh::native::NATIVE_PROTOCOL_VERSION,
+                    server_random: crypto::random_32(),
+                    server_ephemeral_public: server_public,
+                    host_key: host_public_key(&host_signing),
+                    chosen_aead: "chacha20poly1305".to_string(),
+                    server_key_epoch: 1,
+                    auth_challenge: crypto::random_32(),
+                    rate_limit_remaining: Some(locked.config.native_auth_rate_limit_per_minute),
+                    host_signature: Vec::new(),
+                };
+                sign_server_hello(&host_signing, &req.hello, &mut hello)?;
+                let session_key = derive_native_session_key(
+                    &server_secret,
+                    req.hello.client_ephemeral_public,
+                    &req.hello,
+                    &hello,
+                )?;
+                let mut pending_id = [0u8; 16];
+                pending_id.copy_from_slice(&hello.auth_challenge[..16]);
+                locked.pending_native.insert(
+                    pending_id,
+                    PendingNativeAuth {
+                        client: req.hello,
+                        server: hello.clone(),
+                        session_key,
+                        peer,
+                        created_at: Instant::now(),
+                    },
+                );
+                Ok((pending_id, hello))
+            }
+        }
+    };
+    let (pending_id, hello) = match built {
+        Ok(built) => built,
+        Err(err) => return send_reject(socket, peer, &err.to_string()).await,
+    };
+    let body = protocol::to_body(&NativeServerHelloBody { hello })?;
+    let out = protocol::encode_plain(PacketKind::NativeServerHello, pending_id, 1, 0, &body)?;
+    socket.send_to(&out, peer).await?;
+    Ok(())
+}
+
+async fn handle_native_user_auth(
+    state: &Arc<Mutex<ServerState>>,
+    socket: &Arc<UdpSocket>,
+    peer: SocketAddr,
+    packet: &protocol::Packet,
+) -> Result<()> {
+    let pending = {
+        let mut locked = state.lock().expect("server state poisoned");
+        locked.pending_native.remove(&packet.header.conn_id)
+    };
+    let Some(pending) = pending else {
+        return send_reject_to_client(socket, peer, packet.header.conn_id, "unknown native auth")
+            .await;
+    };
+    if pending.peer != peer {
+        return send_reject_to_client(
+            socket,
+            peer,
+            packet.header.conn_id,
+            "native auth peer changed",
+        )
+        .await;
+    }
+    if pending.created_at.elapsed() > Duration::from_secs(30) {
+        return send_reject_to_client(socket, peer, packet.header.conn_id, "native auth expired")
+            .await;
+    }
+    let body = protocol::decrypt_body(packet, &pending.session_key, CLIENT_TO_SERVER)?;
+    let req: NativeUserAuthBody = protocol::from_body(&body)?;
+
+    let attached: Result<(
+        [u8; 16],
+        [u8; 16],
+        Vec<u8>,
+        [u8; 32],
+        String,
+        String,
+        u64,
+        Vec<u8>,
+    )> = {
+        let mut locked = state.lock().expect("server state poisoned");
+        verify_native_user_auth_from_config(
+            &locked.config,
+            &pending.client,
+            &pending.server,
+            &req.auth,
+        )
+        .context("verify native user auth")
+        .and_then(|_| {
+            let session_name = pending.client.requested_session.clone();
+            let mode = pending.client.requested_mode.clone();
+            let cols = pending.client.terminal_size.0;
+            let rows = pending.client.terminal_size.1;
+            if !locked.sessions.contains_key(&session_name) {
+                if locked.config.create_on_attach {
+                    locked.ensure_session(&session_name, cols, rows)?;
+                } else {
+                    return Err(anyhow!("session does not exist"));
+                }
+            }
+            let session_key = pending.session_key;
+            let key_id = {
+                let digest = crypto::sha256(&session_key);
+                let mut out = [0u8; 16];
+                out.copy_from_slice(&digest[..16]);
+                out
+            };
+            let attach_ticket_psk = crypto::random_32();
+            let issued_at = now_secs()?;
+            let attach_ticket = build_attach_ticket(
+                &locked.secret,
+                crypto::sha256(&locked.secret),
+                1,
+                pending.client.requested_user.clone(),
+                session_name.clone(),
+                mode.clone(),
+                issued_at,
+                issued_at + locked.config.attach_ticket_ttl_secs,
+                &attach_ticket_psk,
+            )?;
+            let session = locked
+                .sessions
+                .get_mut(&session_name)
+                .expect("session exists");
+            if mode != "view-only" {
+                session.pty.resize(cols, rows)?;
+                session.parser.set_size(rows, cols);
+            }
+            let client_id = crypto::random_16();
+            let snapshot = screen_snapshot(session.parser.screen());
+            let screen = session.parser.screen().clone();
+            let output_seq = session.output_seq;
+            session.clients.insert(
+                client_id,
+                ClientState {
+                    endpoint: peer,
+                    mode: mode.clone(),
+                    session_key,
+                    last_acked: output_seq,
+                    replay: ReplayWindow::default(),
+                    send_seq: 1,
+                    cols,
+                    rows,
+                    last_seen: Instant::now(),
+                    pending: VecDeque::new(),
+                    last_screen: Some(screen),
+                },
+            );
+            Ok((
+                client_id,
+                key_id,
+                attach_ticket,
+                attach_ticket_psk,
+                session_name,
+                mode,
+                output_seq,
+                snapshot,
+            ))
+        })
+    };
+
+    let (
+        client_id,
+        key_id,
+        attach_ticket,
+        attach_ticket_psk,
+        session_name,
+        mode,
+        output_seq,
+        snapshot,
+    ) = match attached {
+        Ok(attached) => attached,
+        Err(err) => {
+            return send_reject_to_client(socket, peer, packet.header.conn_id, &err.to_string())
+                .await;
+        }
+    };
+
+    let ok = NativeAuthOk {
+        client_id,
+        session: session_name,
+        mode,
+        session_key: pending.session_key,
+        session_key_id: key_id,
+        attach_ticket,
+        attach_ticket_psk,
+        initial_seq: output_seq,
+        snapshot,
+        policy_flags: Vec::new(),
+    };
+    let body = protocol::to_body(&NativeAuthOkBody { ok })?;
+    let out = protocol::encode_encrypted(
+        PacketKind::NativeAuthOk,
+        client_id,
+        1,
+        packet.header.seq,
+        &pending.session_key,
+        SERVER_TO_CLIENT,
+        &body,
+    )?;
+    socket.send_to(&out, peer).await?;
+    Ok(())
+}
+
 async fn handle_bootstrap_attach(
     state: &Arc<Mutex<ServerState>>,
     socket: &Arc<UdpSocket>,