mesh/agent/src/p2p/endpoint.rs

// Created by: Claude
// Date: 2026-01-04
// Purpose: QUIC endpoint management with P2P handshake
// Refs: AGENT.md, signaling_v_2.md

use anyhow::Result;
use quinn::{Endpoint, Connection, RecvStream, SendStream};
use std::sync::Arc;
use std::net::SocketAddr;
use tokio::sync::Mutex;
use std::collections::HashMap;
use tracing::{info, warn, error};

use super::{tls, protocol::*};

pub struct QuicEndpoint {
    endpoint: Endpoint,
    local_port: u16,
    active_sessions: Arc<Mutex<HashMap<String, ActiveSession>>>,
    // Cache local pour validation des session_tokens
    valid_tokens: Arc<Mutex<HashMap<String, SessionTokenCache>>>,
}

struct ActiveSession {
    pub session_id: String,
    pub connection: Connection,
}

struct SessionTokenCache {
    session_id: String,
    session_token: String,
    expires_at: std::time::SystemTime,
}

impl QuicEndpoint {
    pub async fn new(port: u16) -> Result<Self> {
        let rustls_server_config = tls::make_server_config()?;
        let server_config = quinn::ServerConfig::with_crypto(Arc::new(rustls_server_config));
        let addr: SocketAddr = format!("0.0.0.0:{}", port).parse()?;

        let endpoint = Endpoint::server(server_config, addr)?;
        let local_port = endpoint.local_addr()?.port();

        info!("QUIC endpoint listening on port {}", local_port);

        Ok(Self {
            endpoint,
            local_port,
            active_sessions: Arc::new(Mutex::new(HashMap::new())),
            valid_tokens: Arc::new(Mutex::new(HashMap::new())),
        })
    }

    pub fn local_port(&self) -> u16 {
        self.local_port
    }

    /// Ajouter un token au cache (appelé par P2PHandler lors de p2p.session.created)
    pub async fn add_valid_token(&self, session_id: String, session_token: String, ttl_secs: u64) {
        let expires_at = std::time::SystemTime::now() + std::time::Duration::from_secs(ttl_secs);
        let cache_entry = SessionTokenCache {
            session_id: session_id.clone(),
            session_token,
            expires_at,
        };
        self.valid_tokens.lock().await.insert(session_id.clone(), cache_entry);
        info!("Token cached for session: {} (TTL: {}s)", session_id, ttl_secs);
    }

    /// Valider un token depuis le cache local
    async fn validate_token(&self, session_id: &str, session_token: &str) -> Result<()> {
        let tokens = self.valid_tokens.lock().await;

        if let Some(cached) = tokens.get(session_id) {
            if cached.session_token != session_token {
                anyhow::bail!("Token mismatch");
            }

            if std::time::SystemTime::now() > cached.expires_at {
                anyhow::bail!("Token expired");
            }

            Ok(())
        } else {
            anyhow::bail!("Session not found in cache")
        }
    }

    /// Accept loop (spawn dans main)
    pub async fn accept_loop(self: Arc<Self>) -> Result<()> {
        info!("Starting QUIC accept loop");

        loop {
            let incoming = match self.endpoint.accept().await {
                Some(incoming) => incoming,
                None => {
                    info!("QUIC endpoint closed");
                    return Ok(());
                }
            };

            let endpoint_clone = Arc::clone(&self);
            tokio::spawn(async move {
                if let Err(e) = endpoint_clone.handle_incoming(incoming).await {
                    error!("Failed to handle incoming connection: {}", e);
                }
            });
        }
    }

    pub fn close(&self) {
        self.endpoint.close(0u32.into(), b"shutdown");
    }

    async fn handle_incoming(&self, incoming: quinn::Connecting) -> Result<()> {
        let connection = incoming.await?;
        info!("Incoming QUIC connection from {}", connection.remote_address());

        // Wait for P2P_HELLO
        let (send, recv) = connection.accept_bi().await?;
        let hello = self.receive_hello(recv).await?;

        info!("P2P_HELLO received: session_id={}", hello.session_id);

        // Valider session_token via cache local
        if let Err(e) = self.validate_token(&hello.session_id, &hello.session_token).await {
            warn!("Token validation failed: {}", e);
            self.send_response(send, &P2PResponse::Deny {
                reason: format!("Invalid token: {}", e),
            }).await?;
            return Ok(());
        }

        // Send P2P_OK
        self.send_response(send, &P2PResponse::Ok).await?;
        info!("P2P handshake successful for session: {}", hello.session_id);

        // Store session
        let session = ActiveSession {
            session_id: hello.session_id.clone(),
            connection,
        };
        self.active_sessions.lock().await.insert(hello.session_id, session);

        Ok(())
    }

    /// Connect to remote peer
    pub async fn connect_to_peer(
        &self,
        remote_addr: SocketAddr,
        session_id: String,
        session_token: String,
        device_id: String,
    ) -> Result<Connection> {
        let rustls_client_config = tls::make_client_config()?;
        let mut client_config = quinn::ClientConfig::new(Arc::new(rustls_client_config));

        // Configurer transport parameters si nécessaire
        let mut transport_config = quinn::TransportConfig::default();
        transport_config.max_idle_timeout(Some(std::time::Duration::from_secs(30).try_into()?));
        client_config.transport_config(Arc::new(transport_config));

        info!("Connecting to peer at {}", remote_addr);

        let connection = self.endpoint.connect_with(
            client_config,
            remote_addr,
            "mesh-peer",
        )?.await?;

        info!("QUIC connection established to {}", remote_addr);

        // Send P2P_HELLO
        let (mut send, recv) = connection.open_bi().await?;
        self.send_hello(&mut send, session_id.clone(), session_token, device_id).await?;

        // Wait for P2P_OK
        let response = self.receive_response(recv).await?;
        match response {
            P2PResponse::Ok => {
                info!("P2P handshake successful for session: {}", session_id);
                Ok(connection)
            }
            P2PResponse::Deny { reason } => {
                anyhow::bail!("P2P handshake denied: {}", reason)
            }
        }
    }

    async fn send_hello(
        &self,
        stream: &mut SendStream,
        session_id: String,
        session_token: String,
        device_id: String,
    ) -> Result<()> {
        let hello = P2PHello {
            t: "P2P_HELLO".to_string(),
            session_id,
            session_token,
            from_device_id: device_id,
        };

        let json = serde_json::to_vec(&hello)?;
        stream.write_all(&json).await?;
        stream.finish().await?;

        info!("Sent P2P_HELLO");

        Ok(())
    }

    async fn receive_hello(&self, mut stream: RecvStream) -> Result<P2PHello> {
        let data = stream.read_to_end(4096).await?;
        let hello: P2PHello = serde_json::from_slice(&data)?;

        if hello.t != "P2P_HELLO" {
            anyhow::bail!("Expected P2P_HELLO, got {}", hello.t);
        }

        Ok(hello)
    }

    async fn send_response(&self, mut stream: SendStream, response: &P2PResponse) -> Result<()> {
        let json = serde_json::to_vec(response)?;
        stream.write_all(&json).await?;
        stream.finish().await?;
        Ok(())
    }

    async fn receive_response(&self, mut stream: RecvStream) -> Result<P2PResponse> {
        let data = stream.read_to_end(4096).await?;
        Ok(serde_json::from_slice(&data)?)
    }

    /// Get active session by ID
    pub async fn get_session(&self, session_id: &str) -> Option<Connection> {
        self.active_sessions.lock().await.get(session_id).map(|s| s.connection.clone())
    }
}