ant_quic/

quic_node.rs

//! QUIC-based P2P node with NAT traversal
//!
//! This module provides a QUIC-based implementation of the P2P node
//! that integrates with the NAT traversal protocol.

use std::{
    collections::HashMap,
    net::SocketAddr,
    sync::Arc,
    time::{Duration, Instant},
};

use tracing::{debug, error, info};

use crate::{
    auth::{AuthConfig, AuthManager, AuthMessage, AuthProtocol},
    crypto::raw_public_keys::key_utils::{
        derive_peer_id_from_public_key, generate_ed25519_keypair,
    },
    nat_traversal_api::{
        EndpointRole, NatTraversalConfig, NatTraversalEndpoint, NatTraversalError,
        NatTraversalEvent, NatTraversalStatistics, PeerId,
    },
};

/// QUIC-based P2P node with NAT traversal
#[derive(Clone, Debug)]
pub struct QuicP2PNode {
    /// NAT traversal endpoint
    nat_endpoint: Arc<NatTraversalEndpoint>,
    /// Active peer connections (maps peer ID to their socket address)
    connected_peers: Arc<tokio::sync::RwLock<HashMap<PeerId, SocketAddr>>>,
    /// Node statistics
    stats: Arc<tokio::sync::Mutex<NodeStats>>,
    /// Node configuration
    config: QuicNodeConfig,
    /// Authentication manager
    auth_manager: Arc<AuthManager>,
    /// Our peer ID
    peer_id: PeerId,
}

/// Configuration for QUIC P2P node
#[derive(Debug, Clone)]
pub struct QuicNodeConfig {
    /// Role of this node
    pub role: EndpointRole,
    /// Bootstrap nodes
    pub bootstrap_nodes: Vec<SocketAddr>,
    /// Enable coordinator services
    pub enable_coordinator: bool,
    /// Max concurrent connections
    pub max_connections: usize,
    /// Connection timeout
    pub connection_timeout: Duration,
    /// Statistics interval
    pub stats_interval: Duration,
    /// Authentication configuration
    pub auth_config: AuthConfig,
    /// Bind address for the node
    pub bind_addr: Option<SocketAddr>,
}

impl Default for QuicNodeConfig {
    fn default() -> Self {
        Self {
            role: EndpointRole::Client,
            bootstrap_nodes: Vec::new(),
            enable_coordinator: false,
            max_connections: 100,
            connection_timeout: Duration::from_secs(30),
            stats_interval: Duration::from_secs(30),
            auth_config: AuthConfig::default(),
            bind_addr: None,
        }
    }
}

#[derive(Debug, Clone)]
pub struct ConnectionMetrics {
    /// Bytes sent to this peer
    pub bytes_sent: u64,
    /// Bytes received from this peer
    pub bytes_received: u64,
    /// Round-trip time
    pub rtt: Option<Duration>,
    /// Packet loss rate (0.0 to 1.0)
    pub packet_loss: f64,
}

#[derive(Debug, Clone)]
pub struct NodeStats {
    /// Number of active connections
    pub active_connections: usize,
    /// Total successful connections
    pub successful_connections: u64,
    /// Total failed connections
    pub failed_connections: u64,
    /// NAT traversal attempts
    pub nat_traversal_attempts: u64,
    /// Successful NAT traversals
    pub nat_traversal_successes: u64,
    /// Node start time
    pub start_time: Instant,
}

impl Default for NodeStats {
    fn default() -> Self {
        Self {
            active_connections: 0,
            successful_connections: 0,
            failed_connections: 0,
            nat_traversal_attempts: 0,
            nat_traversal_successes: 0,
            start_time: Instant::now(),
        }
    }
}

impl QuicP2PNode {
    /// Create a new QUIC P2P node
    pub async fn new(
        config: QuicNodeConfig,
    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        // Generate Ed25519 keypair for authentication
        let (secret_key, public_key) = generate_ed25519_keypair();
        let peer_id = derive_peer_id_from_public_key(&public_key);

        info!("Creating QUIC P2P node with peer ID: {:?}", peer_id);

        // Create authentication manager
        let auth_manager = Arc::new(AuthManager::new(secret_key, config.auth_config.clone()));

        // Create NAT traversal configuration
        let nat_config = NatTraversalConfig {
            role: config.role,
            bootstrap_nodes: config.bootstrap_nodes.clone(),
            max_candidates: 50,
            coordination_timeout: Duration::from_secs(10),
            enable_symmetric_nat: true,
            // Bootstrap nodes should not enable relay fallback
            enable_relay_fallback: !matches!(config.role, EndpointRole::Bootstrap),
            max_concurrent_attempts: 5,
            bind_addr: config.bind_addr,
            prefer_rfc_nat_traversal: true, // Default to RFC format for standards compliance
            timeouts: crate::config::nat_timeouts::TimeoutConfig::default(),
        };

        // Create event callback for NAT traversal events
        let stats_clone = Arc::new(tokio::sync::Mutex::new(NodeStats {
            start_time: Instant::now(),
            ..Default::default()
        }));
        let stats_for_callback = Arc::clone(&stats_clone);

        let event_callback = Box::new(move |event: NatTraversalEvent| {
            let stats = stats_for_callback.clone();
            tokio::spawn(async move {
                let mut stats = stats.lock().await;
                match event {
                    NatTraversalEvent::CoordinationRequested { .. } => {
                        stats.nat_traversal_attempts += 1;
                    }
                    NatTraversalEvent::ConnectionEstablished { .. } => {
                        stats.nat_traversal_successes += 1;
                    }
                    _ => {}
                }
            });
        });

        // Create NAT traversal endpoint
        let nat_endpoint =
            Arc::new(NatTraversalEndpoint::new(nat_config, Some(event_callback)).await?);

        Ok(Self {
            nat_endpoint,
            connected_peers: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
            stats: stats_clone,
            config,
            auth_manager,
            peer_id,
        })
    }

    /// Get the node configuration
    pub fn get_config(&self) -> &QuicNodeConfig {
        &self.config
    }

    /// Connect directly to a bootstrap node
    pub async fn connect_to_bootstrap(
        &self,
        bootstrap_addr: SocketAddr,
    ) -> Result<PeerId, NatTraversalError> {
        info!("Connecting to bootstrap node at {}", bootstrap_addr);

        // Get the quinn endpoint from NAT traversal endpoint
        let endpoint = self.nat_endpoint.get_quinn_endpoint().ok_or_else(|| {
            NatTraversalError::ConfigError("Quinn endpoint not available".to_string())
        })?;

        // Connect using the QUIC endpoint directly
        match endpoint.connect(bootstrap_addr, "bootstrap-node") {
            Ok(connecting) => {
                match connecting.await {
                    Ok(_connection) => {
                        // Extract peer ID from the connection
                        // For now, we'll generate a temporary peer ID based on the address
                        // In a real implementation, we'd exchange peer IDs during the handshake
                        let peer_id = self.derive_peer_id_from_address(bootstrap_addr);

                        // Store the connection
                        self.connected_peers
                            .write()
                            .await
                            .insert(peer_id, bootstrap_addr);

                        // Update stats
                        {
                            let mut stats = self.stats.lock().await;
                            stats.active_connections += 1;
                            stats.successful_connections += 1;
                        }

                        // Fire connection established event
                        if let Some(ref callback) = self.nat_endpoint.get_event_callback() {
                            callback(NatTraversalEvent::ConnectionEstablished {
                                peer_id,
                                remote_address: bootstrap_addr,
                            });
                        }

                        info!(
                            "Successfully connected to bootstrap node {} with peer ID {:?}",
                            bootstrap_addr, peer_id
                        );
                        Ok(peer_id)
                    }
                    Err(e) => {
                        error!(
                            "Failed to establish connection to bootstrap node {}: {}",
                            bootstrap_addr, e
                        );
                        {
                            let mut stats = self.stats.lock().await;
                            stats.failed_connections += 1;
                        }
                        Err(NatTraversalError::NetworkError(format!(
                            "Connection failed: {e}"
                        )))
                    }
                }
            }
            Err(e) => {
                error!(
                    "Failed to initiate connection to bootstrap node {}: {}",
                    bootstrap_addr, e
                );
                {
                    let mut stats = self.stats.lock().await;
                    stats.failed_connections += 1;
                }
                Err(NatTraversalError::NetworkError(format!(
                    "Connect error: {e}"
                )))
            }
        }
    }

    /// Derive a peer ID from a socket address (temporary solution)
    fn derive_peer_id_from_address(&self, addr: SocketAddr) -> PeerId {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        let mut hasher = DefaultHasher::new();
        addr.hash(&mut hasher);
        let hash = hasher.finish();

        let mut peer_id_bytes = [0u8; 32];
        peer_id_bytes[..8].copy_from_slice(&hash.to_le_bytes());
        let port_bytes = addr.port().to_le_bytes();
        peer_id_bytes[8..10].copy_from_slice(&port_bytes);

        PeerId(peer_id_bytes)
    }

    /// Connect to a peer using NAT traversal
    pub async fn connect_to_peer(
        &self,
        peer_id: PeerId,
        coordinator: SocketAddr,
    ) -> Result<SocketAddr, NatTraversalError> {
        info!(
            "Initiating connection to peer {:?} via coordinator {}",
            peer_id, coordinator
        );

        // Update stats
        {
            let mut stats = self.stats.lock().await;
            stats.nat_traversal_attempts += 1;
        }

        // Initiate NAT traversal
        self.nat_endpoint
            .initiate_nat_traversal(peer_id, coordinator)?;

        // Poll for completion (in production, this would be event-driven)
        let start = Instant::now();
        let timeout = self.config.connection_timeout;

        while start.elapsed() < timeout {
            let events = self.nat_endpoint.poll(Instant::now())?;

            for event in events {
                match event {
                    NatTraversalEvent::ConnectionEstablished {
                        peer_id: evt_peer,
                        remote_address,
                    } => {
                        if evt_peer == peer_id {
                            // Store peer connection
                            {
                                let mut peers = self.connected_peers.write().await;
                                peers.insert(peer_id, remote_address);
                            }

                            // Update stats
                            {
                                let mut stats = self.stats.lock().await;
                                stats.successful_connections += 1;
                                stats.active_connections += 1;
                                stats.nat_traversal_successes += 1;
                            }

                            info!(
                                "Successfully connected to peer {:?} at {}",
                                peer_id, remote_address
                            );

                            // Perform authentication if required
                            if self.config.auth_config.require_authentication {
                                match self.authenticate_as_initiator(&peer_id).await {
                                    Ok(_) => {
                                        info!("Authentication successful with peer {:?}", peer_id);
                                    }
                                    Err(e) => {
                                        error!(
                                            "Authentication failed with peer {:?}: {}",
                                            peer_id, e
                                        );
                                        // Remove from connected peers
                                        self.connected_peers.write().await.remove(&peer_id);
                                        // Update stats
                                        let mut stats = self.stats.lock().await;
                                        stats.active_connections =
                                            stats.active_connections.saturating_sub(1);
                                        stats.failed_connections += 1;
                                        return Err(NatTraversalError::ConfigError(format!(
                                            "Authentication failed: {e}"
                                        )));
                                    }
                                }
                            }

                            return Ok(remote_address);
                        }
                    }
                    NatTraversalEvent::TraversalFailed {
                        peer_id: evt_peer,
                        error,
                        fallback_available: _,
                    } => {
                        if evt_peer == peer_id {
                            // Update stats
                            {
                                let mut stats = self.stats.lock().await;
                                stats.failed_connections += 1;
                            }

                            error!("NAT traversal failed for peer {:?}: {}", peer_id, error);
                            return Err(error);
                        }
                    }
                    _ => {
                        debug!("Received event: {:?}", event);
                    }
                }
            }

            // Brief sleep to avoid busy waiting
            tokio::time::sleep(Duration::from_millis(100)).await;
        }

        // Timeout
        {
            let mut stats = self.stats.lock().await;
            stats.failed_connections += 1;
        }

        Err(NatTraversalError::Timeout)
    }

    /// Accept incoming connections
    pub async fn accept(
        &self,
    ) -> Result<(SocketAddr, PeerId), Box<dyn std::error::Error + Send + Sync>> {
        info!("Waiting for incoming connection...");

        // Accept connection through the NAT traversal endpoint
        match self.nat_endpoint.accept_connection().await {
            Ok((peer_id, connection)) => {
                let remote_addr = connection.remote_address();

                // Store the connection
                {
                    let mut peers = self.connected_peers.write().await;
                    peers.insert(peer_id, remote_addr);
                }

                // Update stats
                {
                    let mut stats = self.stats.lock().await;
                    stats.successful_connections += 1;
                    stats.active_connections += 1;
                }

                info!(
                    "Accepted connection from peer {:?} at {}",
                    peer_id, remote_addr
                );

                // Handle authentication if required
                if self.config.auth_config.require_authentication {
                    // Start a task to handle incoming authentication
                    let self_clone = self.clone();
                    let auth_peer_id = peer_id;
                    tokio::spawn(async move {
                        if let Err(e) = self_clone.handle_incoming_auth(auth_peer_id).await {
                            error!(
                                "Failed to handle authentication for peer {:?}: {}",
                                auth_peer_id, e
                            );
                            // Remove the peer if auth fails
                            self_clone
                                .connected_peers
                                .write()
                                .await
                                .remove(&auth_peer_id);
                            let mut stats = self_clone.stats.lock().await;
                            stats.active_connections = stats.active_connections.saturating_sub(1);
                        }
                    });
                }

                Ok((remote_addr, peer_id))
            }
            Err(e) => {
                // Update stats
                {
                    let mut stats = self.stats.lock().await;
                    stats.failed_connections += 1;
                }

                error!("Failed to accept connection: {}", e);
                Err(Box::new(e))
            }
        }
    }

    /// Send data to a peer
    pub async fn send_to_peer(
        &self,
        peer_id: &PeerId,
        data: &[u8],
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let peers = self.connected_peers.read().await;

        if let Some(remote_addr) = peers.get(peer_id) {
            debug!(
                "Sending {} bytes to peer {:?} at {}",
                data.len(),
                peer_id,
                remote_addr
            );

            // Get the Quinn connection for this peer from the NAT traversal endpoint
            match self.nat_endpoint.get_connection(peer_id) {
                Ok(Some(connection)) => {
                    // Open a unidirectional stream for data transmission
                    let mut send_stream = connection
                        .open_uni()
                        .await
                        .map_err(|e| format!("Failed to open unidirectional stream: {e}"))?;

                    // Send the data
                    send_stream
                        .write_all(data)
                        .await
                        .map_err(|e| format!("Failed to write data: {e}"))?;

                    // Finish the stream
                    send_stream
                        .finish()
                        .map_err(|e| format!("Failed to finish stream: {e}"))?;

                    debug!(
                        "Successfully sent {} bytes to peer {:?}",
                        data.len(),
                        peer_id
                    );
                    Ok(())
                }
                Ok(None) => {
                    error!("No active connection found for peer {:?}", peer_id);
                    Err("No active connection".into())
                }
                Err(e) => {
                    error!("Failed to get connection for peer {:?}: {}", peer_id, e);
                    Err(Box::new(e))
                }
            }
        } else {
            error!("Peer {:?} not connected", peer_id);
            Err("Peer not connected".into())
        }
    }

    /// Receive data from peers
    pub async fn receive(
        &self,
    ) -> Result<(PeerId, Vec<u8>), Box<dyn std::error::Error + Send + Sync>> {
        debug!("Waiting to receive data from any connected peer...");

        // Get all connected peers
        let peers = {
            let peers_guard = self.connected_peers.read().await;
            peers_guard.clone()
        };

        if peers.is_empty() {
            return Err("No connected peers".into());
        }

        // Try to receive data from any connected peer
        // In a real implementation, this would use a more sophisticated approach
        // like select! over multiple connection streams
        for (peer_id, _remote_addr) in peers.iter() {
            match self.nat_endpoint.get_connection(peer_id) {
                Ok(Some(connection)) => {
                    // Try to accept incoming unidirectional streams
                    match tokio::time::timeout(Duration::from_millis(100), connection.accept_uni())
                        .await
                    {
                        Ok(Ok(mut recv_stream)) => {
                            debug!(
                                "Receiving data from unidirectional stream from peer {:?}",
                                peer_id
                            );

                            // Read all data from the stream
                            match recv_stream.read_to_end(1024 * 1024).await {
                                // 1MB limit
                                Ok(buffer) => {
                                    if !buffer.is_empty() {
                                        debug!(
                                            "Received {} bytes from peer {:?}",
                                            buffer.len(),
                                            peer_id
                                        );
                                        return Ok((*peer_id, buffer));
                                    }
                                }
                                Err(e) => {
                                    debug!(
                                        "Failed to read from stream for peer {:?}: {}",
                                        peer_id, e
                                    );
                                }
                            }
                        }
                        Ok(Err(e)) => {
                            debug!("Failed to accept uni stream from peer {:?}: {}", peer_id, e);
                        }
                        Err(_) => {
                            // Timeout - try bidirectional streams
                        }
                    }

                    // Also try to accept bidirectional streams
                    match tokio::time::timeout(Duration::from_millis(100), connection.accept_bi())
                        .await
                    {
                        Ok(Ok((_send_stream, mut recv_stream))) => {
                            debug!(
                                "Receiving data from bidirectional stream from peer {:?}",
                                peer_id
                            );

                            // Read all data from the receive side
                            match recv_stream.read_to_end(1024 * 1024).await {
                                // 1MB limit
                                Ok(buffer) => {
                                    if !buffer.is_empty() {
                                        debug!(
                                            "Received {} bytes from peer {:?} via bidirectional stream",
                                            buffer.len(),
                                            peer_id
                                        );
                                        return Ok((*peer_id, buffer));
                                    }
                                }
                                Err(e) => {
                                    debug!(
                                        "Failed to read from bidirectional stream for peer {:?}: {}",
                                        peer_id, e
                                    );
                                }
                            }
                        }
                        Ok(Err(e)) => {
                            debug!(
                                "Failed to accept bidirectional stream from peer {:?}: {}",
                                peer_id, e
                            );
                        }
                        Err(_) => {
                            // Timeout - continue to next peer
                        }
                    }
                }
                Ok(None) => {
                    debug!("No active connection for peer {:?}", peer_id);
                }
                Err(e) => {
                    debug!("Failed to get connection for peer {:?}: {}", peer_id, e);
                }
            }
        }

        // If we get here, no data was received from any peer
        Err("No data available from any connected peer".into())
    }

    /// Get current statistics
    pub async fn get_stats(&self) -> NodeStats {
        self.stats.lock().await.clone()
    }

    /// Get access to the NAT traversal endpoint
    pub fn get_nat_endpoint(
        &self,
    ) -> Result<&NatTraversalEndpoint, Box<dyn std::error::Error + Send + Sync>> {
        Ok(&*self.nat_endpoint)
    }

    /// Start periodic statistics reporting
    pub fn start_stats_task(&self) -> tokio::task::JoinHandle<()> {
        let stats = Arc::clone(&self.stats);
        let interval_duration = self.config.stats_interval;

        tokio::spawn(async move {
            let mut interval = tokio::time::interval(interval_duration);

            loop {
                interval.tick().await;

                let stats_snapshot = stats.lock().await.clone();

                info!(
                    "Node statistics - Connections: {}/{}, NAT traversal: {}/{}",
                    stats_snapshot.active_connections,
                    stats_snapshot.successful_connections,
                    stats_snapshot.nat_traversal_successes,
                    stats_snapshot.nat_traversal_attempts
                );
            }
        })
    }

    /// Get NAT traversal statistics
    pub async fn get_nat_stats(
        &self,
    ) -> Result<NatTraversalStatistics, Box<dyn std::error::Error + Send + Sync>> {
        self.nat_endpoint.get_nat_stats()
    }

    /// Inject an observed address (temporary workaround for OBSERVED_ADDRESS frame integration)
    pub fn inject_observed_address(
        &self,
        observed_address: SocketAddr,
        from_peer: PeerId,
    ) -> Result<(), NatTraversalError> {
        self.nat_endpoint
            .inject_observed_address(observed_address, from_peer)
    }

    /// Get connection metrics for a specific peer
    pub async fn get_connection_metrics(
        &self,
        peer_id: &PeerId,
    ) -> Result<ConnectionMetrics, Box<dyn std::error::Error + Send + Sync>> {
        match self.nat_endpoint.get_connection(peer_id) {
            Ok(Some(connection)) => {
                // Get basic RTT from the connection
                let rtt = connection.rtt();

                // Get congestion window and other stats
                let stats = connection.stats();

                Ok(ConnectionMetrics {
                    bytes_sent: stats.udp_tx.bytes,
                    bytes_received: stats.udp_rx.bytes,
                    rtt: Some(rtt),
                    packet_loss: stats.path.lost_packets as f64
                        / (stats.path.sent_packets + stats.path.lost_packets).max(1) as f64,
                })
            }
            Ok(None) => Err("Connection not found".into()),
            Err(e) => Err(format!("Failed to get connection: {e}").into()),
        }
    }

    /// Get this node's peer ID
    pub fn peer_id(&self) -> PeerId {
        self.peer_id
    }

    /// Get this node's public key bytes
    pub fn public_key_bytes(&self) -> [u8; 32] {
        self.auth_manager.public_key_bytes()
    }

    /// Send an authentication message to a peer
    async fn send_auth_message(
        &self,
        peer_id: &PeerId,
        message: AuthMessage,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let data = AuthManager::serialize_message(&message)?;
        self.send_to_peer(peer_id, &data).await
    }

    /// Perform authentication handshake as initiator
    async fn authenticate_as_initiator(
        &self,
        peer_id: &PeerId,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        info!("Starting authentication with peer {:?}", peer_id);

        // Send authentication request
        let auth_request = self.auth_manager.create_auth_request();
        self.send_auth_message(peer_id, auth_request).await?;

        // Wait for challenge
        let timeout_duration = self.config.auth_config.auth_timeout;
        let start = Instant::now();

        while start.elapsed() < timeout_duration {
            match tokio::time::timeout(Duration::from_secs(1), self.receive()).await {
                Ok(Ok((recv_peer_id, data))) => {
                    if recv_peer_id == *peer_id {
                        match AuthManager::deserialize_message(&data) {
                            Ok(AuthMessage::Challenge { nonce, .. }) => {
                                // Create and send challenge response
                                let response =
                                    self.auth_manager.create_challenge_response(nonce)?;
                                self.send_auth_message(peer_id, response).await?;
                            }
                            Ok(AuthMessage::AuthSuccess { .. }) => {
                                info!("Authentication successful with peer {:?}", peer_id);
                                return Ok(());
                            }
                            Ok(AuthMessage::AuthFailure { reason }) => {
                                return Err(format!("Authentication failed: {reason}").into());
                            }
                            _ => continue,
                        }
                    }
                }
                _ => continue,
            }
        }

        Err("Authentication timeout".into())
    }

    /// Handle incoming authentication messages
    async fn handle_auth_message(
        &self,
        peer_id: PeerId,
        message: AuthMessage,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let auth_protocol = AuthProtocol::new(Arc::clone(&self.auth_manager));

        match auth_protocol.handle_message(peer_id, message).await {
            Ok(Some(response)) => {
                self.send_auth_message(&peer_id, response).await?;
            }
            Ok(None) => {
                // No response needed
            }
            Err(e) => {
                error!("Authentication error: {}", e);
                let failure = AuthMessage::AuthFailure {
                    reason: e.to_string(),
                };
                self.send_auth_message(&peer_id, failure).await?;
                return Err(Box::new(e));
            }
        }

        Ok(())
    }

    /// Check if a peer is authenticated
    pub async fn is_peer_authenticated(&self, peer_id: &PeerId) -> bool {
        self.auth_manager.is_authenticated(peer_id).await
    }

    /// Get list of authenticated peers
    pub async fn list_authenticated_peers(&self) -> Vec<PeerId> {
        self.auth_manager.list_authenticated_peers().await
    }

    /// Handle incoming authentication from a peer
    async fn handle_incoming_auth(
        &self,
        peer_id: PeerId,
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        info!("Handling incoming authentication from peer {:?}", peer_id);

        let timeout_duration = self.config.auth_config.auth_timeout;
        let start = Instant::now();

        while start.elapsed() < timeout_duration {
            match tokio::time::timeout(Duration::from_secs(1), self.receive()).await {
                Ok(Ok((recv_peer_id, data))) => {
                    if recv_peer_id == peer_id {
                        match AuthManager::deserialize_message(&data) {
                            Ok(auth_msg) => {
                                self.handle_auth_message(peer_id, auth_msg).await?;

                                // Check if authentication is complete
                                if self.auth_manager.is_authenticated(&peer_id).await {
                                    info!("Peer {:?} successfully authenticated", peer_id);
                                    return Ok(());
                                }
                            }
                            Err(_) => {
                                // Not an auth message, ignore
                                continue;
                            }
                        }
                    }
                }
                _ => continue,
            }
        }

        Err("Authentication timeout waiting for peer".into())
    }
}