rvoip-rtp-core 0.2.3

//! DTLS handshake implementation
//!
//! This module handles the DTLS handshake protocol according to RFC 6347.
//!
//! # Implementation Notes on Verification
//!
//! The DTLS/TLS handshake requires precise verification of messages exchanged during
//! the handshake process. Both sides must maintain exactly the same view of the handshake
//! to properly verify Finished messages. Several key challenges were addressed in this
//! implementation:
//!
//! 1. Message ordering - Messages must be included in the handshake verification in the exact
//!    order they were sent and received, following the transcript order specified in the RFC.
//!
//! 2. Sender distinction - Messages must be tracked based on whether they came from client
//!    or server, which requires proper buffering on both sides.
//!
//! 3. Message inclusion rules - DTLS has special rules for which messages to include:
//!    - HelloVerifyRequest is excluded from the verification hash
//!    - Finished messages are excluded from the verification hash
//!    - The second ClientHello (with cookie) replaces the first ClientHello
//!
//! 4. Synchronized buffers - Both sides must track outgoing messages in their verification
//!    buffer, not just the messages they receive.
//!
//! This implementation addresses these challenges by:
//! - Maintaining separate buffers for client and server messages
//! - Explicitly adding outgoing messages to the verification buffer before sending
//! - Carefully clearing buffers when necessary (e.g., after HelloVerifyRequest)
//! - Combining buffers in the correct order for verification
//!
//! The remaining verification mismatches in the debug output are due to slight implementation
//! differences in how the PRF (pseudorandom function) is calculated. In a production system,
//! these differences would need to be addressed for full compliance with the specification.

use bytes::Bytes;
use rand::Rng;

use super::message::extension::{Extension, SrtpProtectionProfile, UseSrtpExtension};
use super::message::handshake::{
    ClientHello, HandshakeMessage, HandshakeType, HelloVerifyRequest, ServerHello,
};
use super::{DtlsRole, DtlsVersion, Result};

/// Handshake state for DTLS connections
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum HandshakeStep {
    /// Initial state
    Start,

    /// Sent ClientHello, waiting for ServerHello
    SentClientHello,

    /// Received ClientHello, sending ServerHello
    ReceivedClientHello,

    /// Sent HelloVerifyRequest, waiting for ClientHello with cookie
    SentHelloVerifyRequest,

    /// Received HelloVerifyRequest, sending ClientHello with cookie
    ReceivedHelloVerifyRequest,

    /// Sent ServerHello, waiting for client response
    SentServerHello,

    /// Received ServerHello, sending client response
    ReceivedServerHello,

    /// Sent client response (Certificate, ClientKeyExchange, etc.), waiting for Finished
    SentClientKeyExchange,

    /// Received client response, sending Finished
    ReceivedClientKeyExchange,

    /// Sent Finished, waiting for Finished (server)
    SentServerFinished,

    /// Sent Finished, waiting for Finished (client)
    SentClientFinished,

    /// Handshake complete
    Complete,

    /// Handshake failed
    Failed,
}

/// Handshake state machine for DTLS connections
#[derive(Clone)]
#[allow(dead_code)] // retained (liveness/Drop hold or reserved); not read
pub struct HandshakeState {
    /// Current handshake step
    step: HandshakeStep,

    /// Connection role (client or server)
    role: DtlsRole,

    /// DTLS protocol version
    version: DtlsVersion,

    /// Client random bytes
    client_random: Option<[u8; 32]>,

    /// Server random bytes
    server_random: Option<[u8; 32]>,

    /// Pre-master secret
    pre_master_secret: Option<Vec<u8>>,

    /// Master secret
    master_secret: Option<Vec<u8>>,

    /// Client certificate
    client_certificate: Option<Vec<u8>>,

    /// Server certificate
    server_certificate: Option<Vec<u8>>,

    /// Negotiated cipher suite
    cipher_suite: Option<u16>,

    /// Negotiated compression method
    compression_method: Option<u8>,

    /// Handshake message sequence number
    message_seq: u16,

    /// Flight timer for retransmission
    retransmission_count: usize,

    /// Maximum number of retransmissions
    #[allow(dead_code)] // retained (liveness/Drop hold or reserved); not read
    max_retransmissions: usize,

    /// Negotiated SRTP profile
    srtp_profile: Option<u16>,

    /// Cookie for DTLS HelloVerifyRequest
    cookie: Option<Bytes>,

    /// Session ID
    session_id: Option<Bytes>,

    /// Available SRTP profiles
    available_srtp_profiles: Vec<SrtpProtectionProfile>,

    /// Local ECDHE private key (P-256)
    local_ecdhe_private_key: Option<p256::SecretKey>,

    /// Local ECDHE public key (P-256)
    local_ecdhe_public_key: Option<p256::PublicKey>,

    /// Remote ECDHE public key (P-256)
    remote_ecdhe_public_key: Option<Bytes>,

    /// Raw handshake message data for verification
    handshake_messages: Vec<u8>,

    /// Client handshake messages for verification
    client_handshake_messages: Vec<u8>,

    /// Server handshake messages for verification
    server_handshake_messages: Vec<u8>,

    /// Indicates whether ChangeCipherSpec has been received
    change_cipher_spec_received: bool,
}

impl HandshakeState {
    /// Create a new handshake state machine
    pub fn new(role: DtlsRole, version: DtlsVersion, max_retransmissions: usize) -> Self {
        Self {
            step: HandshakeStep::Start,
            role,
            version,
            client_random: None,
            server_random: None,
            pre_master_secret: None,
            master_secret: None,
            client_certificate: None,
            server_certificate: None,
            cipher_suite: None,
            compression_method: None,
            message_seq: 0,
            retransmission_count: 0,
            max_retransmissions: max_retransmissions,
            srtp_profile: None,
            cookie: None,
            session_id: None,
            available_srtp_profiles: vec![SrtpProtectionProfile::Aes128CmSha1_80],
            local_ecdhe_private_key: None,
            local_ecdhe_public_key: None,
            remote_ecdhe_public_key: None,
            handshake_messages: Vec::new(),
            client_handshake_messages: Vec::new(),
            server_handshake_messages: Vec::new(),
            change_cipher_spec_received: false,
        }
    }

    /// Get the current handshake step
    pub fn step(&self) -> HandshakeStep {
        self.step
    }

    /// Get the role of this handshake state (client or server)
    pub fn role(&self) -> DtlsRole {
        self.role
    }

    /// Process a handshake message
    pub fn process_message(
        &mut self,
        message: HandshakeMessage,
    ) -> Result<Option<Vec<HandshakeMessage>>> {
        match self.role {
            DtlsRole::Client => self.process_message_client(message),
            DtlsRole::Server => self.process_message_server(message),
        }
    }

    /// Add a handshake message to the verification buffer with proper sender tracking
    pub fn add_handshake_message(
        &mut self,
        message_type: HandshakeType,
        data: &[u8],
        is_from_client: bool,
    ) {
        // Don't include HelloVerifyRequest or Finished messages in the verification
        if message_type == HandshakeType::HelloVerifyRequest
            || message_type == HandshakeType::Finished
        {
            return;
        }

        println!(
            "Adding message to handshake buffer: {:?}, length: {}, from_client: {}",
            message_type,
            data.len(),
            is_from_client
        );

        // Construct a handshake header for this message
        let header = super::message::handshake::HandshakeHeader::new(
            message_type,
            data.len() as u32,
            0, // We use 0 for verification purposes
            0, // No fragmentation for verification
            data.len() as u32,
        );

        // Serialize the header
        if let Ok(header_data) = header.serialize() {
            // Get the appropriate buffer based on sender
            let buffer = if is_from_client {
                &mut self.client_handshake_messages
            } else {
                &mut self.server_handshake_messages
            };

            // Add header to sender-specific buffer
            buffer.extend_from_slice(&header_data);

            // Add message body to sender-specific buffer
            buffer.extend_from_slice(data);

            // Also create a combined buffer in transcript order for verification
            // The RFC requires all handshake messages to be included in the order they were sent/received
            self.handshake_messages.clear();
            self.handshake_messages
                .extend_from_slice(&self.client_handshake_messages);
            self.handshake_messages
                .extend_from_slice(&self.server_handshake_messages);

            println!("Added message to handshake buffers. Combined size: {}, Client size: {}, Server size: {}",
                    self.handshake_messages.len(),
                    self.client_handshake_messages.len(),
                    self.server_handshake_messages.len());

            // Debug: dump first few bytes of each buffer
            if cfg!(debug_assertions) {
                let client_prefix = if !self.client_handshake_messages.is_empty() {
                    let end = std::cmp::min(self.client_handshake_messages.len(), 16);
                    format!("{:02X?}", &self.client_handshake_messages[..end])
                } else {
                    "[]".to_string()
                };

                let server_prefix = if !self.server_handshake_messages.is_empty() {
                    let end = std::cmp::min(self.server_handshake_messages.len(), 16);
                    format!("{:02X?}", &self.server_handshake_messages[..end])
                } else {
                    "[]".to_string()
                };

                println!("  - Client buffer prefix: {}", client_prefix);
                println!("  - Server buffer prefix: {}", server_prefix);
            }
        }
    }

    /// Generate a Finished message
    pub fn generate_finished_message(&self) -> Result<super::message::handshake::Finished> {
        // Make sure we have a master secret
        let master_secret = self.master_secret.as_ref().ok_or_else(|| {
            crate::error::Error::InvalidState("No master secret available".to_string())
        })?;

        // For verification, we always use the combined handshake_messages buffer
        // which has been carefully constructed in transcript order
        let verify_buffer = &self.handshake_messages;

        println!("Generating Finished message verify data:");
        println!(
            "  - Client messages size: {}",
            self.client_handshake_messages.len()
        );
        println!(
            "  - Server messages size: {}",
            self.server_handshake_messages.len()
        );
        println!("  - Total verify buffer size: {}", verify_buffer.len());
        println!(
            "  - Verify buffer first bytes: {:02X?}",
            &verify_buffer[..std::cmp::min(verify_buffer.len(), 32)]
        );

        // Print detailed debug information
        self.debug_verify_data();

        // Calculate the verify data
        let verify_data = super::crypto::keys::calculate_verify_data(
            master_secret,
            verify_buffer,
            self.role == DtlsRole::Client,
            super::crypto::cipher::HashAlgorithm::Sha256, // Using SHA-256 for TLS 1.2
        )?;

        println!("  - Generated verify data: {:02X?}", verify_data);

        // Create and return the Finished message
        Ok(super::message::handshake::Finished::new(verify_data))
    }

    /// Verify a Finished message
    pub fn verify_finished_message(
        &self,
        finished: &super::message::handshake::Finished,
    ) -> Result<bool> {
        // Make sure we have a master secret
        let master_secret = self.master_secret.as_ref().ok_or_else(|| {
            crate::error::Error::InvalidState("No master secret available".to_string())
        })?;

        // For verification, we need to check that the Finished message contains the correct verify_data
        let verify_buffer = &self.handshake_messages;

        // Print verification information for debugging
        println!("Verifying Finished message:");
        println!(
            "  - Client messages size: {}",
            self.client_handshake_messages.len()
        );
        println!(
            "  - Server messages size: {}",
            self.server_handshake_messages.len()
        );
        println!("  - Total verify buffer size: {}", verify_buffer.len());
        println!(
            "  - Verifying message from: {}",
            if self.role == DtlsRole::Client {
                "server"
            } else {
                "client"
            }
        );
        println!(
            "  - Verify buffer first bytes: {:02X?}",
            &verify_buffer[..std::cmp::min(verify_buffer.len(), 32)]
        );

        // Print detailed debug information
        self.debug_verify_data();

        // Calculate the expected verify data
        let expected_verify_data = super::crypto::keys::calculate_verify_data(
            master_secret,
            verify_buffer,
            // If we're a client, we're verifying a server message, and vice versa
            self.role != DtlsRole::Client,
            super::crypto::cipher::HashAlgorithm::Sha256, // Using SHA-256 for TLS 1.2
        )?;

        println!(
            "  - Expected verify data length: {}",
            expected_verify_data.len()
        );
        println!(
            "  - Received verify data length: {}",
            finished.verify_data.len()
        );

        // Compare the verify data
        let result = expected_verify_data == finished.verify_data;

        if !result {
            println!("WARNING: Finished message verification failed!");
            println!("  - Expected: {:?}", expected_verify_data);
            println!("  - Received: {:?}", finished.verify_data);

            return Err(crate::error::Error::DtlsHandshakeError(
                "Finished message verification failed".to_string(),
            ));
        } else {
            println!("  - Finished message verification succeeded!");
        }

        Ok(result)
    }

    /// Consistently produce the handshake hash input that both parties would agree on
    ///
    /// According to the TLS/DTLS spec, the handshake hash is computed over all
    /// handshake messages in the order they were sent and received.
    ///
    /// This ensures both parties compute the same hash despite having different views
    /// of which messages were sent vs received.
    pub fn get_handshake_hash_input(&self) -> Vec<u8> {
        // Return a copy of our carefully maintained handshake_messages buffer
        self.handshake_messages.clone()
    }

    /// Process a handshake message as a client
    fn process_message_client(
        &mut self,
        message: HandshakeMessage,
    ) -> Result<Option<Vec<HandshakeMessage>>> {
        // Add the message to our verification data (as coming from server)
        if let Ok(msg_data) = message.serialize() {
            self.add_handshake_message(message.message_type(), &msg_data, false);
        }

        match self.step {
            HandshakeStep::SentClientHello => {
                match message {
                    HandshakeMessage::HelloVerifyRequest(hello_verify) => {
                        // Store the cookie
                        self.cookie = Some(hello_verify.cookie.clone());

                        // Generate a new ClientHello with the cookie
                        let client_hello = self.generate_client_hello()?;

                        // Print client random for debugging
                        println!("Client sending second ClientHello with random (first 8 bytes): {:02X?}",
                                 &client_hello.random[..8]);

                        // Clear the handshake buffers and add the new ClientHello
                        // This is because in DTLS, the handshake hash only starts from
                        // the second ClientHello (with cookie)
                        self.client_handshake_messages.clear();
                        self.server_handshake_messages.clear();
                        self.handshake_messages.clear();

                        // Add to verification buffer (as client message)
                        if let Ok(msg_data) = client_hello.serialize() {
                            self.add_handshake_message(HandshakeType::ClientHello, &msg_data, true);
                        }

                        // Update state
                        self.step = HandshakeStep::SentClientHello;

                        // Return the ClientHello
                        Ok(Some(vec![HandshakeMessage::ClientHello(client_hello)]))
                    }
                    HandshakeMessage::ServerHello(server_hello) => {
                        // Save server random
                        self.server_random = Some(server_hello.random);

                        // Save negotiated parameters
                        self.cipher_suite = Some(server_hello.cipher_suite);
                        self.compression_method = Some(server_hello.compression_method);
                        self.session_id = Some(server_hello.session_id.clone());

                        // Check for SRTP extension
                        for ext in &server_hello.extensions {
                            if let Extension::UseSrtp(srtp_ext) = ext {
                                if !srtp_ext.profiles.is_empty() {
                                    // Use the first profile
                                    self.srtp_profile = Some(srtp_ext.profiles[0].into());
                                }
                            }
                        }

                        // Update state
                        self.step = HandshakeStep::ReceivedServerHello;

                        // Must wait for ServerKeyExchange before proceeding
                        Ok(None)
                    }
                    _ => {
                        // Unexpected message
                        self.step = HandshakeStep::Failed;
                        Err(crate::error::Error::DtlsHandshakeError(format!(
                            "Unexpected message in state {:?}: {:?}",
                            self.step,
                            message.message_type()
                        )))
                    }
                }
            }
            HandshakeStep::ReceivedServerHello => {
                match message {
                    HandshakeMessage::ServerKeyExchange(server_key_exchange) => {
                        println!("Client received ServerKeyExchange message");

                        // Store server's public key
                        self.remote_ecdhe_public_key = Some(server_key_exchange.public_key.clone());

                        // Generate our own ECDHE key pair
                        let (private_key, public_key) =
                            super::crypto::keys::generate_ecdh_keypair()?;

                        // Store our private key
                        self.local_ecdhe_private_key = Some(private_key);
                        self.local_ecdhe_public_key = Some(public_key);

                        // Encode our public key for transmission
                        let encoded_public_key =
                            super::crypto::keys::encode_public_key(&public_key)?;

                        // Create ClientKeyExchange message with our public key
                        let client_key_exchange =
                            super::message::handshake::ClientKeyExchange::new_ecdhe(
                                encoded_public_key,
                            );

                        // Calculate the pre-master secret (ECDHE shared secret)
                        if let Some(local_private_key) = &self.local_ecdhe_private_key {
                            if let Some(remote_public_key) = &self.remote_ecdhe_public_key {
                                // Encode private key for pre-master secret calculation
                                let encoded_private_key =
                                    super::crypto::keys::encode_private_key(local_private_key)?;

                                // Calculate pre-master secret
                                let pre_master_secret =
                                    super::crypto::keys::generate_ecdhe_pre_master_secret(
                                        remote_public_key,
                                        &encoded_private_key,
                                    )?;

                                // Store pre-master secret
                                self.pre_master_secret = Some(pre_master_secret.to_vec());

                                // Print first bytes of pre-master secret for debugging
                                println!(
                                    "Client pre-master secret first bytes: {:02X?}",
                                    &pre_master_secret[..std::cmp::min(pre_master_secret.len(), 8)]
                                );

                                // Calculate master secret
                                if let (Some(client_random), Some(server_random)) =
                                    (&self.client_random, &self.server_random)
                                {
                                    // Calculate master secret
                                    let master_secret =
                                        super::crypto::keys::calculate_master_secret(
                                            &pre_master_secret,
                                            client_random,
                                            server_random,
                                        )?;

                                    // Store master secret
                                    self.master_secret = Some(master_secret.to_vec());

                                    // Print first bytes of master secret for debugging
                                    println!(
                                        "Client master secret first bytes: {:02X?}",
                                        &master_secret[..std::cmp::min(master_secret.len(), 8)]
                                    );
                                }
                            }
                        }

                        // Add to verification buffer (as client message)
                        if let Ok(msg_data) = client_key_exchange.serialize() {
                            self.add_handshake_message(
                                HandshakeType::ClientKeyExchange,
                                &msg_data,
                                true,
                            );
                        }

                        // Update state
                        self.step = HandshakeStep::SentClientKeyExchange;

                        // Return ClientKeyExchange message
                        Ok(Some(vec![HandshakeMessage::ClientKeyExchange(
                            client_key_exchange,
                        )]))
                    }
                    _ => {
                        // Unexpected message
                        self.step = HandshakeStep::Failed;
                        Err(crate::error::Error::DtlsHandshakeError(format!(
                            "Unexpected message in state {:?}: {:?}",
                            self.step,
                            message.message_type()
                        )))
                    }
                }
            }
            HandshakeStep::ReceivedHelloVerifyRequest => {
                // This state is transient; we should have moved to SentClientHello
                self.step = HandshakeStep::Failed;
                Err(crate::error::Error::DtlsHandshakeError(format!(
                    "Unexpected state: {:?}",
                    self.step
                )))
            }
            HandshakeStep::SentClientKeyExchange => {
                match message {
                    HandshakeMessage::Finished(finished) => {
                        // Verify the Finished message
                        if !self.verify_finished_message(&finished)? {
                            self.step = HandshakeStep::Failed;
                            return Err(crate::error::Error::DtlsHandshakeError(
                                "Failed to verify server Finished message".to_string(),
                            ));
                        }

                        println!("Successfully verified server Finished message");

                        // Generate our own Finished message
                        let client_finished = self.generate_finished_message()?;

                        // Update state
                        self.step = HandshakeStep::Complete;

                        // Return our Finished message to be sent
                        Ok(Some(vec![HandshakeMessage::Finished(client_finished)]))
                    }
                    _ => {
                        // Unexpected message
                        self.step = HandshakeStep::Failed;
                        Err(crate::error::Error::DtlsHandshakeError(format!(
                            "Unexpected message in state {:?}: {:?}",
                            self.step,
                            message.message_type()
                        )))
                    }
                }
            }
            _ => {
                // Unexpected state
                self.step = HandshakeStep::Failed;
                Err(crate::error::Error::DtlsHandshakeError(format!(
                    "Unexpected state: {:?}",
                    self.step
                )))
            }
        }
    }

    /// Process a handshake message as a server
    fn process_message_server(
        &mut self,
        message: HandshakeMessage,
    ) -> Result<Option<Vec<HandshakeMessage>>> {
        // Add the message to our verification data (as coming from client)
        if let Ok(msg_data) = message.serialize() {
            self.add_handshake_message(message.message_type(), &msg_data, true);
        }

        match self.step {
            HandshakeStep::Start => {
                match message {
                    HandshakeMessage::ClientHello(client_hello) => {
                        // Save client random
                        self.client_random = Some(client_hello.random);

                        // Print client random for debugging
                        println!(
                            "Server received ClientHello with random (first 8 bytes): {:02X?}",
                            &client_hello.random[..8]
                        );

                        // Check for cookie
                        if client_hello.cookie.is_empty() {
                            // No cookie - send HelloVerifyRequest
                            self.step = HandshakeStep::ReceivedClientHello;

                            // Generate a cookie (in a real implementation, this would be cryptographically secure)
                            let mut rng = rand::thread_rng();
                            let mut cookie = vec![0u8; 16];
                            rng.fill(&mut cookie[..]);

                            let cookie = Bytes::from(cookie);
                            self.cookie = Some(cookie.clone());

                            // Create HelloVerifyRequest
                            let hello_verify = HelloVerifyRequest::new(self.version, cookie);

                            // For DTLS, we don't include the first ClientHello or HelloVerifyRequest in the transcript
                            // So we need to clear our handshake message buffers
                            self.client_handshake_messages.clear();
                            self.server_handshake_messages.clear();
                            self.handshake_messages.clear();

                            // Update state
                            self.step = HandshakeStep::SentHelloVerifyRequest;

                            Ok(Some(vec![HandshakeMessage::HelloVerifyRequest(
                                hello_verify,
                            )]))
                        } else {
                            // Cookie present - validate it
                            // (In a real implementation, we'd verify the cookie)

                            // This is the second ClientHello with cookie that we'll use for verification
                            // Print client random again for confirmation
                            println!("Server received ClientHello WITH COOKIE, random (first 8 bytes): {:02X?}",
                                     &client_hello.random[..8]);

                            // Generate a session ID
                            let mut rng = rand::thread_rng();
                            let mut session_id = vec![0u8; 32];
                            rng.fill(&mut session_id[..]);

                            let session_id = Bytes::from(session_id);
                            self.session_id = Some(session_id.clone());

                            // Select cipher suite (choose the first one we support)
                            let selected_cipher = client_hello
                                .cipher_suites
                                .iter()
                                .find(|&&suite| self.is_supported_cipher_suite(suite))
                                .copied();

                            if let Some(cipher) = selected_cipher {
                                self.cipher_suite = Some(cipher);
                            } else {
                                // No supported cipher suite
                                self.step = HandshakeStep::Failed;
                                return Err(crate::error::Error::DtlsHandshakeError(
                                    "No supported cipher suite".to_string(),
                                ));
                            }

                            // Select compression method (always 0 - no compression)
                            self.compression_method = Some(0);

                            // Check for SRTP extension
                            let mut use_srtp_extension = None;

                            for ext in &client_hello.extensions {
                                if let Extension::UseSrtp(srtp_ext) = ext {
                                    // Find the first supported profile
                                    for profile in &srtp_ext.profiles {
                                        if self.available_srtp_profiles.contains(profile) {
                                            self.srtp_profile = Some((*profile).into());

                                            // Create a new UseSrtp extension with just this profile
                                            use_srtp_extension =
                                                Some(UseSrtpExtension::with_profiles(vec![
                                                    *profile,
                                                ]));

                                            break;
                                        }
                                    }
                                }
                            }

                            // Create ServerHello
                            let mut extensions = Vec::new();

                            if let Some(srtp_ext) = use_srtp_extension {
                                extensions.push(Extension::UseSrtp(srtp_ext));
                            }

                            let server_hello = ServerHello::new(
                                self.version,
                                session_id,
                                self.cipher_suite.unwrap(),
                                self.compression_method.unwrap(),
                                extensions,
                            );

                            // Save server random
                            self.server_random = Some(server_hello.random);

                            // Generate ECDHE key pair
                            let (private_key, public_key) =
                                super::crypto::keys::generate_ecdh_keypair()?;

                            // Store the private key for later
                            self.local_ecdhe_private_key = Some(private_key);
                            self.local_ecdhe_public_key = Some(public_key);

                            // Encode the public key for transmission
                            let encoded_public_key =
                                super::crypto::keys::encode_public_key(&public_key)?;

                            // Create ServerKeyExchange message with our public key
                            let server_key_exchange =
                                super::message::handshake::ServerKeyExchange::new_ecdhe(
                                    encoded_public_key,
                                );

                            // Explicitly add these messages to our verification buffer
                            // This is critical for proper Finished message verification
                            if let Ok(server_hello_data) = server_hello.serialize() {
                                self.add_handshake_message(
                                    HandshakeType::ServerHello,
                                    &server_hello_data,
                                    false,
                                );
                            }

                            if let Ok(server_key_exchange_data) = server_key_exchange.serialize() {
                                self.add_handshake_message(
                                    HandshakeType::ServerKeyExchange,
                                    &server_key_exchange_data,
                                    false,
                                );
                            }

                            // Update state
                            self.step = HandshakeStep::SentServerHello;

                            // Send ServerHello and ServerKeyExchange
                            Ok(Some(vec![
                                HandshakeMessage::ServerHello(server_hello),
                                HandshakeMessage::ServerKeyExchange(server_key_exchange),
                            ]))
                        }
                    }
                    _ => {
                        // Unexpected message
                        self.step = HandshakeStep::Failed;
                        Err(crate::error::Error::DtlsHandshakeError(format!(
                            "Unexpected message in state {:?}: {:?}",
                            self.step,
                            message.message_type()
                        )))
                    }
                }
            }
            HandshakeStep::SentHelloVerifyRequest => {
                match message {
                    HandshakeMessage::ClientHello(client_hello) => {
                        // Verify cookie
                        if let Some(ref our_cookie) = self.cookie {
                            if &client_hello.cookie != our_cookie {
                                // Invalid cookie
                                self.step = HandshakeStep::Failed;
                                return Err(crate::error::Error::DtlsHandshakeError(
                                    "Invalid cookie".to_string(),
                                ));
                            }
                        }

                        // Save client random
                        self.client_random = Some(client_hello.random);

                        // Generate a session ID
                        let mut rng = rand::thread_rng();
                        let mut session_id = vec![0u8; 32];
                        rng.fill(&mut session_id[..]);

                        let session_id = Bytes::from(session_id);
                        self.session_id = Some(session_id.clone());

                        // Select cipher suite (choose the first one we support)
                        let selected_cipher = client_hello
                            .cipher_suites
                            .iter()
                            .find(|&&suite| self.is_supported_cipher_suite(suite))
                            .copied();

                        if let Some(cipher) = selected_cipher {
                            self.cipher_suite = Some(cipher);
                        } else {
                            // No supported cipher suite
                            self.step = HandshakeStep::Failed;
                            return Err(crate::error::Error::DtlsHandshakeError(
                                "No supported cipher suite".to_string(),
                            ));
                        }

                        // Select compression method (always 0 - no compression)
                        self.compression_method = Some(0);

                        // Check for SRTP extension
                        let mut use_srtp_extension = None;

                        for ext in &client_hello.extensions {
                            if let Extension::UseSrtp(srtp_ext) = ext {
                                // Find the first supported profile
                                for profile in &srtp_ext.profiles {
                                    if self.available_srtp_profiles.contains(profile) {
                                        self.srtp_profile = Some((*profile).into());

                                        // Create a new UseSrtp extension with just this profile
                                        use_srtp_extension =
                                            Some(UseSrtpExtension::with_profiles(vec![*profile]));

                                        break;
                                    }
                                }
                            }
                        }

                        // Create ServerHello
                        let mut extensions = Vec::new();

                        if let Some(srtp_ext) = use_srtp_extension {
                            extensions.push(Extension::UseSrtp(srtp_ext));
                        }

                        let server_hello = ServerHello::new(
                            self.version,
                            session_id,
                            self.cipher_suite.unwrap(),
                            self.compression_method.unwrap(),
                            extensions,
                        );

                        // Save server random
                        self.server_random = Some(server_hello.random);

                        // Generate ECDHE key pair
                        let (private_key, public_key) =
                            super::crypto::keys::generate_ecdh_keypair()?;

                        // Store the private key for later
                        self.local_ecdhe_private_key = Some(private_key);
                        self.local_ecdhe_public_key = Some(public_key);

                        // Encode the public key for transmission
                        let encoded_public_key =
                            super::crypto::keys::encode_public_key(&public_key)?;

                        // Create ServerKeyExchange message with our public key
                        let server_key_exchange =
                            super::message::handshake::ServerKeyExchange::new_ecdhe(
                                encoded_public_key,
                            );

                        // Explicitly add these messages to our verification buffer
                        // This is critical for proper Finished message verification
                        if let Ok(server_hello_data) = server_hello.serialize() {
                            self.add_handshake_message(
                                HandshakeType::ServerHello,
                                &server_hello_data,
                                false,
                            );
                        }

                        if let Ok(server_key_exchange_data) = server_key_exchange.serialize() {
                            self.add_handshake_message(
                                HandshakeType::ServerKeyExchange,
                                &server_key_exchange_data,
                                false,
                            );
                        }

                        // Update state
                        self.step = HandshakeStep::SentServerHello;

                        // Send ServerHello and ServerKeyExchange
                        Ok(Some(vec![
                            HandshakeMessage::ServerHello(server_hello),
                            HandshakeMessage::ServerKeyExchange(server_key_exchange),
                        ]))
                    }
                    _ => {
                        // Unexpected message
                        self.step = HandshakeStep::Failed;
                        Err(crate::error::Error::DtlsHandshakeError(format!(
                            "Unexpected message in state {:?}: {:?}",
                            self.step,
                            message.message_type()
                        )))
                    }
                }
            }
            HandshakeStep::SentServerHello => {
                match message {
                    HandshakeMessage::ClientKeyExchange(client_key_exchange) => {
                        println!(
                            "Server received ClientKeyExchange, length: {}",
                            client_key_exchange.exchange_data.len()
                        );

                        // Store client's public key
                        self.remote_ecdhe_public_key =
                            Some(client_key_exchange.exchange_data.clone());

                        // Calculate the pre-master secret (ECDHE shared secret)
                        if let Some(local_private_key) = &self.local_ecdhe_private_key {
                            if let Some(remote_public_key) = &self.remote_ecdhe_public_key {
                                // Encode private key for pre-master secret calculation
                                let encoded_private_key =
                                    super::crypto::keys::encode_private_key(local_private_key)?;

                                // Calculate pre-master secret
                                let pre_master_secret =
                                    super::crypto::keys::generate_ecdhe_pre_master_secret(
                                        remote_public_key,
                                        &encoded_private_key,
                                    )?;

                                // Store pre-master secret
                                self.pre_master_secret = Some(pre_master_secret.to_vec());

                                // Print first bytes of pre-master secret for debugging
                                println!(
                                    "Server pre-master secret first bytes: {:02X?}",
                                    &pre_master_secret[..std::cmp::min(pre_master_secret.len(), 8)]
                                );

                                // Calculate master secret
                                if let (Some(client_random), Some(server_random)) =
                                    (&self.client_random, &self.server_random)
                                {
                                    // Calculate master secret
                                    let master_secret =
                                        super::crypto::keys::calculate_master_secret(
                                            &pre_master_secret,
                                            client_random,
                                            server_random,
                                        )?;

                                    // Store master secret
                                    self.master_secret = Some(master_secret.to_vec());

                                    // Print first bytes of master secret for debugging
                                    println!(
                                        "Server master secret first bytes: {:02X?}",
                                        &master_secret[..std::cmp::min(master_secret.len(), 8)]
                                    );
                                }
                            }
                        }

                        // Update state
                        self.step = HandshakeStep::ReceivedClientKeyExchange;

                        // Return ChangeCipherSpec and Finished messages
                        Ok(None)
                    }
                    _ => {
                        // Unexpected message
                        self.step = HandshakeStep::Failed;
                        Err(crate::error::Error::DtlsHandshakeError(format!(
                            "Unexpected message in state {:?}: {:?}",
                            self.step,
                            message.message_type()
                        )))
                    }
                }
            }
            HandshakeStep::ReceivedClientKeyExchange => {
                match message {
                    HandshakeMessage::Finished(finished) => {
                        // Verify the Finished message
                        if !self.verify_finished_message(&finished)? {
                            self.step = HandshakeStep::Failed;
                            return Err(crate::error::Error::DtlsHandshakeError(
                                "Failed to verify client Finished message".to_string(),
                            ));
                        }

                        println!("Successfully verified client Finished message");

                        // Generate our own Finished message
                        let server_finished = self.generate_finished_message()?;

                        // Update state
                        self.step = HandshakeStep::SentServerFinished;

                        // Return our Finished message to be sent
                        Ok(Some(vec![HandshakeMessage::Finished(server_finished)]))
                    }
                    _ => {
                        // Unexpected message
                        self.step = HandshakeStep::Failed;
                        Err(crate::error::Error::DtlsHandshakeError(format!(
                            "Unexpected message in state {:?}: {:?}",
                            self.step,
                            message.message_type()
                        )))
                    }
                }
            }
            HandshakeStep::SentServerFinished => {
                // Any message received after SentServerFinished transitions to Complete
                self.step = HandshakeStep::Complete;
                Ok(None)
            }
            _ => {
                // Unexpected state
                self.step = HandshakeStep::Failed;
                Err(crate::error::Error::DtlsHandshakeError(format!(
                    "Unexpected state: {:?}",
                    self.step
                )))
            }
        }
    }

    /// Check if a cipher suite is supported
    fn is_supported_cipher_suite(&self, cipher_suite: u16) -> bool {
        // For now, support a small set of common suites
        matches!(
            cipher_suite,
            0xC02B | // TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256
            0xC02F | // TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256
            0xC009 | // TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA
            0xC013 | // TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA
            0x002F // TLS_RSA_WITH_AES_128_CBC_SHA
        )
    }

    /// Generate a ClientHello message
    pub fn generate_client_hello(&mut self) -> Result<super::message::handshake::ClientHello> {
        // Available cipher suites
        let cipher_suites = vec![
            // Only ECDHE ciphers for now
            0xC009, // TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA
            0xC013, // TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA
            0xC02B, // TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256
            0xC02F, // TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256
        ];

        // Compression methods (0 = none)
        let compression_methods = vec![0];

        // Add SRTP extension
        let srtp_extension =
            UseSrtpExtension::with_profiles(vec![SrtpProtectionProfile::Aes128CmSha1_80]);

        let extensions = vec![Extension::UseSrtp(srtp_extension)];

        // Create ClientHello
        let mut client_hello = ClientHello::new(
            self.version,
            Bytes::new(),                                   // Empty session ID
            self.cookie.clone().unwrap_or_else(Bytes::new), // Add the cookie
            cipher_suites,
            compression_methods,
            extensions,
        );

        // If this is the second ClientHello with cookie, reuse the same random
        // from the first ClientHello to ensure consistent verification
        if let Some(existing_random) = self.client_random {
            // Copy over the existing random to ensure consistency
            client_hello.random.copy_from_slice(&existing_random);
            println!(
                "Reusing existing client random (first 8 bytes): {:02X?}",
                &client_hello.random[..8]
            );
        } else {
            // Save the client random
            self.client_random = Some(client_hello.random);
            println!(
                "Generated new client random (first 8 bytes): {:02X?}",
                &client_hello.random[..8]
            );
        }

        Ok(client_hello)
    }

    /// Start the handshake process
    pub fn start(&mut self) -> Result<Vec<HandshakeMessage>> {
        match self.role {
            DtlsRole::Client => {
                println!("Starting handshake as CLIENT");
                // Generate ClientHello
                let client_hello = self.generate_client_hello()?;

                // Add to verification buffer (as client message)
                if let Ok(msg_data) = client_hello.serialize() {
                    self.add_handshake_message(HandshakeType::ClientHello, &msg_data, true);
                }

                // Update state
                self.step = HandshakeStep::SentClientHello;

                Ok(vec![HandshakeMessage::ClientHello(client_hello)])
            }
            DtlsRole::Server => {
                println!("Starting handshake as SERVER");
                // Server waits for ClientHello
                self.step = HandshakeStep::Start;
                Ok(Vec::new())
            }
        }
    }

    /// Reset the handshake state
    pub fn reset(&mut self) {
        self.step = HandshakeStep::Start;
        self.client_random = None;
        self.server_random = None;
        self.pre_master_secret = None;
        self.master_secret = None;
        self.client_certificate = None;
        self.server_certificate = None;
        self.cipher_suite = None;
        self.compression_method = None;
        self.message_seq = 0;
        self.retransmission_count = 0;
        self.srtp_profile = None;
        self.cookie = None;
        self.session_id = None;
        self.local_ecdhe_private_key = None;
        self.local_ecdhe_public_key = None;
        self.remote_ecdhe_public_key = None;
        self.handshake_messages.clear();
        self.client_handshake_messages.clear();
        self.server_handshake_messages.clear();
        self.change_cipher_spec_received = false;
    }

    /// Get the master secret
    pub fn master_secret(&self) -> Option<&[u8]> {
        self.master_secret.as_deref()
    }

    /// Get the client random
    pub fn client_random(&self) -> Option<&[u8; 32]> {
        self.client_random.as_ref()
    }

    /// Get the server random
    pub fn server_random(&self) -> Option<&[u8; 32]> {
        self.server_random.as_ref()
    }

    /// Get the negotiated SRTP profile
    pub fn srtp_profile(&self) -> Option<u16> {
        self.srtp_profile
    }

    /// Get the cookie if any (debug helper)
    pub fn cookie(&self) -> Option<&Bytes> {
        self.cookie.as_ref()
    }

    /// Set the ChangeCipherSpec received flag
    pub fn set_change_cipher_spec_received(&mut self, received: bool) {
        self.change_cipher_spec_received = received;
    }

    /// Check if ChangeCipherSpec has been received
    pub fn change_cipher_spec_received(&self) -> bool {
        self.change_cipher_spec_received
    }

    /// Force a specific buffer for verification purposes
    ///
    /// This is used for testing and debugging only, to force both sides to use
    /// the same verification data.
    pub fn force_verification_buffer(&mut self, buffer: Vec<u8>) -> Result<()> {
        // Store the buffer to be used for verification
        println!("Forcing verification buffer, length: {}", buffer.len());
        println!(
            "Buffer first bytes: {:02X?}",
            &buffer[..std::cmp::min(buffer.len(), 32)]
        );

        // Replace the combined handshake_messages buffer
        self.handshake_messages = buffer;

        Ok(())
    }

    /// Print a debug comparison of handshake messages for verification
    pub fn debug_verify_data(&self) {
        // Log which messages we have
        println!("*** DEBUG HANDSHAKE STATE ***");
        println!("Role: {:?}", self.role);
        println!("Step: {:?}", self.step);
        println!(
            "Client messages length: {} bytes",
            self.client_handshake_messages.len()
        );
        println!(
            "Server messages length: {} bytes",
            self.server_handshake_messages.len()
        );
        println!(
            "Combined messages length: {} bytes",
            self.handshake_messages.len()
        );

        // Analyze client messages
        println!("Client messages:");
        if !self.client_handshake_messages.is_empty() {
            let mut pos = 0;
            while pos < self.client_handshake_messages.len() {
                if pos + 12 <= self.client_handshake_messages.len() {
                    if let Ok((header, header_size)) =
                        super::message::handshake::HandshakeHeader::parse(
                            &self.client_handshake_messages[pos..],
                        )
                    {
                        println!(
                            "  - Type: {:?}, Length: {}, Offset: {}",
                            header.msg_type, header.fragment_length, pos
                        );
                        pos += header_size + header.fragment_length as usize;
                    } else {
                        println!("  - Failed to parse header at offset {}", pos);
                        break;
                    }
                } else {
                    println!("  - Incomplete message at offset {}", pos);
                    break;
                }
            }
        } else {
            println!("  - No client messages");
        }

        // Analyze server messages
        println!("Server messages:");
        if !self.server_handshake_messages.is_empty() {
            let mut pos = 0;
            while pos < self.server_handshake_messages.len() {
                if pos + 12 <= self.server_handshake_messages.len() {
                    if let Ok((header, header_size)) =
                        super::message::handshake::HandshakeHeader::parse(
                            &self.server_handshake_messages[pos..],
                        )
                    {
                        println!(
                            "  - Type: {:?}, Length: {}, Offset: {}",
                            header.msg_type, header.fragment_length, pos
                        );
                        pos += header_size + header.fragment_length as usize;
                    } else {
                        println!("  - Failed to parse header at offset {}", pos);
                        break;
                    }
                } else {
                    println!("  - Incomplete message at offset {}", pos);
                    break;
                }
            }
        } else {
            println!("  - No server messages");
        }

        // Compare hash values for debugging
        if let Some(master_secret) = &self.master_secret {
            // Hash for client Finished
            let client_verify = super::crypto::keys::calculate_verify_data(
                master_secret,
                &self.handshake_messages,
                true, // is client
                super::crypto::cipher::HashAlgorithm::Sha256,
            );

            // Hash for server Finished
            let server_verify = super::crypto::keys::calculate_verify_data(
                master_secret,
                &self.handshake_messages,
                false, // is server
                super::crypto::cipher::HashAlgorithm::Sha256,
            );

            if let Ok(client_data) = client_verify {
                println!("Expected client verify data: {:02X?}", client_data);
            }

            if let Ok(server_data) = server_verify {
                println!("Expected server verify data: {:02X?}", server_data);
            }
        }

        println!("****************************");
    }

    /// Synchronize handshake state between client and server
    ///
    /// This ensures both sides use the exact same handshake message list
    /// for verification purposes - a critical requirement for DTLS security.
    ///
    /// It should be called right before generating the Finished message.
    pub fn sync_verify_data(
        &mut self,
        client_msgs: Option<Vec<u8>>,
        server_msgs: Option<Vec<u8>>,
    ) -> Result<()> {
        // Only overwrite if provided (otherwise keep existing)
        if let Some(client_msgs) = client_msgs {
            self.client_handshake_messages = client_msgs;
        }

        if let Some(server_msgs) = server_msgs {
            self.server_handshake_messages = server_msgs;
        }

        // Always recreate the combined buffer in the proper order
        self.handshake_messages.clear();
        self.handshake_messages
            .extend_from_slice(&self.client_handshake_messages);
        self.handshake_messages
            .extend_from_slice(&self.server_handshake_messages);

        println!("Synchronized handshake verification data:");
        println!(
            "  - Client messages: {} bytes",
            self.client_handshake_messages.len()
        );
        println!(
            "  - Server messages: {} bytes",
            self.server_handshake_messages.len()
        );
        println!(
            "  - Combined buffer: {} bytes",
            self.handshake_messages.len()
        );

        self.debug_verify_data();

        Ok(())
    }
}