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

This commit is contained in:
2026-03-19 10:44:22 +00:00
parent e4807be00b
commit 6abfd2ff2a
11 changed files with 2051 additions and 19 deletions

View File

@@ -16,3 +16,4 @@ log = "0.4"
rustls-pemfile = "2"
tokio-util = "0.7"
socket2 = "0.5"
quinn = "0.11"

View File

@@ -13,6 +13,8 @@ use serde::{Deserialize, Serialize};
use bytes::Bytes;
use remoteingress_protocol::*;
use crate::transport::TransportMode;
use crate::transport::quic as quic_transport;
type EdgeTlsStream = tokio_rustls::client::TlsStream<TcpStream>;
@@ -47,6 +49,9 @@ pub struct EdgeConfig {
/// Useful for testing on localhost where edge and upstream share the same machine.
#[serde(default)]
pub bind_address: Option<String>,
/// Transport mode for the tunnel connection (defaults to TcpTls).
#[serde(default)]
pub transport_mode: Option<TransportMode>,
}
/// Handshake config received from hub after authentication.
@@ -210,6 +215,8 @@ async fn edge_main_loop(
let mut backoff_ms: u64 = 1000;
let max_backoff_ms: u64 = 30000;
let transport_mode = config.transport_mode.unwrap_or(TransportMode::TcpTls);
// Build TLS config ONCE outside the reconnect loop — preserves session
// cache across reconnections for TLS session resumption (saves 1 RTT).
let tls_config = rustls::ClientConfig::builder()
@@ -218,24 +225,77 @@ async fn edge_main_loop(
.with_no_client_auth();
let connector = TlsConnector::from(Arc::new(tls_config));
// Build QUIC client config ONCE (shares session cache across reconnections).
let quic_client_config = quic_transport::build_quic_client_config();
let quic_endpoint = if matches!(transport_mode, TransportMode::Quic | TransportMode::QuicWithFallback) {
match quinn::Endpoint::client("0.0.0.0:0".parse().unwrap()) {
Ok(mut ep) => {
ep.set_default_client_config(quic_client_config);
Some(ep)
}
Err(e) => {
log::error!("Failed to create QUIC endpoint: {}", e);
None
}
}
} else {
None
};
loop {
// Create a per-connection child token
let connection_token = cancel_token.child_token();
// Try to connect to hub
let result = connect_to_hub_and_run(
&config,
&connected,
&public_ip,
&active_streams,
&next_stream_id,
&event_tx,
&listen_ports,
&mut shutdown_rx,
&connection_token,
&connector,
)
.await;
// Try to connect to hub using the configured transport
let result = match transport_mode {
TransportMode::TcpTls => {
connect_to_hub_and_run(
&config, &connected, &public_ip, &active_streams, &next_stream_id,
&event_tx, &listen_ports, &mut shutdown_rx, &connection_token, &connector,
).await
}
TransportMode::Quic => {
if let Some(ep) = &quic_endpoint {
connect_to_hub_and_run_quic(
&config, &connected, &public_ip, &active_streams, &next_stream_id,
&event_tx, &listen_ports, &mut shutdown_rx, &connection_token, ep,
).await
} else {
EdgeLoopResult::Reconnect("quic_endpoint_unavailable".to_string())
}
}
TransportMode::QuicWithFallback => {
if let Some(ep) = &quic_endpoint {
// Try QUIC first with a 5s timeout
let quic_result = tokio::time::timeout(
Duration::from_secs(5),
connect_to_hub_quic_handshake(&config, ep, &connection_token),
).await;
match quic_result {
Ok(Ok(quic_conn)) => {
connect_to_hub_and_run_quic_with_connection(
&config, &connected, &public_ip, &active_streams, &next_stream_id,
&event_tx, &listen_ports, &mut shutdown_rx, &connection_token,
quic_conn,
).await
}
_ => {
log::info!("QUIC connect failed or timed out, falling back to TCP+TLS");
connect_to_hub_and_run(
&config, &connected, &public_ip, &active_streams, &next_stream_id,
&event_tx, &listen_ports, &mut shutdown_rx, &connection_token, &connector,
).await
}
}
} else {
// No QUIC endpoint, fall back to TCP+TLS
connect_to_hub_and_run(
&config, &connected, &public_ip, &active_streams, &next_stream_id,
&event_tx, &listen_ports, &mut shutdown_rx, &connection_token, &connector,
).await
}
}
};
// Cancel connection token to kill all orphaned tasks from this cycle
connection_token.cancel();
@@ -942,6 +1002,437 @@ async fn handle_client_connection(
let _ = edge_id; // used for logging context
}
// ===== QUIC transport functions =====
/// Perform QUIC handshake only (used by QuicWithFallback to test connectivity).
async fn connect_to_hub_quic_handshake(
config: &EdgeConfig,
endpoint: &quinn::Endpoint,
_connection_token: &CancellationToken,
) -> Result<quinn::Connection, Box<dyn std::error::Error + Send + Sync>> {
let addr = format!("{}:{}", config.hub_host, config.hub_port);
let server_addr: std::net::SocketAddr = tokio::net::lookup_host(&addr)
.await?
.next()
.ok_or("DNS resolution failed")?;
// QUIC/TLS SNI requires a hostname, not an IP address.
// If hub_host is an IP, use the same fallback as the TCP+TLS path.
let server_name = match rustls::pki_types::ServerName::try_from(config.hub_host.as_str()) {
Ok(rustls::pki_types::ServerName::DnsName(_)) => config.hub_host.clone(),
_ => "remoteingress-hub".to_string(),
};
let connection = endpoint.connect(server_addr, &server_name)?.await?;
Ok(connection)
}
/// QUIC edge: connect to hub, authenticate, and run the stream multiplexer.
async fn connect_to_hub_and_run_quic(
config: &EdgeConfig,
connected: &Arc<RwLock<bool>>,
public_ip: &Arc<RwLock<Option<String>>>,
active_streams: &Arc<AtomicU32>,
next_stream_id: &Arc<AtomicU32>,
event_tx: &mpsc::Sender<EdgeEvent>,
listen_ports: &Arc<RwLock<Vec<u16>>>,
shutdown_rx: &mut mpsc::Receiver<()>,
connection_token: &CancellationToken,
endpoint: &quinn::Endpoint,
) -> EdgeLoopResult {
// Establish QUIC connection
let quic_conn = match connect_to_hub_quic_handshake(config, endpoint, connection_token).await {
Ok(c) => c,
Err(e) => {
log::error!("QUIC connect failed: {}", e);
return EdgeLoopResult::Reconnect(format!("quic_connect_failed: {}", e));
}
};
connect_to_hub_and_run_quic_with_connection(
config, connected, public_ip, active_streams, next_stream_id,
event_tx, listen_ports, shutdown_rx, connection_token, quic_conn,
).await
}
/// QUIC edge: run with an already-established QUIC connection.
async fn connect_to_hub_and_run_quic_with_connection(
config: &EdgeConfig,
connected: &Arc<RwLock<bool>>,
public_ip: &Arc<RwLock<Option<String>>>,
active_streams: &Arc<AtomicU32>,
next_stream_id: &Arc<AtomicU32>,
event_tx: &mpsc::Sender<EdgeEvent>,
listen_ports: &Arc<RwLock<Vec<u16>>>,
shutdown_rx: &mut mpsc::Receiver<()>,
connection_token: &CancellationToken,
quic_conn: quinn::Connection,
) -> EdgeLoopResult {
log::info!("QUIC connection established to {}", quic_conn.remote_address());
// Open control stream (first bidirectional stream)
let (mut ctrl_send, mut ctrl_recv) = match quic_conn.open_bi().await {
Ok(s) => s,
Err(e) => {
log::error!("Failed to open QUIC control stream: {}", e);
return EdgeLoopResult::Reconnect(format!("quic_ctrl_open_failed: {}", e));
}
};
// Auth handshake on control stream (same protocol as TCP+TLS)
let auth_line = format!("EDGE {} {}\n", config.edge_id, config.secret);
if let Err(e) = ctrl_send.write_all(auth_line.as_bytes()).await {
return EdgeLoopResult::Reconnect(format!("quic_auth_write_failed: {}", e));
}
// Read handshake response (newline-delimited JSON)
let mut handshake_bytes = Vec::with_capacity(512);
let mut byte = [0u8; 1];
loop {
match ctrl_recv.read_exact(&mut byte).await {
Ok(()) => {
handshake_bytes.push(byte[0]);
if byte[0] == b'\n' { break; }
if handshake_bytes.len() > 8192 {
return EdgeLoopResult::Reconnect("quic_handshake_too_long".to_string());
}
}
Err(e) => {
log::error!("QUIC handshake read failed: {}", e);
return EdgeLoopResult::Reconnect(format!("quic_handshake_read_failed: {}", e));
}
}
}
let handshake_line = String::from_utf8_lossy(&handshake_bytes);
let handshake: HandshakeConfig = match serde_json::from_str(handshake_line.trim()) {
Ok(h) => h,
Err(e) => {
log::error!("Invalid QUIC handshake response: {}", e);
return EdgeLoopResult::Reconnect(format!("quic_handshake_invalid: {}", e));
}
};
log::info!(
"QUIC handshake from hub: ports {:?}, stun_interval {}s",
handshake.listen_ports,
handshake.stun_interval_secs
);
*connected.write().await = true;
let _ = event_tx.try_send(EdgeEvent::TunnelConnected);
log::info!("Connected to hub via QUIC at {}", quic_conn.remote_address());
*listen_ports.write().await = handshake.listen_ports.clone();
let _ = event_tx.try_send(EdgeEvent::PortsAssigned {
listen_ports: handshake.listen_ports.clone(),
});
// Start STUN discovery
let stun_interval = handshake.stun_interval_secs;
let public_ip_clone = public_ip.clone();
let event_tx_clone = event_tx.clone();
let stun_token = connection_token.clone();
let stun_handle = tokio::spawn(async move {
loop {
tokio::select! {
ip_result = crate::stun::discover_public_ip() => {
if let Some(ip) = ip_result {
let mut pip = public_ip_clone.write().await;
let changed = pip.as_ref() != Some(&ip);
*pip = Some(ip.clone());
if changed {
let _ = event_tx_clone.try_send(EdgeEvent::PublicIpDiscovered { ip });
}
}
}
_ = stun_token.cancelled() => break,
}
tokio::select! {
_ = tokio::time::sleep(Duration::from_secs(stun_interval)) => {}
_ = stun_token.cancelled() => break,
}
}
});
// Start TCP listeners for the assigned ports.
// For QUIC, each client connection opens a new QUIC bidirectional stream.
let mut port_listeners: HashMap<u16, JoinHandle<()>> = HashMap::new();
let bind_address = config.bind_address.as_deref().unwrap_or("0.0.0.0");
apply_port_config_quic(
&handshake.listen_ports,
&mut port_listeners,
&quic_conn,
active_streams,
next_stream_id,
&config.edge_id,
connection_token,
bind_address,
);
// Monitor control stream for config updates, and connection health.
// Also handle shutdown signals.
let result = 'quic_loop: loop {
tokio::select! {
// Read control messages from hub
ctrl_msg = quic_transport::read_ctrl_message(&mut ctrl_recv) => {
match ctrl_msg {
Ok(Some((msg_type, payload))) => {
match msg_type {
quic_transport::CTRL_CONFIG => {
if let Ok(update) = serde_json::from_slice::<ConfigUpdate>(&payload) {
log::info!("QUIC config update from hub: ports {:?}", update.listen_ports);
*listen_ports.write().await = update.listen_ports.clone();
let _ = event_tx.try_send(EdgeEvent::PortsUpdated {
listen_ports: update.listen_ports.clone(),
});
apply_port_config_quic(
&update.listen_ports,
&mut port_listeners,
&quic_conn,
active_streams,
next_stream_id,
&config.edge_id,
connection_token,
bind_address,
);
}
}
quic_transport::CTRL_PING => {
// Respond with PONG on control stream
if let Err(e) = quic_transport::write_ctrl_message(
&mut ctrl_send, quic_transport::CTRL_PONG, &[],
).await {
log::error!("Failed to send QUIC PONG: {}", e);
break 'quic_loop EdgeLoopResult::Reconnect(
format!("quic_pong_failed: {}", e),
);
}
}
_ => {
log::warn!("Unknown QUIC control message type: {}", msg_type);
}
}
}
Ok(None) => {
log::info!("Hub closed QUIC control stream (EOF)");
break 'quic_loop EdgeLoopResult::Reconnect("quic_ctrl_eof".to_string());
}
Err(e) => {
log::error!("QUIC control stream read error: {}", e);
break 'quic_loop EdgeLoopResult::Reconnect(
format!("quic_ctrl_error: {}", e),
);
}
}
}
// QUIC connection closed
reason = quic_conn.closed() => {
log::info!("QUIC connection closed: {}", reason);
break 'quic_loop EdgeLoopResult::Reconnect(format!("quic_closed: {}", reason));
}
// Shutdown signal
_ = connection_token.cancelled() => {
if shutdown_rx.try_recv().is_ok() {
break 'quic_loop EdgeLoopResult::Shutdown;
}
break 'quic_loop EdgeLoopResult::Shutdown;
}
}
};
// Cleanup
connection_token.cancel();
stun_handle.abort();
for (_, h) in port_listeners.drain() {
h.abort();
}
// Graceful QUIC close
quic_conn.close(quinn::VarInt::from_u32(0), b"shutdown");
result
}
/// Apply port config for QUIC transport: spawn TCP listeners that open QUIC streams.
fn apply_port_config_quic(
new_ports: &[u16],
port_listeners: &mut HashMap<u16, JoinHandle<()>>,
quic_conn: &quinn::Connection,
active_streams: &Arc<AtomicU32>,
next_stream_id: &Arc<AtomicU32>,
edge_id: &str,
connection_token: &CancellationToken,
bind_address: &str,
) {
let new_set: std::collections::HashSet<u16> = new_ports.iter().copied().collect();
let old_set: std::collections::HashSet<u16> = port_listeners.keys().copied().collect();
// Remove ports no longer needed
for &port in old_set.difference(&new_set) {
if let Some(handle) = port_listeners.remove(&port) {
log::info!("Stopping QUIC listener on port {}", port);
handle.abort();
}
}
// Add new ports
for &port in new_set.difference(&old_set) {
let quic_conn = quic_conn.clone();
let active_streams = active_streams.clone();
let next_stream_id = next_stream_id.clone();
let _edge_id = edge_id.to_string();
let port_token = connection_token.child_token();
let bind_addr = bind_address.to_string();
let handle = tokio::spawn(async move {
let listener = match TcpListener::bind((bind_addr.as_str(), port)).await {
Ok(l) => l,
Err(e) => {
log::error!("Failed to bind port {} (QUIC): {}", port, e);
return;
}
};
log::info!("Listening on port {} (QUIC transport)", port);
loop {
tokio::select! {
accept_result = listener.accept() => {
match accept_result {
Ok((client_stream, client_addr)) => {
let _ = client_stream.set_nodelay(true);
let ka = socket2::TcpKeepalive::new()
.with_time(Duration::from_secs(60));
#[cfg(target_os = "linux")]
let ka = ka.with_interval(Duration::from_secs(60));
let _ = socket2::SockRef::from(&client_stream).set_tcp_keepalive(&ka);
let stream_id = next_stream_id.fetch_add(1, Ordering::Relaxed);
let quic_conn = quic_conn.clone();
let active_streams = active_streams.clone();
let client_token = port_token.child_token();
active_streams.fetch_add(1, Ordering::Relaxed);
tokio::spawn(async move {
handle_client_connection_quic(
client_stream,
client_addr,
stream_id,
port,
quic_conn,
client_token,
).await;
// Saturating decrement
loop {
let current = active_streams.load(Ordering::Relaxed);
if current == 0 { break; }
if active_streams.compare_exchange_weak(
current, current - 1,
Ordering::Relaxed, Ordering::Relaxed,
).is_ok() {
break;
}
}
});
}
Err(e) => {
log::error!("Accept error on port {} (QUIC): {}", port, e);
}
}
}
_ = port_token.cancelled() => {
log::info!("Port {} QUIC listener cancelled", port);
break;
}
}
}
});
port_listeners.insert(port, handle);
}
}
/// Handle a single client connection via QUIC transport.
/// Opens a new QUIC bidirectional stream, sends the PROXY header,
/// then bidirectionally copies data between the client TCP socket and the QUIC stream.
async fn handle_client_connection_quic(
client_stream: TcpStream,
client_addr: std::net::SocketAddr,
stream_id: u32,
dest_port: u16,
quic_conn: quinn::Connection,
client_token: CancellationToken,
) {
let client_ip = client_addr.ip().to_string();
let client_port = client_addr.port();
let edge_ip = "0.0.0.0";
// Open a new QUIC bidirectional stream for this client connection
let (mut quic_send, mut quic_recv) = match quic_conn.open_bi().await {
Ok(s) => s,
Err(e) => {
log::error!("Stream {} failed to open QUIC bi stream: {}", stream_id, e);
return;
}
};
// Send PROXY header as first bytes on the stream
let proxy_header = build_proxy_v1_header(&client_ip, edge_ip, client_port, dest_port);
if let Err(e) = quic_transport::write_proxy_header(&mut quic_send, &proxy_header).await {
log::error!("Stream {} failed to write PROXY header: {}", stream_id, e);
return;
}
let (mut client_read, mut client_write) = client_stream.into_split();
// Task: QUIC -> client (download direction)
let dl_token = client_token.clone();
let mut dl_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)) => {
if client_write.write_all(&buf[..n]).await.is_err() {
break;
}
}
Ok(None) => break, // QUIC stream finished
Err(_) => break,
}
}
_ = dl_token.cancelled() => break,
}
}
let _ = client_write.shutdown().await;
});
// Task: client -> QUIC (upload direction)
let mut buf = vec![0u8; 32768];
loop {
tokio::select! {
read_result = client_read.read(&mut buf) => {
match read_result {
Ok(0) => break, // client EOF
Ok(n) => {
if quic_send.write_all(&buf[..n]).await.is_err() {
break;
}
}
Err(_) => break,
}
}
_ = client_token.cancelled() => break,
}
}
// Wait for download task to finish before closing the QUIC stream
let _ = tokio::time::timeout(Duration::from_secs(300), &mut dl_task).await;
// Gracefully close the QUIC send stream
let _ = quic_send.finish();
dl_task.abort();
}
#[cfg(test)]
mod tests {
use super::*;
@@ -971,6 +1462,7 @@ mod tests {
edge_id: "e1".to_string(),
secret: "sec".to_string(),
bind_address: None,
transport_mode: None,
};
let json = serde_json::to_string(&config).unwrap();
let back: EdgeConfig = serde_json::from_str(&json).unwrap();
@@ -988,6 +1480,44 @@ mod tests {
assert_eq!(hc.stun_interval_secs, 120);
}
#[test]
fn test_edge_config_transport_mode_deserialize() {
let json = r#"{
"hubHost": "hub.test",
"hubPort": 8443,
"edgeId": "e1",
"secret": "s",
"transportMode": "quic"
}"#;
let config: EdgeConfig = serde_json::from_str(json).unwrap();
assert_eq!(config.transport_mode, Some(TransportMode::Quic));
}
#[test]
fn test_edge_config_transport_mode_default() {
let json = r#"{
"hubHost": "hub.test",
"hubPort": 8443,
"edgeId": "e1",
"secret": "s"
}"#;
let config: EdgeConfig = serde_json::from_str(json).unwrap();
assert_eq!(config.transport_mode, None);
}
#[test]
fn test_edge_config_transport_mode_quic_with_fallback() {
let json = r#"{
"hubHost": "hub.test",
"hubPort": 8443,
"edgeId": "e1",
"secret": "s",
"transportMode": "quicWithFallback"
}"#;
let config: EdgeConfig = serde_json::from_str(json).unwrap();
assert_eq!(config.transport_mode, Some(TransportMode::QuicWithFallback));
}
#[test]
fn test_handshake_config_default_stun_interval() {
let json = r#"{"listenPorts": [443]}"#;
@@ -1088,6 +1618,7 @@ mod tests {
edge_id: "test-edge".to_string(),
secret: "test-secret".to_string(),
bind_address: None,
transport_mode: None,
});
let status = edge.get_status().await;
assert!(!status.running);
@@ -1105,6 +1636,7 @@ mod tests {
edge_id: "e".to_string(),
secret: "s".to_string(),
bind_address: None,
transport_mode: None,
});
let rx1 = edge.take_event_rx().await;
assert!(rx1.is_some());
@@ -1120,6 +1652,7 @@ mod tests {
edge_id: "e".to_string(),
secret: "s".to_string(),
bind_address: None,
transport_mode: None,
});
edge.stop().await; // should not panic
let status = edge.get_status().await;

View File

@@ -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::*;

View File

@@ -1,5 +1,6 @@
pub mod hub;
pub mod edge;
pub mod stun;
pub mod transport;
pub use remoteingress_protocol as protocol;

View File

@@ -0,0 +1,22 @@
pub mod quic;
use serde::{Deserialize, Serialize};
/// Transport mode for the tunnel connection between edge and hub.
///
/// - `TcpTls`: TCP + TLS with frame-based multiplexing via TunnelIo (default).
/// - `Quic`: QUIC with native stream multiplexing (one QUIC stream per tunneled connection).
/// - `QuicWithFallback`: Try QUIC first, fall back to TCP+TLS if UDP is blocked.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub enum TransportMode {
TcpTls,
Quic,
QuicWithFallback,
}
impl Default for TransportMode {
fn default() -> Self {
TransportMode::TcpTls
}
}

View File

@@ -0,0 +1,191 @@
use std::sync::Arc;
/// QUIC control stream message types (reuses frame type constants for consistency).
pub const CTRL_CONFIG: u8 = 0x06;
pub const CTRL_PING: u8 = 0x07;
pub const CTRL_PONG: u8 = 0x08;
/// Header size for control stream messages: [type:1][length:4] = 5 bytes.
pub const CTRL_HEADER_SIZE: usize = 5;
/// Build a quinn ClientConfig that skips server certificate verification
/// (auth is via shared secret, same as the TCP+TLS path).
pub fn build_quic_client_config() -> quinn::ClientConfig {
let mut tls_config = rustls::ClientConfig::builder()
.dangerous()
.with_custom_certificate_verifier(Arc::new(NoCertVerifier))
.with_no_client_auth();
// QUIC mandates ALPN negotiation (RFC 9001 §8.1).
// Must match the server's ALPN protocol.
tls_config.alpn_protocols = vec![b"remoteingress".to_vec()];
let quic_config = quinn::crypto::rustls::QuicClientConfig::try_from(tls_config)
.expect("failed to build QUIC client config from rustls config");
let mut transport = quinn::TransportConfig::default();
transport.keep_alive_interval(Some(std::time::Duration::from_secs(15)));
transport.max_idle_timeout(Some(
quinn::IdleTimeout::try_from(std::time::Duration::from_secs(45)).unwrap(),
));
// Match MAX_STREAMS_PER_EDGE (1024) from hub.rs.
// Default is 100 which is too low for high-concurrency tunneling.
transport.max_concurrent_bidi_streams(1024u32.into());
let mut client_config = quinn::ClientConfig::new(Arc::new(quic_config));
client_config.transport_config(Arc::new(transport));
client_config
}
/// Build a quinn ServerConfig from the same TLS server config used for TCP+TLS.
pub fn build_quic_server_config(
tls_server_config: rustls::ServerConfig,
) -> Result<quinn::ServerConfig, Box<dyn std::error::Error + Send + Sync>> {
let quic_config = quinn::crypto::rustls::QuicServerConfig::try_from(tls_server_config)?;
let mut transport = quinn::TransportConfig::default();
transport.keep_alive_interval(Some(std::time::Duration::from_secs(15)));
transport.max_idle_timeout(Some(
quinn::IdleTimeout::try_from(std::time::Duration::from_secs(45)).unwrap(),
));
transport.max_concurrent_bidi_streams(1024u32.into());
let mut server_config = quinn::ServerConfig::with_crypto(Arc::new(quic_config));
server_config.transport_config(Arc::new(transport));
Ok(server_config)
}
/// Write a control message to a QUIC send stream.
/// Format: [type:1][length:4][payload:N]
pub async fn write_ctrl_message(
send: &mut quinn::SendStream,
msg_type: u8,
payload: &[u8],
) -> Result<(), std::io::Error> {
let len = payload.len() as u32;
let mut header = [0u8; CTRL_HEADER_SIZE];
header[0] = msg_type;
header[1..5].copy_from_slice(&len.to_be_bytes());
send.write_all(&header).await?;
if !payload.is_empty() {
send.write_all(payload).await?;
}
Ok(())
}
/// Read a control message from a QUIC recv stream.
/// Returns (msg_type, payload). Returns None on EOF.
pub async fn read_ctrl_message(
recv: &mut quinn::RecvStream,
) -> Result<Option<(u8, Vec<u8>)>, std::io::Error> {
let mut header = [0u8; CTRL_HEADER_SIZE];
match recv.read_exact(&mut header).await {
Ok(()) => {}
Err(e) => {
if let quinn::ReadExactError::FinishedEarly(_) = e {
return Ok(None);
}
return Err(std::io::Error::new(std::io::ErrorKind::Other, e));
}
}
let msg_type = header[0];
let len = u32::from_be_bytes([header[1], header[2], header[3], header[4]]) as usize;
let mut payload = vec![0u8; len];
if len > 0 {
recv.read_exact(&mut payload).await.map_err(|e| {
std::io::Error::new(std::io::ErrorKind::Other, e)
})?;
}
Ok(Some((msg_type, payload)))
}
/// Write the PROXY v1 header as the first bytes on a QUIC data stream.
/// The header is length-prefixed so the receiver knows where it ends and data begins.
/// Format: [header_len:4][proxy_header:N]
pub async fn write_proxy_header(
send: &mut quinn::SendStream,
proxy_header: &str,
) -> Result<(), std::io::Error> {
let header_bytes = proxy_header.as_bytes();
let len = header_bytes.len() as u32;
send.write_all(&len.to_be_bytes()).await?;
send.write_all(header_bytes).await?;
Ok(())
}
/// Read the PROXY v1 header from the first bytes of a QUIC data stream.
/// Returns the header string.
pub async fn read_proxy_header(
recv: &mut quinn::RecvStream,
) -> Result<String, std::io::Error> {
let mut len_buf = [0u8; 4];
recv.read_exact(&mut len_buf).await.map_err(|e| {
std::io::Error::new(std::io::ErrorKind::Other, e)
})?;
let len = u32::from_be_bytes(len_buf) as usize;
if len > 8192 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"proxy header too long",
));
}
let mut header = vec![0u8; len];
recv.read_exact(&mut header).await.map_err(|e| {
std::io::Error::new(std::io::ErrorKind::Other, e)
})?;
String::from_utf8(header).map_err(|_| {
std::io::Error::new(std::io::ErrorKind::InvalidData, "proxy header not UTF-8")
})
}
/// TLS certificate verifier that accepts any certificate (auth is via shared secret).
/// Same as the one in edge.rs but placed here so the QUIC module is self-contained.
#[derive(Debug)]
struct NoCertVerifier;
impl rustls::client::danger::ServerCertVerifier for NoCertVerifier {
fn verify_server_cert(
&self,
_end_entity: &rustls::pki_types::CertificateDer<'_>,
_intermediates: &[rustls::pki_types::CertificateDer<'_>],
_server_name: &rustls::pki_types::ServerName<'_>,
_ocsp_response: &[u8],
_now: rustls::pki_types::UnixTime,
) -> Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
Ok(rustls::client::danger::ServerCertVerified::assertion())
}
fn verify_tls12_signature(
&self,
_message: &[u8],
_cert: &rustls::pki_types::CertificateDer<'_>,
_dss: &rustls::DigitallySignedStruct,
) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
}
fn verify_tls13_signature(
&self,
_message: &[u8],
_cert: &rustls::pki_types::CertificateDer<'_>,
_dss: &rustls::DigitallySignedStruct,
) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
}
fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
vec![
rustls::SignatureScheme::RSA_PKCS1_SHA256,
rustls::SignatureScheme::RSA_PKCS1_SHA384,
rustls::SignatureScheme::RSA_PKCS1_SHA512,
rustls::SignatureScheme::ECDSA_NISTP256_SHA256,
rustls::SignatureScheme::ECDSA_NISTP384_SHA384,
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,
]
}
}