quicproquo/crates/quicprochat-rpc/src/auth_handshake.rs

//! Auth handshake protocol for QUIC connections.
//!
//! The session token is sent as a "connection init" message on the first
//! bi-stream before any RPC calls.
//!
//! ## Protocol
//! ```text
//! Client → Server: [0x01 magic][token_len: u16 BE][token bytes]
//! Server → Client: [0x01 magic][0x00 status OK]
//! ```

use crate::error::RpcError;

/// Magic byte identifying an auth init frame.
pub const AUTH_INIT_MAGIC: u8 = 0x01;

/// Status byte: auth accepted.
const AUTH_STATUS_OK: u8 = 0x00;

/// Maximum token length (64 KiB).
const MAX_TOKEN_LEN: usize = 65535;

/// Write an auth init frame to a QUIC send stream.
pub async fn send_auth_init(
    send: &mut quinn::SendStream,
    token: &[u8],
) -> Result<(), RpcError> {
    if token.len() > MAX_TOKEN_LEN {
        return Err(RpcError::Encode(format!(
            "auth token too large: {} bytes (max {MAX_TOKEN_LEN})",
            token.len()
        )));
    }

    let token_len = token.len() as u16;
    let mut buf = Vec::with_capacity(1 + 2 + token.len());
    buf.push(AUTH_INIT_MAGIC);
    buf.extend_from_slice(&token_len.to_be_bytes());
    buf.extend_from_slice(token);

    send.write_all(&buf)
        .await
        .map_err(|e| RpcError::Connection(format!("send auth init: {e}")))?;

    Ok(())
}

/// Read an auth init frame from a QUIC recv stream.
pub async fn recv_auth_init(
    recv: &mut quinn::RecvStream,
) -> Result<Vec<u8>, RpcError> {
    // Read magic byte.
    let mut header = [0u8; 3];
    recv.read_exact(&mut header)
        .await
        .map_err(|e| RpcError::Connection(format!("read auth init header: {e}")))?;

    if header[0] != AUTH_INIT_MAGIC {
        return Err(RpcError::Decode(format!(
            "bad auth init magic: expected 0x{AUTH_INIT_MAGIC:02x}, got 0x{:02x}",
            header[0]
        )));
    }

    let token_len = u16::from_be_bytes([header[1], header[2]]) as usize;
    if token_len > MAX_TOKEN_LEN {
        return Err(RpcError::Decode(format!(
            "auth token length {token_len} exceeds max {MAX_TOKEN_LEN}"
        )));
    }

    let mut token = vec![0u8; token_len];
    if token_len > 0 {
        recv.read_exact(&mut token)
            .await
            .map_err(|e| RpcError::Connection(format!("read auth token: {e}")))?;
    }

    Ok(token)
}

/// Send auth ack (success response).
pub async fn send_auth_ack(
    send: &mut quinn::SendStream,
) -> Result<(), RpcError> {
    let buf = [AUTH_INIT_MAGIC, AUTH_STATUS_OK];
    send.write_all(&buf)
        .await
        .map_err(|e| RpcError::Connection(format!("send auth ack: {e}")))?;
    Ok(())
}

/// Read auth ack from the server.
pub async fn recv_auth_ack(
    recv: &mut quinn::RecvStream,
) -> Result<(), RpcError> {
    let mut buf = [0u8; 2];
    recv.read_exact(&mut buf)
        .await
        .map_err(|e| RpcError::Connection(format!("read auth ack: {e}")))?;

    if buf[0] != AUTH_INIT_MAGIC {
        return Err(RpcError::Decode(format!(
            "bad auth ack magic: expected 0x{AUTH_INIT_MAGIC:02x}, got 0x{:02x}",
            buf[0]
        )));
    }

    if buf[1] != AUTH_STATUS_OK {
        return Err(RpcError::Decode(format!(
            "auth rejected with status 0x{:02x}",
            buf[1]
        )));
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn auth_init_magic_is_0x01() {
        assert_eq!(AUTH_INIT_MAGIC, 0x01);
    }

    #[test]
    fn token_too_large_returns_encode_error() {
        // We cannot call send_auth_init without a real stream, but we can
        // verify the length check logic by constructing the guard condition.
        let big_token = vec![0u8; MAX_TOKEN_LEN + 1];
        assert!(big_token.len() > MAX_TOKEN_LEN);
    }
}