feat(remoteingress-core): add UDP tunneling over QUIC datagrams and expand transport-specific test coverage
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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");
|
||||
|
||||
{
|
||||
|
||||
@@ -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));
|
||||
|
||||
Reference in New Issue
Block a user