feat(remoteingress-core): add UDP tunneling support between edge and hub
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user