2026-02-16 11:22:23 +00:00
|
|
|
use std::collections::HashMap;
|
|
|
|
|
use std::sync::atomic::{AtomicU32, Ordering};
|
|
|
|
|
use std::sync::Arc;
|
2026-02-18 18:20:53 +00:00
|
|
|
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
|
2026-02-16 11:22:23 +00:00
|
|
|
use tokio::net::{TcpListener, TcpStream};
|
|
|
|
|
use tokio::sync::{mpsc, Mutex, RwLock};
|
2026-02-18 18:20:53 +00:00
|
|
|
use tokio::task::JoinHandle;
|
2026-02-16 11:22:23 +00:00
|
|
|
use tokio_rustls::TlsConnector;
|
|
|
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
|
|
|
|
|
|
use remoteingress_protocol::*;
|
|
|
|
|
|
2026-02-18 18:20:53 +00:00
|
|
|
/// Edge configuration (hub-host + credentials only; ports come from hub).
|
2026-02-16 11:22:23 +00:00
|
|
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
|
|
|
|
#[serde(rename_all = "camelCase")]
|
|
|
|
|
pub struct EdgeConfig {
|
|
|
|
|
pub hub_host: String,
|
|
|
|
|
pub hub_port: u16,
|
|
|
|
|
pub edge_id: String,
|
|
|
|
|
pub secret: String,
|
2026-02-18 18:20:53 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// 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>,
|
2026-02-16 11:22:23 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Events emitted by the edge.
|
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
|
|
|
#[serde(rename_all = "camelCase")]
|
|
|
|
|
#[serde(tag = "type")]
|
|
|
|
|
pub enum EdgeEvent {
|
|
|
|
|
TunnelConnected,
|
|
|
|
|
TunnelDisconnected,
|
|
|
|
|
#[serde(rename_all = "camelCase")]
|
|
|
|
|
PublicIpDiscovered { ip: String },
|
2026-02-18 18:20:53 +00:00
|
|
|
#[serde(rename_all = "camelCase")]
|
|
|
|
|
PortsAssigned { listen_ports: Vec<u16> },
|
|
|
|
|
#[serde(rename_all = "camelCase")]
|
|
|
|
|
PortsUpdated { listen_ports: Vec<u16> },
|
2026-02-16 11:22:23 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Edge status response.
|
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
|
|
|
#[serde(rename_all = "camelCase")]
|
|
|
|
|
pub struct EdgeStatus {
|
|
|
|
|
pub running: bool,
|
|
|
|
|
pub connected: bool,
|
|
|
|
|
pub public_ip: Option<String>,
|
|
|
|
|
pub active_streams: usize,
|
|
|
|
|
pub listen_ports: Vec<u16>,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// The tunnel edge that listens for client connections and multiplexes them to the hub.
|
|
|
|
|
pub struct TunnelEdge {
|
|
|
|
|
config: RwLock<EdgeConfig>,
|
|
|
|
|
event_tx: mpsc::UnboundedSender<EdgeEvent>,
|
|
|
|
|
event_rx: Mutex<Option<mpsc::UnboundedReceiver<EdgeEvent>>>,
|
|
|
|
|
shutdown_tx: Mutex<Option<mpsc::Sender<()>>>,
|
|
|
|
|
running: RwLock<bool>,
|
|
|
|
|
connected: Arc<RwLock<bool>>,
|
|
|
|
|
public_ip: Arc<RwLock<Option<String>>>,
|
|
|
|
|
active_streams: Arc<AtomicU32>,
|
|
|
|
|
next_stream_id: Arc<AtomicU32>,
|
2026-02-18 18:20:53 +00:00
|
|
|
listen_ports: Arc<RwLock<Vec<u16>>>,
|
2026-02-16 11:22:23 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
impl TunnelEdge {
|
|
|
|
|
pub fn new(config: EdgeConfig) -> Self {
|
|
|
|
|
let (event_tx, event_rx) = mpsc::unbounded_channel();
|
|
|
|
|
Self {
|
|
|
|
|
config: RwLock::new(config),
|
|
|
|
|
event_tx,
|
|
|
|
|
event_rx: Mutex::new(Some(event_rx)),
|
|
|
|
|
shutdown_tx: Mutex::new(None),
|
|
|
|
|
running: RwLock::new(false),
|
|
|
|
|
connected: Arc::new(RwLock::new(false)),
|
|
|
|
|
public_ip: Arc::new(RwLock::new(None)),
|
|
|
|
|
active_streams: Arc::new(AtomicU32::new(0)),
|
|
|
|
|
next_stream_id: Arc::new(AtomicU32::new(1)),
|
2026-02-18 18:20:53 +00:00
|
|
|
listen_ports: Arc::new(RwLock::new(Vec::new())),
|
2026-02-16 11:22:23 +00:00
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Take the event receiver (can only be called once).
|
|
|
|
|
pub async fn take_event_rx(&self) -> Option<mpsc::UnboundedReceiver<EdgeEvent>> {
|
|
|
|
|
self.event_rx.lock().await.take()
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Get the current edge status.
|
|
|
|
|
pub async fn get_status(&self) -> EdgeStatus {
|
|
|
|
|
EdgeStatus {
|
|
|
|
|
running: *self.running.read().await,
|
|
|
|
|
connected: *self.connected.read().await,
|
|
|
|
|
public_ip: self.public_ip.read().await.clone(),
|
|
|
|
|
active_streams: self.active_streams.load(Ordering::Relaxed) as usize,
|
2026-02-18 18:20:53 +00:00
|
|
|
listen_ports: self.listen_ports.read().await.clone(),
|
2026-02-16 11:22:23 +00:00
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Start the edge: connect to hub, start listeners.
|
|
|
|
|
pub async fn start(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
|
|
|
|
let config = self.config.read().await.clone();
|
|
|
|
|
let (shutdown_tx, shutdown_rx) = mpsc::channel::<()>(1);
|
|
|
|
|
*self.shutdown_tx.lock().await = Some(shutdown_tx);
|
|
|
|
|
*self.running.write().await = true;
|
|
|
|
|
|
|
|
|
|
let connected = self.connected.clone();
|
|
|
|
|
let public_ip = self.public_ip.clone();
|
|
|
|
|
let active_streams = self.active_streams.clone();
|
|
|
|
|
let next_stream_id = self.next_stream_id.clone();
|
|
|
|
|
let event_tx = self.event_tx.clone();
|
2026-02-18 18:20:53 +00:00
|
|
|
let listen_ports = self.listen_ports.clone();
|
2026-02-16 11:22:23 +00:00
|
|
|
|
|
|
|
|
tokio::spawn(async move {
|
|
|
|
|
edge_main_loop(
|
|
|
|
|
config,
|
|
|
|
|
connected,
|
|
|
|
|
public_ip,
|
|
|
|
|
active_streams,
|
|
|
|
|
next_stream_id,
|
|
|
|
|
event_tx,
|
2026-02-18 18:20:53 +00:00
|
|
|
listen_ports,
|
2026-02-16 11:22:23 +00:00
|
|
|
shutdown_rx,
|
|
|
|
|
)
|
|
|
|
|
.await;
|
|
|
|
|
});
|
|
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Stop the edge.
|
|
|
|
|
pub async fn stop(&self) {
|
|
|
|
|
if let Some(tx) = self.shutdown_tx.lock().await.take() {
|
|
|
|
|
let _ = tx.send(()).await;
|
|
|
|
|
}
|
|
|
|
|
*self.running.write().await = false;
|
|
|
|
|
*self.connected.write().await = false;
|
2026-02-18 18:20:53 +00:00
|
|
|
self.listen_ports.write().await.clear();
|
2026-02-16 11:22:23 +00:00
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
async fn edge_main_loop(
|
|
|
|
|
config: EdgeConfig,
|
|
|
|
|
connected: Arc<RwLock<bool>>,
|
|
|
|
|
public_ip: Arc<RwLock<Option<String>>>,
|
|
|
|
|
active_streams: Arc<AtomicU32>,
|
|
|
|
|
next_stream_id: Arc<AtomicU32>,
|
|
|
|
|
event_tx: mpsc::UnboundedSender<EdgeEvent>,
|
2026-02-18 18:20:53 +00:00
|
|
|
listen_ports: Arc<RwLock<Vec<u16>>>,
|
2026-02-16 11:22:23 +00:00
|
|
|
mut shutdown_rx: mpsc::Receiver<()>,
|
|
|
|
|
) {
|
|
|
|
|
let mut backoff_ms: u64 = 1000;
|
|
|
|
|
let max_backoff_ms: u64 = 30000;
|
|
|
|
|
|
|
|
|
|
loop {
|
|
|
|
|
// Try to connect to hub
|
|
|
|
|
let result = connect_to_hub_and_run(
|
|
|
|
|
&config,
|
|
|
|
|
&connected,
|
|
|
|
|
&public_ip,
|
|
|
|
|
&active_streams,
|
|
|
|
|
&next_stream_id,
|
|
|
|
|
&event_tx,
|
2026-02-18 18:20:53 +00:00
|
|
|
&listen_ports,
|
2026-02-16 11:22:23 +00:00
|
|
|
&mut shutdown_rx,
|
|
|
|
|
)
|
|
|
|
|
.await;
|
|
|
|
|
|
|
|
|
|
*connected.write().await = false;
|
|
|
|
|
let _ = event_tx.send(EdgeEvent::TunnelDisconnected);
|
|
|
|
|
active_streams.store(0, Ordering::Relaxed);
|
2026-02-18 18:20:53 +00:00
|
|
|
listen_ports.write().await.clear();
|
2026-02-16 11:22:23 +00:00
|
|
|
|
|
|
|
|
match result {
|
|
|
|
|
EdgeLoopResult::Shutdown => break,
|
|
|
|
|
EdgeLoopResult::Reconnect => {
|
|
|
|
|
log::info!("Reconnecting in {}ms...", backoff_ms);
|
|
|
|
|
tokio::select! {
|
|
|
|
|
_ = tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)) => {}
|
|
|
|
|
_ = shutdown_rx.recv() => break,
|
|
|
|
|
}
|
|
|
|
|
backoff_ms = (backoff_ms * 2).min(max_backoff_ms);
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
enum EdgeLoopResult {
|
|
|
|
|
Shutdown,
|
|
|
|
|
Reconnect,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
async fn connect_to_hub_and_run(
|
|
|
|
|
config: &EdgeConfig,
|
|
|
|
|
connected: &Arc<RwLock<bool>>,
|
|
|
|
|
public_ip: &Arc<RwLock<Option<String>>>,
|
|
|
|
|
active_streams: &Arc<AtomicU32>,
|
|
|
|
|
next_stream_id: &Arc<AtomicU32>,
|
|
|
|
|
event_tx: &mpsc::UnboundedSender<EdgeEvent>,
|
2026-02-18 18:20:53 +00:00
|
|
|
listen_ports: &Arc<RwLock<Vec<u16>>>,
|
2026-02-16 11:22:23 +00:00
|
|
|
shutdown_rx: &mut mpsc::Receiver<()>,
|
|
|
|
|
) -> EdgeLoopResult {
|
|
|
|
|
// Build TLS connector that skips cert verification (auth is via secret)
|
|
|
|
|
let tls_config = rustls::ClientConfig::builder()
|
|
|
|
|
.dangerous()
|
|
|
|
|
.with_custom_certificate_verifier(Arc::new(NoCertVerifier))
|
|
|
|
|
.with_no_client_auth();
|
|
|
|
|
|
|
|
|
|
let connector = TlsConnector::from(Arc::new(tls_config));
|
|
|
|
|
|
|
|
|
|
let addr = format!("{}:{}", config.hub_host, config.hub_port);
|
|
|
|
|
let tcp = match TcpStream::connect(&addr).await {
|
|
|
|
|
Ok(s) => s,
|
|
|
|
|
Err(e) => {
|
|
|
|
|
log::error!("Failed to connect to hub at {}: {}", addr, e);
|
|
|
|
|
return EdgeLoopResult::Reconnect;
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
let server_name = rustls::pki_types::ServerName::try_from(config.hub_host.clone())
|
|
|
|
|
.unwrap_or_else(|_| rustls::pki_types::ServerName::try_from("remoteingress-hub".to_string()).unwrap());
|
|
|
|
|
|
|
|
|
|
let tls_stream = match connector.connect(server_name, tcp).await {
|
|
|
|
|
Ok(s) => s,
|
|
|
|
|
Err(e) => {
|
|
|
|
|
log::error!("TLS handshake failed: {}", e);
|
|
|
|
|
return EdgeLoopResult::Reconnect;
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
let (read_half, mut write_half) = tokio::io::split(tls_stream);
|
|
|
|
|
|
|
|
|
|
// Send auth line
|
|
|
|
|
let auth_line = format!("EDGE {} {}\n", config.edge_id, config.secret);
|
|
|
|
|
if write_half.write_all(auth_line.as_bytes()).await.is_err() {
|
|
|
|
|
return EdgeLoopResult::Reconnect;
|
|
|
|
|
}
|
|
|
|
|
|
2026-02-18 18:20:53 +00:00
|
|
|
// 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
|
|
|
|
|
);
|
|
|
|
|
|
2026-02-16 11:22:23 +00:00
|
|
|
*connected.write().await = true;
|
|
|
|
|
let _ = event_tx.send(EdgeEvent::TunnelConnected);
|
|
|
|
|
log::info!("Connected to hub at {}", addr);
|
|
|
|
|
|
2026-02-18 18:20:53 +00:00
|
|
|
// 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(),
|
|
|
|
|
});
|
|
|
|
|
|
2026-02-16 11:22:23 +00:00
|
|
|
// Start STUN discovery
|
2026-02-18 18:20:53 +00:00
|
|
|
let stun_interval = handshake.stun_interval_secs;
|
2026-02-16 11:22:23 +00:00
|
|
|
let public_ip_clone = public_ip.clone();
|
|
|
|
|
let event_tx_clone = event_tx.clone();
|
|
|
|
|
let stun_handle = tokio::spawn(async move {
|
|
|
|
|
loop {
|
|
|
|
|
if let Some(ip) = crate::stun::discover_public_ip().await {
|
|
|
|
|
let mut pip = public_ip_clone.write().await;
|
|
|
|
|
let changed = pip.as_ref() != Some(&ip);
|
|
|
|
|
*pip = Some(ip.clone());
|
|
|
|
|
if changed {
|
|
|
|
|
let _ = event_tx_clone.send(EdgeEvent::PublicIpDiscovered { ip });
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
tokio::time::sleep(std::time::Duration::from_secs(stun_interval)).await;
|
|
|
|
|
}
|
|
|
|
|
});
|
|
|
|
|
|
|
|
|
|
// Client socket map: stream_id -> sender for writing data back to client
|
|
|
|
|
let client_writers: Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>> =
|
|
|
|
|
Arc::new(Mutex::new(HashMap::new()));
|
|
|
|
|
|
|
|
|
|
// Shared tunnel writer
|
|
|
|
|
let tunnel_writer = Arc::new(Mutex::new(write_half));
|
|
|
|
|
|
2026-02-18 18:20:53 +00:00
|
|
|
// 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) {
|
2026-02-16 11:22:23 +00:00
|
|
|
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();
|
2026-02-18 18:20:53 +00:00
|
|
|
let edge_id = edge_id.to_string();
|
2026-02-16 11:22:23 +00:00
|
|
|
|
|
|
|
|
let handle = tokio::spawn(async move {
|
|
|
|
|
let listener = match TcpListener::bind(("0.0.0.0", port)).await {
|
|
|
|
|
Ok(l) => l,
|
|
|
|
|
Err(e) => {
|
|
|
|
|
log::error!("Failed to bind port {}: {}", port, e);
|
|
|
|
|
return;
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|
log::info!("Listening on port {}", port);
|
|
|
|
|
|
|
|
|
|
loop {
|
|
|
|
|
match listener.accept().await {
|
|
|
|
|
Ok((client_stream, client_addr)) => {
|
|
|
|
|
let stream_id = next_stream_id.fetch_add(1, Ordering::Relaxed);
|
|
|
|
|
let tunnel_writer = tunnel_writer.clone();
|
|
|
|
|
let client_writers = client_writers.clone();
|
|
|
|
|
let active_streams = active_streams.clone();
|
|
|
|
|
let edge_id = edge_id.clone();
|
|
|
|
|
|
|
|
|
|
active_streams.fetch_add(1, Ordering::Relaxed);
|
|
|
|
|
|
|
|
|
|
tokio::spawn(async move {
|
|
|
|
|
handle_client_connection(
|
|
|
|
|
client_stream,
|
|
|
|
|
client_addr,
|
|
|
|
|
stream_id,
|
|
|
|
|
port,
|
|
|
|
|
&edge_id,
|
|
|
|
|
tunnel_writer,
|
|
|
|
|
client_writers,
|
|
|
|
|
)
|
|
|
|
|
.await;
|
|
|
|
|
active_streams.fetch_sub(1, Ordering::Relaxed);
|
|
|
|
|
});
|
|
|
|
|
}
|
|
|
|
|
Err(e) => {
|
|
|
|
|
log::error!("Accept error on port {}: {}", port, e);
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
});
|
2026-02-18 18:20:53 +00:00
|
|
|
port_listeners.insert(port, handle);
|
2026-02-16 11:22:23 +00:00
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
async fn handle_client_connection(
|
|
|
|
|
client_stream: TcpStream,
|
|
|
|
|
client_addr: std::net::SocketAddr,
|
|
|
|
|
stream_id: u32,
|
|
|
|
|
dest_port: u16,
|
|
|
|
|
edge_id: &str,
|
|
|
|
|
tunnel_writer: Arc<Mutex<tokio::io::WriteHalf<tokio_rustls::client::TlsStream<TcpStream>>>>,
|
|
|
|
|
client_writers: Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>>,
|
|
|
|
|
) {
|
|
|
|
|
let client_ip = client_addr.ip().to_string();
|
|
|
|
|
let client_port = client_addr.port();
|
|
|
|
|
|
|
|
|
|
// Determine edge IP (use 0.0.0.0 as placeholder — hub doesn't use it for routing)
|
|
|
|
|
let edge_ip = "0.0.0.0";
|
|
|
|
|
|
|
|
|
|
// Send OPEN frame with PROXY v1 header
|
|
|
|
|
let proxy_header = build_proxy_v1_header(&client_ip, edge_ip, client_port, dest_port);
|
|
|
|
|
let open_frame = encode_frame(stream_id, FRAME_OPEN, proxy_header.as_bytes());
|
|
|
|
|
{
|
|
|
|
|
let mut w = tunnel_writer.lock().await;
|
|
|
|
|
if w.write_all(&open_frame).await.is_err() {
|
|
|
|
|
return;
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Set up channel for data coming back from hub
|
|
|
|
|
let (back_tx, mut back_rx) = mpsc::channel::<Vec<u8>>(256);
|
|
|
|
|
{
|
|
|
|
|
let mut writers = client_writers.lock().await;
|
|
|
|
|
writers.insert(stream_id, back_tx);
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
let (mut client_read, mut client_write) = client_stream.into_split();
|
|
|
|
|
|
|
|
|
|
// Task: hub -> client
|
|
|
|
|
let hub_to_client = tokio::spawn(async move {
|
|
|
|
|
while let Some(data) = back_rx.recv().await {
|
|
|
|
|
if client_write.write_all(&data).await.is_err() {
|
|
|
|
|
break;
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
let _ = client_write.shutdown().await;
|
|
|
|
|
});
|
|
|
|
|
|
|
|
|
|
// Task: client -> hub
|
|
|
|
|
let mut buf = vec![0u8; 32768];
|
|
|
|
|
loop {
|
|
|
|
|
match client_read.read(&mut buf).await {
|
|
|
|
|
Ok(0) => break,
|
|
|
|
|
Ok(n) => {
|
|
|
|
|
let data_frame = encode_frame(stream_id, FRAME_DATA, &buf[..n]);
|
|
|
|
|
let mut w = tunnel_writer.lock().await;
|
|
|
|
|
if w.write_all(&data_frame).await.is_err() {
|
|
|
|
|
break;
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
Err(_) => break,
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Send CLOSE frame
|
|
|
|
|
let close_frame = encode_frame(stream_id, FRAME_CLOSE, &[]);
|
|
|
|
|
{
|
|
|
|
|
let mut w = tunnel_writer.lock().await;
|
|
|
|
|
let _ = w.write_all(&close_frame).await;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Cleanup
|
|
|
|
|
{
|
|
|
|
|
let mut writers = client_writers.lock().await;
|
|
|
|
|
writers.remove(&stream_id);
|
|
|
|
|
}
|
|
|
|
|
hub_to_client.abort();
|
|
|
|
|
let _ = edge_id; // used for logging context
|
|
|
|
|
}
|
|
|
|
|
|
2026-02-18 18:35:53 +00:00
|
|
|
#[cfg(test)]
|
|
|
|
|
mod tests {
|
|
|
|
|
use super::*;
|
|
|
|
|
|
|
|
|
|
// --- Serde tests ---
|
|
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
|
fn test_edge_config_deserialize_camel_case() {
|
|
|
|
|
let json = r#"{
|
|
|
|
|
"hubHost": "hub.example.com",
|
|
|
|
|
"hubPort": 8443,
|
|
|
|
|
"edgeId": "edge-1",
|
|
|
|
|
"secret": "my-secret"
|
|
|
|
|
}"#;
|
|
|
|
|
let config: EdgeConfig = serde_json::from_str(json).unwrap();
|
|
|
|
|
assert_eq!(config.hub_host, "hub.example.com");
|
|
|
|
|
assert_eq!(config.hub_port, 8443);
|
|
|
|
|
assert_eq!(config.edge_id, "edge-1");
|
|
|
|
|
assert_eq!(config.secret, "my-secret");
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
|
fn test_edge_config_serialize_roundtrip() {
|
|
|
|
|
let config = EdgeConfig {
|
|
|
|
|
hub_host: "host.test".to_string(),
|
|
|
|
|
hub_port: 9999,
|
|
|
|
|
edge_id: "e1".to_string(),
|
|
|
|
|
secret: "sec".to_string(),
|
|
|
|
|
};
|
|
|
|
|
let json = serde_json::to_string(&config).unwrap();
|
|
|
|
|
let back: EdgeConfig = serde_json::from_str(&json).unwrap();
|
|
|
|
|
assert_eq!(back.hub_host, config.hub_host);
|
|
|
|
|
assert_eq!(back.hub_port, config.hub_port);
|
|
|
|
|
assert_eq!(back.edge_id, config.edge_id);
|
|
|
|
|
assert_eq!(back.secret, config.secret);
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
|
fn test_handshake_config_deserialize_all_fields() {
|
|
|
|
|
let json = r#"{"listenPorts": [80, 443], "stunIntervalSecs": 120}"#;
|
|
|
|
|
let hc: HandshakeConfig = serde_json::from_str(json).unwrap();
|
|
|
|
|
assert_eq!(hc.listen_ports, vec![80, 443]);
|
|
|
|
|
assert_eq!(hc.stun_interval_secs, 120);
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
|
fn test_handshake_config_default_stun_interval() {
|
|
|
|
|
let json = r#"{"listenPorts": [443]}"#;
|
|
|
|
|
let hc: HandshakeConfig = serde_json::from_str(json).unwrap();
|
|
|
|
|
assert_eq!(hc.listen_ports, vec![443]);
|
|
|
|
|
assert_eq!(hc.stun_interval_secs, 300);
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
|
fn test_config_update_deserialize() {
|
|
|
|
|
let json = r#"{"listenPorts": [8080, 9090]}"#;
|
|
|
|
|
let update: ConfigUpdate = serde_json::from_str(json).unwrap();
|
|
|
|
|
assert_eq!(update.listen_ports, vec![8080, 9090]);
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
|
fn test_edge_status_serialize() {
|
|
|
|
|
let status = EdgeStatus {
|
|
|
|
|
running: true,
|
|
|
|
|
connected: true,
|
|
|
|
|
public_ip: Some("1.2.3.4".to_string()),
|
|
|
|
|
active_streams: 5,
|
|
|
|
|
listen_ports: vec![443],
|
|
|
|
|
};
|
|
|
|
|
let json = serde_json::to_value(&status).unwrap();
|
|
|
|
|
assert_eq!(json["running"], true);
|
|
|
|
|
assert_eq!(json["connected"], true);
|
|
|
|
|
assert_eq!(json["publicIp"], "1.2.3.4");
|
|
|
|
|
assert_eq!(json["activeStreams"], 5);
|
|
|
|
|
assert_eq!(json["listenPorts"], serde_json::json!([443]));
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
|
fn test_edge_status_serialize_none_ip() {
|
|
|
|
|
let status = EdgeStatus {
|
|
|
|
|
running: false,
|
|
|
|
|
connected: false,
|
|
|
|
|
public_ip: None,
|
|
|
|
|
active_streams: 0,
|
|
|
|
|
listen_ports: vec![],
|
|
|
|
|
};
|
|
|
|
|
let json = serde_json::to_value(&status).unwrap();
|
|
|
|
|
assert!(json["publicIp"].is_null());
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
|
fn test_edge_event_tunnel_connected() {
|
|
|
|
|
let event = EdgeEvent::TunnelConnected;
|
|
|
|
|
let json = serde_json::to_value(&event).unwrap();
|
|
|
|
|
assert_eq!(json["type"], "tunnelConnected");
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
|
fn test_edge_event_tunnel_disconnected() {
|
|
|
|
|
let event = EdgeEvent::TunnelDisconnected;
|
|
|
|
|
let json = serde_json::to_value(&event).unwrap();
|
|
|
|
|
assert_eq!(json["type"], "tunnelDisconnected");
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
|
fn test_edge_event_public_ip_discovered() {
|
|
|
|
|
let event = EdgeEvent::PublicIpDiscovered {
|
|
|
|
|
ip: "203.0.113.1".to_string(),
|
|
|
|
|
};
|
|
|
|
|
let json = serde_json::to_value(&event).unwrap();
|
|
|
|
|
assert_eq!(json["type"], "publicIpDiscovered");
|
|
|
|
|
assert_eq!(json["ip"], "203.0.113.1");
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
|
fn test_edge_event_ports_assigned() {
|
|
|
|
|
let event = EdgeEvent::PortsAssigned {
|
|
|
|
|
listen_ports: vec![443, 8080],
|
|
|
|
|
};
|
|
|
|
|
let json = serde_json::to_value(&event).unwrap();
|
|
|
|
|
assert_eq!(json["type"], "portsAssigned");
|
|
|
|
|
assert_eq!(json["listenPorts"], serde_json::json!([443, 8080]));
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
|
fn test_edge_event_ports_updated() {
|
|
|
|
|
let event = EdgeEvent::PortsUpdated {
|
|
|
|
|
listen_ports: vec![9090],
|
|
|
|
|
};
|
|
|
|
|
let json = serde_json::to_value(&event).unwrap();
|
|
|
|
|
assert_eq!(json["type"], "portsUpdated");
|
|
|
|
|
assert_eq!(json["listenPorts"], serde_json::json!([9090]));
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// --- Async tests ---
|
|
|
|
|
|
|
|
|
|
#[tokio::test]
|
|
|
|
|
async fn test_tunnel_edge_new_get_status() {
|
|
|
|
|
let edge = TunnelEdge::new(EdgeConfig {
|
|
|
|
|
hub_host: "localhost".to_string(),
|
|
|
|
|
hub_port: 8443,
|
|
|
|
|
edge_id: "test-edge".to_string(),
|
|
|
|
|
secret: "test-secret".to_string(),
|
|
|
|
|
});
|
|
|
|
|
let status = edge.get_status().await;
|
|
|
|
|
assert!(!status.running);
|
|
|
|
|
assert!(!status.connected);
|
|
|
|
|
assert!(status.public_ip.is_none());
|
|
|
|
|
assert_eq!(status.active_streams, 0);
|
|
|
|
|
assert!(status.listen_ports.is_empty());
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[tokio::test]
|
|
|
|
|
async fn test_tunnel_edge_take_event_rx() {
|
|
|
|
|
let edge = TunnelEdge::new(EdgeConfig {
|
|
|
|
|
hub_host: "localhost".to_string(),
|
|
|
|
|
hub_port: 8443,
|
|
|
|
|
edge_id: "e".to_string(),
|
|
|
|
|
secret: "s".to_string(),
|
|
|
|
|
});
|
|
|
|
|
let rx1 = edge.take_event_rx().await;
|
|
|
|
|
assert!(rx1.is_some());
|
|
|
|
|
let rx2 = edge.take_event_rx().await;
|
|
|
|
|
assert!(rx2.is_none());
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[tokio::test]
|
|
|
|
|
async fn test_tunnel_edge_stop_without_start() {
|
|
|
|
|
let edge = TunnelEdge::new(EdgeConfig {
|
|
|
|
|
hub_host: "localhost".to_string(),
|
|
|
|
|
hub_port: 8443,
|
|
|
|
|
edge_id: "e".to_string(),
|
|
|
|
|
secret: "s".to_string(),
|
|
|
|
|
});
|
|
|
|
|
edge.stop().await; // should not panic
|
|
|
|
|
let status = edge.get_status().await;
|
|
|
|
|
assert!(!status.running);
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
2026-02-16 11:22:23 +00:00
|
|
|
/// TLS certificate verifier that accepts any certificate (auth is via shared secret).
|
|
|
|
|
#[derive(Debug)]
|
|
|
|
|
struct NoCertVerifier;
|
|
|
|
|
|
|
|
|
|
impl rustls::client::danger::ServerCertVerifier for NoCertVerifier {
|
|
|
|
|
fn verify_server_cert(
|
|
|
|
|
&self,
|
|
|
|
|
_end_entity: &rustls::pki_types::CertificateDer<'_>,
|
|
|
|
|
_intermediates: &[rustls::pki_types::CertificateDer<'_>],
|
|
|
|
|
_server_name: &rustls::pki_types::ServerName<'_>,
|
|
|
|
|
_ocsp_response: &[u8],
|
|
|
|
|
_now: rustls::pki_types::UnixTime,
|
|
|
|
|
) -> Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
|
|
|
|
|
Ok(rustls::client::danger::ServerCertVerified::assertion())
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
fn verify_tls12_signature(
|
|
|
|
|
&self,
|
|
|
|
|
_message: &[u8],
|
|
|
|
|
_cert: &rustls::pki_types::CertificateDer<'_>,
|
|
|
|
|
_dss: &rustls::DigitallySignedStruct,
|
|
|
|
|
) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
|
|
|
|
|
Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
fn verify_tls13_signature(
|
|
|
|
|
&self,
|
|
|
|
|
_message: &[u8],
|
|
|
|
|
_cert: &rustls::pki_types::CertificateDer<'_>,
|
|
|
|
|
_dss: &rustls::DigitallySignedStruct,
|
|
|
|
|
) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
|
|
|
|
|
Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
|
|
|
|
|
vec![
|
|
|
|
|
rustls::SignatureScheme::RSA_PKCS1_SHA256,
|
|
|
|
|
rustls::SignatureScheme::RSA_PKCS1_SHA384,
|
|
|
|
|
rustls::SignatureScheme::RSA_PKCS1_SHA512,
|
|
|
|
|
rustls::SignatureScheme::ECDSA_NISTP256_SHA256,
|
|
|
|
|
rustls::SignatureScheme::ECDSA_NISTP384_SHA384,
|
|
|
|
|
rustls::SignatureScheme::ECDSA_NISTP521_SHA512,
|
|
|
|
|
rustls::SignatureScheme::RSA_PSS_SHA256,
|
|
|
|
|
rustls::SignatureScheme::RSA_PSS_SHA384,
|
|
|
|
|
rustls::SignatureScheme::RSA_PSS_SHA512,
|
|
|
|
|
rustls::SignatureScheme::ED25519,
|
|
|
|
|
rustls::SignatureScheme::ED448,
|
|
|
|
|
]
|
|
|
|
|
}
|
|
|
|
|
}
|