serve.zone/remoteingress
3
0
Fork 0
Code Issues Pull Requests Actions 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

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::sync::{mpsc, Mutex, RwLock};
use tokio_rustls::TlsAcceptor;
use tokio_util::sync::CancellationToken;
use serde::{Deserialize, Serialize};
use remoteingress_protocol::*;
@@ -95,21 +96,24 @@ pub struct TunnelHub {
config: RwLock<HubConfig>,
allowed_edges: Arc<RwLock<HashMap<String, AllowedEdge>>>,
connected_edges: Arc<Mutex<HashMap<String, ConnectedEdgeInfo>>>,
event_tx: mpsc::UnboundedSender<HubEvent>,
event_rx: Mutex<Option<mpsc::UnboundedReceiver<HubEvent>>>,
event_tx: mpsc::Sender<HubEvent>,
event_rx: Mutex<Option<mpsc::Receiver<HubEvent>>>,
shutdown_tx: Mutex<Option<mpsc::Sender<()>>>,
running: RwLock<bool>,
cancel_token: CancellationToken,
}
struct ConnectedEdgeInfo {
connected_at: u64,
active_streams: Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>>,
config_tx: mpsc::Sender<EdgeConfigUpdate>,
#[allow(dead_code)] // kept alive for Drop — cancels child tokens when edge is removed
cancel_token: CancellationToken,
}
impl TunnelHub {
pub fn new(config: HubConfig) -> Self {
let (event_tx, event_rx) = mpsc::unbounded_channel();
let (event_tx, event_rx) = mpsc::channel(1024);
Self {
config: RwLock::new(config),
allowed_edges: Arc::new(RwLock::new(HashMap::new())),
@@ -118,11 +122,12 @@ impl TunnelHub {
event_rx: Mutex::new(Some(event_rx)),
shutdown_tx: Mutex::new(None),
running: RwLock::new(false),
cancel_token: CancellationToken::new(),
}
}
/// 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()
}
@@ -198,6 +203,7 @@ impl TunnelHub {
let connected = self.connected_edges.clone();
let event_tx = self.event_tx.clone();
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 {
loop {
@@ -211,9 +217,10 @@ impl TunnelHub {
let connected = connected.clone();
let event_tx = event_tx.clone();
let target = target_host.clone();
let edge_token = hub_token.child_token();
tokio::spawn(async move {
if let Err(e) = handle_edge_connection(
stream, acceptor, allowed, connected, event_tx, target,
stream, acceptor, allowed, connected, event_tx, target, edge_token,
).await {
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() => {
log::info!("Hub shutting down");
break;
@@ -237,6 +248,7 @@ impl TunnelHub {
/// Stop the hub.
pub async fn stop(&self) {
self.cancel_token.cancel();
if let Some(tx) = self.shutdown_tx.lock().await.take() {
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.
async fn handle_edge_connection(
stream: TcpStream,
acceptor: TlsAcceptor,
allowed: Arc<RwLock<HashMap<String, AllowedEdge>>>,
connected: Arc<Mutex<HashMap<String, ConnectedEdgeInfo>>>,
event_tx: mpsc::UnboundedSender<HubEvent>,
event_tx: mpsc::Sender<HubEvent>,
target_host: String,
edge_token: CancellationToken,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let tls_stream = acceptor.accept(stream).await?;
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);
let _ = event_tx.send(HubEvent::EdgeConnected {
let _ = event_tx.try_send(HubEvent::EdgeConnected {
edge_id: edge_id.clone(),
});
@@ -321,6 +340,7 @@ async fn handle_edge_connection(
connected_at: now,
active_streams: streams.clone(),
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
let config_writer = write_half.clone();
let config_edge_id = edge_id.clone();
let config_token = edge_token.clone();
let config_handle = tokio::spawn(async move {
while let Some(update) = config_rx.recv().await {
if let Ok(payload) = serde_json::to_vec(&update) {
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;
loop {
tokio::select! {
update = config_rx.recv() => {
match update {
Some(update) => {
if let Ok(payload) = serde_json::to_vec(&update) {
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);
loop {
match frame_reader.next_frame().await {
Ok(Some(frame)) => {
match frame.frame_type {
FRAME_OPEN => {
// Payload is PROXY v1 header line
let proxy_header = String::from_utf8_lossy(&frame.payload).to_string();
tokio::select! {
frame_result = frame_reader.next_frame() => {
match frame_result {
Ok(Some(frame)) => {
match frame.frame_type {
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
let dest_port = parse_dest_port_from_proxy(&proxy_header).unwrap_or(443);
// Parse destination port from PROXY header
let dest_port = parse_dest_port_from_proxy(&proxy_header).unwrap_or(443);
let stream_id = frame.stream_id;
let edge_id_clone = edge_id.clone();
let event_tx_clone = event_tx.clone();
let streams_clone = streams.clone();
let writer_clone = write_half.clone();
let target = target_host.clone();
let stream_id = frame.stream_id;
let edge_id_clone = edge_id.clone();
let event_tx_clone = event_tx.clone();
let streams_clone = streams.clone();
let writer_clone = write_half.clone();
let target = target_host.clone();
let stream_token = edge_token.child_token();
let _ = event_tx.send(HubEvent::StreamOpened {
edge_id: edge_id.clone(),
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;
let _ = event_tx.try_send(HubEvent::StreamOpened {
edge_id: edge_id.clone(),
stream_id,
});
// Forward data from SmartProxy back to edge
let mut buf = vec![0u8; 32768];
loop {
match up_read.read(&mut buf).await {
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,
}
// 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);
}
// Send CLOSE_BACK to edge
let close_frame = encode_frame(stream_id, FRAME_CLOSE_BACK, &[]);
let mut w = writer_clone.lock().await;
let _ = w.write_all(&close_frame).await;
// 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?;
writer_for_edge_data.abort();
Ok::<(), Box<dyn std::error::Error + Send + Sync>>(())
}
.await;
let (mut up_read, mut up_write) =
upstream.into_split();
if let Err(e) = result {
log::error!("Stream {} error: {}", stream_id, e);
// Send CLOSE_BACK on error
let close_frame = encode_frame(stream_id, FRAME_CLOSE_BACK, &[]);
let mut w = writer_clone.lock().await;
let _ = w.write_all(&close_frame).await;
}
// Forward data from edge (via channel) to SmartProxy
let writer_token = stream_token.clone();
let writer_for_edge_data = tokio::spawn(async move {
loop {
tokio::select! {
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
{
let mut s = streams_clone.lock().await;
s.remove(&stream_id);
// Forward data from SmartProxy back to edge
let mut buf = vec![0u8; 32768];
loop {
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 => {
let mut s = streams.lock().await;
s.remove(&frame.stream_id);
Ok(None) => {
log::info!("Edge {} disconnected (EOF)", edge_id);
break;
}
_ => {
log::warn!("Unexpected frame type {} from edge", frame.frame_type);
Err(e) => {
log::error!("Edge {} frame error: {}", edge_id, e);
break;
}
}
}
Ok(None) => {
log::info!("Edge {} disconnected (EOF)", edge_id);
break;
}
Err(e) => {
log::error!("Edge {} frame error: {}", edge_id, e);
_ = edge_token.cancelled() => {
log::info!("Edge {} cancelled by hub", edge_id);
break;
}
}
}
// Cleanup
// Cleanup: cancel edge token to propagate to all child tasks
edge_token.cancel();
config_handle.abort();
{
let mut edges = connected.lock().await;
edges.remove(&edge_id);
}
let _ = event_tx.send(HubEvent::EdgeDisconnected {
let _ = event_tx.try_send(HubEvent::EdgeDisconnected {
edge_id: edge_id.clone(),
});