feat(remoteingress (edge/hub/protocol)): add dynamic port configuration: handshake, FRAME_CONFIG frames, and hot-reloadable listeners
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user