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,
|
||||
|
||||
Reference in New Issue
Block a user