first
This commit is contained in:
Gilles Soulier
241
agent/src/p2p/endpoint.rs
Normal file
241
agent/src/p2p/endpoint.rs
Normal file
@@ -0,0 +1,241 @@
|
||||
// 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())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user