huskies/server/src/crdt_sync/handshake/mod.rs

//! Auth handshake for the server-side `/crdt-sync` WebSocket.
//!
//! # Extended mutual-auth handshake protocol
//!
//! ```text
//! Connecting peer (client):          Responding node (server):
//!   hello(nonce) ──────────────────► receive hello
//!                ◄──────────────────  server_auth(pubkey, sign("huskies-v1:{nonce}"))
//!   verify server sig + trusted_keys
//!                ◄──────────────────  challenge(server_nonce)
//!   auth(pubkey, sign(server_nonce)) ►
//!                                     verify client sig + trusted_keys
//! ```
//!
//! Both sides verify the peer's pubkey against `trusted_keys`.  A peer whose
//! pubkey is absent from the allow-list is rejected with close code 4002.

#![allow(unused_imports, dead_code)]
use futures::{SinkExt, StreamExt};
use poem::web::websocket::Message as WsMessage;

use crate::crdt_state;
use crate::node_identity;
use crate::slog;

use super::AUTH_TIMEOUT_SECS;
use super::auth::trusted_keys;
use super::wire::{AuthMessage, ChallengeMessage, HelloMessage, ServerAuthMessage};

/// Perform the extended mutual-auth handshake for a freshly-upgraded WebSocket
/// connection.
///
/// **Protocol (server/responding-node side):**
/// 1. Receive `hello` from the connecting peer (contains client nonce).
/// 2. Sign `"huskies-v1:{nonce}"` and send `server_auth` (this node's pubkey +
///    signature) back to the connecting peer.
/// 3. Send a fresh challenge nonce to the connecting peer.
/// 4. Wait up to [`AUTH_TIMEOUT_SECS`] for a signed `auth` reply.
/// 5. Verify the connecting peer's signature and check its pubkey against the
///    trusted-key allow-list.
///
/// Returns `Some(AuthMessage)` on success.  On failure the connection has
/// already been closed with the appropriate close code (`auth_timeout` or
/// `auth_failed`); the caller should simply return.
pub(super) async fn perform_auth_handshake(
    sink: &mut futures::stream::SplitSink<poem::web::websocket::WebSocketStream, WsMessage>,
    stream: &mut futures::stream::SplitStream<poem::web::websocket::WebSocketStream>,
) -> Option<AuthMessage> {
    // ── Step 1: Receive hello from connecting peer ───────────────────
    let hello_result = tokio::time::timeout(
        std::time::Duration::from_secs(AUTH_TIMEOUT_SECS),
        stream.next(),
    )
    .await;

    let hello_text = match hello_result {
        Ok(Some(Ok(WsMessage::Text(text)))) => text,
        Ok(_) | Err(_) => {
            slog!("[crdt-sync] No hello from peer — closing");
            close_with_auth_failed(sink).await;
            return None;
        }
    };

    let hello: HelloMessage = match serde_json::from_str::<HelloMessage>(&hello_text) {
        Ok(m) if m.r#type == "hello" => m,
        _ => {
            slog!("[crdt-sync] Invalid hello message from peer");
            close_with_auth_failed(sink).await;
            return None;
        }
    };

    // ── Step 2: Sign versioned challenge and send server_auth ────────
    let (server_pubkey_hex, server_sig_hex) =
        match crdt_state::sign_versioned_challenge(&hello.nonce) {
            Some(v) => v,
            None => {
                slog!("[crdt-sync] CRDT not initialised — cannot produce server_auth");
                close_with_auth_failed(sink).await;
                return None;
            }
        };

    let server_auth = ServerAuthMessage {
        r#type: "server_auth".to_string(),
        pubkey_hex: server_pubkey_hex,
        signature_hex: server_sig_hex,
    };
    let server_auth_json = serde_json::to_string(&server_auth).ok()?;
    if sink.send(WsMessage::Text(server_auth_json)).await.is_err() {
        return None;
    }

    // ── Step 3: Send challenge nonce to connecting peer ──────────────
    let challenge = node_identity::generate_challenge();
    let challenge_msg = ChallengeMessage {
        r#type: "challenge".to_string(),
        nonce: challenge.clone(),
    };
    let challenge_json = serde_json::to_string(&challenge_msg).ok()?;
    if sink.send(WsMessage::Text(challenge_json)).await.is_err() {
        return None;
    }

    // ── Step 4: Await signed auth reply from connecting peer ─────────
    let auth_result = tokio::time::timeout(
        std::time::Duration::from_secs(AUTH_TIMEOUT_SECS),
        stream.next(),
    )
    .await;

    let auth_text = match auth_result {
        Ok(Some(Ok(WsMessage::Text(text)))) => text,
        Ok(_) | Err(_) => {
            slog!("[crdt-sync] Auth timeout or connection lost during handshake");
            let _ = sink
                .send(WsMessage::Close(Some((
                    poem::web::websocket::CloseCode::from(4001),
                    "auth_timeout".to_string(),
                ))))
                .await;
            let _ = sink.close().await;
            return None;
        }
    };

    let auth_msg: AuthMessage = match serde_json::from_str(&auth_text) {
        Ok(m) => m,
        Err(_) => {
            slog!("[crdt-sync] Invalid auth message from peer");
            close_with_auth_failed(sink).await;
            return None;
        }
    };

    // ── Step 5: Verify signature and check trusted-key allow-list ────
    let key_trusted = trusted_keys().iter().any(|k| k == &auth_msg.pubkey_hex);
    if !key_trusted {
        slog!(
            "[crdt-sync] Auth rejected: peer pubkey not in trusted_keys: {}",
            auth_msg.pubkey_hex
        );
        close_with_auth_failed(sink).await;
        return None;
    }

    let sig_valid =
        node_identity::verify_challenge(&auth_msg.pubkey_hex, &challenge, &auth_msg.signature_hex);
    if !sig_valid {
        slog!(
            "[crdt-sync] Auth rejected: invalid signature from peer {:.12}…",
            &auth_msg.pubkey_hex
        );
        close_with_auth_failed(sink).await;
        return None;
    }

    slog!(
        "[crdt-sync] Peer authenticated: {:.12}…",
        &auth_msg.pubkey_hex
    );

    Some(auth_msg)
}

/// Close the WebSocket with a generic `auth_failed` reason.
///
/// The close reason is intentionally the same for all auth failures
/// (bad signature, untrusted key, malformed message) to avoid leaking
/// which check failed.
async fn close_with_auth_failed(
    sink: &mut futures::stream::SplitSink<poem::web::websocket::WebSocketStream, WsMessage>,
) {
    let _ = sink
        .send(WsMessage::Close(Some((
            poem::web::websocket::CloseCode::from(4002),
            "auth_failed".to_string(),
        ))))
        .await;
    let _ = sink.close().await;
}

#[cfg(test)]
mod tests;