asupersync 0.3.6

//! QUIC Version Negotiation
//!
//! Implements QUIC version negotiation packet handling as specified in RFC 9000.

use crate::bytes::{Buf, BufMut, Bytes, BytesMut};
use crate::net::atp::handshake::state_machine::{HandshakeError, QuicVersion};
use crate::types::outcome::Outcome;

const VERSION_NEGOTIATION_MARKER: u32 = QuicVersion::Negotiation as u32;

/// Version negotiation packet
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct VersionNegotiationPacket {
    /// Source connection ID from server
    pub source_cid: Bytes,
    /// Destination connection ID from original client packet
    pub dest_cid: Bytes,
    /// List of supported versions
    pub supported_versions: Vec<u32>,
}

impl VersionNegotiationPacket {
    /// Create a new version negotiation packet
    pub fn new(source_cid: Bytes, dest_cid: Bytes, supported_versions: Vec<u32>) -> Self {
        Self {
            source_cid,
            dest_cid,
            supported_versions,
        }
    }

    /// Encode packet to wire format
    pub fn encode(&self) -> Outcome<Bytes, HandshakeError> {
        if let Outcome::Err(error) = validate_supported_versions(&self.supported_versions) {
            return Outcome::err(error);
        }

        let mut buf = BytesMut::new();

        // Long header form with version = 0 for version negotiation
        let first_byte = 0x80; // Long header
        buf.put_u8(first_byte);
        buf.put_u32(VERSION_NEGOTIATION_MARKER);

        // Destination Connection ID
        if self.dest_cid.len() > 255 {
            return Outcome::err(HandshakeError::ConnectionIdError {
                reason: "destination CID too long".to_string(),
            });
        }
        buf.put_u8(self.dest_cid.len() as u8);
        buf.put_slice(&self.dest_cid);

        // Source Connection ID
        if self.source_cid.len() > 255 {
            return Outcome::err(HandshakeError::ConnectionIdError {
                reason: "source CID too long".to_string(),
            });
        }
        buf.put_u8(self.source_cid.len() as u8);
        buf.put_slice(&self.source_cid);

        // Supported versions (4 bytes each)
        for &version in &self.supported_versions {
            buf.put_u32(version);
        }

        Outcome::ok(buf.freeze())
    }

    /// Decode packet from wire format
    pub fn decode(data: &[u8]) -> Outcome<Self, HandshakeError> {
        if data.len() < 7 {
            return Outcome::err(HandshakeError::InvalidPacket {
                reason: "version negotiation packet too short".to_string(),
            });
        }

        let mut buf = data;

        // Check long header form
        let first_byte = buf.get_u8();
        if first_byte & 0x80 == 0 {
            return Outcome::err(HandshakeError::InvalidPacket {
                reason: "not a long header packet".to_string(),
            });
        }

        // Check version is 0
        let version = buf.get_u32();
        if version != VERSION_NEGOTIATION_MARKER {
            return Outcome::err(HandshakeError::InvalidPacket {
                reason: "version negotiation must have version 0".to_string(),
            });
        }

        // Destination Connection ID
        let dest_cid_len = buf.get_u8() as usize;
        if buf.remaining() < dest_cid_len {
            return Outcome::err(HandshakeError::InvalidPacket {
                reason: "insufficient data for destination CID".to_string(),
            });
        }
        let dest_cid = Bytes::copy_from_slice(&buf[..dest_cid_len]);
        buf.advance(dest_cid_len);

        // Source Connection ID
        if buf.is_empty() {
            return Outcome::err(HandshakeError::InvalidPacket {
                reason: "missing source CID length".to_string(),
            });
        }
        let source_cid_len = buf.get_u8() as usize;
        if buf.remaining() < source_cid_len {
            return Outcome::err(HandshakeError::InvalidPacket {
                reason: "insufficient data for source CID".to_string(),
            });
        }
        let source_cid = Bytes::copy_from_slice(&buf[..source_cid_len]);
        buf.advance(source_cid_len);

        // Supported versions
        if buf.remaining() % 4 != 0 {
            return Outcome::err(HandshakeError::InvalidPacket {
                reason: "invalid version list length".to_string(),
            });
        }

        let mut supported_versions = Vec::new();
        while buf.remaining() >= 4 {
            supported_versions.push(buf.get_u32());
        }

        if let Outcome::Err(error) = validate_supported_versions(&supported_versions) {
            return Outcome::err(error);
        }

        Outcome::ok(Self {
            source_cid,
            dest_cid,
            supported_versions,
        })
    }

    /// Check if a version is supported
    pub fn supports_version(&self, version: u32) -> bool {
        self.supported_versions.contains(&version)
    }

    /// Select the best supported version from client's attempted version
    pub fn select_version(&self, attempted_version: u32) -> Option<u32> {
        // A Version Negotiation packet listing the attempted version is invalid
        // for the active connection attempt and must be discarded.
        if self.supports_version(attempted_version) {
            return None;
        }

        self.supported_versions
            .iter()
            .copied()
            .filter(|&version| QuicVersion::is_supported(version))
            .max()
    }
}

/// Version negotiation utilities
pub struct VersionNegotiation;

impl VersionNegotiation {
    /// Check if version negotiation is needed
    pub fn is_negotiation_needed(client_version: u32, server_versions: &[u32]) -> bool {
        !server_versions.contains(&client_version)
    }

    /// Create server response for unsupported client version
    pub fn create_server_response(
        client_dest_cid: Bytes,
        server_source_cid: Bytes,
    ) -> VersionNegotiationPacket {
        VersionNegotiationPacket::new(
            server_source_cid,
            client_dest_cid,
            QuicVersion::supported_versions(),
        )
    }

    /// Validate version negotiation packet from server
    pub fn validate_server_response(
        packet: &VersionNegotiationPacket,
        original_dest_cid: &[u8],
        original_source_cid: &[u8],
        attempted_version: u32,
    ) -> Outcome<(), HandshakeError> {
        // Destination CID must match original source CID
        if packet.dest_cid.as_ref() != original_source_cid {
            return Outcome::err(HandshakeError::InvalidPacket {
                reason: "version negotiation destination CID mismatch".to_string(),
            });
        }

        // Source CID must match original destination CID
        if packet.source_cid.as_ref() != original_dest_cid {
            return Outcome::err(HandshakeError::InvalidPacket {
                reason: "version negotiation source CID mismatch".to_string(),
            });
        }

        if let Outcome::Err(error) = validate_supported_versions(&packet.supported_versions) {
            return Outcome::err(error);
        }

        if packet.supports_version(attempted_version) {
            return Outcome::err(HandshakeError::InvalidPacket {
                reason: "version negotiation lists attempted version".to_string(),
            });
        }

        if packet.select_version(attempted_version).is_none() {
            return Outcome::err(HandshakeError::UnsupportedVersion {
                version: attempted_version,
            });
        }

        Outcome::ok(())
    }
}

fn validate_supported_versions(versions: &[u32]) -> Outcome<(), HandshakeError> {
    if versions.is_empty() {
        return Outcome::err(HandshakeError::InvalidPacket {
            reason: "no supported versions".to_string(),
        });
    }

    if versions.contains(&VERSION_NEGOTIATION_MARKER) {
        return Outcome::err(HandshakeError::InvalidPacket {
            reason: "supported versions include version negotiation marker".to_string(),
        });
    }

    Outcome::ok(())
}

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

    #[test]
    fn test_version_negotiation_packet_roundtrip() {
        let source_cid = Bytes::from_static(b"server_cid");
        let dest_cid = Bytes::from_static(b"client_cid");
        let supported_versions = vec![0x00000001, 0x12345678];

        let packet = VersionNegotiationPacket::new(
            source_cid.clone(),
            dest_cid.clone(),
            supported_versions.clone(),
        );

        let encoded = packet.encode().unwrap();
        let decoded = VersionNegotiationPacket::decode(&encoded).unwrap();

        assert_eq!(decoded.source_cid, source_cid);
        assert_eq!(decoded.dest_cid, dest_cid);
        assert_eq!(decoded.supported_versions, supported_versions);
    }

    #[test]
    fn test_version_support_check() {
        let packet = VersionNegotiationPacket::new(
            Bytes::from_static(b"src"),
            Bytes::from_static(b"dst"),
            vec![0x00000001, 0x12345678],
        );

        assert!(packet.supports_version(0x00000001));
        assert!(packet.supports_version(0x12345678));
        assert!(!packet.supports_version(0xabcdef00));
    }

    #[test]
    fn test_version_selection() {
        let packet = VersionNegotiationPacket::new(
            Bytes::from_static(b"src"),
            Bytes::from_static(b"dst"),
            vec![0x00000001, 0x12345678],
        );

        // A server that lists the attempted version did not need negotiation.
        assert_eq!(packet.select_version(0x00000001), None);

        // Unsupported attempted versions select the best version supported locally.
        assert_eq!(packet.select_version(0xabcdef00), Some(0x00000001));

        let unsupported_only = VersionNegotiationPacket::new(
            Bytes::from_static(b"src"),
            Bytes::from_static(b"dst"),
            vec![0x12345678],
        );
        assert_eq!(unsupported_only.select_version(0xabcdef00), None);
    }

    #[test]
    fn test_negotiation_needed() {
        let server_versions = vec![0x00000001, 0x12345678];

        assert!(!VersionNegotiation::is_negotiation_needed(
            0x00000001,
            &server_versions
        ));
        assert!(VersionNegotiation::is_negotiation_needed(
            0xabcdef00,
            &server_versions
        ));
    }

    #[test]
    fn test_invalid_packet_decode() {
        // Too short
        let result = VersionNegotiationPacket::decode(&[0x80, 0x00]);
        assert!(result.is_err());

        // Short header (not long header)
        let result = VersionNegotiationPacket::decode(&[0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]);
        assert!(result.is_err());

        // Non-zero version
        let result = VersionNegotiationPacket::decode(&[0x80, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00]);
        assert!(result.is_err());
    }

    #[test]
    fn test_version_zero_rejected_in_supported_versions() {
        let packet = VersionNegotiationPacket::new(
            Bytes::from_static(b"src"),
            Bytes::from_static(b"dst"),
            vec![VERSION_NEGOTIATION_MARKER],
        );
        assert!(packet.encode().is_err());

        let encoded_with_zero = [
            0x80, 0, 0, 0, 0, // Long header, Version Negotiation marker
            0, // Empty destination CID
            0, // Empty source CID
            0, 0, 0, 0, // Invalid supported version 0
        ];
        assert!(VersionNegotiationPacket::decode(&encoded_with_zero).is_err());
    }

    #[test]
    fn test_validate_server_response_discards_attempted_version() {
        let packet = VersionNegotiationPacket::new(
            Bytes::from_static(b"original_dest"),
            Bytes::from_static(b"original_source"),
            vec![QuicVersion::V1 as u32],
        );

        let result = VersionNegotiation::validate_server_response(
            &packet,
            b"original_dest",
            b"original_source",
            QuicVersion::V1 as u32,
        );

        assert!(result.is_err());
    }

    #[test]
    fn test_validate_server_response_accepts_supported_alternative() {
        let packet = VersionNegotiationPacket::new(
            Bytes::from_static(b"original_dest"),
            Bytes::from_static(b"original_source"),
            vec![QuicVersion::V1 as u32],
        );

        let result = VersionNegotiation::validate_server_response(
            &packet,
            b"original_dest",
            b"original_source",
            0xabcdef00,
        );

        assert!(result.is_ok());
    }
}