diff --git a/Cargo.lock b/Cargo.lock index f52ae973..c5bb140c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4269,7 +4269,9 @@ dependencies = [ "pqcrypto-traits", "pretty_assertions", "pty-process", + "quinn", "rand 0.8.5", + "rcgen", "readabilityrs", "refinery", "regex", @@ -4456,6 +4458,16 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099" +[[package]] +name = "pem" +version = "3.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be" +dependencies = [ + "base64 0.22.1", + "serde_core", +] + [[package]] name = "percent-encoding" version = "2.3.2" @@ -5151,6 +5163,19 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "rcgen" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75e669e5202259b5314d1ea5397316ad400819437857b90861765f24c4cf80a2" +dependencies = [ + "pem", + "ring", + "rustls-pki-types", + "time", + "yasna", +] + [[package]] name = "readabilityrs" version = "0.1.2" @@ -8636,6 +8661,15 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cfe53a6657fd280eaa890a3bc59152892ffa3e30101319d168b781ed6529b049" +[[package]] +name = "yasna" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17bb3549cc1321ae1296b9cdc2698e2b6cb1992adfa19a8c72e5b7a738f44cd" +dependencies = [ + "time", +] + [[package]] name = "yoke" version = "0.8.1" diff --git a/Cargo.toml b/Cargo.toml index 46e2250e..6c0dfa57 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -191,7 +191,8 @@ pqcrypto-kyber = { version = "0.8", optional = true } pqcrypto-traits = { version = "0.3", optional = true } sys-info = { version = "0.9", optional = true } hostname = { version = "0.4", optional = true } -tokio-tungstenite = { version = "0.26", optional = true } +quinn = { version = "0.11", default-features = false, features = ["runtime-tokio", "rustls-ring"], optional = true } +rcgen = { version = "0.13", optional = true } # macOS keychain [target.'cfg(target_os = "macos")'.dependencies] @@ -244,7 +245,7 @@ integration = [] html-to-markdown = ["dep:html-to-markdown-rs", "dep:readabilityrs"] bedrock = ["dep:aws-config", "dep:aws-sdk-bedrockruntime", "dep:aws-smithy-types"] import = ["dep:json5", "libsql"] -cluster = ["dep:pqcrypto-kyber", "dep:pqcrypto-traits", "dep:sys-info", "dep:hostname", "dep:tokio-tungstenite"] +cluster = ["dep:pqcrypto-kyber", "dep:pqcrypto-traits", "dep:sys-info", "dep:hostname", "dep:quinn", "dep:rcgen", "dep:rustls"] [[test]] name = "e2e_thread_scheduling" diff --git a/providers.json b/providers.json index 517e2a26..1b70aa62 100644 --- a/providers.json +++ b/providers.json @@ -55,6 +55,24 @@ "can_list_models": true } }, + { + "id": "vllm", + "aliases": ["vllm-server", "vllm_server"], + "protocol": "open_ai_completions", + "default_base_url": "http://localhost:8000/v1", + "base_url_env": "VLLM_BASE_URL", + "api_key_env": "VLLM_API_KEY", + "api_key_required": false, + "model_env": "VLLM_MODEL", + "default_model": "default", + "description": "vLLM inference server — local or remote, OpenAI-compatible. Set VLLM_MODEL to the model name loaded in the server (e.g. meta-llama/Llama-3-8B-Instruct). Increase LLM_REQUEST_TIMEOUT_SECS for large models.", + "setup": { + "kind": "open_ai_compatible", + "secret_name": "llm_vllm_api_key", + "display_name": "vLLM", + "can_list_models": true + } + }, { "id": "openai_compatible", "aliases": [ @@ -69,7 +87,7 @@ "model_env": "LLM_MODEL", "default_model": "default", "extra_headers_env": "LLM_EXTRA_HEADERS", - "description": "Custom OpenAI-compatible endpoint (vLLM, LiteLLM, etc.)", + "description": "Custom OpenAI-compatible endpoint (LiteLLM, LM Studio, etc.)", "setup": { "kind": "open_ai_compatible", "secret_name": "llm_compatible_api_key", diff --git a/src/cluster/config.rs b/src/cluster/config.rs index 8aeba8c5..ecd68640 100644 --- a/src/cluster/config.rs +++ b/src/cluster/config.rs @@ -32,6 +32,20 @@ pub struct ClusterConfig { pub broadcast_addr: Ipv4Addr, /// Path to persist mesh keypair. pub keys_path: String, + /// Static remote peers to connect to (host:port, may include hostnames). + /// Resolved and retried periodically. LAN discovery via UDP broadcast is unaffected. + /// Set via CLUSTER_STATIC_PEERS (comma-separated). + pub static_peers: Vec, + /// Externally-reachable address advertised to static/remote peers (host:port). + /// Needed when this node is behind NAT with a port forward or public IP. + /// Set via CLUSTER_ADVERTISE_ADDR. + pub advertise_addr: Option, + /// Enable VLESS proxy inbound on the mesh QUIC port (ALPN "oproxy/1"). + /// Set via CLUSTER_PROXY_ENABLED=1. + pub proxy_enabled: bool, + /// UUIDs allowed to authenticate as VLESS proxy clients. + /// Set via CLUSTER_PROXY_UUIDS (comma-separated standard UUID strings). + pub proxy_uuids: Vec, } impl ClusterConfig { @@ -66,6 +80,24 @@ impl ClusterConfig { .unwrap_or(Ipv4Addr::BROADCAST), keys_path: std::env::var("CLUSTER_KEYS_PATH") .unwrap_or_else(|_| format!("{}/.optimclaw/mesh_keys.json", home)), + static_peers: std::env::var("CLUSTER_STATIC_PEERS") + .unwrap_or_default() + .split(',') + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(String::from) + .collect(), + advertise_addr: std::env::var("CLUSTER_ADVERTISE_ADDR").ok(), + proxy_enabled: std::env::var("CLUSTER_PROXY_ENABLED") + .map(|v| v == "1" || v == "true") + .unwrap_or(false), + proxy_uuids: std::env::var("CLUSTER_PROXY_UUIDS") + .unwrap_or_default() + .split(',') + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(String::from) + .collect(), } } } diff --git a/src/cluster/mod.rs b/src/cluster/mod.rs index 8627c37e..c05f5a5d 100644 --- a/src/cluster/mod.rs +++ b/src/cluster/mod.rs @@ -1,8 +1,9 @@ //! Autonomous AI mesh network for OptimClaw. //! //! Nodes self-discover via UDP beacons, form a peer-to-peer overlay mesh -//! with post-quantum encrypted WebSocket channels, and route tasks -//! intelligently based on GPU capability, model availability, and load. +//! with post-quantum encrypted QUIC channels (UDP, NAT hole-punch capable), +//! and route tasks intelligently based on GPU capability, model availability, +//! and load. A VLESS proxy inbound is optionally available on the same port. pub mod api; pub mod beacon; @@ -10,7 +11,9 @@ pub mod config; pub mod crypto; pub mod gossip; pub mod overlay; +pub mod proxy; pub mod router; +pub mod transport; pub mod types; use std::collections::HashMap; @@ -55,13 +58,13 @@ impl MeshNode { // Message channel for incoming mesh messages let (incoming_tx, mut incoming_rx) = mpsc::channel::<(NodeId, MeshMessage)>(256); - // Initialize overlay mesh + // Initialize overlay mesh (binds the QUIC endpoint) let overlay = Arc::new(OverlayMesh::new( identity.clone(), gossip.clone(), config.clone(), incoming_tx, - )); + )?); let (shutdown_tx, _) = tokio::sync::broadcast::channel(1); let (event_tx, _) = tokio::sync::broadcast::channel(128); @@ -82,10 +85,25 @@ impl MeshNode { let mut beacon_rx = beacon::start_beacon(&config, identity.clone()).await?; tracing::info!("Beacon broadcasting on port {}", config.beacon_port); - // Start the overlay WebSocket listener - overlay.start_listener().await?; + // Parse proxy UUIDs (only matters when CLUSTER_PROXY_ENABLED=1) + let proxy_uuids: Arc> = Arc::new( + config + .proxy_uuids + .iter() + .filter_map(|s| match proxy::parse_uuid(s) { + Ok(u) => Some(u), + Err(e) => { + tracing::warn!("Ignoring invalid proxy UUID '{}': {}", s, e); + None + } + }) + .collect(), + ); - // Task: process discovered peers -- connect via PQ-encrypted WebSocket + // Start the QUIC overlay listener (also dispatches VLESS proxy if enabled) + overlay.start_listener(proxy_uuids).await?; + + // Task: process discovered peers — connect via PQ-encrypted QUIC (UDP hole-punch) let overlay_disc = overlay.clone(); let gossip_disc = gossip.clone(); tokio::spawn(async move { @@ -148,6 +166,55 @@ impl MeshNode { } }); + // Task: static peer connector — only active when CLUSTER_STATIC_PEERS is set. + // LAN peers continue to be discovered via UDP broadcast above; this task handles + // nodes on remote/NAT'd networks that are reachable by hostname or public IP. + if !config.static_peers.is_empty() { + let overlay_sp = overlay.clone(); + let static_peers = config.static_peers.clone(); + // Retry every 3× the beacon interval so we don't spam unreachable hosts. + let retry_interval = config.beacon_interval * 3; + let mut shutdown_rx_sp = shutdown_tx.subscribe(); + tokio::spawn(async move { + loop { + for peer_str in &static_peers { + match tokio::net::lookup_host(peer_str.as_str()).await { + Ok(mut addrs) => { + if let Some(addr) = addrs.next() { + if overlay_sp.is_connected_by_addr(&addr).await { + continue; + } + let overlay_conn = overlay_sp.clone(); + let label = peer_str.clone(); + tokio::spawn(async move { + if let Err(e) = overlay_conn.connect_to_addr(addr).await { + tracing::debug!( + "Static peer {} unreachable: {}", + label, + e + ); + } + }); + } + } + Err(e) => { + tracing::debug!( + "Failed to resolve static peer {}: {}", + peer_str, + e + ); + } + } + } + + tokio::select! { + _ = tokio::time::sleep(retry_interval) => {} + _ = shutdown_rx_sp.recv() => break, + } + } + }); + } + // Task: periodic gossip / failure detection let gossip_tick = gossip.clone(); let overlay_tick = overlay.clone(); @@ -424,7 +491,6 @@ async fn handle_mesh_message( // Execute the task via optimclaw single-message mode let overlay_exec = overlay.identity.node_id; let from_id = from; - let overlay_ref = overlay.incoming_tx.clone(); let task_id_exec = task_id.clone(); let content = envelope.content.clone(); diff --git a/src/cluster/overlay.rs b/src/cluster/overlay.rs index f56275d1..e9019093 100644 --- a/src/cluster/overlay.rs +++ b/src/cluster/overlay.rs @@ -1,30 +1,35 @@ -//! WebSocket overlay mesh with post-quantum encrypted channels. +//! QUIC overlay mesh with post-quantum encrypted channels. //! -//! Each peer connection is a WebSocket tunnel encrypted with AES-256-GCM -//! after an ML-KEM-768 key exchange handshake. +//! Transport: QUIC over UDP (NAT hole-punching capable, like Nebula). +//! Security: ML-KEM-768 key encapsulation + AES-256-GCM per-session, over the +//! first QUIC bidirectional stream. The QUIC TLS layer uses an ephemeral +//! self-signed cert for transport only — all authentication is PQ. +//! +//! Dedup rule: lower NodeId is always the connector (initiator). The higher-ID +//! side is the responder. Static peers (unknown remote ID) skip this check and +//! let the post-handshake is_connected guard prevent duplicate channels. use std::collections::HashMap; use std::net::SocketAddr; use std::sync::Arc; -use tokio::net::{TcpListener, TcpStream}; -use tokio::sync::{mpsc, RwLock}; -use tokio_tungstenite::tungstenite::Message as WsMessage; -use tokio_tungstenite::{accept_async, connect_async}; -use futures::{SinkExt, StreamExt}; + use aes_gcm::Aes256Gcm; use serde::{Deserialize, Serialize}; +use tokio::sync::{mpsc, RwLock}; use super::config::ClusterConfig; use super::crypto::{self, MeshIdentity}; use super::gossip::GossipState; +use super::transport::{self, ALPN_MESH, ALPN_PROXY}; use super::types::*; -/// A connected peer with its encrypted channel. +// ── Peer connection ─────────────────────────────────────────────────────────── + +/// A connected peer with its outbound message channel. pub struct PeerConnection { pub node_id: NodeId, pub addr: SocketAddr, pub send_tx: mpsc::Sender>, - pub send_nonce: std::sync::atomic::AtomicU64, } impl std::fmt::Debug for PeerConnection { @@ -36,7 +41,9 @@ impl std::fmt::Debug for PeerConnection { } } -/// Handshake message sent during PQ key exchange. +// ── Handshake messages ──────────────────────────────────────────────────────── + +/// Sent by the initiator: our KEM public key + signing key + node ID. #[derive(Serialize, Deserialize)] struct HandshakeInit { node_id: NodeId, @@ -44,6 +51,7 @@ struct HandshakeInit { verify_pk: Vec, } +/// Sent by the responder: KEM ciphertext + signing key + Ed25519 signature. #[derive(Serialize, Deserialize)] struct HandshakeReply { node_id: NodeId, @@ -52,13 +60,19 @@ struct HandshakeReply { signature: Vec, } -/// Manages the overlay mesh connections. +// ── OverlayMesh ─────────────────────────────────────────────────────────────── + +/// Manages the QUIC overlay mesh connections. pub struct OverlayMesh { pub identity: Arc, pub peers: RwLock>, pub gossip: Arc, pub config: ClusterConfig, pub incoming_tx: mpsc::Sender<(NodeId, MeshMessage)>, + /// Shared QUIC endpoint — used for both listening and outbound dials. + endpoint: quinn::Endpoint, + /// Client config used when dialing peers (skips cert verification). + client_cfg: quinn::ClientConfig, } impl OverlayMesh { @@ -67,69 +81,99 @@ impl OverlayMesh { gossip: Arc, config: ClusterConfig, incoming_tx: mpsc::Sender<(NodeId, MeshMessage)>, - ) -> Self { - Self { + ) -> anyhow::Result { + let bind_addr = SocketAddr::new(config.bind_addr, config.mesh_port); + let endpoint = transport::make_server_endpoint(bind_addr)?; + let client_cfg = transport::make_client_config()?; + + Ok(Self { identity, peers: RwLock::new(HashMap::new()), gossip, config, incoming_tx, - } + endpoint, + client_cfg, + }) } - /// Start the WebSocket listener for incoming peer connections. - pub async fn start_listener(self: &Arc) -> anyhow::Result<()> { - let listen_addr = format!("{}:{}", self.config.bind_addr, self.config.mesh_port); - let listener = TcpListener::bind(&listen_addr).await?; - tracing::info!("Mesh overlay listening on {}", listen_addr); + // ── Listener ────────────────────────────────────────────────────────────── + + /// Accept incoming QUIC connections and dispatch by ALPN. + /// + /// `"omesh/1"` → PQ mesh handshake. + /// `"oproxy/1"` → VLESS proxy (handled by proxy::handle_proxy_connection). + pub async fn start_listener( + self: &Arc, + proxy_uuids: Arc>, + ) -> anyhow::Result<()> { + let listen_addr = SocketAddr::new(self.config.bind_addr, self.config.mesh_port); + tracing::info!("Mesh overlay listening on UDP {}", listen_addr); let mesh = self.clone(); tokio::spawn(async move { - loop { - match listener.accept().await { - Ok((stream, addr)) => { - let mesh = mesh.clone(); - tokio::spawn(async move { - if let Err(e) = mesh.handle_incoming(stream, addr).await { - tracing::debug!("Incoming peer {} handshake failed: {}", addr, e); + while let Some(incoming) = mesh.endpoint.accept().await { + let mesh = mesh.clone(); + let proxy_uuids = proxy_uuids.clone(); + tokio::spawn(async move { + let conn = match incoming.await { + Ok(c) => c, + Err(e) => { + tracing::debug!("QUIC accept error: {}", e); + return; + } + }; + + let alpn = conn + .handshake_data() + .and_then(|d| { + d.downcast::() + .ok() + }) + .and_then(|d| d.protocol.clone()); + + match alpn.as_deref() { + Some(p) if p == ALPN_MESH => { + if let Err(e) = mesh.handle_incoming(conn).await { + tracing::debug!("Mesh handshake failed: {}", e); } - }); + } + Some(p) if p == ALPN_PROXY => { + if let Err(e) = + super::proxy::handle_proxy_connection(conn, proxy_uuids).await + { + tracing::debug!("Proxy connection failed: {}", e); + } + } + _ => tracing::debug!("Incoming QUIC conn with unknown ALPN, dropping"), } - Err(e) => { - tracing::warn!("Mesh accept error: {}", e); - tokio::time::sleep(std::time::Duration::from_secs(1)).await; - } - } + }); } }); Ok(()) } - /// Handle an incoming WebSocket connection (responder side of PQ handshake). - async fn handle_incoming(&self, stream: TcpStream, addr: SocketAddr) -> anyhow::Result<()> { - let ws = accept_async(stream).await?; - let (mut ws_tx, mut ws_rx) = ws.split(); + // ── Incoming PQ handshake (responder) ───────────────────────────────────── - // Step 1: Receive initiator's handshake (kem_pk + verify_pk + node_id) - let init_msg = ws_rx - .next() - .await - .ok_or_else(|| anyhow::anyhow!("No handshake received"))??; + async fn handle_incoming(self: &Arc, conn: quinn::Connection) -> anyhow::Result<()> { + let peer_addr = conn.remote_address(); + let (mut send, mut recv) = conn.accept_bi().await?; - let init: HandshakeInit = serde_json::from_slice(&init_msg.into_data())?; + // Receive initiator's HandshakeInit + let init_bytes = transport::recv_frame(&mut recv).await?; + let init: HandshakeInit = serde_json::from_slice(&init_bytes)?; - // Dedup: if lower ID should be connector, reject + // Dedup: if the initiator has a higher ID than us, they should be the + // connector, not us — drop this and let the ordering play out. if init.node_id > self.identity.node_id { - // We have the lower ID -- we should be the connector, not the responder - // Drop this connection; we'll connect outbound instead return Ok(()); } - // Step 2: Encapsulate shared secret with initiator's KEM public key + // Encapsulate shared secret with initiator's KEM public key let (shared_secret, ciphertext) = crypto::encapsulate(&init.kem_pk)?; - // Step 3: Sign the ciphertext + our node_id + // Sign ciphertext || our node_id so the initiator can verify us let mut sign_data = Vec::new(); sign_data.extend_from_slice(&ciphertext); sign_data.extend_from_slice(&self.identity.node_id); @@ -141,176 +185,217 @@ impl OverlayMesh { verify_pk: self.identity.verify_key_bytes(), signature, }; + transport::send_frame(&mut send, &serde_json::to_vec(&reply)?).await?; - ws_tx - .send(WsMessage::Binary(serde_json::to_vec(&reply)?.into())) - .await?; - - // Step 4: Derive symmetric keys + // Derive symmetric keys and start the encrypted peer loop let (send_cipher, recv_cipher) = - crypto::derive_keys(&shared_secret, &self.identity.node_id, &init.node_id); + crypto::derive_keys(&shared_secret, &init.node_id, &self.identity.node_id); - // Connection established -- run the encrypted message loop - self.run_peer_loop(init.node_id, addr, ws_tx, ws_rx, send_cipher, recv_cipher) + self.run_peer_loop(init.node_id, peer_addr, send, recv, send_cipher, recv_cipher) .await; Ok(()) } - /// Initiate an outbound connection to a discovered peer (initiator side). - pub async fn connect_to_peer(self: &Arc, peer_addr: SocketAddr, peer_node_id: &NodeId) -> anyhow::Result<()> { - // Dedup: lower NodeId is always the connector + // ── Outbound dial — known peer ID (LAN beacon discovery) ───────────────── + + /// Connect to a peer discovered via UDP beacon (node ID known). + /// + /// Lower-ID rule: only the node with the lower ID initiates. Higher-ID + /// nodes wait for the remote to connect to them. + pub async fn connect_to_peer( + self: &Arc, + peer_addr: SocketAddr, + peer_node_id: &NodeId, + ) -> anyhow::Result<()> { if self.identity.node_id >= *peer_node_id { - // We have the higher ID -- wait for the other side to connect to us return Ok(()); } - if self.is_connected(peer_node_id).await { return Ok(()); } + self.dial(peer_addr, Some(*peer_node_id)).await + } - let url = format!("ws://{}", peer_addr); - let (ws, _) = connect_async(&url).await?; - let (mut ws_tx, mut ws_rx) = ws.split(); + // ── Outbound dial — unknown peer ID (static peers / hole-punch) ────────── - // Step 1: Send our handshake (kem_pk + verify_pk + node_id) + /// Connect to a static peer whose NodeId is not known until the PQ handshake. + /// + /// Both sides of a static peer pair attempt this simultaneously. Sending the + /// QUIC Initial packet from each side opens a NAT pinhole (UDP hole punching), + /// so at least one direction succeeds. The post-handshake `is_connected` guard + /// prevents duplicate channels if both sides finish the handshake concurrently. + pub async fn connect_to_addr(self: &Arc, peer_addr: SocketAddr) -> anyhow::Result<()> { + if self.is_connected_by_addr(&peer_addr).await { + return Ok(()); + } + self.dial(peer_addr, None).await + } + + // ── Core dial logic ─────────────────────────────────────────────────────── + + async fn dial( + self: &Arc, + peer_addr: SocketAddr, + known_peer_id: Option, + ) -> anyhow::Result<()> { + let conn = self + .endpoint + .connect_with(self.client_cfg.clone(), peer_addr, "omesh")? + .await?; + + let (mut send, mut recv) = conn.open_bi().await?; + + // Send our HandshakeInit let init = HandshakeInit { node_id: self.identity.node_id, kem_pk: self.identity.kem_pk_bytes(), verify_pk: self.identity.verify_key_bytes(), }; + transport::send_frame(&mut send, &serde_json::to_vec(&init)?).await?; - ws_tx - .send(WsMessage::Binary(serde_json::to_vec(&init)?.into())) - .await?; + // Receive HandshakeReply + let reply_bytes = transport::recv_frame(&mut recv).await?; + let reply: HandshakeReply = serde_json::from_slice(&reply_bytes)?; - // Step 2: Receive responder's reply (ciphertext + verify_pk + signature) - let reply_msg = ws_rx - .next() - .await - .ok_or_else(|| anyhow::anyhow!("No handshake reply"))??; - - let reply: HandshakeReply = serde_json::from_slice(&reply_msg.into_data())?; - - // Step 3: Verify signature + // Verify Ed25519 signature over (ciphertext || responder node_id) let peer_verify_key = ed25519_dalek::VerifyingKey::from_bytes( - reply.verify_pk.as_slice().try_into() + reply + .verify_pk + .as_slice() + .try_into() .map_err(|_| anyhow::anyhow!("Invalid verify key length"))?, )?; - let mut sign_data = Vec::new(); sign_data.extend_from_slice(&reply.ciphertext); sign_data.extend_from_slice(&reply.node_id); - if !MeshIdentity::verify(&peer_verify_key, &sign_data, &reply.signature) { - anyhow::bail!("Invalid handshake signature from peer"); + anyhow::bail!("Invalid handshake signature from {}", peer_addr); } - // Step 4: Decapsulate shared secret - let shared_secret = self.identity.decapsulate(&reply.ciphertext)?; + // If known peer ID was supplied, verify it matches + if let Some(expected) = known_peer_id { + if reply.node_id != expected { + anyhow::bail!( + "Peer ID mismatch: expected {} got {}", + node_id_hex(&expected), + node_id_hex(&reply.node_id) + ); + } + } - // Step 5: Derive symmetric keys (we are the initiator) + // Guard against races (both sides completed simultaneously) + if self.is_connected(&reply.node_id).await { + return Ok(()); + } + + let shared_secret = self.identity.decapsulate(&reply.ciphertext)?; let (send_cipher, recv_cipher) = crypto::derive_keys(&shared_secret, &self.identity.node_id, &reply.node_id); tracing::info!( - "PQ handshake complete with {} (ML-KEM-768 + AES-256-GCM)", - node_id_hex(&reply.node_id) + "PQ handshake complete with {} at {} (ML-KEM-768 + AES-256-GCM)", + node_id_hex(&reply.node_id), + peer_addr, ); - // Run encrypted message loop let mesh = self.clone(); let peer_id = reply.node_id; tokio::spawn(async move { - mesh.run_peer_loop(peer_id, peer_addr, ws_tx, ws_rx, send_cipher, recv_cipher) + mesh.run_peer_loop(peer_id, peer_addr, send, recv, send_cipher, recv_cipher) .await; }); Ok(()) } - /// Run the encrypted message loop for a connected peer. - async fn run_peer_loop( + // ── Encrypted peer loop ─────────────────────────────────────────────────── + + /// Run the AES-256-GCM encrypted send/receive loop for a connected peer. + /// + /// Frames are length-prefixed over the QUIC bidirectional stream. The QUIC + /// layer provides reliability and ordering; we provide confidentiality and + /// integrity via our PQ-derived symmetric keys. + async fn run_peer_loop( &self, peer_id: NodeId, addr: SocketAddr, - mut ws_tx: S, - mut ws_rx: R, + mut quic_tx: quinn::SendStream, + mut quic_rx: quinn::RecvStream, send_cipher: Aes256Gcm, recv_cipher: Aes256Gcm, - ) where - S: futures::Sink + Unpin + Send + 'static, - R: futures::Stream> + Unpin + Send + 'static, - { - // Channel for outgoing messages (plaintext bytes, encrypted before sending) + ) { let (out_tx, mut out_rx) = mpsc::channel::>(64); - - // Register peer self.add_peer(peer_id, addr, out_tx).await; let incoming_tx = self.incoming_tx.clone(); let send_cipher = Arc::new(send_cipher); let recv_cipher = Arc::new(recv_cipher); - // Sender task: encrypt and send + // Sender: encrypt plaintext → length-prefixed frame → QUIC stream let send_c = send_cipher.clone(); let sender = tokio::spawn(async move { let mut nonce: u64 = 0; while let Some(plaintext) = out_rx.recv().await { nonce += 1; let frame = crypto::encrypt(&send_c, nonce, &plaintext); - if ws_tx.send(WsMessage::Binary(frame.into())).await.is_err() { + if transport::send_frame(&mut quic_tx, &frame).await.is_err() { break; } } }); - // Receiver task: decrypt and dispatch + // Receiver: QUIC stream → decrypt frame → dispatch MeshMessage let recv_c = recv_cipher.clone(); let receiver = tokio::spawn(async move { - while let Some(Ok(msg)) = ws_rx.next().await { - let data = match msg { - WsMessage::Binary(b) => b.to_vec(), - WsMessage::Ping(_) | WsMessage::Pong(_) => continue, - WsMessage::Close(_) => break, - _ => continue, + loop { + let frame = match transport::recv_frame(&mut quic_rx).await { + Ok(f) => f, + Err(_) => break, }; - - match crypto::decrypt(&recv_c, &data) { - Ok(plaintext) => { - match serde_json::from_slice::(&plaintext) { - Ok(mesh_msg) => { - if incoming_tx.send((peer_id, mesh_msg)).await.is_err() { - break; - } - } - Err(e) => { - tracing::warn!("Invalid mesh message from {}: {}", node_id_hex(&peer_id), e); + match crypto::decrypt(&recv_c, &frame) { + Ok(plaintext) => match serde_json::from_slice::(&plaintext) { + Ok(msg) => { + if incoming_tx.send((peer_id, msg)).await.is_err() { + break; } } - } + Err(e) => { + tracing::warn!( + "Invalid mesh message from {}: {}", + node_id_hex(&peer_id), + e + ); + } + }, Err(e) => { tracing::warn!("Decrypt failed from {}: {}", node_id_hex(&peer_id), e); - break; // Crypto failure = drop connection + break; } } } }); - // Wait for either task to finish, then clean up tokio::select! { - _ = sender => {}, - _ = receiver => {}, + _ = sender => {} + _ = receiver => {} } self.remove_peer(&peer_id).await; } - /// Check if we're connected to a peer. + // ── Peer registry ───────────────────────────────────────────────────────── + + /// Check if we're connected to a peer by node ID. pub async fn is_connected(&self, node_id: &NodeId) -> bool { self.peers.read().await.contains_key(node_id) } + /// Check if we already have an outbound connection to a given socket address. + pub async fn is_connected_by_addr(&self, addr: &SocketAddr) -> bool { + self.peers.read().await.values().any(|p| p.addr == *addr) + } + /// Number of active peer connections. pub async fn peer_count(&self) -> usize { self.peers.read().await.len() @@ -321,13 +406,12 @@ impl OverlayMesh { let peers = self.peers.read().await; let peer = peers .get(node_id) - .ok_or_else(|| anyhow::anyhow!("Peer not connected"))?; - + .ok_or_else(|| anyhow::anyhow!("Peer not connected: {}", node_id_hex(node_id)))?; let json = serde_json::to_vec(msg)?; peer.send_tx .send(json) .await - .map_err(|_| anyhow::anyhow!("Peer channel closed")) + .map_err(|_| anyhow::anyhow!("Peer channel closed: {}", node_id_hex(node_id))) } /// Broadcast a message to all connected peers. @@ -339,7 +423,6 @@ impl OverlayMesh { return; } }; - let peers = self.peers.read().await; for (id, peer) in peers.iter() { if peer.send_tx.send(json.clone()).await.is_err() { @@ -349,9 +432,13 @@ impl OverlayMesh { } /// Register a new peer connection. - pub async fn add_peer(&self, node_id: NodeId, addr: SocketAddr, send_tx: mpsc::Sender>) { + pub async fn add_peer( + &self, + node_id: NodeId, + addr: SocketAddr, + send_tx: mpsc::Sender>, + ) { let mut peers = self.peers.write().await; - if peers.len() >= self.config.max_peers && !peers.contains_key(&node_id) { tracing::debug!( "Max peers ({}) reached, rejecting {}", @@ -360,24 +447,13 @@ impl OverlayMesh { ); return; } - tracing::info!("Peer connected: {} ({})", node_id_hex(&node_id), addr); - - peers.insert( - node_id, - PeerConnection { - node_id, - addr, - send_tx, - send_nonce: std::sync::atomic::AtomicU64::new(0), - }, - ); + peers.insert(node_id, PeerConnection { node_id, addr, send_tx }); } /// Remove a peer connection. pub async fn remove_peer(&self, node_id: &NodeId) { - let mut peers = self.peers.write().await; - if peers.remove(node_id).is_some() { + if self.peers.write().await.remove(node_id).is_some() { tracing::info!("Peer disconnected: {}", node_id_hex(node_id)); } } diff --git a/src/cluster/proxy.rs b/src/cluster/proxy.rs new file mode 100644 index 00000000..b97d8815 --- /dev/null +++ b/src/cluster/proxy.rs @@ -0,0 +1,152 @@ +//! VLESS proxy inbound over QUIC. +//! +//! Accepts QUIC connections with ALPN `"oproxy/1"` and speaks the VLESS v0 +//! protocol, making the relay node usable as an Xray-compatible proxy. +//! +//! ## PQ note +//! The QUIC TLS layer uses a classical ephemeral cert (transport confidentiality +//! only). For 100% post-quantum on the proxy path, clients need a PQ-capable +//! QUIC stack (e.g. a PQ-enabled Xray build with ML-KEM support). The mesh +//! node-to-node path is always 100% PQ regardless. +//! +//! ## VLESS v0 header format +//! ``` +//! [version: 1 byte = 0x00] +//! [UUID: 16 bytes ] +//! [addon_len: 1 byte ] +//! [addons: addon_len bytes] +//! [command: 1 byte (1=TCP, 2=UDP)] +//! [port: 2 bytes BE ] +//! [addr_type: 1 byte (1=IPv4, 2=domain, 3=IPv6)] +//! [addr: 4 / N+1 / 16 bytes] +//! ``` +//! Response: `[0x00][addon_len=0x00]` then raw data. + +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; +use std::sync::Arc; + +use anyhow::{bail, Result}; +use tokio::io::AsyncReadExt; +use tokio::net::TcpStream; + +const VLESS_VERSION: u8 = 0x00; +const CMD_TCP: u8 = 0x01; +const CMD_UDP: u8 = 0x02; +const ATYPE_IPV4: u8 = 0x01; +const ATYPE_DOMAIN: u8 = 0x02; +const ATYPE_IPV6: u8 = 0x03; + +/// Handle a QUIC connection that arrived with ALPN `"oproxy/1"`. +/// +/// Validates the VLESS UUID against `allowed_uuids`, connects to the +/// requested target, and splices the QUIC stream bidirectionally with TCP. +pub async fn handle_proxy_connection( + conn: quinn::Connection, + allowed_uuids: Arc>, +) -> Result<()> { + let (mut send, mut recv) = conn.accept_bi().await?; + + // Parse the VLESS request header + let (target_addr, target_port) = parse_vless_header(&mut recv, &allowed_uuids).await?; + + // Send VLESS response header: version(0x00) + addon_len(0x00) + send.write_all(&[VLESS_VERSION, 0x00]).await?; + + // Connect to the target + let target_addr_str = format!("{}:{}", target_addr, target_port); + tracing::debug!("VLESS proxy → {}", target_addr_str); + let tcp = TcpStream::connect(&target_addr_str).await?; + let (mut tcp_rx, mut tcp_tx) = tcp.into_split(); + + // Splice: QUIC recv → TCP, TCP → QUIC send + let quic_to_tcp = tokio::io::copy(&mut recv, &mut tcp_tx); + let tcp_to_quic = tokio::io::copy(&mut tcp_rx, &mut send); + + tokio::select! { + r = quic_to_tcp => { r.ok(); } + r = tcp_to_quic => { r.ok(); } + } + + Ok(()) +} + +/// Parse the VLESS v0 request header from a QUIC receive stream. +/// +/// Returns `(target_host_string, port)` on success, or an error if the +/// UUID is not in `allowed_uuids` or the header is malformed. +async fn parse_vless_header( + recv: &mut quinn::RecvStream, + allowed_uuids: &[[u8; 16]], +) -> Result<(String, u16)> { + // Version + let version = recv.read_u8().await?; + if version != VLESS_VERSION { + bail!("Unsupported VLESS version: {}", version); + } + + // UUID (16 bytes) + let mut uuid = [0u8; 16]; + recv.read_exact(&mut uuid).await?; + if !allowed_uuids.iter().any(|u| u == &uuid) { + bail!("VLESS UUID not authorized"); + } + + // Additional info (skip) + let addon_len = recv.read_u8().await? as usize; + if addon_len > 0 { + let mut skip = vec![0u8; addon_len]; + recv.read_exact(&mut skip).await?; + } + + // Command + let cmd = recv.read_u8().await?; + if cmd != CMD_TCP && cmd != CMD_UDP { + bail!("Unsupported VLESS command: {}", cmd); + } + + // Port (2 bytes BE) + let port = recv.read_u16().await?; + + // Address + let addr_type = recv.read_u8().await?; + let host = match addr_type { + ATYPE_IPV4 => { + let mut octets = [0u8; 4]; + recv.read_exact(&mut octets).await?; + IpAddr::V4(Ipv4Addr::from(octets)).to_string() + } + ATYPE_IPV6 => { + let mut octets = [0u8; 16]; + recv.read_exact(&mut octets).await?; + IpAddr::V6(Ipv6Addr::from(octets)).to_string() + } + ATYPE_DOMAIN => { + let domain_len = recv.read_u8().await? as usize; + if domain_len == 0 || domain_len > 253 { + bail!("Invalid VLESS domain length: {}", domain_len); + } + let mut domain_bytes = vec![0u8; domain_len]; + recv.read_exact(&mut domain_bytes).await?; + String::from_utf8(domain_bytes) + .map_err(|_| anyhow::anyhow!("VLESS domain is not valid UTF-8"))? + } + _ => bail!("Unknown VLESS address type: {}", addr_type), + }; + + Ok((host, port)) +} + +/// Parse a UUID string (e.g. `"550e8400-e29b-41d4-a716-446655440000"`) into 16 bytes. +pub fn parse_uuid(s: &str) -> Result<[u8; 16]> { + let hex: String = s.chars().filter(|c| c.is_ascii_hexdigit()).collect(); + if hex.len() != 32 { + bail!("Invalid UUID: {}", s); + } + let bytes = (0..16) + .map(|i| u8::from_str_radix(&hex[i * 2..i * 2 + 2], 16)) + .collect::, _>>() + .map_err(|_| anyhow::anyhow!("UUID parse failed: {}", s))?; + let mut out = [0u8; 16]; + out.copy_from_slice(&bytes); + Ok(out) +} diff --git a/src/cluster/transport.rs b/src/cluster/transport.rs new file mode 100644 index 00000000..e03b1180 --- /dev/null +++ b/src/cluster/transport.rs @@ -0,0 +1,156 @@ +//! QUIC transport layer for the mesh overlay. +//! +//! Provides UDP-based connectivity (NAT hole-punching capable, like Nebula) via +//! the quinn QUIC implementation. An ephemeral self-signed TLS certificate is used +//! for the QUIC handshake — it provides transport confidentiality only. All real +//! authentication and post-quantum security comes from the ML-KEM-768 application +//! handshake in overlay.rs. +//! +//! Two ALPN values share the single UDP port: +//! "omesh/1" — mesh overlay protocol (PQ handshake + MeshMessage frames) +//! "oproxy/1" — VLESS proxy inbound (for Xray-compatible clients) + +use std::net::SocketAddr; +use std::sync::Arc; + +use anyhow::Result; +use quinn::{ClientConfig, Endpoint, ServerConfig}; +use rustls::pki_types::{CertificateDer, PrivatePkcs8KeyDer, PrivateKeyDer}; + +/// ALPN token for mesh overlay traffic. +pub const ALPN_MESH: &[u8] = b"omesh/1"; +/// ALPN token for VLESS proxy traffic. +pub const ALPN_PROXY: &[u8] = b"oproxy/1"; + +// ── Endpoint construction ───────────────────────────────────────────────────── + +/// Build a QUIC endpoint that listens for incoming connections (server role) +/// and can also initiate outbound connections (client role). +/// +/// Accepts connections with ALPN `"omesh/1"` (mesh) and `"oproxy/1"` (proxy). +/// Uses an ephemeral self-signed certificate — authentication is done by the PQ +/// application handshake, not by TLS certificate verification. +pub fn make_server_endpoint(bind_addr: SocketAddr) -> Result { + let (cert_chain, priv_key) = generate_ephemeral_cert()?; + + let mut tls = rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(cert_chain, priv_key)?; + tls.alpn_protocols = vec![ALPN_MESH.to_vec(), ALPN_PROXY.to_vec()]; + + let server_cfg = ServerConfig::with_crypto(Arc::new( + quinn::crypto::rustls::QuicServerConfig::try_from(tls) + .map_err(|e| anyhow::anyhow!("QUIC server crypto: {}", e))?, + )); + + let endpoint = Endpoint::server(server_cfg, bind_addr)?; + Ok(endpoint) +} + +/// Build a QUIC client config that skips TLS certificate verification. +/// +/// Security does NOT depend on certificate validity — our ML-KEM-768 handshake +/// running on the first QUIC stream provides authentication and PQ confidentiality. +pub fn make_client_config() -> Result { + let mut tls = rustls::ClientConfig::builder() + .dangerous() + .with_custom_certificate_verifier(Arc::new(SkipServerVerification)) + .with_no_client_auth(); + tls.alpn_protocols = vec![ALPN_MESH.to_vec()]; + + Ok(ClientConfig::new(Arc::new( + quinn::crypto::rustls::QuicClientConfig::try_from(tls) + .map_err(|e| anyhow::anyhow!("QUIC client crypto: {}", e))?, + ))) +} + +// ── Length-prefix framing over QUIC streams ─────────────────────────────────── +// +// QUIC streams are byte streams, not message streams. We use 4-byte big-endian +// length prefixes so both handshake JSON and encrypted MeshMessage frames can +// be exchanged as discrete messages. + +/// Write a length-prefixed frame to a QUIC send stream. +pub async fn send_frame(stream: &mut quinn::SendStream, data: &[u8]) -> Result<()> { + let len = u32::try_from(data.len()) + .map_err(|_| anyhow::anyhow!("Frame too large: {} bytes", data.len()))?; + stream.write_all(&len.to_be_bytes()).await?; + stream.write_all(data).await?; + Ok(()) +} + +/// Read a length-prefixed frame from a QUIC receive stream. +pub async fn recv_frame(stream: &mut quinn::RecvStream) -> Result> { + let mut len_buf = [0u8; 4]; + stream.read_exact(&mut len_buf).await?; + let len = u32::from_be_bytes(len_buf) as usize; + if len > 16 * 1024 * 1024 { + anyhow::bail!("Frame too large: {} bytes", len); + } + let mut buf = vec![0u8; len]; + stream.read_exact(&mut buf).await?; + Ok(buf) +} + +// ── Ephemeral certificate ───────────────────────────────────────────────────── + +fn generate_ephemeral_cert() -> Result<(Vec>, PrivateKeyDer<'static>)> { + let cert = rcgen::generate_simple_self_signed(vec!["omesh".to_string()])?; + let cert_der = CertificateDer::from(cert.cert.der().to_vec()); + let key_der = PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(cert.key_pair.serialize_der())); + Ok((vec![cert_der], key_der)) +} + +// ── TLS certificate verifier that accepts any cert ─────────────────────────── + +#[derive(Debug)] +struct SkipServerVerification; + +impl rustls::client::danger::ServerCertVerifier for SkipServerVerification { + fn verify_server_cert( + &self, + _end_entity: &CertificateDer<'_>, + _intermediates: &[CertificateDer<'_>], + _server_name: &rustls::pki_types::ServerName<'_>, + _ocsp_response: &[u8], + _now: rustls::pki_types::UnixTime, + ) -> std::result::Result { + Ok(rustls::client::danger::ServerCertVerified::assertion()) + } + + fn verify_tls12_signature( + &self, + _message: &[u8], + _cert: &CertificateDer<'_>, + _dss: &rustls::DigitallySignedStruct, + ) -> std::result::Result { + Ok(rustls::client::danger::HandshakeSignatureValid::assertion()) + } + + fn verify_tls13_signature( + &self, + _message: &[u8], + _cert: &CertificateDer<'_>, + _dss: &rustls::DigitallySignedStruct, + ) -> std::result::Result { + Ok(rustls::client::danger::HandshakeSignatureValid::assertion()) + } + + fn supported_verify_schemes(&self) -> Vec { + vec![ + rustls::SignatureScheme::RSA_PKCS1_SHA1, + rustls::SignatureScheme::ECDSA_SHA1_Legacy, + rustls::SignatureScheme::RSA_PKCS1_SHA256, + rustls::SignatureScheme::ECDSA_NISTP256_SHA256, + rustls::SignatureScheme::RSA_PKCS1_SHA384, + rustls::SignatureScheme::ECDSA_NISTP384_SHA384, + rustls::SignatureScheme::RSA_PKCS1_SHA512, + rustls::SignatureScheme::ECDSA_NISTP521_SHA512, + rustls::SignatureScheme::RSA_PSS_SHA256, + rustls::SignatureScheme::RSA_PSS_SHA384, + rustls::SignatureScheme::RSA_PSS_SHA512, + rustls::SignatureScheme::ED25519, + rustls::SignatureScheme::ED448, + ] + } +} diff --git a/src/llm/CLAUDE.md b/src/llm/CLAUDE.md index ac36bd88..a795447c 100644 --- a/src/llm/CLAUDE.md +++ b/src/llm/CLAUDE.md @@ -39,6 +39,7 @@ Set via `LLM_BACKEND` env var: | `anthropic` | Anthropic | `ANTHROPIC_API_KEY` | | `github_copilot` | GitHub Copilot Chat API | `GITHUB_COPILOT_TOKEN`, `GITHUB_COPILOT_MODEL` | | `ollama` | Ollama local | `OLLAMA_BASE_URL` | +| `vllm` | vLLM inference server | `VLLM_BASE_URL` (default: `http://localhost:8000/v1`), `VLLM_MODEL`, `VLLM_API_KEY` (optional) | | `openai_compatible` | Any OpenAI-compatible endpoint | `LLM_BASE_URL`, `LLM_API_KEY`, `LLM_MODEL` | | `tinfoil` | Tinfoil TEE inference | `TINFOIL_API_KEY`, `TINFOIL_MODEL` | | `bedrock` | AWS Bedrock (requires `--features bedrock`) | `BEDROCK_REGION`, `BEDROCK_MODEL`, `AWS_PROFILE` | diff --git a/src/llm/mod.rs b/src/llm/mod.rs index df776823..77accad1 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -5,6 +5,7 @@ //! - **OpenAI**: Direct API access with your own key //! - **Anthropic**: Direct API access with your own key //! - **Ollama**: Local model inference +//! - **vLLM**: High-throughput local/remote inference server (OpenAI-compatible) //! - **OpenAI-compatible**: Any endpoint that speaks the OpenAI API //! - **AWS Bedrock**: Native Converse API via aws-sdk-bedrockruntime