feat(remoteingress-core): add UDP tunneling support between edge and hub

This commit is contained in:
2026-03-19 12:02:41 +00:00
parent 61fa69f108
commit a96b4ba84a
9 changed files with 849 additions and 16 deletions

View File

@@ -3,7 +3,7 @@ use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::net::{TcpListener, TcpStream, UdpSocket};
use tokio::sync::{mpsc, Mutex, Notify, RwLock};
use tokio::task::JoinHandle;
use tokio::time::{Instant, sleep_until};
@@ -15,6 +15,7 @@ use bytes::Bytes;
use remoteingress_protocol::*;
use crate::transport::TransportMode;
use crate::transport::quic as quic_transport;
use crate::udp_session::{UdpSessionKey, UdpSessionManager};
type EdgeTlsStream = tokio_rustls::client::TlsStream<TcpStream>;
@@ -59,6 +60,8 @@ pub struct EdgeConfig {
#[serde(rename_all = "camelCase")]
struct HandshakeConfig {
listen_ports: Vec<u16>,
#[serde(default)]
listen_ports_udp: Vec<u16>,
#[serde(default = "default_stun_interval")]
stun_interval_secs: u64,
}
@@ -72,6 +75,8 @@ fn default_stun_interval() -> u64 {
#[serde(rename_all = "camelCase")]
struct ConfigUpdate {
listen_ports: Vec<u16>,
#[serde(default)]
listen_ports_udp: Vec<u16>,
}
/// Events emitted by the edge.
@@ -344,7 +349,8 @@ enum EdgeLoopResult {
}
/// Process a single frame received from the hub side of the tunnel.
/// Handles FRAME_DATA_BACK, FRAME_WINDOW_UPDATE_BACK, FRAME_CLOSE_BACK, FRAME_CONFIG, FRAME_PING.
/// Handles FRAME_DATA_BACK, FRAME_WINDOW_UPDATE_BACK, FRAME_CLOSE_BACK, FRAME_CONFIG, FRAME_PING,
/// and UDP frames: FRAME_UDP_DATA_BACK, FRAME_UDP_CLOSE.
async fn handle_edge_frame(
frame: Frame,
tunnel_io: &mut remoteingress_protocol::TunnelIo<EdgeTlsStream>,
@@ -355,11 +361,14 @@ async fn handle_edge_frame(
tunnel_data_tx: &mpsc::Sender<Bytes>,
tunnel_sustained_tx: &mpsc::Sender<Bytes>,
port_listeners: &mut HashMap<u16, JoinHandle<()>>,
udp_listeners: &mut HashMap<u16, JoinHandle<()>>,
active_streams: &Arc<AtomicU32>,
next_stream_id: &Arc<AtomicU32>,
edge_id: &str,
connection_token: &CancellationToken,
bind_address: &str,
udp_sessions: &Arc<Mutex<UdpSessionManager>>,
udp_sockets: &Arc<Mutex<HashMap<u16, Arc<UdpSocket>>>>,
) -> EdgeFrameAction {
match frame.frame_type {
FRAME_DATA_BACK => {
@@ -394,7 +403,7 @@ async fn handle_edge_frame(
}
FRAME_CONFIG => {
if let Ok(update) = serde_json::from_slice::<ConfigUpdate>(&frame.payload) {
log::info!("Config update from hub: ports {:?}", update.listen_ports);
log::info!("Config update from hub: ports {:?}, udp {:?}", update.listen_ports, update.listen_ports_udp);
*listen_ports.write().await = update.listen_ports.clone();
let _ = event_tx.try_send(EdgeEvent::PortsUpdated {
listen_ports: update.listen_ports.clone(),
@@ -412,12 +421,39 @@ async fn handle_edge_frame(
connection_token,
bind_address,
);
apply_udp_port_config(
&update.listen_ports_udp,
udp_listeners,
tunnel_writer_tx,
tunnel_data_tx,
udp_sessions,
udp_sockets,
next_stream_id,
connection_token,
bind_address,
);
}
}
FRAME_PING => {
// Queue PONG directly — no channel round-trip, guaranteed delivery
tunnel_io.queue_ctrl(encode_frame(0, FRAME_PONG, &[]));
}
FRAME_UDP_DATA_BACK => {
// Dispatch return UDP datagram to the original client
let mut sessions = udp_sessions.lock().await;
if let Some(session) = sessions.get_by_stream_id(frame.stream_id) {
let client_addr = session.client_addr;
let dest_port = session.dest_port;
let sockets = udp_sockets.lock().await;
if let Some(socket) = sockets.get(&dest_port) {
let _ = socket.send_to(&frame.payload, client_addr).await;
}
}
}
FRAME_UDP_CLOSE => {
let mut sessions = udp_sessions.lock().await;
sessions.remove_by_stream_id(frame.stream_id);
}
_ => {
log::warn!("Unexpected frame type {} from hub", frame.frame_type);
}
@@ -581,6 +617,24 @@ async fn connect_to_hub_and_run(
bind_address,
);
// UDP session manager + listeners
let udp_sessions: Arc<Mutex<UdpSessionManager>> =
Arc::new(Mutex::new(UdpSessionManager::new(Duration::from_secs(60))));
let udp_sockets: Arc<Mutex<HashMap<u16, Arc<UdpSocket>>>> =
Arc::new(Mutex::new(HashMap::new()));
let mut udp_listeners: HashMap<u16, JoinHandle<()>> = HashMap::new();
apply_udp_port_config(
&handshake.listen_ports_udp,
&mut udp_listeners,
&tunnel_ctrl_tx,
&tunnel_data_tx,
&udp_sessions,
&udp_sockets,
next_stream_id,
connection_token,
bind_address,
);
// Single-owner I/O engine — no tokio::io::split, no mutex
let mut tunnel_io = remoteingress_protocol::TunnelIo::new(tls_stream, Vec::new());
@@ -605,7 +659,8 @@ async fn connect_to_hub_and_run(
if let EdgeFrameAction::Disconnect(reason) = handle_edge_frame(
frame, &mut tunnel_io, &client_writers, listen_ports, event_tx,
&tunnel_writer_tx, &tunnel_data_tx, &tunnel_sustained_tx, &mut port_listeners,
active_streams, next_stream_id, &config.edge_id, connection_token, bind_address,
&mut udp_listeners, active_streams, next_stream_id, &config.edge_id,
connection_token, bind_address, &udp_sessions, &udp_sockets,
).await {
break 'io_loop EdgeLoopResult::Reconnect(reason);
}
@@ -623,7 +678,8 @@ async fn connect_to_hub_and_run(
if let EdgeFrameAction::Disconnect(reason) = handle_edge_frame(
frame, &mut tunnel_io, &client_writers, listen_ports, event_tx,
&tunnel_writer_tx, &tunnel_data_tx, &tunnel_sustained_tx, &mut port_listeners,
active_streams, next_stream_id, &config.edge_id, connection_token, bind_address,
&mut udp_listeners, active_streams, next_stream_id, &config.edge_id,
connection_token, bind_address, &udp_sessions, &udp_sockets,
).await {
break EdgeLoopResult::Reconnect(reason);
}
@@ -661,6 +717,9 @@ async fn connect_to_hub_and_run(
for (_, h) in port_listeners.drain() {
h.abort();
}
for (_, h) in udp_listeners.drain() {
h.abort();
}
// Graceful TLS shutdown: send close_notify so the hub sees a clean disconnect.
// Stream handlers are already cancelled, so no new data is being produced.
@@ -790,6 +849,107 @@ fn apply_port_config(
}
}
/// Apply UDP port configuration: bind UdpSockets for added ports, abort removed ports.
fn apply_udp_port_config(
new_ports: &[u16],
udp_listeners: &mut HashMap<u16, JoinHandle<()>>,
tunnel_ctrl_tx: &mpsc::Sender<Bytes>,
tunnel_data_tx: &mpsc::Sender<Bytes>,
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();
// Remove ports no longer needed
for &port in old_set.difference(&new_set) {
if let Some(handle) = udp_listeners.remove(&port) {
log::info!("Stopping UDP listener on port {}", port);
handle.abort();
}
// Remove socket from shared map
let sockets = udp_sockets.clone();
tokio::spawn(async move {
sockets.lock().await.remove(&port);
});
}
// Add new ports
for &port in new_set.difference(&old_set) {
let tunnel_ctrl_tx = tunnel_ctrl_tx.clone();
let tunnel_data_tx = tunnel_data_tx.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 UDP port {}: {}", port, e);
return;
}
};
log::info!("Listening on UDP port {}", port);
// Register socket in shared map for return traffic
udp_sockets.lock().await.insert(port, socket.clone());
let mut buf = vec![0u8; 65536]; // max UDP datagram size
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 — allocate stream_id and send UDP_OPEN
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,
);
let open_frame = encode_frame(sid, FRAME_UDP_OPEN, &proxy_header);
let _ = tunnel_ctrl_tx.try_send(open_frame);
log::debug!("New UDP session {} from {} -> port {}", sid, client_addr, port);
sid
};
drop(sessions); // release lock before sending
// Send datagram through tunnel
let data_frame = encode_frame(stream_id, FRAME_UDP_DATA, &buf[..len]);
let _ = tunnel_data_tx.try_send(data_frame);
}
Err(e) => {
log::error!("UDP recv error on port {}: {}", port, e);
}
}
}
_ = port_token.cancelled() => {
log::info!("UDP port {} listener cancelled", port);
break;
}
}
}
});
udp_listeners.insert(port, handle);
}
}
async fn handle_client_connection(
client_stream: TcpStream,
client_addr: std::net::SocketAddr,