feat(remoteingress-core): add UDP tunneling over QUIC datagrams and expand transport-specific test coverage

This commit is contained in:
2026-03-19 12:19:58 +00:00
parent bfa88f8d76
commit 2087567f15
9 changed files with 393 additions and 35 deletions

View File

@@ -1328,8 +1328,24 @@ async fn connect_to_hub_and_run_quic_with_connection(
bind_address,
);
// Monitor control stream for config updates, and connection health.
// Also handle shutdown signals.
// UDP listeners for QUIC transport — uses QUIC datagrams for low-latency forwarding.
let udp_sessions_quic: Arc<Mutex<UdpSessionManager>> =
Arc::new(Mutex::new(UdpSessionManager::new(Duration::from_secs(60))));
let udp_sockets_quic: Arc<Mutex<HashMap<u16, Arc<UdpSocket>>>> =
Arc::new(Mutex::new(HashMap::new()));
let mut udp_listeners_quic: HashMap<u16, JoinHandle<()>> = HashMap::new();
apply_udp_port_config_quic(
&handshake.listen_ports_udp,
&mut udp_listeners_quic,
&quic_conn,
&udp_sessions_quic,
&udp_sockets_quic,
next_stream_id,
connection_token,
bind_address,
);
// Monitor control stream for config updates, connection health, and QUIC datagrams.
let result = 'quic_loop: loop {
tokio::select! {
// Read control messages from hub
@@ -1384,6 +1400,30 @@ async fn connect_to_hub_and_run_quic_with_connection(
}
}
}
// Receive QUIC datagrams (UDP return traffic from hub)
datagram = quic_conn.read_datagram() => {
match datagram {
Ok(data) => {
// Format: [session_id:4][payload:N]
if data.len() >= 4 {
let session_id = u32::from_be_bytes([data[0], data[1], data[2], data[3]]);
let payload = &data[4..];
let mut sessions = udp_sessions_quic.lock().await;
if let Some(session) = sessions.get_by_stream_id(session_id) {
let client_addr = session.client_addr;
let dest_port = session.dest_port;
let sockets = udp_sockets_quic.lock().await;
if let Some(socket) = sockets.get(&dest_port) {
let _ = socket.send_to(payload, client_addr).await;
}
}
}
}
Err(e) => {
log::debug!("QUIC datagram recv error: {}", e);
}
}
}
// QUIC connection closed
reason = quic_conn.closed() => {
log::info!("QUIC connection closed: {}", reason);
@@ -1405,6 +1445,9 @@ async fn connect_to_hub_and_run_quic_with_connection(
for (_, h) in port_listeners.drain() {
h.abort();
}
for (_, h) in udp_listeners_quic.drain() {
h.abort();
}
// Graceful QUIC close
quic_conn.close(quinn::VarInt::from_u32(0), b"shutdown");
@@ -1513,6 +1556,104 @@ fn apply_port_config_quic(
/// 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.
/// Apply UDP port config for QUIC transport: bind UdpSockets that send via QUIC datagrams.
fn apply_udp_port_config_quic(
new_ports: &[u16],
udp_listeners: &mut HashMap<u16, JoinHandle<()>>,
quic_conn: &quinn::Connection,
udp_sessions: &Arc<Mutex<UdpSessionManager>>,
udp_sockets: &Arc<Mutex<HashMap<u16, Arc<UdpSocket>>>>,
next_stream_id: &Arc<AtomicU32>,
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> = udp_listeners.keys().copied().collect();
for &port in old_set.difference(&new_set) {
if let Some(handle) = udp_listeners.remove(&port) {
log::info!("Stopping QUIC UDP listener on port {}", port);
handle.abort();
}
let sockets = udp_sockets.clone();
tokio::spawn(async move { sockets.lock().await.remove(&port); });
}
for &port in new_set.difference(&old_set) {
let quic_conn = quic_conn.clone();
let udp_sessions = udp_sessions.clone();
let udp_sockets = udp_sockets.clone();
let next_stream_id = next_stream_id.clone();
let port_token = connection_token.child_token();
let bind_addr = bind_address.to_string();
let handle = tokio::spawn(async move {
let socket = match UdpSocket::bind((bind_addr.as_str(), port)).await {
Ok(s) => Arc::new(s),
Err(e) => {
log::error!("Failed to bind QUIC UDP port {}: {}", port, e);
return;
}
};
log::info!("Listening on UDP port {} (QUIC datagram transport)", port);
udp_sockets.lock().await.insert(port, socket.clone());
let mut buf = vec![0u8; 65536];
loop {
tokio::select! {
recv_result = socket.recv_from(&mut buf) => {
match recv_result {
Ok((len, client_addr)) => {
let key = UdpSessionKey { client_addr, dest_port: port };
let mut sessions = udp_sessions.lock().await;
let stream_id = if let Some(session) = sessions.get_mut(&key) {
session.stream_id
} else {
// New session — send PROXY v2 header via control-style datagram
let sid = next_stream_id.fetch_add(1, Ordering::Relaxed);
sessions.insert(key, sid);
let client_ip = client_addr.ip().to_string();
let client_port = client_addr.port();
let proxy_header = build_proxy_v2_header_from_str(
&client_ip, "0.0.0.0", client_port, port,
ProxyV2Transport::Udp,
);
// Send OPEN as a QUIC datagram: [session_id:4][0xFF magic:1][proxy_header:28]
let mut open_buf = Vec::with_capacity(4 + 1 + proxy_header.len());
open_buf.extend_from_slice(&sid.to_be_bytes());
open_buf.push(0xFF); // magic byte to distinguish OPEN from DATA
open_buf.extend_from_slice(&proxy_header);
let _ = quic_conn.send_datagram(open_buf.into());
log::debug!("New QUIC UDP session {} from {} -> port {}", sid, client_addr, port);
sid
};
drop(sessions);
// Send datagram: [session_id:4][payload:N]
let mut dgram = Vec::with_capacity(4 + len);
dgram.extend_from_slice(&stream_id.to_be_bytes());
dgram.extend_from_slice(&buf[..len]);
let _ = quic_conn.send_datagram(dgram.into());
}
Err(e) => {
log::error!("QUIC UDP recv error on port {}: {}", port, e);
}
}
}
_ = port_token.cancelled() => {
log::info!("QUIC UDP port {} listener cancelled", port);
break;
}
}
}
});
udp_listeners.insert(port, handle);
}
}
async fn handle_client_connection_quic(
client_stream: TcpStream,
client_addr: std::net::SocketAddr,

View File

@@ -1331,6 +1331,122 @@ async fn handle_edge_connection_quic(
}
});
// UDP sessions for QUIC datagram transport
let quic_udp_sessions: Arc<Mutex<HashMap<u32, mpsc::Sender<Bytes>>>> =
Arc::new(Mutex::new(HashMap::new()));
// Spawn QUIC datagram receiver task
let dgram_conn = quic_conn.clone();
let dgram_sessions = quic_udp_sessions.clone();
let dgram_target = target_host.clone();
let dgram_edge_id = edge_id.clone();
let dgram_token = edge_token.clone();
let dgram_handle = tokio::spawn(async move {
loop {
tokio::select! {
datagram = dgram_conn.read_datagram() => {
match datagram {
Ok(data) => {
if data.len() < 4 { continue; }
let session_id = u32::from_be_bytes([data[0], data[1], data[2], data[3]]);
let payload = &data[4..];
// Check for OPEN magic byte (0xFF)
if !payload.is_empty() && payload[0] == 0xFF {
// This is a session OPEN: [0xFF][proxy_v2_header:28]
let proxy_data = &payload[1..];
let dest_port = if proxy_data.len() >= 28 {
u16::from_be_bytes([proxy_data[26], proxy_data[27]])
} else {
53 // fallback
};
// Create upstream UDP socket
let target = dgram_target.clone();
let conn = dgram_conn.clone();
let sessions = dgram_sessions.clone();
let session_token = dgram_token.child_token();
let (tx, mut rx) = mpsc::channel::<Bytes>(256);
{
let mut s = sessions.lock().await;
s.insert(session_id, tx);
}
tokio::spawn(async move {
let upstream = match UdpSocket::bind("0.0.0.0:0").await {
Ok(s) => Arc::new(s),
Err(e) => {
log::error!("QUIC UDP session {} bind failed: {}", session_id, e);
return;
}
};
if let Err(e) = upstream.connect((target.as_str(), dest_port)).await {
log::error!("QUIC UDP session {} connect failed: {}", session_id, e);
return;
}
// Upstream recv → QUIC datagram back to edge
let upstream_recv = upstream.clone();
let recv_conn = conn.clone();
let recv_token = session_token.clone();
let recv_handle = tokio::spawn(async move {
let mut buf = vec![0u8; 65536];
loop {
tokio::select! {
result = upstream_recv.recv(&mut buf) => {
match result {
Ok(len) => {
let mut dgram = Vec::with_capacity(4 + len);
dgram.extend_from_slice(&session_id.to_be_bytes());
dgram.extend_from_slice(&buf[..len]);
let _ = recv_conn.send_datagram(dgram.into());
}
Err(_) => break,
}
}
_ = recv_token.cancelled() => break,
}
}
});
// Edge datagrams → upstream
loop {
tokio::select! {
data = rx.recv() => {
match data {
Some(datagram) => {
let _ = upstream.send(&datagram).await;
}
None => break,
}
}
_ = session_token.cancelled() => break,
}
}
recv_handle.abort();
});
continue;
}
// Regular data datagram — forward to upstream
let sessions = dgram_sessions.lock().await;
if let Some(tx) = sessions.get(&session_id) {
let _ = tx.try_send(Bytes::copy_from_slice(payload));
}
}
Err(e) => {
log::debug!("QUIC datagram recv error from edge {}: {}", dgram_edge_id, e);
break;
}
}
}
_ = dgram_token.cancelled() => break,
}
}
});
// Control stream loop: forward config updates and handle PONG
let disconnect_reason;
loop {
@@ -1399,6 +1515,7 @@ async fn handle_edge_connection_quic(
// Cleanup
edge_token.cancel();
data_handle.abort();
dgram_handle.abort();
quic_conn.close(quinn::VarInt::from_u32(0), b"hub_shutdown");
{

View File

@@ -31,6 +31,8 @@ pub fn build_quic_client_config() -> quinn::ClientConfig {
// 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());
// Enable QUIC datagrams (RFC 9221) for low-latency UDP tunneling.
transport.datagram_receive_buffer_size(Some(65536));
let mut client_config = quinn::ClientConfig::new(Arc::new(quic_config));
client_config.transport_config(Arc::new(transport));
@@ -49,6 +51,7 @@ pub fn build_quic_server_config(
quinn::IdleTimeout::try_from(std::time::Duration::from_secs(45)).unwrap(),
));
transport.max_concurrent_bidi_streams(1024u32.into());
transport.datagram_receive_buffer_size(Some(65536));
let mut server_config = quinn::ServerConfig::with_crypto(Arc::new(quic_config));
server_config.transport_config(Arc::new(transport));