feat(core,edge,hub,transport): add QUIC tunnel transport support with optional edge transport selection

rust/crates/remoteingress-core/src/hub.rs

@@ -12,6 +12,7 @@ use serde::{Deserialize, Serialize};
 
 use bytes::Bytes;
 use remoteingress_protocol::*;
+use crate::transport::quic as quic_transport;
 
 type HubTlsStream = tokio_rustls::server::TlsStream<TcpStream>;
 
@@ -216,14 +217,35 @@ impl TunnelHub {
         }
     }
 
-    /// Start the hub — listen for TLS connections from edges.
+    /// Start the hub — listen for TLS connections (TCP) and QUIC connections (UDP) from edges.
     pub async fn start(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
         let config = self.config.read().await.clone();
         let tls_config = build_tls_config(&config)?;
-        let acceptor = TlsAcceptor::from(Arc::new(tls_config));
+        let acceptor = TlsAcceptor::from(Arc::new(tls_config.clone()));
 
         let listener = TcpListener::bind(("0.0.0.0", config.tunnel_port)).await?;
-        log::info!("Hub listening on port {}", config.tunnel_port);
+        log::info!("Hub listening on TCP port {}", config.tunnel_port);
+
+        // Start QUIC endpoint on the same port (UDP)
+        let quic_endpoint = match quic_transport::build_quic_server_config(tls_config) {
+            Ok(quic_server_config) => {
+                let bind_addr: std::net::SocketAddr = ([0, 0, 0, 0], config.tunnel_port).into();
+                match quinn::Endpoint::server(quic_server_config, bind_addr) {
+                    Ok(ep) => {
+                        log::info!("Hub listening on QUIC/UDP port {}", config.tunnel_port);
+                        Some(ep)
+                    }
+                    Err(e) => {
+                        log::warn!("Failed to start QUIC endpoint: {} (QUIC disabled)", e);
+                        None
+                    }
+                }
+            }
+            Err(e) => {
+                log::warn!("Failed to build QUIC server config: {} (QUIC disabled)", e);
+                None
+            }
+        };
 
         let (shutdown_tx, mut shutdown_rx) = mpsc::channel::<()>(1);
         *self.shutdown_tx.lock().await = Some(shutdown_tx);
@@ -236,12 +258,62 @@ impl TunnelHub {
         let hub_token = self.cancel_token.clone();
 
         tokio::spawn(async move {
+            // Spawn QUIC acceptor as a separate task
+            let quic_handle = if let Some(quic_ep) = quic_endpoint {
+                let allowed_q = allowed.clone();
+                let connected_q = connected.clone();
+                let event_tx_q = event_tx.clone();
+                let target_q = target_host.clone();
+                let hub_token_q = hub_token.clone();
+                Some(tokio::spawn(async move {
+                    loop {
+                        tokio::select! {
+                            incoming = quic_ep.accept() => {
+                                match incoming {
+                                    Some(incoming) => {
+                                        let allowed = allowed_q.clone();
+                                        let connected = connected_q.clone();
+                                        let event_tx = event_tx_q.clone();
+                                        let target = target_q.clone();
+                                        let edge_token = hub_token_q.child_token();
+                                        let peer_addr = incoming.remote_address().ip().to_string();
+                                        tokio::spawn(async move {
+                                            // Accept the QUIC connection
+                                            let quic_conn = match incoming.await {
+                                                Ok(c) => c,
+                                                Err(e) => {
+                                                    log::error!("QUIC connection error: {}", e);
+                                                    return;
+                                                }
+                                            };
+                                            if let Err(e) = handle_edge_connection_quic(
+                                                quic_conn, allowed, connected, event_tx, target, edge_token, peer_addr,
+                                            ).await {
+                                                log::error!("QUIC edge connection error: {}", e);
+                                            }
+                                        });
+                                    }
+                                    None => {
+                                        log::info!("QUIC endpoint closed");
+                                        break;
+                                    }
+                                }
+                            }
+                            _ = hub_token_q.cancelled() => break,
+                        }
+                    }
+                }))
+            } else {
+                None
+            };
+
+            // TCP+TLS acceptor loop
             loop {
                 tokio::select! {
                     result = listener.accept() => {
                         match result {
                             Ok((stream, addr)) => {
-                                log::info!("Edge connection from {}", addr);
+                                log::info!("Edge connection from {} (TCP+TLS)", addr);
                                 let acceptor = acceptor.clone();
                                 let allowed = allowed.clone();
                                 let connected = connected.clone();
@@ -272,6 +344,11 @@ impl TunnelHub {
                     }
                 }
             }
+
+            // Abort QUIC acceptor if running
+            if let Some(h) = quic_handle {
+                h.abort();
+            }
         });
 
         Ok(())
@@ -956,6 +1033,363 @@ fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
     diff == 0
 }
 
+// ===== QUIC transport functions for hub =====
+
+/// Handle an edge connection arriving via QUIC.
+/// The first bidirectional stream is the control stream (auth + config).
+/// Subsequent bidirectional streams are tunneled client connections.
+async fn handle_edge_connection_quic(
+    quic_conn: quinn::Connection,
+    allowed: Arc<RwLock<HashMap<String, AllowedEdge>>>,
+    connected: Arc<Mutex<HashMap<String, ConnectedEdgeInfo>>>,
+    event_tx: mpsc::Sender<HubEvent>,
+    target_host: String,
+    edge_token: CancellationToken,
+    peer_addr: String,
+) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    log::info!("QUIC edge connection from {}", peer_addr);
+
+    // Accept the control stream (first bidirectional stream from edge)
+    let (mut ctrl_send, mut ctrl_recv) = match quic_conn.accept_bi().await {
+        Ok(s) => s,
+        Err(e) => return Err(format!("QUIC control stream accept failed: {}", e).into()),
+    };
+
+    // Read auth line from control stream
+    let mut auth_buf = Vec::with_capacity(512);
+    loop {
+        let mut byte = [0u8; 1];
+        match ctrl_recv.read_exact(&mut byte).await {
+            Ok(()) => {
+                if byte[0] == b'\n' { break; }
+                auth_buf.push(byte[0]);
+                if auth_buf.len() > 4096 {
+                    return Err("QUIC auth line too long".into());
+                }
+            }
+            Err(e) => return Err(format!("QUIC auth read failed: {}", e).into()),
+        }
+    }
+    let auth_line = String::from_utf8(auth_buf)
+        .map_err(|_| "QUIC auth line not valid UTF-8")?;
+    let auth_line = auth_line.trim();
+
+    let parts: Vec<&str> = auth_line.splitn(3, ' ').collect();
+    if parts.len() != 3 || parts[0] != "EDGE" {
+        return Err("invalid QUIC auth line".into());
+    }
+
+    let edge_id = parts[1].to_string();
+    let secret = parts[2];
+
+    // Verify credentials
+    let (listen_ports, stun_interval_secs) = {
+        let edges = allowed.read().await;
+        match edges.get(&edge_id) {
+            Some(edge) => {
+                if !constant_time_eq(secret.as_bytes(), edge.secret.as_bytes()) {
+                    return Err(format!("invalid secret for edge {}", edge_id).into());
+                }
+                (edge.listen_ports.clone(), edge.stun_interval_secs.unwrap_or(300))
+            }
+            None => return Err(format!("unknown edge {}", edge_id).into()),
+        }
+    };
+
+    log::info!("QUIC edge {} authenticated from {}", edge_id, peer_addr);
+    let _ = event_tx.try_send(HubEvent::EdgeConnected {
+        edge_id: edge_id.clone(),
+        peer_addr: peer_addr.clone(),
+    });
+
+    // Send handshake response on control stream
+    let handshake = HandshakeResponse {
+        listen_ports: listen_ports.clone(),
+        stun_interval_secs,
+    };
+    let mut handshake_json = serde_json::to_string(&handshake)?;
+    handshake_json.push('\n');
+    ctrl_send.write_all(handshake_json.as_bytes()).await
+        .map_err(|e| format!("QUIC handshake write failed: {}", e))?;
+
+    // Track this edge
+    let edge_stream_count = Arc::new(AtomicU32::new(0));
+    let now = std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .unwrap_or_default()
+        .as_secs();
+
+    let (config_tx, mut config_rx) = mpsc::channel::<EdgeConfigUpdate>(16);
+
+    {
+        let mut edges = connected.lock().await;
+        if let Some(old) = edges.remove(&edge_id) {
+            log::info!("QUIC edge {} reconnected, cancelling old connection", edge_id);
+            old.cancel_token.cancel();
+        }
+        edges.insert(
+            edge_id.clone(),
+            ConnectedEdgeInfo {
+                connected_at: now,
+                peer_addr,
+                edge_stream_count: edge_stream_count.clone(),
+                config_tx,
+                cancel_token: edge_token.clone(),
+            },
+        );
+    }
+
+    let stream_semaphore = Arc::new(Semaphore::new(MAX_STREAMS_PER_EDGE));
+
+    // Spawn task to accept data streams (tunneled client connections)
+    let data_stream_conn = quic_conn.clone();
+    let data_target = target_host.clone();
+    let data_edge_id = edge_id.clone();
+    let data_event_tx = event_tx.clone();
+    let data_semaphore = stream_semaphore.clone();
+    let data_stream_count = edge_stream_count.clone();
+    let data_token = edge_token.clone();
+    let data_handle = tokio::spawn(async move {
+        let mut stream_id_counter: u32 = 0;
+        loop {
+            tokio::select! {
+                bi_result = data_stream_conn.accept_bi() => {
+                    match bi_result {
+                        Ok((quic_send, quic_recv)) => {
+                            // Check stream limit
+                            let permit = match data_semaphore.clone().try_acquire_owned() {
+                                Ok(p) => p,
+                                Err(_) => {
+                                    log::warn!("QUIC edge {} exceeded max streams, rejecting", data_edge_id);
+                                    // Drop the streams to reject
+                                    drop(quic_send);
+                                    drop(quic_recv);
+                                    continue;
+                                }
+                            };
+
+                            stream_id_counter += 1;
+                            let stream_id = stream_id_counter;
+                            let target = data_target.clone();
+                            let edge_id = data_edge_id.clone();
+                            let event_tx = data_event_tx.clone();
+                            let stream_count = data_stream_count.clone();
+                            let stream_token = data_token.child_token();
+
+                            let _ = event_tx.try_send(HubEvent::StreamOpened {
+                                edge_id: edge_id.clone(),
+                                stream_id,
+                            });
+
+                            stream_count.fetch_add(1, Ordering::Relaxed);
+                            tokio::spawn(async move {
+                                let _permit = permit;
+                                handle_quic_stream(
+                                    quic_send, quic_recv, stream_id,
+                                    &target, &edge_id, stream_token,
+                                ).await;
+                                stream_count.fetch_sub(1, Ordering::Relaxed);
+                                let _ = event_tx.try_send(HubEvent::StreamClosed {
+                                    edge_id,
+                                    stream_id,
+                                });
+                            });
+                        }
+                        Err(e) => {
+                            log::info!("QUIC edge {} accept_bi ended: {}", data_edge_id, e);
+                            break;
+                        }
+                    }
+                }
+                _ = data_token.cancelled() => break,
+            }
+        }
+    });
+
+    // Control stream loop: forward config updates and handle PONG
+    let disconnect_reason;
+    loop {
+        tokio::select! {
+            // Send config updates from hub to edge
+            update = config_rx.recv() => {
+                match update {
+                    Some(update) => {
+                        if let Ok(payload) = serde_json::to_vec(&update) {
+                            if let Err(e) = quic_transport::write_ctrl_message(
+                                &mut ctrl_send, quic_transport::CTRL_CONFIG, &payload,
+                            ).await {
+                                log::error!("QUIC config send to edge {} failed: {}", edge_id, e);
+                                disconnect_reason = format!("quic_config_send_failed: {}", e);
+                                break;
+                            }
+                            log::info!("Sent QUIC config update to edge {}: ports {:?}", edge_id, update.listen_ports);
+                        }
+                    }
+                    None => {
+                        disconnect_reason = "config_channel_closed".to_string();
+                        break;
+                    }
+                }
+            }
+            // Read control messages from edge (mainly PONG responses)
+            ctrl_msg = quic_transport::read_ctrl_message(&mut ctrl_recv) => {
+                match ctrl_msg {
+                    Ok(Some((msg_type, _payload))) => {
+                        match msg_type {
+                            quic_transport::CTRL_PONG => {
+                                log::debug!("Received QUIC PONG from edge {}", edge_id);
+                            }
+                            _ => {
+                                log::warn!("Unexpected QUIC control message type {} from edge {}", msg_type, edge_id);
+                            }
+                        }
+                    }
+                    Ok(None) => {
+                        log::info!("QUIC edge {} control stream EOF", edge_id);
+                        disconnect_reason = "quic_ctrl_eof".to_string();
+                        break;
+                    }
+                    Err(e) => {
+                        log::error!("QUIC edge {} control stream error: {}", edge_id, e);
+                        disconnect_reason = format!("quic_ctrl_error: {}", e);
+                        break;
+                    }
+                }
+            }
+            // QUIC connection closed
+            reason = quic_conn.closed() => {
+                log::info!("QUIC connection to edge {} closed: {}", edge_id, reason);
+                disconnect_reason = format!("quic_closed: {}", reason);
+                break;
+            }
+            // Hub-initiated cancellation
+            _ = edge_token.cancelled() => {
+                log::info!("QUIC edge {} cancelled by hub", edge_id);
+                disconnect_reason = "cancelled_by_hub".to_string();
+                break;
+            }
+        }
+    }
+
+    // Cleanup
+    edge_token.cancel();
+    data_handle.abort();
+    quic_conn.close(quinn::VarInt::from_u32(0), b"hub_shutdown");
+
+    {
+        let mut edges = connected.lock().await;
+        edges.remove(&edge_id);
+    }
+    let _ = event_tx.try_send(HubEvent::EdgeDisconnected {
+        edge_id,
+        reason: disconnect_reason,
+    });
+
+    Ok(())
+}
+
+/// Handle a single tunneled client connection arriving via a QUIC bidirectional stream.
+/// Reads the PROXY header, connects to SmartProxy, and pipes data bidirectionally.
+async fn handle_quic_stream(
+    mut quic_send: quinn::SendStream,
+    mut quic_recv: quinn::RecvStream,
+    stream_id: u32,
+    target_host: &str,
+    _edge_id: &str,
+    stream_token: CancellationToken,
+) {
+    // Read PROXY header from the beginning of the stream
+    let proxy_header = match quic_transport::read_proxy_header(&mut quic_recv).await {
+        Ok(h) => h,
+        Err(e) => {
+            log::error!("QUIC stream {} failed to read PROXY header: {}", stream_id, e);
+            return;
+        }
+    };
+
+    let dest_port = parse_dest_port_from_proxy(&proxy_header).unwrap_or(443);
+
+    // Connect to SmartProxy
+    let mut upstream = match tokio::time::timeout(
+        Duration::from_secs(10),
+        TcpStream::connect((target_host, dest_port)),
+    ).await {
+        Ok(Ok(s)) => s,
+        Ok(Err(e)) => {
+            log::error!("QUIC stream {} connect to {}:{} failed: {}", stream_id, target_host, dest_port, e);
+            return;
+        }
+        Err(_) => {
+            log::error!("QUIC stream {} connect to {}:{} timed out", stream_id, target_host, dest_port);
+            return;
+        }
+    };
+
+    let _ = upstream.set_nodelay(true);
+    // Send PROXY header to SmartProxy
+    if let Err(e) = upstream.write_all(proxy_header.as_bytes()).await {
+        log::error!("QUIC stream {} failed to write PROXY header to upstream: {}", stream_id, e);
+        return;
+    }
+
+    let (mut up_read, mut up_write) = upstream.into_split();
+
+    // Task: QUIC -> upstream (edge data to SmartProxy)
+    let writer_token = stream_token.clone();
+    let writer_task = tokio::spawn(async move {
+        let mut buf = vec![0u8; 32768];
+        loop {
+            tokio::select! {
+                read_result = quic_recv.read(&mut buf) => {
+                    match read_result {
+                        Ok(Some(n)) => {
+                            let write_result = tokio::select! {
+                                r = tokio::time::timeout(
+                                    Duration::from_secs(60),
+                                    up_write.write_all(&buf[..n]),
+                                ) => r,
+                                _ = writer_token.cancelled() => break,
+                            };
+                            match write_result {
+                                Ok(Ok(())) => {}
+                                Ok(Err(_)) => break,
+                                Err(_) => break,
+                            }
+                        }
+                        Ok(None) => break, // QUIC stream finished
+                        Err(_) => break,
+                    }
+                }
+                _ = writer_token.cancelled() => break,
+            }
+        }
+        let _ = up_write.shutdown().await;
+    });
+
+    // Task: upstream -> QUIC (SmartProxy data to edge)
+    let mut buf = vec![0u8; 32768];
+    loop {
+        tokio::select! {
+            read_result = up_read.read(&mut buf) => {
+                match read_result {
+                    Ok(0) => break,
+                    Ok(n) => {
+                        if quic_send.write_all(&buf[..n]).await.is_err() {
+                            break;
+                        }
+                    }
+                    Err(_) => break,
+                }
+            }
+            _ = stream_token.cancelled() => break,
+        }
+    }
+
+    // Gracefully close the QUIC send stream
+    let _ = quic_send.finish();
+    writer_task.abort();
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;