subc-transport 0.5.0

use std::{error::Error, fmt, future::Future, io, time::Duration};

use hmac::{Hmac, Mac};
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use sha2::Sha256;
use subtle::ConstantTimeEq;
use tokio::{
    io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
    time,
};

use crate::connection_file::{ConnectionInfo, DAEMON_ID_LEN, MIN_KEY_LEN};

pub const NONCE_LEN: usize = 32;
pub const PROOF_LEN: usize = 32;
pub const MAX_AUTH_MESSAGE_LEN: u32 = 4096;
pub const SERVER_PROOF_DOMAIN: &str = "subc-server-v1";
pub const CLIENT_AUTH_DOMAIN: &str = "subc-client-v1";
pub const DEFAULT_CLIENT_ROLE: &str = "client";
pub const WATCHDOG_CLIENT_ROLE: &str = "watchdog";

type HmacSha256 = Hmac<Sha256>;

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ClientHello {
    pub client_nonce: [u8; NONCE_LEN],
    pub role: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ServerProof {
    pub daemon_id: [u8; DAEMON_ID_LEN],
    pub server_nonce: [u8; NONCE_LEN],
    pub daemon_ver: String,
    pub server_proof: [u8; PROOF_LEN],
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ClientAuth {
    pub client_auth: [u8; PROOF_LEN],
}

/// The outcome of a successful handshake.
///
/// WHAT THIS PROVES: the peer possesses the connection key, and (client side)
/// that the daemon does too. Nothing more.
///
/// WHAT `role` IS NOT: it is a string the CLIENT SENT, echoed back unverified.
/// The handshake never checks it against anything, so it carries no authority --
/// any peer holding the key can claim any role. It exists so a caller can tell
/// self-issued traffic (the daemon's own watchdog probe) from real clients when
/// REPORTING, and it must never decide admission, capacity, or privilege.
///
/// A type called `Authenticated` invites reading every field as attested. Only
/// the possession of the key is. Module identity, which IS attested, travels a
/// different path entirely (spawn nonces validated at route.open).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Authenticated {
    pub role: String,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AuthStage {
    ClientHello,
    ServerProof,
    ClientAuth,
}

#[derive(Debug)]
pub enum AuthError {
    Io {
        stage: AuthStage,
        source: io::Error,
    },
    Timeout {
        stage: AuthStage,
        deadline: Duration,
    },
    UnexpectedEof {
        stage: AuthStage,
        expected: usize,
        actual: usize,
    },
    MessageTooLarge {
        stage: AuthStage,
        len: u32,
        max: u32,
    },
    JsonEncode {
        stage: AuthStage,
        source: serde_json::Error,
    },
    JsonDecode {
        stage: AuthStage,
        source: serde_json::Error,
    },
    Random(getrandom::Error),
    KeyTooShort {
        len: usize,
        min: usize,
    },
    InvalidServerProof,
    DaemonIdMismatch,
    InvalidClientAuth,
}

pub fn compute_proof(
    key: &[u8],
    domain: &str,
    client_nonce: &[u8; NONCE_LEN],
    server_nonce: &[u8; NONCE_LEN],
    daemon_id: &[u8],
) -> [u8; PROOF_LEN] {
    let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts keys of any length");
    mac.update(domain.as_bytes());
    mac.update(client_nonce);
    mac.update(server_nonce);
    mac.update(daemon_id);
    mac.finalize().into_bytes().into()
}

/// An absolute handshake deadline. Every per-stage read/write recomputes the time
/// remaining until `at`, so the WHOLE handshake (length byte + body, across all
/// stages, plus error teardown) is bounded by a single wall-clock budget. Passing
/// a bare `Duration` to each step instead would let a slow peer spend the full
/// budget on every length read AND every body read — multiplying the real bound.
#[derive(Clone, Copy)]
struct Deadline {
    at: time::Instant,
    total: Duration,
}

impl Deadline {
    fn starting_now(total: Duration) -> Self {
        Self {
            at: time::Instant::now() + total,
            total,
        }
    }

    /// Time left until the deadline, or `Timeout` if it has already elapsed.
    fn remaining(&self, stage: AuthStage) -> Result<Duration, AuthError> {
        let remaining = self.at.saturating_duration_since(time::Instant::now());
        if remaining.is_zero() {
            Err(AuthError::Timeout {
                stage,
                deadline: self.total,
            })
        } else {
            Ok(remaining)
        }
    }

    /// Time left until the deadline, clamped to zero — for best-effort teardown
    /// that must not outlive the handshake budget.
    fn remaining_or_zero(&self) -> Duration {
        self.at.saturating_duration_since(time::Instant::now())
    }
}

pub async fn authenticate_server<S>(
    stream: &mut S,
    key: &[u8],
    daemon_id: &[u8; DAEMON_ID_LEN],
    daemon_ver: &str,
    deadline: Duration,
) -> Result<Authenticated, AuthError>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    let deadline = Deadline::starting_now(deadline);
    let result = authenticate_server_inner(stream, key, daemon_id, daemon_ver, deadline).await;
    if result.is_err() {
        // Bound teardown by the SAME absolute deadline so a failed handshake (and
        // the unauthenticated-handshake slot it holds) is released promptly instead
        // of waiting out another full budget.
        let _ = time::timeout(deadline.remaining_or_zero(), stream.shutdown()).await;
    }
    result
}

async fn authenticate_server_inner<S>(
    stream: &mut S,
    key: &[u8],
    daemon_id: &[u8; DAEMON_ID_LEN],
    daemon_ver: &str,
    deadline: Deadline,
) -> Result<Authenticated, AuthError>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    validate_key(key)?;

    let hello: ClientHello = read_message(stream, AuthStage::ClientHello, deadline).await?;
    let server_nonce = random_nonce()?;
    let server_proof = compute_proof(
        key,
        SERVER_PROOF_DOMAIN,
        &hello.client_nonce,
        &server_nonce,
        daemon_id,
    );

    write_message(
        stream,
        AuthStage::ServerProof,
        &ServerProof {
            daemon_id: *daemon_id,
            server_nonce,
            daemon_ver: daemon_ver.to_owned(),
            server_proof,
        },
        deadline,
    )
    .await?;

    let client_auth: ClientAuth = read_message(stream, AuthStage::ClientAuth, deadline).await?;
    let expected_client_auth = compute_proof(
        key,
        CLIENT_AUTH_DOMAIN,
        &hello.client_nonce,
        &server_nonce,
        daemon_id,
    );
    if !constant_time_eq(&expected_client_auth, &client_auth.client_auth) {
        return Err(AuthError::InvalidClientAuth);
    }

    Ok(Authenticated { role: hello.role })
}

pub async fn authenticate_client<S>(
    stream: &mut S,
    conn: &ConnectionInfo,
    deadline: Duration,
) -> Result<(), AuthError>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    authenticate_client_with_role(stream, conn, deadline, DEFAULT_CLIENT_ROLE).await
}

pub async fn authenticate_client_with_role<S>(
    stream: &mut S,
    conn: &ConnectionInfo,
    deadline: Duration,
    role: &str,
) -> Result<(), AuthError>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    let deadline = Deadline::starting_now(deadline);
    let result = authenticate_client_inner(stream, conn, deadline, role).await;
    if result.is_err() {
        let _ = time::timeout(deadline.remaining_or_zero(), stream.shutdown()).await;
    }
    result
}

async fn authenticate_client_inner<S>(
    stream: &mut S,
    conn: &ConnectionInfo,
    deadline: Deadline,
    role: &str,
) -> Result<(), AuthError>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    validate_key(&conn.key)?;

    let client_nonce = random_nonce()?;
    write_message(
        stream,
        AuthStage::ClientHello,
        &ClientHello {
            client_nonce,
            role: role.to_owned(),
        },
        deadline,
    )
    .await?;

    let server_proof: ServerProof = read_message(stream, AuthStage::ServerProof, deadline).await?;
    let expected_server_proof = compute_proof(
        &conn.key,
        SERVER_PROOF_DOMAIN,
        &client_nonce,
        &server_proof.server_nonce,
        &server_proof.daemon_id,
    );
    if !constant_time_eq(&expected_server_proof, &server_proof.server_proof) {
        return Err(AuthError::InvalidServerProof);
    }
    if server_proof.daemon_id != conn.daemon_id {
        return Err(AuthError::DaemonIdMismatch);
    }

    let client_auth = compute_proof(
        &conn.key,
        CLIENT_AUTH_DOMAIN,
        &client_nonce,
        &server_proof.server_nonce,
        &server_proof.daemon_id,
    );
    write_message(
        stream,
        AuthStage::ClientAuth,
        &ClientAuth { client_auth },
        deadline,
    )
    .await
}

fn validate_key(key: &[u8]) -> Result<(), AuthError> {
    if key.len() < MIN_KEY_LEN {
        return Err(AuthError::KeyTooShort {
            len: key.len(),
            min: MIN_KEY_LEN,
        });
    }
    Ok(())
}

fn random_nonce() -> Result<[u8; NONCE_LEN], AuthError> {
    let mut nonce = [0u8; NONCE_LEN];
    getrandom::getrandom(&mut nonce).map_err(AuthError::Random)?;
    Ok(nonce)
}

/// Both directions of this comparison are fenced, verified by mutation rather than
/// assumed, because a proof check has the failure mode where a suite proves only
/// that it can say NO.
///
/// ALWAYS-FALSE (no proof ever verifies, every connection in the fleet refused) is
/// caught by the handshake integration tests and by two bootstrap tests -- named
/// for key rotation and singleton probing, so this is coverage carried by tests
/// about something else. Narrowing either would remove it silently.
///
/// ALWAYS-TRUE is caught by `foreign_server_reused_port_never_receives_client_auth`
/// -- the case where a client must refuse a server that cannot produce the proof.
/// Named for the refusal, and it holds that direction directly.
///
/// The construction feeding it is pinned separately by
/// `committed_wire_vectors_pin_the_proof_construction`, which reddens for a constant
/// proof AND for one that folds only part of its input -- the second matters because
/// a partial-input proof still produces different outputs for different inputs, so
/// any distinctness assertion passes it while a proof minted against one daemon
/// verifies against another.
fn constant_time_eq(expected: &[u8; PROOF_LEN], actual: &[u8; PROOF_LEN]) -> bool {
    expected.as_slice().ct_eq(actual.as_slice()).into()
}

async fn read_message<S, T>(
    stream: &mut S,
    stage: AuthStage,
    deadline: Deadline,
) -> Result<T, AuthError>
where
    S: AsyncRead + Unpin,
    T: DeserializeOwned,
{
    // Both the length read and the body read recompute the time remaining against
    // the same absolute deadline, so the two together cannot exceed the budget.
    let mut len_bytes = [0u8; 4];
    read_exact_deadline(stream, &mut len_bytes, stage, deadline).await?;
    let len = u32::from_le_bytes(len_bytes);
    if len > MAX_AUTH_MESSAGE_LEN {
        return Err(AuthError::MessageTooLarge {
            stage,
            len,
            max: MAX_AUTH_MESSAGE_LEN,
        });
    }

    let mut json = vec![0u8; len as usize];
    if !json.is_empty() {
        read_exact_deadline(stream, &mut json, stage, deadline).await?;
    }
    serde_json::from_slice(&json).map_err(|source| AuthError::JsonDecode { stage, source })
}

async fn write_message<S, T>(
    stream: &mut S,
    stage: AuthStage,
    value: &T,
    deadline: Deadline,
) -> Result<(), AuthError>
where
    S: AsyncWrite + Unpin,
    T: Serialize,
{
    let json =
        serde_json::to_vec(value).map_err(|source| AuthError::JsonEncode { stage, source })?;
    let len = u32::try_from(json.len()).map_err(|_| AuthError::MessageTooLarge {
        stage,
        len: u32::MAX,
        max: MAX_AUTH_MESSAGE_LEN,
    })?;
    if len > MAX_AUTH_MESSAGE_LEN {
        return Err(AuthError::MessageTooLarge {
            stage,
            len,
            max: MAX_AUTH_MESSAGE_LEN,
        });
    }

    write_all_deadline(stream, &len.to_le_bytes(), stage, deadline).await?;
    write_all_deadline(stream, &json, stage, deadline).await
}

async fn read_exact_deadline<S>(
    stream: &mut S,
    buf: &mut [u8],
    stage: AuthStage,
    deadline: Deadline,
) -> Result<(), AuthError>
where
    S: AsyncRead + Unpin,
{
    let remaining = deadline.remaining(stage)?;
    let expected = buf.len();
    with_timeout(stage, remaining, async {
        let mut actual = 0;
        while actual < expected {
            let read = stream.read(&mut buf[actual..]).await?;
            if read == 0 {
                return Err(ReadExactError::UnexpectedEof { actual });
            }
            actual += read;
        }
        Ok(())
    })
    .await
    .map_err(|err| match err {
        DeadlineIoError::Io(source) => AuthError::Io { stage, source },
        DeadlineIoError::Timeout => AuthError::Timeout {
            stage,
            deadline: deadline.total,
        },
        DeadlineIoError::UnexpectedEof { actual } => AuthError::UnexpectedEof {
            stage,
            expected,
            actual,
        },
    })
}

async fn write_all_deadline<S>(
    stream: &mut S,
    buf: &[u8],
    stage: AuthStage,
    deadline: Deadline,
) -> Result<(), AuthError>
where
    S: AsyncWrite + Unpin,
{
    let remaining = deadline.remaining(stage)?;
    timeout_io(stage, remaining, deadline.total, stream.write_all(buf)).await
}

async fn timeout_io<T, F>(
    stage: AuthStage,
    remaining: Duration,
    total: Duration,
    future: F,
) -> Result<T, AuthError>
where
    F: Future<Output = io::Result<T>>,
{
    match time::timeout(remaining, future).await {
        Ok(Ok(value)) => Ok(value),
        Ok(Err(source)) => Err(AuthError::Io { stage, source }),
        Err(_) => Err(AuthError::Timeout {
            stage,
            deadline: total,
        }),
    }
}

async fn with_timeout<F>(
    _stage: AuthStage,
    deadline: Duration,
    future: F,
) -> Result<(), DeadlineIoError>
where
    F: Future<Output = Result<(), ReadExactError>>,
{
    match time::timeout(deadline, future).await {
        Ok(Ok(())) => Ok(()),
        Ok(Err(ReadExactError::Io(source))) => Err(DeadlineIoError::Io(source)),
        Ok(Err(ReadExactError::UnexpectedEof { actual })) => {
            Err(DeadlineIoError::UnexpectedEof { actual })
        }
        Err(_) => Err(DeadlineIoError::Timeout),
    }
}

#[derive(Debug)]
enum ReadExactError {
    Io(io::Error),
    UnexpectedEof { actual: usize },
}

impl From<io::Error> for ReadExactError {
    fn from(source: io::Error) -> Self {
        Self::Io(source)
    }
}

#[derive(Debug)]
enum DeadlineIoError {
    Io(io::Error),
    Timeout,
    UnexpectedEof { actual: usize },
}

impl fmt::Display for AuthError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Io { stage, source } => write!(f, "auth {stage:?} I/O error: {source}"),
            Self::Timeout { stage, deadline } => {
                write!(f, "auth {stage:?} timed out after {deadline:?}")
            }
            Self::UnexpectedEof {
                stage,
                expected,
                actual,
            } => write!(
                f,
                "auth {stage:?} ended early: expected {expected} bytes, got {actual}"
            ),
            Self::MessageTooLarge { stage, len, max } => write!(
                f,
                "auth {stage:?} message length {len} exceeds hard cap {max}"
            ),
            Self::JsonEncode { stage, source } => {
                write!(f, "auth {stage:?} JSON encode error: {source}")
            }
            Self::JsonDecode { stage, source } => {
                write!(f, "auth {stage:?} JSON decode error: {source}")
            }
            Self::Random(source) => write!(f, "auth random generation failed: {source}"),
            Self::KeyTooShort { len, min } => {
                write!(f, "auth key is too short: {len} bytes, need at least {min}")
            }
            Self::InvalidServerProof => write!(f, "invalid server auth proof"),
            Self::DaemonIdMismatch => write!(f, "server daemon_id did not match connection file"),
            Self::InvalidClientAuth => write!(f, "invalid client auth proof"),
        }
    }
}

impl Error for AuthError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::Io { source, .. } => Some(source),
            Self::JsonEncode { source, .. } | Self::JsonDecode { source, .. } => Some(source),
            Self::Random(_) => None,
            Self::Timeout { .. }
            | Self::UnexpectedEof { .. }
            | Self::MessageTooLarge { .. }
            | Self::KeyTooShort { .. }
            | Self::InvalidServerProof
            | Self::DaemonIdMismatch
            | Self::InvalidClientAuth => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Endpoint;
    use tokio::{
        io::{duplex, AsyncReadExt, AsyncWriteExt, DuplexStream},
        task::yield_now,
        time::advance,
    };

    const TEST_DAEMON_VER: &str = "subc-auth-test-1";

    /// The Swift client asserts its handshake against a committed vector file, so
    /// that file is a cross-language contract with two authors and no shared
    /// compiler between them. Until this test existed it was only a SNAPSHOT: it
    /// recorded what `compute_proof` produced on the day someone generated it, and
    /// nothing here would fail if the proof construction later changed. The Swift
    /// side would keep passing against the stale bytes and the two implementations
    /// would agree with the file while disagreeing with each other on the wire.
    ///
    /// Asserting the vectors from THIS side turns the snapshot into a pin: changing
    /// the domain separator, the field order, or the MAC breaks the build here,
    /// where the change is being made, instead of surfacing as a handshake failure
    /// against a peer that has not been rebuilt.
    #[test]
    fn committed_wire_vectors_pin_the_proof_construction() {
        // Kept in step with clients/subc-client-swift/Tests/SubcClientTests/
        // Fixtures/wire_vectors.json. Hex is written out rather than read from the
        // file so that regenerating the fixture cannot silently satisfy this test:
        // the expected value has to be changed by a person, in a diff a reviewer
        // sees, in the same commit as the construction change that motivated it.
        struct ProofVector {
            name: &'static str,
            key_hex: String,
            domain: &'static str,
            client_nonce_hex: String,
            server_nonce_hex: String,
            daemon_id_hex: String,
            expected_proof_hex: &'static str,
        }

        fn unhex(hex: &str, what: &str) -> Vec<u8> {
            assert!(hex.len().is_multiple_of(2), "{what}: odd-length hex");
            (0..hex.len())
                .step_by(2)
                .map(|i| {
                    u8::from_str_radix(&hex[i..i + 2], 16)
                        .unwrap_or_else(|_| panic!("{what}: non-hex byte at {i}"))
                })
                .collect()
        }

        let key_hex = "000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f";
        let client_nonce_hex = "ab".repeat(NONCE_LEN);
        let server_nonce_hex = "cd".repeat(NONCE_LEN);
        let daemon_id_hex = "000102030405060708090a0b0c0d0e0f";

        let vectors = [
            ProofVector {
                name: "server_proof_fixed_inputs",
                key_hex: key_hex.to_owned(),
                domain: SERVER_PROOF_DOMAIN,
                client_nonce_hex: client_nonce_hex.clone(),
                server_nonce_hex: server_nonce_hex.clone(),
                daemon_id_hex: daemon_id_hex.to_owned(),
                expected_proof_hex:
                    "ea06076a980bc7558e45017df86de89f3d2fc09861f8460795dea31eadf40527",
            },
            ProofVector {
                name: "client_auth_fixed_inputs",
                key_hex: key_hex.to_owned(),
                domain: CLIENT_AUTH_DOMAIN,
                client_nonce_hex,
                server_nonce_hex,
                daemon_id_hex: daemon_id_hex.to_owned(),
                expected_proof_hex:
                    "a3bc64784dbd94c4f52799e0c66f7b2b8183aa4655594630245c5d4a2fa387a9",
            },
        ];

        for vector in &vectors {
            let client_nonce: [u8; NONCE_LEN] = unhex(&vector.client_nonce_hex, vector.name)
                .try_into()
                .expect("client nonce length");
            let server_nonce: [u8; NONCE_LEN] = unhex(&vector.server_nonce_hex, vector.name)
                .try_into()
                .expect("server nonce length");
            let proof = compute_proof(
                &unhex(&vector.key_hex, vector.name),
                vector.domain,
                &client_nonce,
                &server_nonce,
                &unhex(&vector.daemon_id_hex, vector.name),
            );
            let actual: String = proof.iter().map(|byte| format!("{byte:02x}")).collect();
            if vector.expected_proof_hex.is_empty() {
                panic!(
                    "vector {} needs its expected proof recorded: {actual}",
                    vector.name
                );
            }
            assert_eq!(
                actual, vector.expected_proof_hex,
                "proof construction changed for vector {}: the committed Swift wire \
                 vectors no longer describe this implementation",
                vector.name
            );
        }
    }

    #[tokio::test(start_paused = true)]
    async fn authenticate_server_deadline_is_absolute_across_handshake() {
        let key = vec![0x5a; MIN_KEY_LEN];
        let daemon_id = [0x6b; DAEMON_ID_LEN];
        let deadline = Duration::from_millis(100);
        let stage_delay = Duration::from_millis(60);
        let (mut client, mut server) = duplex(4096);

        let server_task = tokio::spawn(async move {
            authenticate_server(&mut server, &key, &daemon_id, TEST_DAEMON_VER, deadline).await
        });

        yield_now().await;
        assert!(!server_task.is_finished());

        advance(stage_delay).await;
        write_auth_json(
            &mut client,
            &ClientHello {
                client_nonce: [0x11; NONCE_LEN],
                role: DEFAULT_CLIENT_ROLE.to_owned(),
            },
        )
        .await;
        yield_now().await;

        let server_proof: ServerProof = read_auth_json(&mut client).await;
        assert_eq!(server_proof.daemon_id, daemon_id);
        assert_eq!(server_proof.daemon_ver, TEST_DAEMON_VER);
        assert!(!server_task.is_finished());

        advance(stage_delay).await;
        yield_now().await;
        assert!(server_task.is_finished());

        let err = server_task
            .await
            .expect("server task should join")
            .expect_err("server handshake should time out once the total deadline elapses");
        assert!(matches!(
            err,
            AuthError::Timeout {
                stage: AuthStage::ClientAuth,
                ..
            }
        ));
    }

    /// Write only the 4-byte length prefix of an auth message, withholding the
    /// body — used to stall a peer mid-message so the within-stage deadline can be
    /// exercised (the length read and body read must share one absolute budget).
    async fn write_auth_len_only<T>(stream: &mut DuplexStream, value: &T)
    where
        T: Serialize,
    {
        let body = serde_json::to_vec(value).expect("encode auth json");
        stream
            .write_all(&(body.len() as u32).to_le_bytes())
            .await
            .expect("write auth length");
    }

    #[tokio::test(start_paused = true)]
    async fn server_deadline_spans_length_and_body_within_one_stage() {
        // The bug this guards: applying the timeout independently to the length
        // read and the body read lets a single stage consume ~2x the budget. Here
        // the client sends the ClientHello length prefix late, then withholds the
        // body until the absolute deadline has passed. With one shared deadline the
        // server times out at the deadline; with per-read deadlines the body read
        // would get a fresh full window and not time out yet.
        let key = vec![0x5a; MIN_KEY_LEN];
        let daemon_id = [0x6b; DAEMON_ID_LEN];
        let deadline = Duration::from_millis(100);
        let (mut client, mut server) = duplex(4096);

        let server_task = tokio::spawn(async move {
            authenticate_server(&mut server, &key, &daemon_id, TEST_DAEMON_VER, deadline).await
        });

        yield_now().await;
        // Burn most of the budget, then send only the length prefix.
        advance(Duration::from_millis(60)).await;
        write_auth_len_only(
            &mut client,
            &ClientHello {
                client_nonce: [0x11; NONCE_LEN],
                role: DEFAULT_CLIENT_ROLE.to_owned(),
            },
        )
        .await;
        yield_now().await;
        assert!(!server_task.is_finished());

        // Cross the absolute deadline (60 + 50 > 100) without sending the body.
        advance(Duration::from_millis(50)).await;
        yield_now().await;
        assert!(
            server_task.is_finished(),
            "body read must share the handshake deadline, not get a fresh window"
        );
        let err = server_task
            .await
            .expect("join")
            .expect_err("must time out at ClientHello body");
        assert!(matches!(
            err,
            AuthError::Timeout {
                stage: AuthStage::ClientHello,
                ..
            }
        ));
    }

    #[tokio::test(start_paused = true)]
    async fn client_deadline_is_absolute() {
        // The client previously passed the full deadline to each stage with no
        // absolute bound. Here the server stalls the ServerProof body past the
        // deadline; the client must time out at the budget, not wait a fresh window.
        let key = vec![0x5a; MIN_KEY_LEN];
        let daemon_id = [0x6b; DAEMON_ID_LEN];
        let deadline = Duration::from_millis(100);
        let conn = ConnectionInfo {
            schema: 1,
            wire_version: None,
            endpoints: vec![Endpoint {
                host: "127.0.0.1".to_owned(),
                port: 1,
            }],
            key: key.clone(),
            daemon_id,
            pid: 1,
            daemon_ver: TEST_DAEMON_VER.to_owned(),
        };

        let (mut server, mut client) = duplex(4096);
        let client_task =
            tokio::spawn(async move { authenticate_client(&mut client, &conn, deadline).await });

        // Read the client's ClientHello so its write completes.
        let _hello: ClientHello = read_auth_json(&mut server).await;
        yield_now().await;
        assert!(!client_task.is_finished());

        // Send only the ServerProof length prefix, then stall the body past the
        // absolute deadline.
        advance(Duration::from_millis(60)).await;
        let server_nonce = [0x22; NONCE_LEN];
        let server_proof = compute_proof(
            &key,
            SERVER_PROOF_DOMAIN,
            &[0u8; NONCE_LEN],
            &server_nonce,
            &daemon_id,
        );
        write_auth_len_only(
            &mut server,
            &ServerProof {
                daemon_id,
                server_nonce,
                daemon_ver: TEST_DAEMON_VER.to_owned(),
                server_proof,
            },
        )
        .await;
        yield_now().await;
        advance(Duration::from_millis(50)).await;
        yield_now().await;
        assert!(
            client_task.is_finished(),
            "client must bound the whole handshake by one absolute deadline"
        );
        let err = client_task
            .await
            .expect("join")
            .expect_err("client must time out");
        assert!(matches!(err, AuthError::Timeout { .. }));
    }

    async fn read_auth_json<T>(stream: &mut DuplexStream) -> T
    where
        T: DeserializeOwned,
    {
        let mut len_bytes = [0u8; 4];
        stream
            .read_exact(&mut len_bytes)
            .await
            .expect("read auth length");
        let len = u32::from_le_bytes(len_bytes);
        assert!(
            len <= MAX_AUTH_MESSAGE_LEN,
            "test helper received auth message over cap"
        );
        let mut body = vec![0u8; len as usize];
        stream.read_exact(&mut body).await.expect("read auth body");
        serde_json::from_slice(&body).expect("decode auth json")
    }

    async fn write_auth_json<T>(stream: &mut DuplexStream, value: &T)
    where
        T: Serialize,
    {
        let body = serde_json::to_vec(value).expect("encode auth json");
        assert!(
            body.len() <= MAX_AUTH_MESSAGE_LEN as usize,
            "test helper auth message over cap"
        );
        stream
            .write_all(&(body.len() as u32).to_le_bytes())
            .await
            .expect("write auth length");
        stream.write_all(&body).await.expect("write auth body");
    }
}