feat(remoteingress (edge/hub/protocol)): add dynamic port configuration: handshake, FRAME_CONFIG frames, and hot-reloadable listeners

This commit is contained in:
2026-02-18 18:20:53 +00:00
parent 072362a8e6
commit d869589663
8 changed files with 314 additions and 72 deletions

View File

@@ -1,15 +1,16 @@
use std::collections::HashMap;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{mpsc, Mutex, RwLock};
use tokio::task::JoinHandle;
use tokio_rustls::TlsConnector;
use serde::{Deserialize, Serialize};
use remoteingress_protocol::*;
/// Edge configuration.
/// Edge configuration (hub-host + credentials only; ports come from hub).
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct EdgeConfig {
@@ -17,8 +18,26 @@ pub struct EdgeConfig {
pub hub_port: u16,
pub edge_id: String,
pub secret: String,
pub listen_ports: Vec<u16>,
pub stun_interval_secs: Option<u64>,
}
/// Handshake config received from hub after authentication.
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct HandshakeConfig {
listen_ports: Vec<u16>,
#[serde(default = "default_stun_interval")]
stun_interval_secs: u64,
}
fn default_stun_interval() -> u64 {
300
}
/// Runtime config update received from hub via FRAME_CONFIG.
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct ConfigUpdate {
listen_ports: Vec<u16>,
}
/// Events emitted by the edge.
@@ -30,6 +49,10 @@ pub enum EdgeEvent {
TunnelDisconnected,
#[serde(rename_all = "camelCase")]
PublicIpDiscovered { ip: String },
#[serde(rename_all = "camelCase")]
PortsAssigned { listen_ports: Vec<u16> },
#[serde(rename_all = "camelCase")]
PortsUpdated { listen_ports: Vec<u16> },
}
/// Edge status response.
@@ -54,6 +77,7 @@ pub struct TunnelEdge {
public_ip: Arc<RwLock<Option<String>>>,
active_streams: Arc<AtomicU32>,
next_stream_id: Arc<AtomicU32>,
listen_ports: Arc<RwLock<Vec<u16>>>,
}
impl TunnelEdge {
@@ -69,6 +93,7 @@ impl TunnelEdge {
public_ip: Arc::new(RwLock::new(None)),
active_streams: Arc::new(AtomicU32::new(0)),
next_stream_id: Arc::new(AtomicU32::new(1)),
listen_ports: Arc::new(RwLock::new(Vec::new())),
}
}
@@ -84,7 +109,7 @@ impl TunnelEdge {
connected: *self.connected.read().await,
public_ip: self.public_ip.read().await.clone(),
active_streams: self.active_streams.load(Ordering::Relaxed) as usize,
listen_ports: self.config.read().await.listen_ports.clone(),
listen_ports: self.listen_ports.read().await.clone(),
}
}
@@ -100,6 +125,7 @@ impl TunnelEdge {
let active_streams = self.active_streams.clone();
let next_stream_id = self.next_stream_id.clone();
let event_tx = self.event_tx.clone();
let listen_ports = self.listen_ports.clone();
tokio::spawn(async move {
edge_main_loop(
@@ -109,6 +135,7 @@ impl TunnelEdge {
active_streams,
next_stream_id,
event_tx,
listen_ports,
shutdown_rx,
)
.await;
@@ -124,6 +151,7 @@ impl TunnelEdge {
}
*self.running.write().await = false;
*self.connected.write().await = false;
self.listen_ports.write().await.clear();
}
}
@@ -134,6 +162,7 @@ async fn edge_main_loop(
active_streams: Arc<AtomicU32>,
next_stream_id: Arc<AtomicU32>,
event_tx: mpsc::UnboundedSender<EdgeEvent>,
listen_ports: Arc<RwLock<Vec<u16>>>,
mut shutdown_rx: mpsc::Receiver<()>,
) {
let mut backoff_ms: u64 = 1000;
@@ -148,6 +177,7 @@ async fn edge_main_loop(
&active_streams,
&next_stream_id,
&event_tx,
&listen_ports,
&mut shutdown_rx,
)
.await;
@@ -155,6 +185,7 @@ async fn edge_main_loop(
*connected.write().await = false;
let _ = event_tx.send(EdgeEvent::TunnelDisconnected);
active_streams.store(0, Ordering::Relaxed);
listen_ports.write().await.clear();
match result {
EdgeLoopResult::Shutdown => break,
@@ -182,6 +213,7 @@ async fn connect_to_hub_and_run(
active_streams: &Arc<AtomicU32>,
next_stream_id: &Arc<AtomicU32>,
event_tx: &mpsc::UnboundedSender<EdgeEvent>,
listen_ports: &Arc<RwLock<Vec<u16>>>,
shutdown_rx: &mut mpsc::Receiver<()>,
) -> EdgeLoopResult {
// Build TLS connector that skips cert verification (auth is via secret)
@@ -220,12 +252,47 @@ async fn connect_to_hub_and_run(
return EdgeLoopResult::Reconnect;
}
// Read handshake response line from hub (JSON with initial config)
let mut buf_reader = BufReader::new(read_half);
let mut handshake_line = String::new();
match buf_reader.read_line(&mut handshake_line).await {
Ok(0) => {
log::error!("Hub rejected connection (EOF before handshake)");
return EdgeLoopResult::Reconnect;
}
Ok(_) => {}
Err(e) => {
log::error!("Failed to read handshake response: {}", e);
return EdgeLoopResult::Reconnect;
}
}
let handshake: HandshakeConfig = match serde_json::from_str(handshake_line.trim()) {
Ok(h) => h,
Err(e) => {
log::error!("Invalid handshake response: {}", e);
return EdgeLoopResult::Reconnect;
}
};
log::info!(
"Handshake from hub: ports {:?}, stun_interval {}s",
handshake.listen_ports,
handshake.stun_interval_secs
);
*connected.write().await = true;
let _ = event_tx.send(EdgeEvent::TunnelConnected);
log::info!("Connected to hub at {}", addr);
// Store initial ports and emit event
*listen_ports.write().await = handshake.listen_ports.clone();
let _ = event_tx.send(EdgeEvent::PortsAssigned {
listen_ports: handshake.listen_ports.clone(),
});
// Start STUN discovery
let stun_interval = config.stun_interval_secs.unwrap_or(300);
let stun_interval = handshake.stun_interval_secs;
let public_ip_clone = public_ip.clone();
let event_tx_clone = event_tx.clone();
let stun_handle = tokio::spawn(async move {
@@ -249,14 +316,112 @@ async fn connect_to_hub_and_run(
// Shared tunnel writer
let tunnel_writer = Arc::new(Mutex::new(write_half));
// Start TCP listeners for each port
let mut listener_handles = Vec::new();
for &port in &config.listen_ports {
// Start TCP listeners for initial ports (hot-reloadable)
let mut port_listeners: HashMap<u16, JoinHandle<()>> = HashMap::new();
apply_port_config(
&handshake.listen_ports,
&mut port_listeners,
&tunnel_writer,
&client_writers,
active_streams,
next_stream_id,
&config.edge_id,
);
// Read frames from hub
let mut frame_reader = FrameReader::new(buf_reader);
let result = loop {
tokio::select! {
frame_result = frame_reader.next_frame() => {
match frame_result {
Ok(Some(frame)) => {
match frame.frame_type {
FRAME_DATA_BACK => {
let writers = client_writers.lock().await;
if let Some(tx) = writers.get(&frame.stream_id) {
let _ = tx.send(frame.payload).await;
}
}
FRAME_CLOSE_BACK => {
let mut writers = client_writers.lock().await;
writers.remove(&frame.stream_id);
}
FRAME_CONFIG => {
if let Ok(update) = serde_json::from_slice::<ConfigUpdate>(&frame.payload) {
log::info!("Config update from hub: ports {:?}", update.listen_ports);
*listen_ports.write().await = update.listen_ports.clone();
let _ = event_tx.send(EdgeEvent::PortsUpdated {
listen_ports: update.listen_ports.clone(),
});
apply_port_config(
&update.listen_ports,
&mut port_listeners,
&tunnel_writer,
&client_writers,
active_streams,
next_stream_id,
&config.edge_id,
);
}
}
_ => {
log::warn!("Unexpected frame type {} from hub", frame.frame_type);
}
}
}
Ok(None) => {
log::info!("Hub disconnected (EOF)");
break EdgeLoopResult::Reconnect;
}
Err(e) => {
log::error!("Hub frame error: {}", e);
break EdgeLoopResult::Reconnect;
}
}
}
_ = shutdown_rx.recv() => {
break EdgeLoopResult::Shutdown;
}
}
};
// Cleanup
stun_handle.abort();
for (_, h) in port_listeners.drain() {
h.abort();
}
result
}
/// Apply a new port configuration: spawn listeners for added ports, abort removed ports.
fn apply_port_config(
new_ports: &[u16],
port_listeners: &mut HashMap<u16, JoinHandle<()>>,
tunnel_writer: &Arc<Mutex<tokio::io::WriteHalf<tokio_rustls::client::TlsStream<TcpStream>>>>,
client_writers: &Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>>,
active_streams: &Arc<AtomicU32>,
next_stream_id: &Arc<AtomicU32>,
edge_id: &str,
) {
let new_set: std::collections::HashSet<u16> = new_ports.iter().copied().collect();
let old_set: std::collections::HashSet<u16> = port_listeners.keys().copied().collect();
// Remove ports no longer needed
for &port in old_set.difference(&new_set) {
if let Some(handle) = port_listeners.remove(&port) {
log::info!("Stopping listener on port {}", port);
handle.abort();
}
}
// Add new ports
for &port in new_set.difference(&old_set) {
let tunnel_writer = tunnel_writer.clone();
let client_writers = client_writers.clone();
let active_streams = active_streams.clone();
let next_stream_id = next_stream_id.clone();
let edge_id = config.edge_id.clone();
let edge_id = edge_id.to_string();
let handle = tokio::spawn(async move {
let listener = match TcpListener::bind(("0.0.0.0", port)).await {
@@ -299,55 +464,8 @@ async fn connect_to_hub_and_run(
}
}
});
listener_handles.push(handle);
port_listeners.insert(port, handle);
}
// Read frames from hub
let mut frame_reader = FrameReader::new(read_half);
let result = loop {
tokio::select! {
frame_result = frame_reader.next_frame() => {
match frame_result {
Ok(Some(frame)) => {
match frame.frame_type {
FRAME_DATA_BACK => {
let writers = client_writers.lock().await;
if let Some(tx) = writers.get(&frame.stream_id) {
let _ = tx.send(frame.payload).await;
}
}
FRAME_CLOSE_BACK => {
let mut writers = client_writers.lock().await;
writers.remove(&frame.stream_id);
}
_ => {
log::warn!("Unexpected frame type {} from hub", frame.frame_type);
}
}
}
Ok(None) => {
log::info!("Hub disconnected (EOF)");
break EdgeLoopResult::Reconnect;
}
Err(e) => {
log::error!("Hub frame error: {}", e);
break EdgeLoopResult::Reconnect;
}
}
}
_ = shutdown_rx.recv() => {
break EdgeLoopResult::Shutdown;
}
}
};
// Cleanup
stun_handle.abort();
for h in listener_handles {
h.abort();
}
result
}
async fn handle_client_connection(