serve.zone/remoteingress
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

BREAKING CHANGE(remoteingress-core): add cancellation tokens and cooperative shutdown; switch event channels to bounded mpsc and improve cleanup

Browse Source
This commit is contained in:
Jürgen Kunz
2026-02-19 08:45:32 +00:00
parent a3af2487b7
commit 4e511b3350
6 changed files with 339 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
changelog.md
View File

@@ -1,5 +1,15 @@
# Changelog # Changelog
## 2026-02-19 - 4.0.0 - BREAKING CHANGE(remoteingress-core)
add cancellation tokens and cooperative shutdown; switch event channels to bounded mpsc and improve cleanup
- Introduce tokio-util::sync::CancellationToken for hub/edge and per-connection/stream cancellation, enabling cooperative shutdown of spawned tasks.
- Replace unbounded mpsc channels with bounded mpsc::channel(1024) and switch from UnboundedSender/Receiver to Sender/Receiver; use try_send where non-blocking sends are appropriate.
- Wire cancellation tokens through edge and hub codepaths: child tokens per connection, per-port, per-stream; cancel tokens in stop() and Drop impls to ensure deterministic task termination and cleanup.
- Reset stream id counters and clear listener state on reconnect; improved error handling around accept/read loops using tokio::select! and cancellation checks.
- Update Cargo.toml and Cargo.lock to add tokio-util (and related futures entries) as dependencies.
- BREAKING: public API/types changed — take_event_rx return types and event_tx/event_rx fields now use bounded mpsc::Sender/mpsc::Receiver instead of the unbounded variants; callers must adapt to the new types and bounded behavior.
## 2026-02-18 - 3.3.0 - feat(readme) ## 2026-02-18 - 3.3.0 - feat(readme)
document dynamic port assignment and runtime port updates; clarify TLS multiplexing, frame format, and handshake sequence document dynamic port assignment and runtime port updates; clarify TLS multiplexing, frame format, and handshake sequence

26
rust/Cargo.lock generated
View File

@@ -234,6 +234,18 @@ version = "1.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
[[package]]
name = "futures-core"
version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d"
[[package]]
name = "futures-sink"
version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893"
[[package]] [[package]]
name = "getrandom" name = "getrandom"
version = "0.2.17" version = "0.2.17"
@@ -528,6 +540,7 @@ dependencies = [
"serde_json", "serde_json",
"tokio", "tokio",
"tokio-rustls", "tokio-rustls",
"tokio-util",
] ]
[[package]] [[package]]
@@ -758,6 +771,19 @@ dependencies = [
"tokio", "tokio",
] ]
[[package]]
name = "tokio-util"
version = "0.7.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098"
dependencies = [
"bytes",
"futures-core",
"futures-sink",
"pin-project-lite",
"tokio",
]
[[package]] [[package]]
name = "unicode-ident" name = "unicode-ident"
version = "1.0.24" version = "1.0.24"

1
rust/crates/remoteingress-core/Cargo.toml
View File

@@ -13,3 +13,4 @@ serde = { version = "1", features = ["derive"] }
serde_json = "1" serde_json = "1"
log = "0.4" log = "0.4"
rustls-pemfile = "2" rustls-pemfile = "2"
tokio-util = "0.7"

179
rust/crates/remoteingress-core/src/edge.rs
View File

@@ -6,6 +6,7 @@ use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{mpsc, Mutex, RwLock}; use tokio::sync::{mpsc, Mutex, RwLock};
use tokio::task::JoinHandle; use tokio::task::JoinHandle;
use tokio_rustls::TlsConnector; use tokio_rustls::TlsConnector;
use tokio_util::sync::CancellationToken;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use remoteingress_protocol::*; use remoteingress_protocol::*;
@@ -69,8 +70,8 @@ pub struct EdgeStatus {
/// The tunnel edge that listens for client connections and multiplexes them to the hub. /// The tunnel edge that listens for client connections and multiplexes them to the hub.
pub struct TunnelEdge { pub struct TunnelEdge {
config: RwLock<EdgeConfig>, config: RwLock<EdgeConfig>,
event_tx: mpsc::UnboundedSender<EdgeEvent>, event_tx: mpsc::Sender<EdgeEvent>,
event_rx: Mutex<Option<mpsc::UnboundedReceiver<EdgeEvent>>>, event_rx: Mutex<Option<mpsc::Receiver<EdgeEvent>>>,
shutdown_tx: Mutex<Option<mpsc::Sender<()>>>, shutdown_tx: Mutex<Option<mpsc::Sender<()>>>,
running: RwLock<bool>, running: RwLock<bool>,
connected: Arc<RwLock<bool>>, connected: Arc<RwLock<bool>>,
@@ -78,11 +79,12 @@ pub struct TunnelEdge {
active_streams: Arc<AtomicU32>, active_streams: Arc<AtomicU32>,
next_stream_id: Arc<AtomicU32>, next_stream_id: Arc<AtomicU32>,
listen_ports: Arc<RwLock<Vec<u16>>>, listen_ports: Arc<RwLock<Vec<u16>>>,
cancel_token: CancellationToken,
} }
impl TunnelEdge { impl TunnelEdge {
pub fn new(config: EdgeConfig) -> Self { pub fn new(config: EdgeConfig) -> Self {
let (event_tx, event_rx) = mpsc::unbounded_channel(); let (event_tx, event_rx) = mpsc::channel(1024);
Self { Self {
config: RwLock::new(config), config: RwLock::new(config),
event_tx, event_tx,
@@ -94,11 +96,12 @@ impl TunnelEdge {
active_streams: Arc::new(AtomicU32::new(0)), active_streams: Arc::new(AtomicU32::new(0)),
next_stream_id: Arc::new(AtomicU32::new(1)), next_stream_id: Arc::new(AtomicU32::new(1)),
listen_ports: Arc::new(RwLock::new(Vec::new())), listen_ports: Arc::new(RwLock::new(Vec::new())),
cancel_token: CancellationToken::new(),
} }
} }
/// Take the event receiver (can only be called once). /// Take the event receiver (can only be called once).
pub async fn take_event_rx(&self) -> Option<mpsc::UnboundedReceiver<EdgeEvent>> { pub async fn take_event_rx(&self) -> Option<mpsc::Receiver<EdgeEvent>> {
self.event_rx.lock().await.take() self.event_rx.lock().await.take()
} }
@@ -126,6 +129,7 @@ impl TunnelEdge {
let next_stream_id = self.next_stream_id.clone(); let next_stream_id = self.next_stream_id.clone();
let event_tx = self.event_tx.clone(); let event_tx = self.event_tx.clone();
let listen_ports = self.listen_ports.clone(); let listen_ports = self.listen_ports.clone();
let cancel_token = self.cancel_token.clone();
tokio::spawn(async move { tokio::spawn(async move {
edge_main_loop( edge_main_loop(
@@ -137,6 +141,7 @@ impl TunnelEdge {
event_tx, event_tx,
listen_ports, listen_ports,
shutdown_rx, shutdown_rx,
cancel_token,
) )
.await; .await;
}); });
@@ -146,6 +151,7 @@ impl TunnelEdge {
/// Stop the edge. /// Stop the edge.
pub async fn stop(&self) { pub async fn stop(&self) {
self.cancel_token.cancel();
if let Some(tx) = self.shutdown_tx.lock().await.take() { if let Some(tx) = self.shutdown_tx.lock().await.take() {
let _ = tx.send(()).await; let _ = tx.send(()).await;
} }
@@ -155,20 +161,30 @@ impl TunnelEdge {
} }
} }
impl Drop for TunnelEdge {
fn drop(&mut self) {
self.cancel_token.cancel();
}
}
async fn edge_main_loop( async fn edge_main_loop(
config: EdgeConfig, config: EdgeConfig,
connected: Arc<RwLock<bool>>, connected: Arc<RwLock<bool>>,
public_ip: Arc<RwLock<Option<String>>>, public_ip: Arc<RwLock<Option<String>>>,
active_streams: Arc<AtomicU32>, active_streams: Arc<AtomicU32>,
next_stream_id: Arc<AtomicU32>, next_stream_id: Arc<AtomicU32>,
event_tx: mpsc::UnboundedSender<EdgeEvent>, event_tx: mpsc::Sender<EdgeEvent>,
listen_ports: Arc<RwLock<Vec<u16>>>, listen_ports: Arc<RwLock<Vec<u16>>>,
mut shutdown_rx: mpsc::Receiver<()>, mut shutdown_rx: mpsc::Receiver<()>,
cancel_token: CancellationToken,
) { ) {
let mut backoff_ms: u64 = 1000; let mut backoff_ms: u64 = 1000;
let max_backoff_ms: u64 = 30000; let max_backoff_ms: u64 = 30000;
loop { loop {
// Create a per-connection child token
let connection_token = cancel_token.child_token();
// Try to connect to hub // Try to connect to hub
let result = connect_to_hub_and_run( let result = connect_to_hub_and_run(
&config, &config,
@@ -179,12 +195,18 @@ async fn edge_main_loop(
&event_tx, &event_tx,
&listen_ports, &listen_ports,
&mut shutdown_rx, &mut shutdown_rx,
&connection_token,
) )
.await; .await;
// Cancel connection token to kill all orphaned tasks from this cycle
connection_token.cancel();
*connected.write().await = false; *connected.write().await = false;
let _ = event_tx.send(EdgeEvent::TunnelDisconnected); let _ = event_tx.try_send(EdgeEvent::TunnelDisconnected);
active_streams.store(0, Ordering::Relaxed); active_streams.store(0, Ordering::Relaxed);
// Reset stream ID counter for next connection cycle
next_stream_id.store(1, Ordering::Relaxed);
listen_ports.write().await.clear(); listen_ports.write().await.clear();
match result { match result {
@@ -193,6 +215,7 @@ async fn edge_main_loop(
log::info!("Reconnecting in {}ms...", backoff_ms); log::info!("Reconnecting in {}ms...", backoff_ms);
tokio::select! { tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)) => {} _ = tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)) => {}
_ = cancel_token.cancelled() => break,
_ = shutdown_rx.recv() => break, _ = shutdown_rx.recv() => break,
} }
backoff_ms = (backoff_ms * 2).min(max_backoff_ms); backoff_ms = (backoff_ms * 2).min(max_backoff_ms);
@@ -212,9 +235,10 @@ async fn connect_to_hub_and_run(
public_ip: &Arc<RwLock<Option<String>>>, public_ip: &Arc<RwLock<Option<String>>>,
active_streams: &Arc<AtomicU32>, active_streams: &Arc<AtomicU32>,
next_stream_id: &Arc<AtomicU32>, next_stream_id: &Arc<AtomicU32>,
event_tx: &mpsc::UnboundedSender<EdgeEvent>, event_tx: &mpsc::Sender<EdgeEvent>,
listen_ports: &Arc<RwLock<Vec<u16>>>, listen_ports: &Arc<RwLock<Vec<u16>>>,
shutdown_rx: &mut mpsc::Receiver<()>, shutdown_rx: &mut mpsc::Receiver<()>,
connection_token: &CancellationToken,
) -> EdgeLoopResult { ) -> EdgeLoopResult {
// Build TLS connector that skips cert verification (auth is via secret) // Build TLS connector that skips cert verification (auth is via secret)
let tls_config = rustls::ClientConfig::builder() let tls_config = rustls::ClientConfig::builder()
@@ -282,12 +306,12 @@ async fn connect_to_hub_and_run(
); );
*connected.write().await = true; *connected.write().await = true;
let _ = event_tx.send(EdgeEvent::TunnelConnected); let _ = event_tx.try_send(EdgeEvent::TunnelConnected);
log::info!("Connected to hub at {}", addr); log::info!("Connected to hub at {}", addr);
// Store initial ports and emit event // Store initial ports and emit event
*listen_ports.write().await = handshake.listen_ports.clone(); *listen_ports.write().await = handshake.listen_ports.clone();
let _ = event_tx.send(EdgeEvent::PortsAssigned { let _ = event_tx.try_send(EdgeEvent::PortsAssigned {
listen_ports: handshake.listen_ports.clone(), listen_ports: handshake.listen_ports.clone(),
}); });
@@ -295,17 +319,26 @@ async fn connect_to_hub_and_run(
let stun_interval = handshake.stun_interval_secs; let stun_interval = handshake.stun_interval_secs;
let public_ip_clone = public_ip.clone(); let public_ip_clone = public_ip.clone();
let event_tx_clone = event_tx.clone(); let event_tx_clone = event_tx.clone();
let stun_token = connection_token.clone();
let stun_handle = tokio::spawn(async move { let stun_handle = tokio::spawn(async move {
loop { loop {
if let Some(ip) = crate::stun::discover_public_ip().await { tokio::select! {
let mut pip = public_ip_clone.write().await; ip_result = crate::stun::discover_public_ip() => {
let changed = pip.as_ref() != Some(&ip); if let Some(ip) = ip_result {
*pip = Some(ip.clone()); let mut pip = public_ip_clone.write().await;
if changed { let changed = pip.as_ref() != Some(&ip);
let _ = event_tx_clone.send(EdgeEvent::PublicIpDiscovered { ip }); *pip = Some(ip.clone());
if changed {
let _ = event_tx_clone.try_send(EdgeEvent::PublicIpDiscovered { ip });
}
}
} }
_ = stun_token.cancelled() => break,
}
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_secs(stun_interval)) => {}
_ = stun_token.cancelled() => break,
} }
tokio::time::sleep(std::time::Duration::from_secs(stun_interval)).await;
} }
}); });
@@ -326,6 +359,7 @@ async fn connect_to_hub_and_run(
active_streams, active_streams,
next_stream_id, next_stream_id,
&config.edge_id, &config.edge_id,
connection_token,
); );
// Read frames from hub // Read frames from hub
@@ -350,7 +384,7 @@ async fn connect_to_hub_and_run(
if let Ok(update) = serde_json::from_slice::<ConfigUpdate>(&frame.payload) { 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 {:?}", update.listen_ports);
*listen_ports.write().await = update.listen_ports.clone(); *listen_ports.write().await = update.listen_ports.clone();
let _ = event_tx.send(EdgeEvent::PortsUpdated { let _ = event_tx.try_send(EdgeEvent::PortsUpdated {
listen_ports: update.listen_ports.clone(), listen_ports: update.listen_ports.clone(),
}); });
apply_port_config( apply_port_config(
@@ -361,6 +395,7 @@ async fn connect_to_hub_and_run(
active_streams, active_streams,
next_stream_id, next_stream_id,
&config.edge_id, &config.edge_id,
connection_token,
); );
} }
} }
@@ -379,13 +414,18 @@ async fn connect_to_hub_and_run(
} }
} }
} }
_ = connection_token.cancelled() => {
log::info!("Connection cancelled");
break EdgeLoopResult::Shutdown;
}
_ = shutdown_rx.recv() => { _ = shutdown_rx.recv() => {
break EdgeLoopResult::Shutdown; break EdgeLoopResult::Shutdown;
} }
} }
}; };
// Cleanup // Cancel connection token to propagate to all child tasks BEFORE aborting
connection_token.cancel();
stun_handle.abort(); stun_handle.abort();
for (_, h) in port_listeners.drain() { for (_, h) in port_listeners.drain() {
h.abort(); h.abort();
@@ -403,6 +443,7 @@ fn apply_port_config(
active_streams: &Arc<AtomicU32>, active_streams: &Arc<AtomicU32>,
next_stream_id: &Arc<AtomicU32>, next_stream_id: &Arc<AtomicU32>,
edge_id: &str, edge_id: &str,
connection_token: &CancellationToken,
) { ) {
let new_set: std::collections::HashSet<u16> = new_ports.iter().copied().collect(); let new_set: std::collections::HashSet<u16> = new_ports.iter().copied().collect();
let old_set: std::collections::HashSet<u16> = port_listeners.keys().copied().collect(); let old_set: std::collections::HashSet<u16> = port_listeners.keys().copied().collect();
@@ -422,6 +463,7 @@ fn apply_port_config(
let active_streams = active_streams.clone(); let active_streams = active_streams.clone();
let next_stream_id = next_stream_id.clone(); let next_stream_id = next_stream_id.clone();
let edge_id = edge_id.to_string(); let edge_id = edge_id.to_string();
let port_token = connection_token.child_token();
let handle = tokio::spawn(async move { let handle = tokio::spawn(async move {
let listener = match TcpListener::bind(("0.0.0.0", port)).await { let listener = match TcpListener::bind(("0.0.0.0", port)).await {
@@ -434,32 +476,42 @@ fn apply_port_config(
log::info!("Listening on port {}", port); log::info!("Listening on port {}", port);
loop { loop {
match listener.accept().await { tokio::select! {
Ok((client_stream, client_addr)) => { accept_result = listener.accept() => {
let stream_id = next_stream_id.fetch_add(1, Ordering::Relaxed); match accept_result {
let tunnel_writer = tunnel_writer.clone(); Ok((client_stream, client_addr)) => {
let client_writers = client_writers.clone(); let stream_id = next_stream_id.fetch_add(1, Ordering::Relaxed);
let active_streams = active_streams.clone(); let tunnel_writer = tunnel_writer.clone();
let edge_id = edge_id.clone(); let client_writers = client_writers.clone();
let active_streams = active_streams.clone();
let edge_id = edge_id.clone();
let client_token = port_token.child_token();
active_streams.fetch_add(1, Ordering::Relaxed); active_streams.fetch_add(1, Ordering::Relaxed);
tokio::spawn(async move { tokio::spawn(async move {
handle_client_connection( handle_client_connection(
client_stream, client_stream,
client_addr, client_addr,
stream_id, stream_id,
port, port,
&edge_id, &edge_id,
tunnel_writer, tunnel_writer,
client_writers, client_writers,
) client_token,
.await; )
active_streams.fetch_sub(1, Ordering::Relaxed); .await;
}); active_streams.fetch_sub(1, Ordering::Relaxed);
});
}
Err(e) => {
log::error!("Accept error on port {}: {}", port, e);
}
}
} }
Err(e) => { _ = port_token.cancelled() => {
log::error!("Accept error on port {}: {}", port, e); log::info!("Port {} listener cancelled", port);
break;
} }
} }
} }
@@ -476,6 +528,7 @@ async fn handle_client_connection(
edge_id: &str, edge_id: &str,
tunnel_writer: Arc<Mutex<tokio::io::WriteHalf<tokio_rustls::client::TlsStream<TcpStream>>>>, tunnel_writer: Arc<Mutex<tokio::io::WriteHalf<tokio_rustls::client::TlsStream<TcpStream>>>>,
client_writers: Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>>, client_writers: Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>>,
client_token: CancellationToken,
) { ) {
let client_ip = client_addr.ip().to_string(); let client_ip = client_addr.ip().to_string();
let client_port = client_addr.port(); let client_port = client_addr.port();
@@ -503,10 +556,21 @@ async fn handle_client_connection(
let (mut client_read, mut client_write) = client_stream.into_split(); let (mut client_read, mut client_write) = client_stream.into_split();
// Task: hub -> client // Task: hub -> client
let hub_to_client_token = client_token.clone();
let hub_to_client = tokio::spawn(async move { let hub_to_client = tokio::spawn(async move {
while let Some(data) = back_rx.recv().await { loop {
if client_write.write_all(&data).await.is_err() { tokio::select! {
break; data = back_rx.recv() => {
match data {
Some(data) => {
if client_write.write_all(&data).await.is_err() {
break;
}
}
None => break,
}
}
_ = hub_to_client_token.cancelled() => break,
} }
} }
let _ = client_write.shutdown().await; let _ = client_write.shutdown().await;
@@ -515,22 +579,27 @@ async fn handle_client_connection(
// Task: client -> hub // Task: client -> hub
let mut buf = vec![0u8; 32768]; let mut buf = vec![0u8; 32768];
loop { loop {
match client_read.read(&mut buf).await { tokio::select! {
Ok(0) => break, read_result = client_read.read(&mut buf) => {
Ok(n) => { match read_result {
let data_frame = encode_frame(stream_id, FRAME_DATA, &buf[..n]); Ok(0) => break,
let mut w = tunnel_writer.lock().await; Ok(n) => {
if w.write_all(&data_frame).await.is_err() { let data_frame = encode_frame(stream_id, FRAME_DATA, &buf[..n]);
break; let mut w = tunnel_writer.lock().await;
if w.write_all(&data_frame).await.is_err() {
break;
}
}
Err(_) => break,
} }
} }
Err(_) => break, _ = client_token.cancelled() => break,
} }
} }
// Send CLOSE frame // Send CLOSE frame (only if not cancelled)
let close_frame = encode_frame(stream_id, FRAME_CLOSE, &[]); if !client_token.is_cancelled() {
{ let close_frame = encode_frame(stream_id, FRAME_CLOSE, &[]);
let mut w = tunnel_writer.lock().await; let mut w = tunnel_writer.lock().await;
let _ = w.write_all(&close_frame).await; let _ = w.write_all(&close_frame).await;
} }

293
rust/crates/remoteingress-core/src/hub.rs
View File

@@ -4,6 +4,7 @@ use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::net::{TcpListener, TcpStream}; use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{mpsc, Mutex, RwLock}; use tokio::sync::{mpsc, Mutex, RwLock};
use tokio_rustls::TlsAcceptor; use tokio_rustls::TlsAcceptor;
use tokio_util::sync::CancellationToken;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use remoteingress_protocol::*; use remoteingress_protocol::*;
@@ -95,21 +96,24 @@ pub struct TunnelHub {
config: RwLock<HubConfig>, config: RwLock<HubConfig>,
allowed_edges: Arc<RwLock<HashMap<String, AllowedEdge>>>, allowed_edges: Arc<RwLock<HashMap<String, AllowedEdge>>>,
connected_edges: Arc<Mutex<HashMap<String, ConnectedEdgeInfo>>>, connected_edges: Arc<Mutex<HashMap<String, ConnectedEdgeInfo>>>,
event_tx: mpsc::UnboundedSender<HubEvent>, event_tx: mpsc::Sender<HubEvent>,
event_rx: Mutex<Option<mpsc::UnboundedReceiver<HubEvent>>>, event_rx: Mutex<Option<mpsc::Receiver<HubEvent>>>,
shutdown_tx: Mutex<Option<mpsc::Sender<()>>>, shutdown_tx: Mutex<Option<mpsc::Sender<()>>>,
running: RwLock<bool>, running: RwLock<bool>,
cancel_token: CancellationToken,
} }
struct ConnectedEdgeInfo { struct ConnectedEdgeInfo {
connected_at: u64, connected_at: u64,
active_streams: Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>>, active_streams: Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>>,
config_tx: mpsc::Sender<EdgeConfigUpdate>, config_tx: mpsc::Sender<EdgeConfigUpdate>,
#[allow(dead_code)] // kept alive for Drop — cancels child tokens when edge is removed
cancel_token: CancellationToken,
} }
impl TunnelHub { impl TunnelHub {
pub fn new(config: HubConfig) -> Self { pub fn new(config: HubConfig) -> Self {
let (event_tx, event_rx) = mpsc::unbounded_channel(); let (event_tx, event_rx) = mpsc::channel(1024);
Self { Self {
config: RwLock::new(config), config: RwLock::new(config),
allowed_edges: Arc::new(RwLock::new(HashMap::new())), allowed_edges: Arc::new(RwLock::new(HashMap::new())),
@@ -118,11 +122,12 @@ impl TunnelHub {
event_rx: Mutex::new(Some(event_rx)), event_rx: Mutex::new(Some(event_rx)),
shutdown_tx: Mutex::new(None), shutdown_tx: Mutex::new(None),
running: RwLock::new(false), running: RwLock::new(false),
cancel_token: CancellationToken::new(),
} }
} }
/// Take the event receiver (can only be called once). /// Take the event receiver (can only be called once).
pub async fn take_event_rx(&self) -> Option<mpsc::UnboundedReceiver<HubEvent>> { pub async fn take_event_rx(&self) -> Option<mpsc::Receiver<HubEvent>> {
self.event_rx.lock().await.take() self.event_rx.lock().await.take()
} }
@@ -198,6 +203,7 @@ impl TunnelHub {
let connected = self.connected_edges.clone(); let connected = self.connected_edges.clone();
let event_tx = self.event_tx.clone(); let event_tx = self.event_tx.clone();
let target_host = config.target_host.unwrap_or_else(|| "127.0.0.1".to_string()); let target_host = config.target_host.unwrap_or_else(|| "127.0.0.1".to_string());
let hub_token = self.cancel_token.clone();
tokio::spawn(async move { tokio::spawn(async move {
loop { loop {
@@ -211,9 +217,10 @@ impl TunnelHub {
let connected = connected.clone(); let connected = connected.clone();
let event_tx = event_tx.clone(); let event_tx = event_tx.clone();
let target = target_host.clone(); let target = target_host.clone();
let edge_token = hub_token.child_token();
tokio::spawn(async move { tokio::spawn(async move {
if let Err(e) = handle_edge_connection( if let Err(e) = handle_edge_connection(
stream, acceptor, allowed, connected, event_tx, target, stream, acceptor, allowed, connected, event_tx, target, edge_token,
).await { ).await {
log::error!("Edge connection error: {}", e); log::error!("Edge connection error: {}", e);
} }
@@ -224,6 +231,10 @@ impl TunnelHub {
} }
} }
} }
_ = hub_token.cancelled() => {
log::info!("Hub shutting down (token cancelled)");
break;
}
_ = shutdown_rx.recv() => { _ = shutdown_rx.recv() => {
log::info!("Hub shutting down"); log::info!("Hub shutting down");
break; break;
@@ -237,6 +248,7 @@ impl TunnelHub {
/// Stop the hub. /// Stop the hub.
pub async fn stop(&self) { pub async fn stop(&self) {
self.cancel_token.cancel();
if let Some(tx) = self.shutdown_tx.lock().await.take() { if let Some(tx) = self.shutdown_tx.lock().await.take() {
let _ = tx.send(()).await; let _ = tx.send(()).await;
} }
@@ -246,14 +258,21 @@ impl TunnelHub {
} }
} }
impl Drop for TunnelHub {
fn drop(&mut self) {
self.cancel_token.cancel();
}
}
/// Handle a single edge connection: authenticate, then enter frame loop. /// Handle a single edge connection: authenticate, then enter frame loop.
async fn handle_edge_connection( async fn handle_edge_connection(
stream: TcpStream, stream: TcpStream,
acceptor: TlsAcceptor, acceptor: TlsAcceptor,
allowed: Arc<RwLock<HashMap<String, AllowedEdge>>>, allowed: Arc<RwLock<HashMap<String, AllowedEdge>>>,
connected: Arc<Mutex<HashMap<String, ConnectedEdgeInfo>>>, connected: Arc<Mutex<HashMap<String, ConnectedEdgeInfo>>>,
event_tx: mpsc::UnboundedSender<HubEvent>, event_tx: mpsc::Sender<HubEvent>,
target_host: String, target_host: String,
edge_token: CancellationToken,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> { ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let tls_stream = acceptor.accept(stream).await?; let tls_stream = acceptor.accept(stream).await?;
let (read_half, mut write_half) = tokio::io::split(tls_stream); let (read_half, mut write_half) = tokio::io::split(tls_stream);
@@ -289,7 +308,7 @@ async fn handle_edge_connection(
}; };
log::info!("Edge {} authenticated", edge_id); log::info!("Edge {} authenticated", edge_id);
let _ = event_tx.send(HubEvent::EdgeConnected { let _ = event_tx.try_send(HubEvent::EdgeConnected {
edge_id: edge_id.clone(), edge_id: edge_id.clone(),
}); });
@@ -321,6 +340,7 @@ async fn handle_edge_connection(
connected_at: now, connected_at: now,
active_streams: streams.clone(), active_streams: streams.clone(),
config_tx, config_tx,
cancel_token: edge_token.clone(),
}, },
); );
} }
@@ -331,16 +351,27 @@ async fn handle_edge_connection(
// Spawn task to forward config updates as FRAME_CONFIG frames // Spawn task to forward config updates as FRAME_CONFIG frames
let config_writer = write_half.clone(); let config_writer = write_half.clone();
let config_edge_id = edge_id.clone(); let config_edge_id = edge_id.clone();
let config_token = edge_token.clone();
let config_handle = tokio::spawn(async move { let config_handle = tokio::spawn(async move {
while let Some(update) = config_rx.recv().await { loop {
if let Ok(payload) = serde_json::to_vec(&update) { tokio::select! {
let frame = encode_frame(0, FRAME_CONFIG, &payload); update = config_rx.recv() => {
let mut w = config_writer.lock().await; match update {
if w.write_all(&frame).await.is_err() { Some(update) => {
log::error!("Failed to send config update to edge {}", config_edge_id); if let Ok(payload) = serde_json::to_vec(&update) {
break; let frame = encode_frame(0, FRAME_CONFIG, &payload);
let mut w = config_writer.lock().await;
if w.write_all(&frame).await.is_err() {
log::error!("Failed to send config update to edge {}", config_edge_id);
break;
}
log::info!("Sent config update to edge {}: ports {:?}", config_edge_id, update.listen_ports);
}
}
None => break,
}
} }
log::info!("Sent config update to edge {}: ports {:?}", config_edge_id, update.listen_ports); _ = config_token.cancelled() => break,
} }
} }
}); });
@@ -349,134 +380,164 @@ async fn handle_edge_connection(
let mut frame_reader = FrameReader::new(buf_reader); let mut frame_reader = FrameReader::new(buf_reader);
loop { loop {
match frame_reader.next_frame().await { tokio::select! {
Ok(Some(frame)) => { frame_result = frame_reader.next_frame() => {
match frame.frame_type { match frame_result {
FRAME_OPEN => { Ok(Some(frame)) => {
// Payload is PROXY v1 header line match frame.frame_type {
let proxy_header = String::from_utf8_lossy(&frame.payload).to_string(); FRAME_OPEN => {
// Payload is PROXY v1 header line
let proxy_header = String::from_utf8_lossy(&frame.payload).to_string();
// Parse destination port from PROXY header // Parse destination port from PROXY header
let dest_port = parse_dest_port_from_proxy(&proxy_header).unwrap_or(443); let dest_port = parse_dest_port_from_proxy(&proxy_header).unwrap_or(443);
let stream_id = frame.stream_id; let stream_id = frame.stream_id;
let edge_id_clone = edge_id.clone(); let edge_id_clone = edge_id.clone();
let event_tx_clone = event_tx.clone(); let event_tx_clone = event_tx.clone();
let streams_clone = streams.clone(); let streams_clone = streams.clone();
let writer_clone = write_half.clone(); let writer_clone = write_half.clone();
let target = target_host.clone(); let target = target_host.clone();
let stream_token = edge_token.child_token();
let _ = event_tx.send(HubEvent::StreamOpened { let _ = event_tx.try_send(HubEvent::StreamOpened {
edge_id: edge_id.clone(), edge_id: edge_id.clone(),
stream_id, stream_id,
});
// Create channel for data from edge to this stream
let (data_tx, mut data_rx) = mpsc::channel::<Vec<u8>>(256);
{
let mut s = streams.lock().await;
s.insert(stream_id, data_tx);
}
// Spawn task: connect to SmartProxy, send PROXY header, pipe data
tokio::spawn(async move {
let result = async {
let mut upstream =
TcpStream::connect((target.as_str(), dest_port)).await?;
upstream.write_all(proxy_header.as_bytes()).await?;
let (mut up_read, mut up_write) =
upstream.into_split();
// Forward data from edge (via channel) to SmartProxy
let writer_for_edge_data = tokio::spawn(async move {
while let Some(data) = data_rx.recv().await {
if up_write.write_all(&data).await.is_err() {
break;
}
}
let _ = up_write.shutdown().await;
}); });
// Forward data from SmartProxy back to edge // Create channel for data from edge to this stream
let mut buf = vec![0u8; 32768]; let (data_tx, mut data_rx) = mpsc::channel::<Vec<u8>>(256);
loop { {
match up_read.read(&mut buf).await { let mut s = streams.lock().await;
Ok(0) => break, s.insert(stream_id, data_tx);
Ok(n) => {
let frame =
encode_frame(stream_id, FRAME_DATA_BACK, &buf[..n]);
let mut w = writer_clone.lock().await;
if w.write_all(&frame).await.is_err() {
break;
}
}
Err(_) => break,
}
} }
// Send CLOSE_BACK to edge // Spawn task: connect to SmartProxy, send PROXY header, pipe data
let close_frame = encode_frame(stream_id, FRAME_CLOSE_BACK, &[]); tokio::spawn(async move {
let mut w = writer_clone.lock().await; let result = async {
let _ = w.write_all(&close_frame).await; let mut upstream =
TcpStream::connect((target.as_str(), dest_port)).await?;
upstream.write_all(proxy_header.as_bytes()).await?;
writer_for_edge_data.abort(); let (mut up_read, mut up_write) =
Ok::<(), Box<dyn std::error::Error + Send + Sync>>(()) upstream.into_split();
}
.await;
if let Err(e) = result { // Forward data from edge (via channel) to SmartProxy
log::error!("Stream {} error: {}", stream_id, e); let writer_token = stream_token.clone();
// Send CLOSE_BACK on error let writer_for_edge_data = tokio::spawn(async move {
let close_frame = encode_frame(stream_id, FRAME_CLOSE_BACK, &[]); loop {
let mut w = writer_clone.lock().await; tokio::select! {
let _ = w.write_all(&close_frame).await; data = data_rx.recv() => {
} match data {
Some(data) => {
if up_write.write_all(&data).await.is_err() {
break;
}
}
None => break,
}
}
_ = writer_token.cancelled() => break,
}
}
let _ = up_write.shutdown().await;
});
// Clean up stream // Forward data from SmartProxy back to edge
{ let mut buf = vec![0u8; 32768];
let mut s = streams_clone.lock().await; loop {
s.remove(&stream_id); tokio::select! {
read_result = up_read.read(&mut buf) => {
match read_result {
Ok(0) => break,
Ok(n) => {
let frame =
encode_frame(stream_id, FRAME_DATA_BACK, &buf[..n]);
let mut w = writer_clone.lock().await;
if w.write_all(&frame).await.is_err() {
break;
}
}
Err(_) => break,
}
}
_ = stream_token.cancelled() => break,
}
}
// Send CLOSE_BACK to edge (only if not cancelled)
if !stream_token.is_cancelled() {
let close_frame = encode_frame(stream_id, FRAME_CLOSE_BACK, &[]);
let mut w = writer_clone.lock().await;
let _ = w.write_all(&close_frame).await;
}
writer_for_edge_data.abort();
Ok::<(), Box<dyn std::error::Error + Send + Sync>>(())
}
.await;
if let Err(e) = result {
log::error!("Stream {} error: {}", stream_id, e);
// Send CLOSE_BACK on error (only if not cancelled)
if !stream_token.is_cancelled() {
let close_frame = encode_frame(stream_id, FRAME_CLOSE_BACK, &[]);
let mut w = writer_clone.lock().await;
let _ = w.write_all(&close_frame).await;
}
}
// Clean up stream
{
let mut s = streams_clone.lock().await;
s.remove(&stream_id);
}
let _ = event_tx_clone.try_send(HubEvent::StreamClosed {
edge_id: edge_id_clone,
stream_id,
});
});
}
FRAME_DATA => {
let s = streams.lock().await;
if let Some(tx) = s.get(&frame.stream_id) {
let _ = tx.send(frame.payload).await;
}
}
FRAME_CLOSE => {
let mut s = streams.lock().await;
s.remove(&frame.stream_id);
}
_ => {
log::warn!("Unexpected frame type {} from edge", frame.frame_type);
} }
let _ = event_tx_clone.send(HubEvent::StreamClosed {
edge_id: edge_id_clone,
stream_id,
});
});
}
FRAME_DATA => {
let s = streams.lock().await;
if let Some(tx) = s.get(&frame.stream_id) {
let _ = tx.send(frame.payload).await;
} }
} }
FRAME_CLOSE => { Ok(None) => {
let mut s = streams.lock().await; log::info!("Edge {} disconnected (EOF)", edge_id);
s.remove(&frame.stream_id); break;
} }
_ => { Err(e) => {
log::warn!("Unexpected frame type {} from edge", frame.frame_type); log::error!("Edge {} frame error: {}", edge_id, e);
break;
} }
} }
} }
Ok(None) => { _ = edge_token.cancelled() => {
log::info!("Edge {} disconnected (EOF)", edge_id); log::info!("Edge {} cancelled by hub", edge_id);
break;
}
Err(e) => {
log::error!("Edge {} frame error: {}", edge_id, e);
break; break;
} }
} }
} }
// Cleanup // Cleanup: cancel edge token to propagate to all child tasks
edge_token.cancel();
config_handle.abort(); config_handle.abort();
{ {
let mut edges = connected.lock().await; let mut edges = connected.lock().await;
edges.remove(&edge_id); edges.remove(&edge_id);
} }
let _ = event_tx.send(HubEvent::EdgeDisconnected { let _ = event_tx.try_send(HubEvent::EdgeDisconnected {
edge_id: edge_id.clone(), edge_id: edge_id.clone(),
}); });

2
ts/00_commitinfo_data.ts
View File

@@ -3,6 +3,6 @@
*/ */
export const commitinfo = { export const commitinfo = {
name: '@serve.zone/remoteingress', name: '@serve.zone/remoteingress',
version: '3.3.0', version: '4.0.0',
description: 'Edge ingress tunnel for DcRouter - accepts incoming TCP connections at network edge and tunnels them to DcRouter SmartProxy preserving client IP via PROXY protocol v1.' description: 'Edge ingress tunnel for DcRouter - accepts incoming TCP connections at network edge and tunnels them to DcRouter SmartProxy preserving client IP via PROXY protocol v1.'
} }