maestro/crates/maestro-registry/src/handler.rs

//! Per-connection handler for the `WS /daemon` endpoint.
//!
//! The handler is split from the axum-specific glue so it can be exercised
//! directly with an in-process pair of channels in tests. The axum side just
//! adapts a `WebSocket` into the trait below.

use std::time::Duration;

use axum::extract::ws::{Message, WebSocket};
use futures_util::{SinkExt, StreamExt};
use maestro_protocol::{DaemonInbound, DaemonOutbound, ErrorInfo};
use tokio::sync::mpsc;

use crate::state::Registry;

/// Bound the registry's per-daemon outbound queue. 64 in-flight R→D requests
/// is plenty for a single daemon and bounds memory use under back-pressure.
const OUTBOUND_BUFFER: usize = 64;

/// Run the handshake + message loop on a freshly-accepted WebSocket. Returns
/// the `daemon_id` once the connection terminates (for logging).
///
/// `heartbeat_timeout` is the maximum gap allowed between inbound frames —
/// once exceeded, the connection is dropped and the daemon's sessions are
/// removed from the catalog. Production code passes
/// [`crate::state::HEARTBEAT_TIMEOUT`]; tests use a short value.
pub async fn handle_socket(
    socket: WebSocket,
    registry: Registry,
    expected_token: &str,
    heartbeat_timeout: Duration,
) -> Option<String> {
    let (mut sink, mut stream) = socket.split();

    let hello = match read_hello(&mut stream).await {
        Some(h) => h,
        None => return None,
    };

    let (daemon_id, hostname, version, token, daemon_url) = match hello {
        DaemonOutbound::Hello {
            daemon_id,
            hostname,
            version,
            token,
            daemon_url,
        } => (daemon_id, hostname, version, token, daemon_url),
        _ => {
            let _ = send_error(&mut sink, "expected_hello", "first message must be Hello").await;
            return None;
        }
    };

    if token != expected_token {
        let _ = send_error(&mut sink, "auth_failed", "invalid token").await;
        return None;
    }

    let (out_tx, mut out_rx) = mpsc::channel::<DaemonInbound>(OUTBOUND_BUFFER);
    registry
        .register(daemon_id.clone(), hostname, version, daemon_url, out_tx)
        .await;

    let _ = sink
        .send(Message::Text(
            serde_json::to_string(&DaemonInbound::Ack).expect("ack serialize"),
        ))
        .await;

    loop {
        tokio::select! {
            inbound = tokio::time::timeout(heartbeat_timeout, stream.next()) => {
                match inbound {
                    Err(_) => {
                        tracing::info!(%daemon_id, "heartbeat timeout — dropping");
                        break;
                    }
                    Ok(None) => break,
                    Ok(Some(Err(e))) => {
                        tracing::info!(%daemon_id, ?e, "ws read error");
                        break;
                    }
                    Ok(Some(Ok(Message::Close(_)))) => break,
                    Ok(Some(Ok(Message::Text(text)))) => {
                        if !dispatch_text(&registry, &daemon_id, &text).await {
                            // malformed; close the connection
                            break;
                        }
                    }
                    Ok(Some(Ok(Message::Binary(_) | Message::Ping(_) | Message::Pong(_)))) => {
                        // Daemons only send text frames on this endpoint.
                    }
                }
            }
            Some(msg) = out_rx.recv() => {
                let text = serde_json::to_string(&msg).expect("serialize daemon-inbound");
                if sink.send(Message::Text(text)).await.is_err() {
                    break;
                }
            }
        }
    }

    registry.deregister(&daemon_id).await;
    Some(daemon_id)
}

/// Reads the first text frame, parses it as a `DaemonOutbound`, returns it.
async fn read_hello(
    stream: &mut futures_util::stream::SplitStream<WebSocket>,
) -> Option<DaemonOutbound> {
    match tokio::time::timeout(Duration::from_secs(5), stream.next()).await {
        Ok(Some(Ok(Message::Text(t)))) => serde_json::from_str(&t).ok(),
        _ => None,
    }
}

async fn dispatch_text(registry: &Registry, daemon_id: &str, text: &str) -> bool {
    let parsed: DaemonOutbound = match serde_json::from_str(text) {
        Ok(p) => p,
        Err(e) => {
            tracing::warn!(%daemon_id, error = %e, "malformed message");
            return false;
        }
    };
    match parsed {
        DaemonOutbound::Hello { .. } => {
            tracing::warn!(%daemon_id, "second Hello on live connection — ignoring");
        }
        DaemonOutbound::SessionsSnapshot { sessions } => {
            registry.apply_snapshot(daemon_id, sessions).await;
        }
        DaemonOutbound::SessionAdded { session } => {
            registry.add_session(daemon_id, session).await;
        }
        DaemonOutbound::SessionRemoved { id } => {
            registry.remove_session(daemon_id, id).await;
        }
        DaemonOutbound::SessionUpdated { session } => {
            registry.update_session(daemon_id, session).await;
        }
        DaemonOutbound::Heartbeat { .. } => {
            // Per-read timeout above is the heartbeat watchdog; no extra work.
            // perimeter_id / private_addr advertised here are persisted by
            // NAP-243 for the T2 transport.
        }
        DaemonOutbound::CreateSessionResult { reply_to, result } => {
            registry.complete_create(reply_to, result).await;
        }
        DaemonOutbound::KillSessionResult { reply_to, result } => {
            registry.complete_kill(reply_to, result).await;
        }
        DaemonOutbound::ResolveAttach { reply_to, .. }
        | DaemonOutbound::RelayAck { reply_to, .. } => {
            // Handled by NAP-243 (broker mode); type already lands here so old
            // daemons sending it don't kill the connection on parse failure.
            tracing::debug!(%daemon_id, ?reply_to, "broker message received but unhandled (NAP-243 pending)");
        }
    }
    true
}

async fn send_error(
    sink: &mut futures_util::stream::SplitSink<WebSocket, Message>,
    code: &str,
    message: &str,
) -> Result<(), axum::Error> {
    let err = DaemonInbound::Error(ErrorInfo {
        code: code.into(),
        message: message.into(),
    });
    let text = serde_json::to_string(&err).expect("serialize error");
    sink.send(Message::Text(text)).await
}