monocoque-rs-zmtp 0.4.0

//! Synchronous ZMTP handshake that completes before spawning background tasks.
//!
//! This eliminates race conditions by ensuring both peers complete the handshake
//! protocol before any application data can be sent.
//!
//! ## Memory Allocation Strategy
//!
//! This module uses **stack arrays** for all fixed-size protocol buffers:
//! - Greeting: 64-byte stack array
//! - Frame header: 2-byte stack array
//! - Length field: 8-byte stack array
//!
//! The READY body uses a small `Vec` allocation (typically ~27 bytes) because:
//! 1. compio's ownership-passing API requires owned buffers (can't use &mut slice)
//! 2. Size is dynamic but bounded (max 512 bytes enforced)
//! 3. Handshake happens once per connection (not in hot path)
//! 4. Total allocation overhead: ~93 bytes one-time per connection
//!
//! After handshake completes, the main data path uses the `core::io` read slab for zero-copy IO.

use crate::codec::ZmtpError;
use crate::security::curve::CurveHandshakeResult;
use crate::session::SocketType;
use crate::utils::{FLAG_COMMAND, build_ready, encode_frame};
use bytes::{BufMut, Bytes, BytesMut};
use compio_buf::BufResult;
use compio_io::{AsyncRead, AsyncWrite};
use monocoque_core::options::SocketOptions;
use monocoque_core::timeout::{read_exact_with_timeout, write_all_with_timeout};
use std::time::Duration;
use tracing::{debug, warn};

/// Result of a successful handshake
#[derive(Debug)]
pub struct HandshakeResult {
    pub peer_identity: Option<Bytes>,
    pub peer_socket_type: SocketType,
    pub curve_cipher: Option<crate::security::curve::CurveMessageCipher>,
}

/// Security mechanism to use for the ZMTP handshake.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SecurityMechanism {
    /// No authentication (default).
    Null,
    /// Username/password authentication (PLAIN).
    Plain,
    /// Public-key encryption (CURVE).
    Curve,
}

impl SecurityMechanism {
    /// Detect the mechanism from socket options.
    ///
    /// Priority: CURVE > PLAIN > NULL.
    pub fn from_options(options: &SocketOptions) -> Self {
        if options.curve_secretkey.is_some() || options.curve_server {
            Self::Curve
        } else if options.plain_server || options.plain_username.is_some() {
            Self::Plain
        } else {
            Self::Null
        }
    }

    /// The ASCII mechanism name used in ZMTP greetings (20-byte field).
    pub fn as_greeting_bytes(self) -> &'static [u8] {
        match self {
            Self::Null => b"NULL",
            Self::Plain => b"PLAIN",
            Self::Curve => b"CURVE",
        }
    }
}

/// Performs the complete ZMTP handshake, selecting the security mechanism from options.
///
/// This is the primary handshake entry point for sockets that have security configured.
///
/// A single total deadline (`timeout`, the socket's `handshake_timeout`) bounds
/// the WHOLE handshake, not just each step. Per-step timeouts alone let a slow
/// peer reset the clock on every read/write and hold the handshake open
/// indefinitely (slowloris); the total budget here caps the entire greeting,
/// security exchange, and READY at one wall-clock value. A legitimate handshake
/// completes in well under it; a trickling peer is dropped when it elapses.
pub async fn perform_handshake_with_options<S>(
    stream: &mut S,
    local_socket_type: SocketType,
    identity: Option<&[u8]>,
    timeout: Option<Duration>,
    options: &SocketOptions,
) -> Result<HandshakeResult, ZmtpError>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    // The handshake future is large; box it at the deadline boundary so the
    // combinator (and every caller's future) stays small.
    let inner = Box::pin(perform_handshake_inner(
        stream,
        local_socket_type,
        identity,
        timeout,
        options,
    ));
    match timeout {
        Some(budget) => match monocoque_core::rt::timeout(budget, inner).await {
            Ok(result) => result,
            Err(_elapsed) => {
                warn!(
                    "[HANDSHAKE] total handshake deadline of {:?} exceeded; closing (slowloris defense)",
                    budget
                );
                Err(ZmtpError::Protocol)
            }
        },
        // No configured timeout: preserve the prior unbounded behavior.
        None => inner.await,
    }
}

/// The handshake body, bounded by the total deadline in
/// [`perform_handshake_with_options`]. `timeout` is still passed through as the
/// per-step read/write bound.
#[allow(clippy::too_many_lines)]
async fn perform_handshake_inner<S>(
    stream: &mut S,
    local_socket_type: SocketType,
    identity: Option<&[u8]>,
    timeout: Option<Duration>,
    options: &SocketOptions,
) -> Result<HandshakeResult, ZmtpError>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    let mechanism = SecurityMechanism::from_options(options);

    debug!(
        "[HANDSHAKE] Starting handshake for {} (timeout: {:?}, mechanism: {:?})",
        local_socket_type.as_str(),
        timeout,
        mechanism
    );

    // Step 1: Send our greeting
    debug!("[HANDSHAKE] Step 1: Sending greeting...");
    let greeting_bytes = build_greeting_with_mechanism(mechanism, options);
    let BufResult(write_res, _) = write_all_with_timeout(stream, greeting_bytes.clone(), timeout)
        .await
        .map_err(|e| {
            warn!("[HANDSHAKE] Step 1: Failed to send ZMTP greeting: {}", e);
            ZmtpError::Protocol
        })?;
    write_res.map_err(|e| {
        warn!(
            "[HANDSHAKE] Step 1: Failed to write ZMTP greeting bytes: {}",
            e
        );
        ZmtpError::Protocol
    })?;
    debug!(
        "[HANDSHAKE] Step 1 DONE: Sent greeting ({} bytes)",
        greeting_bytes.len()
    );

    // Step 2: Receive peer greeting
    debug!("[HANDSHAKE] Step 2: Receiving peer greeting...");
    let greeting_buf = [0u8; 64];
    let BufResult(read_res, greeting_buf) = read_exact_with_timeout(stream, greeting_buf, timeout)
        .await
        .map_err(|e| {
            warn!("[HANDSHAKE] Step 2: Failed to receive ZMTP greeting: {}", e);
            ZmtpError::Protocol
        })?;
    read_res.map_err(|e| {
        warn!(
            "[HANDSHAKE] Step 2: Failed to read ZMTP greeting bytes: {}",
            e
        );
        ZmtpError::Protocol
    })?;
    debug!("[HANDSHAKE] Step 2 DONE: Received peer greeting (64 bytes)");

    // Validate greeting signature
    if greeting_buf[0] != 0xFF || greeting_buf[9] != 0x7F {
        warn!(
            "[HANDSHAKE] ZMTP greeting: invalid signature bytes (expected [0]=0xff [9]=0x7f, got [0]=0x{:02x} [9]=0x{:02x})",
            greeting_buf[0], greeting_buf[9]
        );
        return Err(ZmtpError::Protocol);
    }

    // Parse peer greeting to check mechanism compatibility
    use crate::greeting::ZmtpGreeting;
    let peer_greeting = ZmtpGreeting::parse(&Bytes::copy_from_slice(&greeting_buf[..]))
        .map_err(|_| ZmtpError::Protocol)?;
    let expected_mech = mechanism.as_greeting_bytes();
    let peer_mech_str = peer_greeting.mechanism_str();
    let our_mech_name = std::str::from_utf8(expected_mech).unwrap_or("NULL");
    if !peer_mech_str.eq_ignore_ascii_case(our_mech_name) {
        warn!(
            "[HANDSHAKE] Security mechanism mismatch: we advertise {:?}, peer advertises {:?}",
            our_mech_name, peer_mech_str
        );
        return Err(ZmtpError::Protocol);
    }
    if greeting_buf[9] != 0x7F {
        warn!(
            "[HANDSHAKE] ZMTP greeting: expected signature byte 0x7f at offset 9, got 0x{:02x}",
            greeting_buf[9]
        );
        return Err(ZmtpError::Protocol);
    }

    let peer_major = greeting_buf[10];
    if mechanism != SecurityMechanism::Null && peer_major != 3 {
        warn!(
            "[HANDSHAKE] non-NULL mechanism {:?} cannot negotiate with ZMTP major version {}",
            mechanism, peer_major
        );
        return Err(ZmtpError::Protocol);
    }

    let peer_mechanism = parse_greeting_mechanism(&greeting_buf[12..32])?;
    if peer_mechanism != mechanism {
        warn!(
            "[HANDSHAKE] security mechanism mismatch: local {:?}, peer {:?}",
            mechanism, peer_mechanism
        );
        return Err(ZmtpError::Protocol);
    }

    // Step 3: Run security-mechanism-specific exchange (between greeting and READY)
    let curve_cipher: Option<crate::security::curve::CurveMessageCipher> = None;
    match mechanism {
        SecurityMechanism::Null => {
            // No mechanism-level exchange for NULL; proceed directly to READY.
        }
        SecurityMechanism::Plain => {
            run_plain_exchange(stream, options, timeout).await?;
        }
        SecurityMechanism::Curve => {
            // CURVE handshake carries all metadata internally (no separate READY needed).
            let cr =
                run_curve_exchange(stream, options, timeout, local_socket_type, identity).await?;
            let peer_socket_type = parse_socket_type(cr.peer_socket_type.as_ref())?;
            check_socket_type_compatibility(local_socket_type, peer_socket_type)?;
            return Ok(HandshakeResult {
                peer_identity: cr.peer_identity,
                peer_socket_type,
                curve_cipher: cr.cipher,
            });
        }
    }

    // Step 4: Send the metadata command (NULL and PLAIN only). ZMTP PLAIN has
    // the client send INITIATE (the server then replies READY); NULL and the
    // PLAIN server send READY.
    let plain_client = mechanism == SecurityMechanism::Plain && !options.plain_server;
    debug!(
        "[HANDSHAKE] Step 4: Sending {} command...",
        if plain_client { "INITIATE" } else { "READY" }
    );
    let ready_body = if plain_client {
        crate::utils::build_initiate(local_socket_type.as_str(), identity)
    } else {
        build_ready(local_socket_type.as_str(), identity)
    };
    let ready_frame = encode_frame(FLAG_COMMAND, &ready_body);
    let BufResult(write_res, _) = write_all_with_timeout(stream, ready_frame.clone(), timeout)
        .await
        .map_err(|e| {
            warn!(
                "[HANDSHAKE] Step 4: Failed to send ZMTP READY command: {}",
                e
            );
            ZmtpError::Protocol
        })?;
    write_res.map_err(|e| {
        warn!(
            "[HANDSHAKE] Step 4: Failed to write ZMTP READY command bytes: {}",
            e
        );
        ZmtpError::Protocol
    })?;
    debug!(
        "[HANDSHAKE] Step 4 DONE: Sent READY command ({} bytes)",
        ready_frame.len()
    );

    // Step 5: Receive peer READY command
    debug!("[HANDSHAKE] Step 5: Receiving peer READY command...");
    let header_buf = [0u8; 2];
    let BufResult(read_res, header_buf) = read_exact_with_timeout(stream, header_buf, timeout)
        .await
        .map_err(|e| {
            warn!(
                "[HANDSHAKE] Step 5: Failed to receive ZMTP READY frame header: {}",
                e
            );
            ZmtpError::Protocol
        })?;
    read_res.map_err(|e| {
        warn!(
            "[HANDSHAKE] Step 5: Failed to read ZMTP READY frame header bytes: {}",
            e
        );
        ZmtpError::Protocol
    })?;
    debug!(
        "[HANDSHAKE] Step 5a DONE: Read header [{:02x}, {:02x}]",
        header_buf[0], header_buf[1]
    );

    let flags = header_buf[0];
    let is_command = (flags & FLAG_COMMAND) != 0;
    let is_long = (flags & 0x02) != 0;

    if !is_command {
        warn!(
            "[HANDSHAKE] ZMTP READY step: expected COMMAND frame (flags & 0x04 != 0), \
             got flags=0x{:02x}  -  peer sent a data frame instead of READY",
            flags
        );
        return Err(ZmtpError::Protocol);
    }

    // Read body length
    let body_len = if is_long {
        let len_buf = [0u8; 8];
        let BufResult(read_res, len_buf) = read_exact_with_timeout(stream, len_buf, timeout)
            .await
            .map_err(|e| {
                warn!(
                    "[HANDSHAKE] Step 5: Failed to receive ZMTP READY long-frame length: {}",
                    e
                );
                ZmtpError::Protocol
            })?;
        read_res.map_err(|e| {
            warn!(
                "[HANDSHAKE] Step 5: Failed to read ZMTP READY long-frame length bytes: {}",
                e
            );
            ZmtpError::Protocol
        })?;
        let raw_len = u64::from_be_bytes(len_buf);
        if raw_len > usize::MAX as u64 {
            warn!(
                "[HANDSHAKE] ZMTP READY long-frame length overflows usize: {}",
                raw_len
            );
            return Err(ZmtpError::Protocol);
        }
        raw_len as usize
    } else {
        header_buf[1] as usize
    };
    debug!("[HANDSHAKE] Step 5b DONE: body_len={}", body_len);

    // Read body
    const MAX_READY_SIZE: usize = 512;
    if body_len > MAX_READY_SIZE {
        warn!(
            "[HANDSHAKE] ZMTP READY body too large: got {} bytes, maximum allowed is {} bytes",
            body_len, MAX_READY_SIZE
        );
        return Err(ZmtpError::Protocol);
    }
    let body_buf = vec![0u8; body_len];
    let BufResult(read_res, body_buf) = read_exact_with_timeout(stream, body_buf, timeout)
        .await
        .map_err(|e| {
            warn!(
                "[HANDSHAKE] Step 5: Failed to receive ZMTP READY body ({} bytes): {}",
                body_len, e
            );
            ZmtpError::Protocol
        })?;
    read_res.map_err(|e| {
        warn!(
            "[HANDSHAKE] Step 5: Failed to read ZMTP READY body bytes: {}",
            e
        );
        ZmtpError::Protocol
    })?;
    debug!("[HANDSHAKE] Step 5c DONE: Read {} bytes of body", body_len);

    // Parse READY command
    let ready_bytes = Bytes::from(body_buf);
    let (peer_socket_type, peer_identity) = parse_ready_command(&ready_bytes)?;
    check_socket_type_compatibility(local_socket_type, peer_socket_type)?;

    debug!(
        "[HANDSHAKE] Handshake complete! Peer is {}",
        peer_socket_type.as_str()
    );

    Ok(HandshakeResult {
        peer_identity,
        peer_socket_type,
        curve_cipher,
    })
}

// ---------------------------------------------------------------------------
// Per-mechanism security exchanges
// ---------------------------------------------------------------------------

/// Run the PLAIN authentication exchange.
///
/// - Client mode: send HELLO, receive WELCOME/ERROR.
/// - Server mode: receive HELLO, validate, send WELCOME/ERROR.
async fn run_plain_exchange<S>(
    stream: &mut S,
    options: &SocketOptions,
    timeout: Option<Duration>,
) -> Result<(), ZmtpError>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    use crate::security::plain::{PlainCredentials, plain_client_handshake};

    if options.plain_server {
        debug!("[HANDSHAKE] Running PLAIN server exchange");
        let domain = options.zap_domain.as_str();
        crate::security::plain::plain_server_handshake_zap(stream, domain, "unknown", timeout)
            .await
            .map(|_| ())
    } else if let Some(ref username) = options.plain_username {
        debug!("[HANDSHAKE] Running PLAIN client exchange");
        let password = options.plain_password().unwrap_or("");
        let credentials = PlainCredentials::new(username.clone(), password);
        plain_client_handshake(stream, &credentials, timeout).await
    } else {
        // Should not happen (mechanism detection guards this), but be safe.
        Ok(())
    }
}

/// Run the CURVE key-exchange handshake.
///
/// Returns a `CurveHandshakeResult` containing the peer's socket type, identity,
/// and (server-side only) the authenticated client public key.
///
/// - Client mode: `curve_secretkey` + `curve_serverkey` must be set.
/// - Server mode: `curve_server` flag + `curve_secretkey` must be set.
async fn run_curve_exchange<S>(
    stream: &mut S,
    options: &SocketOptions,
    timeout: Option<Duration>,
    local_socket_type: SocketType,
    local_identity: Option<&[u8]>,
) -> Result<CurveHandshakeResult, ZmtpError>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    use crate::security::curve::{
        CurveClient, CurveKeyPair, CurvePublicKey, CurveSecretKey, CurveServer,
    };

    if options.curve_server {
        debug!("[HANDSHAKE] Running CURVE server exchange");
        let secret_bytes = *options.curve_secretkey().ok_or_else(|| {
            warn!("[HANDSHAKE] CURVE server mode requires curve_secretkey to be set, but it is missing");
            ZmtpError::Protocol
        })?;
        let server_secret = CurveSecretKey::from_bytes(secret_bytes);
        let server_public = server_secret.public_key();
        let server_keypair = CurveKeyPair::from_keys(server_public, server_secret);

        if options.zap_domain.is_empty() {
            // No zap_domain means no ZAP authorization runs at all: every client
            // that completes the CURVE handshake is admitted (encryption only, no
            // access control). Warn loudly so this is never a silent default. Set
            // a zap_domain and run a ZAP handler to authorize client keys.
            warn!(
                "[HANDSHAKE] CURVE server has no zap_domain: accepting any client key \
                 without authorization (encryption only). Configure ZAP to authenticate peers."
            );
            let mut curve_server = CurveServer::new(server_keypair, local_socket_type.as_str());
            curve_server.handshake(stream, timeout).await
        } else {
            use crate::security::curve::curve_server_handshake_zap;
            curve_server_handshake_zap(
                stream,
                server_keypair,
                options.zap_domain.clone(),
                timeout,
                "unknown",
                local_socket_type.as_str(),
            )
            .await
        }
    } else if let (Some(secret_bytes), Some(server_key_bytes)) =
        (options.curve_secretkey(), options.curve_serverkey)
    {
        debug!("[HANDSHAKE] Running CURVE client exchange");
        let client_secret = CurveSecretKey::from_bytes(*secret_bytes);
        let client_public = client_secret.public_key();
        let client_keypair = CurveKeyPair::from_keys(client_public, client_secret);
        let server_public = CurvePublicKey::from_bytes(server_key_bytes);

        let local_id = local_identity.map(Bytes::copy_from_slice);
        let mut curve_client = CurveClient::new(
            client_keypair,
            server_public,
            local_socket_type.as_str(),
            local_id,
        );
        curve_client.handshake(stream, timeout).await
    } else if options.curve_secretkey.is_some() {
        // curve_secretkey set for a client but curve_serverkey is absent.
        warn!(
            "[HANDSHAKE] CURVE client mode requires both curve_secretkey and curve_serverkey, \
             but curve_serverkey (server's public key) is missing"
        );
        Err(ZmtpError::Protocol)
    } else {
        Err(ZmtpError::Protocol)
    }
}

// ---------------------------------------------------------------------------
// Greeting helpers
// ---------------------------------------------------------------------------

/// Build a ZMTP 3.0 greeting (64 bytes) advertising the given security mechanism.
fn build_greeting_with_mechanism(mechanism: SecurityMechanism, options: &SocketOptions) -> Bytes {
    let mut b = BytesMut::with_capacity(64);

    // Signature
    b.extend_from_slice(&[0xFF]);
    b.extend_from_slice(&[0u8; 8]);
    b.extend_from_slice(&[0x7F]);

    // Version 3.0
    b.extend_from_slice(&[0x03, 0x00]);

    // Mechanism field: 20 bytes, ASCII name padded with NUL
    let mech_name = mechanism.as_greeting_bytes();
    b.extend_from_slice(mech_name);
    let padding = 20usize.saturating_sub(mech_name.len());
    b.put_bytes(0, padding);

    // As-server flag (byte 32): 1 if this side acts as CURVE/PLAIN server
    let as_server = match mechanism {
        SecurityMechanism::Curve => options.curve_server,
        SecurityMechanism::Plain => options.plain_server,
        SecurityMechanism::Null => false,
    };
    b.extend_from_slice(&[u8::from(as_server)]);

    // Padding to reach 64 bytes total
    b.extend_from_slice(&[0u8; 31]);

    b.freeze()
}

fn parse_greeting_mechanism(field: &[u8]) -> Result<SecurityMechanism, ZmtpError> {
    let len = field
        .iter()
        .position(|&byte| byte == 0)
        .unwrap_or(field.len());
    match &field[..len] {
        b"NULL" => Ok(SecurityMechanism::Null),
        b"PLAIN" => Ok(SecurityMechanism::Plain),
        b"CURVE" => Ok(SecurityMechanism::Curve),
        _ => Err(ZmtpError::Protocol),
    }
}

/// Parse READY command to extract socket type and identity
pub fn parse_ready_command(body: &Bytes) -> Result<(SocketType, Option<Bytes>), ZmtpError> {
    // READY format:
    // - 1 byte: command name length
    // - N bytes: "READY"
    // - Properties as key-value pairs

    if body.is_empty() {
        warn!("[HANDSHAKE] ZMTP READY parse: empty body");
        return Err(ZmtpError::Protocol);
    }

    // Accept READY or INITIATE (the client's PLAIN metadata command); both carry
    // identical Socket-Type/Identity properties, only the name differs.
    let name_len = body[0] as usize;
    let name = body.get(1..1 + name_len);
    if name != Some(b"READY".as_ref()) && name != Some(b"INITIATE".as_ref()) {
        warn!(
            "[HANDSHAKE] ZMTP READY parse: expected command name \"READY\" or \"INITIATE\", \
             got length={} name={:?}",
            name_len,
            name.map(|b| String::from_utf8_lossy(b).into_owned())
                .unwrap_or_default()
        );
        return Err(ZmtpError::Protocol);
    }

    // Parse properties, starting after the (variable-length) command name.
    let mut offset = 1 + name_len;
    let mut socket_type = None;
    let mut identity = None;

    while offset < body.len() {
        if offset + 1 > body.len() {
            warn!("[HANDSHAKE] READY property truncated at key-length byte");
            return Err(ZmtpError::Protocol);
        }

        let key_len = body[offset] as usize;
        offset += 1;

        if offset + key_len > body.len() {
            warn!(
                "[HANDSHAKE] READY property key truncated (key_len={})",
                key_len
            );
            return Err(ZmtpError::Protocol);
        }

        let key = &body[offset..offset + key_len];
        offset += key_len;

        if offset + 4 > body.len() {
            warn!("[HANDSHAKE] READY property value-length truncated");
            return Err(ZmtpError::Protocol);
        }

        let value_len = u32::from_be_bytes([
            body[offset],
            body[offset + 1],
            body[offset + 2],
            body[offset + 3],
        ]) as usize;
        offset += 4;

        // value_len is an attacker-controlled 32-bit length. `offset + value_len`
        // can wrap on a 32-bit target, pass a naive `> body.len()` check, and
        // then panic slicing out of bounds. checked_add rejects the overflow.
        let value_start = offset;
        let value_end = offset
            .checked_add(value_len)
            .filter(|&end| end <= body.len());
        let Some(value_end) = value_end else {
            warn!(
                "[HANDSHAKE] READY property value truncated or length overflow (value_len={})",
                value_len
            );
            return Err(ZmtpError::Protocol);
        };
        offset = value_end;

        match key {
            b"Socket-Type" => {
                if socket_type.is_some() {
                    return Err(ZmtpError::Protocol);
                }
                // The longest valid socket-type name is "STREAM" (6 bytes). Cap
                // the value so a bogus Socket-Type cannot force a large slice
                // before parse_socket_type rejects it.
                if value_len > 16 {
                    warn!(
                        "[HANDSHAKE] READY Socket-Type property too long: {} bytes",
                        value_len
                    );
                    return Err(ZmtpError::Protocol);
                }
                socket_type = Some(parse_socket_type(&body[value_start..value_end])?);
            }
            b"Identity" => {
                if identity.is_some() {
                    return Err(ZmtpError::Protocol);
                }
                // ZMQ spec limits identities to 255 bytes.
                if value_len > 255 {
                    warn!(
                        "[HANDSHAKE] READY Identity property too long: {} bytes (max 255)",
                        value_len
                    );
                    return Err(ZmtpError::Protocol);
                }
                // Zero-copy: slice the existing Bytes instead of copying
                identity = Some(body.slice(value_start..value_end));
            }
            _ => {
                // Ignore unknown properties
            }
        }
    }

    let socket_type = socket_type.ok_or_else(|| {
        warn!("[HANDSHAKE] ZMTP READY parse: peer READY command is missing the required \"Socket-Type\" property");
        ZmtpError::Protocol
    })?;
    Ok((socket_type, identity))
}

/// Parse socket type from bytes
fn parse_socket_type(value: &[u8]) -> Result<SocketType, ZmtpError> {
    match value {
        b"PAIR" => Ok(SocketType::Pair),
        b"DEALER" => Ok(SocketType::Dealer),
        b"ROUTER" => Ok(SocketType::Router),
        b"PUB" => Ok(SocketType::Pub),
        b"SUB" => Ok(SocketType::Sub),
        b"XPUB" => Ok(SocketType::Xpub),
        b"XSUB" => Ok(SocketType::Xsub),
        b"REQ" => Ok(SocketType::Req),
        b"REP" => Ok(SocketType::Rep),
        b"PUSH" => Ok(SocketType::Push),
        b"PULL" => Ok(SocketType::Pull),
        _ => {
            warn!(
                "[HANDSHAKE] ZMTP READY parse: unknown Socket-Type value {:?}",
                String::from_utf8_lossy(value)
            );
            Err(ZmtpError::Protocol)
        }
    }
}

/// Reject socket-type pairings that ZMQ does not allow (e.g. PUB with REQ).
///
/// libzmq refuses an incompatible pairing at the ZMTP layer; without this a
/// mismatched peer completes the handshake and then silently misbehaves. The
/// rules mirror libzmq: each type lists the peer types it may talk to.
fn check_socket_type_compatibility(local: SocketType, peer: SocketType) -> Result<(), ZmtpError> {
    use SocketType::{Dealer, Pair, Pub, Pull, Push, Rep, Req, Router, Sub, Xpub, Xsub};

    let compatible_peers: &[SocketType] = match local {
        Pair => &[Pair],
        Pub | Xpub => &[Sub, Xsub],
        Sub | Xsub => &[Pub, Xpub],
        Req => &[Rep, Router],
        Rep => &[Req, Dealer],
        Dealer => &[Rep, Dealer, Router],
        Router => &[Req, Dealer, Router],
        Push => &[Pull],
        Pull => &[Push],
    };

    if compatible_peers.contains(&peer) {
        Ok(())
    } else {
        warn!(
            "[HANDSHAKE] incompatible socket types: local {} cannot talk to peer {}",
            local.as_str(),
            peer.as_str()
        );
        Err(ZmtpError::Protocol)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bytes::BytesMut;
    use compio_buf::BufResult;
    use monocoque_core::options::SocketOptions;
    use monocoque_core::rt::{LocalRuntime, TcpListener, TcpStream};
    use monocoque_core::timeout::{read_exact_with_timeout, write_all_with_timeout};

    const TEST_TIMEOUT: Duration = Duration::from_secs(1);

    fn ready_body(properties: &[(&[u8], &[u8])]) -> Bytes {
        let mut body = Vec::new();
        body.extend_from_slice(b"\x05READY");

        for (key, value) in properties {
            body.push(key.len() as u8);
            body.extend_from_slice(key);
            body.extend_from_slice(&(value.len() as u32).to_be_bytes());
            body.extend_from_slice(value);
        }

        Bytes::from(body)
    }

    #[test]
    fn parse_ready_rejects_duplicate_socket_type_property() {
        let body = ready_body(&[(b"Socket-Type", b"DEALER"), (b"Socket-Type", b"ROUTER")]);

        assert!(matches!(
            parse_ready_command(&body),
            Err(ZmtpError::Protocol)
        ));
    }

    #[test]
    fn parse_ready_rejects_duplicate_identity_property() {
        let body = ready_body(&[
            (b"Socket-Type", b"DEALER"),
            (b"Identity", b"trusted"),
            (b"Identity", b"shadow"),
        ]);

        assert!(matches!(
            parse_ready_command(&body),
            Err(ZmtpError::Protocol)
        ));
    }

    async fn read_client_greeting(stream: &mut TcpStream) {
        let greeting = [0u8; 64];
        let BufResult(read_res, _) = read_exact_with_timeout(stream, greeting, Some(TEST_TIMEOUT))
            .await
            .unwrap();
        read_res.unwrap();
    }

    async fn write_greeting(stream: &mut TcpStream, greeting: Vec<u8>) {
        let BufResult(write_res, _) = write_all_with_timeout(stream, greeting, Some(TEST_TIMEOUT))
            .await
            .unwrap();
        write_res.unwrap();
    }

    async fn maybe_read_plain_hello(stream: &mut TcpStream) -> Option<[u8; 6]> {
        let header = [0u8; 6];
        let Ok(BufResult(read_res, header)) =
            read_exact_with_timeout(stream, header, Some(TEST_TIMEOUT)).await
        else {
            return None;
        };
        read_res.ok()?;
        Some(header)
    }

    async fn maybe_complete_ready_exchange(stream: &mut TcpStream) {
        let header = [0u8; 2];
        let Ok(BufResult(read_res, header)) =
            read_exact_with_timeout(stream, header, Some(TEST_TIMEOUT)).await
        else {
            return;
        };
        if read_res.is_err() {
            return;
        }

        let body = vec![0u8; header[1] as usize];
        let Ok(BufResult(read_res, _)) =
            read_exact_with_timeout(stream, body, Some(TEST_TIMEOUT)).await
        else {
            return;
        };
        if read_res.is_err() {
            return;
        }

        let ready_body = crate::utils::build_ready("PAIR", None);
        let ready_frame = crate::utils::encode_frame(crate::utils::FLAG_COMMAND, &ready_body);
        let Ok(BufResult(write_res, _)) =
            write_all_with_timeout(stream, ready_frame.to_vec(), Some(TEST_TIMEOUT)).await
        else {
            return;
        };
        let _ = write_res;
    }

    #[test]
    fn ready_parser_rejects_truncated_property_after_socket_type() {
        let mut body = BytesMut::from(crate::utils::build_ready("PAIR", None).as_ref());
        body.extend_from_slice(&[8]);
        body.extend_from_slice(b"Identity");
        body.extend_from_slice(&5u32.to_be_bytes());
        body.extend_from_slice(b"a");

        assert!(
            parse_ready_command(&body.freeze()).is_err(),
            "READY parser accepted a command with truncated trailing identity metadata"
        );
    }

    #[test]
    fn handshake_rejects_peer_greeting_with_invalid_signature_tail() {
        LocalRuntime::new().unwrap().block_on(async {
            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let addr = listener.local_addr().unwrap();

            let peer_task = monocoque_core::rt::spawn(async move {
                let (mut stream, _) = listener.accept().await.unwrap();
                read_client_greeting(&mut stream).await;

                let mut bad_greeting =
                    build_greeting_with_mechanism(SecurityMechanism::Null, &SocketOptions::new())
                        .to_vec();
                bad_greeting[9] = 0x00;
                write_greeting(&mut stream, bad_greeting).await;
                maybe_complete_ready_exchange(&mut stream).await;
            });

            let mut stream = TcpStream::connect(addr).await.unwrap();
            let result = perform_handshake_with_options(
                &mut stream,
                SocketType::Req,
                None,
                Some(TEST_TIMEOUT),
                &SocketOptions::new(),
            )
            .await;

            assert!(
                result.is_err(),
                "handshake accepted a peer greeting with an invalid ZMTP signature tail"
            );

            monocoque_core::rt::join(peer_task).await;
        });
    }

    #[test]
    fn non_null_handshake_rejects_peer_greeting_with_unsupported_major_version() {
        LocalRuntime::new().unwrap().block_on(async {
            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let addr = listener.local_addr().unwrap();

            let peer_task = monocoque_core::rt::spawn(async move {
                let (mut stream, _) = listener.accept().await.unwrap();

                read_client_greeting(&mut stream).await;

                let mut bad_greeting = build_greeting_with_mechanism(
                    SecurityMechanism::Plain,
                    &SocketOptions::new().with_plain_server(true),
                )
                .to_vec();
                bad_greeting[10] = 2;
                write_greeting(&mut stream, bad_greeting).await;

                assert!(
                    maybe_read_plain_hello(&mut stream).await.as_ref() != Some(b"\x05HELLO"),
                    "PLAIN client sent security commands to a ZMTP 2.x peer"
                );
            });

            let mut stream = TcpStream::connect(addr).await.unwrap();
            let options = SocketOptions::new().with_plain_credentials("alice", "secret");
            let result = perform_handshake_with_options(
                &mut stream,
                SocketType::Req,
                None,
                Some(TEST_TIMEOUT),
                &options,
            )
            .await;

            assert!(
                result.is_err(),
                "non-NULL handshake accepted an unsupported ZMTP major version during security negotiation"
            );

            monocoque_core::rt::join(peer_task).await;
        });
    }

    #[test]
    fn plain_client_does_not_send_credentials_to_peer_advertising_null() {
        LocalRuntime::new().unwrap().block_on(async {
            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let addr = listener.local_addr().unwrap();

            let peer_task = monocoque_core::rt::spawn(async move {
                let (mut stream, _) = listener.accept().await.unwrap();

                read_client_greeting(&mut stream).await;

                let peer_greeting =
                    build_greeting_with_mechanism(SecurityMechanism::Null, &SocketOptions::new());
                write_greeting(&mut stream, peer_greeting.to_vec()).await;

                assert!(
                    maybe_read_plain_hello(&mut stream).await.as_ref() != Some(b"\x05HELLO"),
                    "PLAIN client sent credentials to a peer that advertised NULL security"
                );
            });

            let mut stream = TcpStream::connect(addr).await.unwrap();
            let options = SocketOptions::new().with_plain_credentials("alice", "secret");
            let _ = perform_handshake_with_options(
                &mut stream,
                SocketType::Req,
                None,
                Some(TEST_TIMEOUT),
                &options,
            )
            .await;

            monocoque_core::rt::join(peer_task).await;
        });
    }

    #[test]
    fn total_deadline_bounds_a_handshake_that_stalls_after_the_greeting() {
        LocalRuntime::new().unwrap().block_on(async {
            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let addr = listener.local_addr().unwrap();

            // The peer completes the greeting - so the handshake gets past step 2
            // - then stalls with the connection held open, never sending the
            // PLAIN WELCOME the client waits for. Without a bound the client would
            // wait forever; the total handshake deadline must close it.
            let peer_task = monocoque_core::rt::spawn(async move {
                let (mut stream, _) = listener.accept().await.unwrap();
                read_client_greeting(&mut stream).await;
                let peer_greeting = build_greeting_with_mechanism(
                    SecurityMechanism::Plain,
                    &SocketOptions::new().with_plain_server(true),
                );
                write_greeting(&mut stream, peer_greeting.to_vec()).await;
                // Hold the connection open, sending nothing further.
                std::future::pending::<()>().await;
            });

            let mut stream = TcpStream::connect(addr).await.unwrap();
            let options = SocketOptions::new().with_plain_credentials("alice", "secret");
            let deadline = Duration::from_millis(400);
            let start = std::time::Instant::now();
            let result = perform_handshake_with_options(
                &mut stream,
                SocketType::Req,
                None,
                Some(deadline),
                &options,
            )
            .await;
            let elapsed = start.elapsed();

            assert!(
                result.is_err(),
                "handshake must fail against a peer that stalls after the greeting"
            );
            assert!(
                elapsed < deadline * 8,
                "handshake ran {elapsed:?}, far beyond the {deadline:?} deadline; a stalling peer was not bounded"
            );

            // The peer stalls forever; drop its task rather than joining it.
            drop(peer_task);
        });
    }

    #[test]
    fn parse_ready_rejects_overflowing_value_length_without_panicking() {
        // READY body: name "READY", one property keyed "Socket-Type" whose
        // 4-byte value length is u32::MAX. A naive `offset + value_len` wraps on
        // a 32-bit target and slices out of bounds; the checked add must reject
        // it (and it also exceeds the body on 64-bit). Either way: Err, no panic.
        let mut body = Vec::new();
        body.push(5u8);
        body.extend_from_slice(b"READY");
        body.push(11u8);
        body.extend_from_slice(b"Socket-Type");
        body.extend_from_slice(&u32::MAX.to_be_bytes()); // value_len = 0xFFFFFFFF
        // No value bytes follow.

        let result = parse_ready_command(&Bytes::from(body));
        assert!(
            matches!(result, Err(ZmtpError::Protocol)),
            "oversized value length must be rejected, not panic or accepted"
        );
    }

    #[test]
    fn parse_ready_rejects_oversized_socket_type() {
        // A Socket-Type value longer than any real name (here 20 bytes of 'A')
        // must be rejected by the length cap before parse_socket_type runs.
        let value = [b'A'; 20];
        let mut body = Vec::new();
        body.push(5u8);
        body.extend_from_slice(b"READY");
        body.push(11u8);
        body.extend_from_slice(b"Socket-Type");
        body.extend_from_slice(&(value.len() as u32).to_be_bytes());
        body.extend_from_slice(&value);

        assert!(matches!(
            parse_ready_command(&Bytes::from(body)),
            Err(ZmtpError::Protocol)
        ));
    }

    #[test]
    fn socket_type_compatibility_matches_zmq_rules() {
        use SocketType::{Dealer, Pair, Pub, Pull, Push, Rep, Req, Router, Sub};

        // Valid pairings.
        for (a, b) in [
            (Req, Rep),
            (Rep, Req),
            (Req, Router),
            (Dealer, Router),
            (Dealer, Dealer),
            (Pub, Sub),
            (Push, Pull),
            (Pair, Pair),
        ] {
            assert!(
                check_socket_type_compatibility(a, b).is_ok(),
                "{} should be compatible with {}",
                a.as_str(),
                b.as_str()
            );
        }

        // Invalid pairings.
        for (a, b) in [
            (Pub, Req),
            (Req, Pub),
            (Push, Sub),
            (Pair, Dealer),
            (Req, Req),
        ] {
            assert!(
                check_socket_type_compatibility(a, b).is_err(),
                "{} should be incompatible with {}",
                a.as_str(),
                b.as_str()
            );
        }
    }
}