refactor: 移除独立 hub/proxy/executor/gateway crate,统一为 gateway tunnel 架构

- 删除 aether-hub、aether-proxy 独立项目及其 Dockerfile/配置
- 删除 crates/aether-executor 和 crates/aether-gateway 全部模块
- 新增 apps/ 目录作为应用入口
- 将 hub 概念重构为 gateway tunnel transport
- 将 executor 重构为 execution runtime
- 新增 tunnel.rs 合约定义和 testkit tunnel/execution_runtime 模块
- 更新 Python 服务层和测试适配新架构命名
This commit is contained in:
fawney19
2026-04-03 14:59:58 +08:00
parent ddf18fed9a
commit 8f26e1a31f
983 changed files with 103098 additions and 105837 deletions

View File

@@ -0,0 +1,238 @@
//! WebSocket tunnel client: connect, authenticate, and run the tunnel.
use std::sync::Arc;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::sync::watch;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http;
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
use tracing::{debug, info, warn};
use crate::state::{AppState, ServerContext};
use super::{dispatcher, heartbeat, writer};
/// Outcome of a tunnel session.
pub enum TunnelOutcome {
/// Graceful shutdown requested by the local process.
Shutdown,
/// Remote side disconnected or connection lost — should reconnect.
Disconnected,
}
/// Connect to Aether's WebSocket tunnel endpoint and run until disconnected.
///
/// `conn_idx` identifies which connection in the pool this is (0-based).
/// Only connection 0 sends heartbeats to avoid resetting shared metrics.
pub async fn connect_and_run(
state: &Arc<AppState>,
server: &Arc<ServerContext>,
conn_idx: usize,
shutdown: &mut watch::Receiver<bool>,
) -> Result<TunnelOutcome, anyhow::Error> {
let ws_url = build_tunnel_url(server);
info!(url = %ws_url, conn = conn_idx, "connecting tunnel");
// Build WebSocket request with auth headers
let mut request = ws_url.clone().into_client_request()?;
let headers = request.headers_mut();
headers.insert(
"Authorization",
http::HeaderValue::from_str(&format!("Bearer {}", server.management_token))?,
);
let node_id = server.node_id.read().unwrap().clone();
headers.insert("X-Node-Id", http::HeaderValue::from_str(&node_id)?);
// Use dynamic node_name (may be updated by remote config) instead of
// the static server.node_name, so that remote name changes take effect
// on the next reconnect.
let dynamic_node_name = server.dynamic.load().node_name.clone();
headers.insert(
"X-Node-Name",
http::HeaderValue::from_str(&dynamic_node_name)?,
);
// Advertise per-connection max concurrent streams so the backend can
// respect the proxy's capacity limit (backward-compatible: old backends
// ignore this header).
let max_streams = state.config.tunnel_max_streams.unwrap_or(128);
headers.insert("X-Tunnel-Max-Streams", http::HeaderValue::from(max_streams));
// Parse host:port from URL
let uri: http::Uri = ws_url.parse()?;
let host = uri
.host()
.ok_or_else(|| anyhow::anyhow!("missing host in tunnel URL"))?;
let is_tls = uri.scheme_str() == Some("wss");
let port = uri.port_u16().unwrap_or(if is_tls { 443 } else { 80 });
// TCP connect with timeout
let connect_timeout = Duration::from_secs(state.config.tunnel_connect_timeout_secs);
let tcp_stream = tokio::time::timeout(connect_timeout, TcpStream::connect((host, port)))
.await
.map_err(|_| {
anyhow::anyhow!(
"tunnel TCP connect timeout ({}s)",
connect_timeout.as_secs()
)
})??;
// Configure TCP parameters via socket2
configure_tcp_socket(&tcp_stream, state);
// WebSocket upgrade (with TLS if wss://)
let connector = if is_tls {
Some(tokio_tungstenite::Connector::Rustls(Arc::clone(
&state.tunnel_tls_config,
)))
} else {
None
};
// Match Python-side _MAX_FRAME_SIZE (64 MiB) to prevent tungstenite's
// default 16 MiB limit from rejecting large AI API payloads (multi-image
// base64 requests can exceed 16 MiB).
let ws_config = WebSocketConfig {
max_frame_size: Some(64 << 20),
max_message_size: Some(64 << 20),
..Default::default()
};
let handshake_timeout = Duration::from_secs(state.config.tunnel_connect_timeout_secs);
let (ws_stream, _response) = tokio::time::timeout(
handshake_timeout,
tokio_tungstenite::client_async_tls_with_config(
request,
tcp_stream,
Some(ws_config),
connector,
),
)
.await
.map_err(|_| {
anyhow::anyhow!(
"tunnel WebSocket handshake timeout ({}s)",
handshake_timeout.as_secs()
)
})??;
info!(
conn = conn_idx,
tcp_keepalive_secs = state.config.tunnel_tcp_keepalive_secs,
tcp_nodelay = state.config.tunnel_tcp_nodelay,
connect_timeout_secs = state.config.tunnel_connect_timeout_secs,
stale_timeout_secs = state.config.tunnel_stale_timeout_secs,
"tunnel connected"
);
// NOTE: reconnect_attempts reset is handled by the caller (mod.rs)
// based on how long the connection stayed alive.
// Split into read/write halves
let (ws_sink, ws_read) = futures_util::StreamExt::split(ws_stream);
// Spawn writer task (with WebSocket ping keepalive)
let ping_interval = Duration::from_secs(state.config.tunnel_ping_interval_secs);
let (frame_tx, mut writer_handle) = writer::spawn_writer(ws_sink, ping_interval);
// Spawn heartbeat task (only for primary connection to avoid
// resetting shared atomic metrics via swap(0))
let hb_handle = if conn_idx == 0 {
heartbeat::spawn(
Arc::clone(state),
Arc::clone(server),
frame_tx.clone(),
shutdown.clone(),
)
} else {
heartbeat::spawn_noop()
};
// Run dispatcher (blocks until disconnect or shutdown).
// Also watch for writer exit — if the write half dies (e.g. the peer
// closed the connection) but the read half stays open, dispatcher would
// block forever on `ws_stream.next()`. Monitoring `writer_handle`
// ensures we detect this and trigger a reconnect promptly.
let state_clone = Arc::clone(state);
let server_clone = Arc::clone(server);
let outcome = tokio::select! {
result = dispatcher::run(state_clone, server_clone, ws_read, frame_tx.clone(), hb_handle) => {
match result {
Ok(()) => TunnelOutcome::Disconnected,
Err(e) => return Err(e),
}
}
writer_result = &mut writer_handle => {
match writer_result {
Ok(()) => warn!("writer task exited normally, triggering reconnect"),
Err(e) => {
if e.is_panic() {
tracing::error!(error = %e, "writer task panicked, triggering reconnect");
} else {
warn!(error = %e, "writer task cancelled, triggering reconnect");
}
}
}
TunnelOutcome::Disconnected
}
_ = shutdown.changed() => {
debug!("shutdown during tunnel dispatch");
TunnelOutcome::Shutdown
}
};
// Drop our sender; the writer will exit once all stream handler clones
// are also dropped (i.e. after they finish their in-flight work).
drop(frame_tx);
// Wait for the writer task to finish with a generous timeout — the
// dispatcher already waits up to 30s for stream handlers, so 35s here
// covers that plus a small margin.
// Skip if the writer already exited (the select branch that fired).
if !writer_handle.is_finished() {
let _ = tokio::time::timeout(Duration::from_secs(35), writer_handle).await;
}
info!("tunnel disconnected");
Ok(outcome)
}
/// Configure TCP keepalive and NODELAY on an established socket.
fn configure_tcp_socket(stream: &TcpStream, state: &Arc<AppState>) {
let sock_ref = socket2::SockRef::from(stream);
if state.config.tunnel_tcp_keepalive_secs > 0 {
let keepalive = socket2::TcpKeepalive::new()
.with_time(Duration::from_secs(state.config.tunnel_tcp_keepalive_secs))
.with_interval(Duration::from_secs(5));
#[cfg(not(target_os = "windows"))]
let keepalive = keepalive.with_retries(3);
if let Err(e) = sock_ref.set_tcp_keepalive(&keepalive) {
warn!(error = %e, "failed to set TCP keepalive on tunnel socket");
}
}
if state.config.tunnel_tcp_nodelay {
if let Err(e) = sock_ref.set_nodelay(true) {
warn!(error = %e, "failed to set TCP_NODELAY on tunnel socket");
}
}
}
/// Build rustls ClientConfig with system root certificates.
pub fn build_tls_config() -> rustls::ClientConfig {
let root_store =
rustls::RootCertStore::from_iter(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
rustls::ClientConfig::builder()
.with_root_certificates(root_store)
.with_no_client_auth()
}
fn build_tunnel_url(server: &ServerContext) -> String {
let base = server.aether_url.trim_end_matches('/');
let ws_base = if base.starts_with("https://") {
base.replacen("https://", "wss://", 1)
} else if base.starts_with("http://") {
base.replacen("http://", "ws://", 1)
} else {
format!("wss://{}", base)
};
format!("{}/api/internal/proxy-tunnel", ws_base)
}

View File

@@ -0,0 +1,249 @@
//! Frame dispatcher: reads incoming WebSocket frames and routes them.
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use bytes::Bytes;
use futures_util::StreamExt;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use tokio_tungstenite::tungstenite::Message;
use tracing::{debug, error, info, warn};
use crate::state::{AppState, ServerContext};
use super::heartbeat::HeartbeatHandle;
use super::protocol::{decompress_if_gzip, Frame, MsgType, RequestMeta};
use super::stream_handler;
use super::writer::FrameSender;
/// Run the dispatcher loop, reading from the WebSocket stream.
pub async fn run<S>(
state: Arc<AppState>,
server: Arc<ServerContext>,
mut ws_stream: S,
frame_tx: FrameSender,
heartbeat: HeartbeatHandle,
) -> Result<(), anyhow::Error>
where
S: StreamExt<Item = Result<Message, tokio_tungstenite::tungstenite::Error>>
+ Unpin
+ Send
+ 'static,
{
// Active streams: stream_id -> body sender
let mut streams: HashMap<u32, mpsc::Sender<Frame>> = HashMap::new();
// Track spawned stream handlers so we can wait for them on shutdown
let mut handler_handles: Vec<JoinHandle<()>> = Vec::new();
let max_streams = state.config.tunnel_max_streams.unwrap_or(128) as usize;
let mut frames_since_cleanup: u32 = 0;
let stale_timeout = Duration::from_secs(state.config.tunnel_stale_timeout_secs);
// Track last time we received any data to detect stale connections
let mut last_data_at = tokio::time::Instant::now();
let read_err = loop {
let msg_result = tokio::select! {
msg = ws_stream.next() => {
match msg {
Some(r) => r,
None => break None,
}
}
_ = tokio::time::sleep_until(last_data_at + stale_timeout) => {
warn!(
stale_secs = stale_timeout.as_secs(),
"tunnel connection stale, no data received"
);
break None;
}
};
let msg = match msg_result {
Ok(m) => m,
Err(e) => {
error!(error = %e, "WebSocket read error");
break Some(e);
}
};
// Any successfully received message proves the connection is alive
last_data_at = tokio::time::Instant::now();
let data = match msg {
Message::Binary(data) => Bytes::from(data),
Message::Ping(_) => continue,
Message::Pong(_) => continue,
Message::Close(_) => {
info!("received WebSocket close");
break None;
}
_ => continue,
};
let frame = match Frame::decode(data) {
Ok(f) => f,
Err(e) => {
warn!(error = %e, "failed to decode frame");
continue;
}
};
match frame.msg_type {
MsgType::RequestHeaders => {
// Decompress if the frame is gzip-compressed, then parse metadata
let payload = match decompress_if_gzip(&frame) {
Ok(p) => p,
Err(e) => {
warn!(stream_id = frame.stream_id, error = %e, "frame decompress failed");
continue;
}
};
let meta: RequestMeta = match serde_json::from_slice(&payload) {
Ok(m) => m,
Err(e) => {
warn!(stream_id = frame.stream_id, error = %e, "invalid request metadata");
// Use try_send to avoid blocking the read loop
if frame_tx
.try_send(Frame::new(
frame.stream_id,
MsgType::StreamError,
0,
Bytes::from(format!("invalid request metadata: {e}")),
))
.is_err()
{
warn!(
stream_id = frame.stream_id,
"writer channel full, StreamError dropped"
);
}
continue;
}
};
if streams.len() >= max_streams {
warn!(
stream_id = frame.stream_id,
"max concurrent streams reached"
);
if frame_tx
.try_send(Frame::new(
frame.stream_id,
MsgType::StreamError,
0,
Bytes::from("max concurrent streams reached"),
))
.is_err()
{
warn!(
stream_id = frame.stream_id,
"writer channel full, StreamError dropped"
);
}
continue;
}
// Create body channel and spawn handler
let (body_tx, body_rx) = mpsc::channel::<Frame>(64);
streams.insert(frame.stream_id, body_tx);
let state_clone = Arc::clone(&state);
let server_clone = Arc::clone(&server);
let tx_clone = frame_tx.clone();
let sid = frame.stream_id;
let handle = tokio::spawn(async move {
stream_handler::handle_stream(
state_clone,
server_clone,
sid,
meta,
body_rx,
tx_clone,
)
.await;
});
handler_handles.push(handle);
debug!(stream_id = frame.stream_id, "new stream started");
}
MsgType::RequestBody => {
if let Some(tx) = streams.get(&frame.stream_id) {
let is_end = frame.is_end_stream();
let sid = frame.stream_id;
let _ = tx.send(frame).await;
if is_end {
streams.remove(&sid);
}
}
}
MsgType::StreamEnd | MsgType::StreamError => {
// Client-side cancellation or end
if let Some(tx) = streams.remove(&frame.stream_id) {
let _ = tx.send(frame).await;
}
}
MsgType::Ping => {
// Use try_send to avoid blocking the read loop when writer is congested
if frame_tx
.try_send(Frame::control(MsgType::Pong, frame.payload))
.is_err()
{
warn!("writer channel full, Pong dropped");
}
}
MsgType::HeartbeatAck => {
heartbeat.on_ack(frame.payload).await;
}
MsgType::GoAway => {
info!("received GOAWAY");
break None;
}
_ => {
debug!(msg_type = ?frame.msg_type, "ignoring unexpected frame type");
}
}
// Periodically clean up finished handles to avoid unbounded growth.
// Trigger every 64 frames OR when the count exceeds max_streams.
frames_since_cleanup += 1;
if frames_since_cleanup >= 64 || handler_handles.len() > max_streams {
handler_handles.retain(|h| !h.is_finished());
frames_since_cleanup = 0;
}
};
// Drop body senders so stream handlers waiting on body_rx will unblock
streams.clear();
// Wait for active stream handlers to finish so their frame_tx clones
// are dropped before the writer closes the sink.
drain_handlers(handler_handles).await;
match read_err {
Some(e) => Err(e.into()),
None => Ok(()),
}
}
/// Wait for all active stream handlers to finish (with a timeout).
async fn drain_handlers(handles: Vec<JoinHandle<()>>) {
if handles.is_empty() {
return;
}
let count = handles.len();
debug!(count, "waiting for active stream handlers to finish");
let _ = tokio::time::timeout(Duration::from_secs(30), async {
for h in handles {
let _ = h.await;
}
})
.await;
}

View File

@@ -0,0 +1,373 @@
//! Tunnel heartbeat: sends metrics over the tunnel, processes ACKs.
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;
use std::time::SystemTime;
use std::time::UNIX_EPOCH;
use bytes::Bytes;
use tokio::sync::watch;
use tracing::{debug, info, warn};
use crate::registration::client::RemoteConfig;
use crate::runtime;
use crate::state::AppState;
use crate::state::ServerContext;
use super::protocol::{Frame, MsgType};
use super::writer::FrameSender;
const CURRENT_VERSION: &str = env!("CARGO_PKG_VERSION");
static UPGRADE_IN_PROGRESS: AtomicBool = AtomicBool::new(false);
static NON_ROOT_UPGRADE_WARNED: AtomicBool = AtomicBool::new(false);
enum AckDecision {
Accept {
heartbeat_id: Option<u64>,
upgrade_to: Option<String>,
},
Ignore,
}
/// Handle for the dispatcher to forward HeartbeatAck frames.
#[derive(Clone)]
pub struct HeartbeatHandle {
ack_tx: tokio::sync::mpsc::Sender<Bytes>,
}
impl HeartbeatHandle {
pub async fn on_ack(&self, payload: Bytes) {
let _ = self.ack_tx.send(payload).await;
}
}
/// Create a no-op heartbeat handle that silently discards ACKs.
/// Used for non-primary tunnel connections (conn_idx > 0) to avoid
/// resetting shared atomic metrics via `swap(0)`.
pub fn spawn_noop() -> HeartbeatHandle {
let (ack_tx, _) = tokio::sync::mpsc::channel::<Bytes>(1);
// receiver is immediately dropped; on_ack() calls will silently fail
HeartbeatHandle { ack_tx }
}
#[derive(Debug, Clone, Copy, Default)]
struct HeartbeatSnapshot {
requests: u64,
latency_ns: u64,
failed: u64,
dns_failures: u64,
stream_errors: u64,
}
/// Spawn the heartbeat task. Returns a handle for forwarding ACKs.
pub fn spawn(
state: Arc<AppState>,
server: Arc<ServerContext>,
frame_tx: FrameSender,
mut shutdown: watch::Receiver<bool>,
) -> HeartbeatHandle {
let (ack_tx, mut ack_rx) = tokio::sync::mpsc::channel::<Bytes>(4);
tokio::spawn(async move {
// Read initial interval from dynamic config (may be updated by remote config).
let initial_interval = Duration::from_secs(server.dynamic.load().heartbeat_interval);
let mut current_interval = initial_interval;
// At most one in-flight heartbeat snapshot is tracked at a time.
// Snapshot is only cleared after receiving an ACK, which avoids losing
// interval counters when ACK/frame delivery is temporarily unstable.
let mut pending: Option<(u64, HeartbeatSnapshot)> = None;
let mut next_heartbeat_id: u64 = 1;
let heartbeat_session_id = format!(
"{}-{}",
std::process::id(),
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_nanos()
);
// Skip first immediate tick by sleeping first.
tokio::time::sleep(current_interval).await;
loop {
tokio::select! {
_ = tokio::time::sleep(current_interval) => {
let (heartbeat_id, snapshot) = if let Some((id, snap)) = pending {
(id, snap)
} else {
let snap = collect_snapshot(&server);
let id = next_heartbeat_id;
next_heartbeat_id = next_heartbeat_id.wrapping_add(1);
if next_heartbeat_id == 0 {
next_heartbeat_id = 1;
}
pending = Some((id, snap));
(id, snap)
};
let payload = build_heartbeat_payload(
&state,
&server,
&heartbeat_session_id,
heartbeat_id,
snapshot
).await;
let frame = Frame::control(MsgType::HeartbeatData, payload);
if frame_tx.send(frame).await.is_err() {
if let Some((_, snap)) = pending.take() {
restore_snapshot(&server, snap);
}
break; // Writer closed
}
debug!("sent heartbeat data");
// Re-read interval from dynamic config (remote config may have
// updated it since the last heartbeat).
let new_interval = Duration::from_secs(
server.dynamic.load().heartbeat_interval
);
if new_interval != current_interval {
debug!(
old_secs = current_interval.as_secs(),
new_secs = new_interval.as_secs(),
"heartbeat interval updated from dynamic config"
);
current_interval = new_interval;
}
}
Some(ack_payload) = ack_rx.recv() => {
match handle_ack(&server, &ack_payload) {
AckDecision::Accept {
heartbeat_id: ack_id,
upgrade_to,
} => {
if let Some((pending_id, _)) = pending {
match ack_id {
Some(id) if id == pending_id => {
pending = None;
}
None => {
// Backward-compatible with servers that don't echo
// heartbeat_id in ACK payload yet.
pending = None;
}
_ => {}
}
}
maybe_trigger_upgrade(upgrade_to);
}
AckDecision::Ignore => {}
}
}
_ = shutdown.changed() => {
debug!("heartbeat task shutting down");
if let Some((_, snap)) = pending.take() {
restore_snapshot(&server, snap);
}
break;
}
}
}
});
HeartbeatHandle { ack_tx }
}
fn collect_snapshot(server: &ServerContext) -> HeartbeatSnapshot {
HeartbeatSnapshot {
requests: server.metrics.total_requests.swap(0, Ordering::AcqRel),
latency_ns: server.metrics.total_latency_ns.swap(0, Ordering::AcqRel),
failed: server.metrics.failed_requests.swap(0, Ordering::AcqRel),
dns_failures: server.metrics.dns_failures.swap(0, Ordering::AcqRel),
stream_errors: server.metrics.stream_errors.swap(0, Ordering::AcqRel),
}
}
fn restore_snapshot(server: &ServerContext, snap: HeartbeatSnapshot) {
if snap.requests > 0 {
server
.metrics
.total_requests
.fetch_add(snap.requests, Ordering::Release);
}
if snap.latency_ns > 0 {
server
.metrics
.total_latency_ns
.fetch_add(snap.latency_ns, Ordering::Release);
}
if snap.failed > 0 {
server
.metrics
.failed_requests
.fetch_add(snap.failed, Ordering::Release);
}
if snap.dns_failures > 0 {
server
.metrics
.dns_failures
.fetch_add(snap.dns_failures, Ordering::Release);
}
if snap.stream_errors > 0 {
server
.metrics
.stream_errors
.fetch_add(snap.stream_errors, Ordering::Release);
}
}
async fn build_heartbeat_payload(
state: &AppState,
server: &ServerContext,
heartbeat_session_id: &str,
heartbeat_id: u64,
snapshot: HeartbeatSnapshot,
) -> Bytes {
let node_id = server.node_id.read().unwrap().clone();
let avg_latency_ms = if snapshot.requests > 0 {
Some(snapshot.latency_ns as f64 / snapshot.requests as f64 / 1_000_000.0)
} else {
None
};
let local_admission = state.stream_concurrency_snapshot().map(|snapshot| {
serde_json::json!({
"limit": snapshot.limit,
"in_flight": snapshot.in_flight,
"available_permits": snapshot.available_permits,
"high_watermark": snapshot.high_watermark,
"rejected_total": snapshot.rejected,
})
});
let distributed_admission = match state.distributed_stream_concurrency_snapshot().await {
Ok(Some(snapshot)) => Some(serde_json::json!({
"limit": snapshot.limit,
"in_flight": snapshot.in_flight,
"available_permits": snapshot.available_permits,
"high_watermark": snapshot.high_watermark,
"rejected_total": snapshot.rejected,
})),
Ok(None) => None,
Err(err) => Some(serde_json::json!({
"error": err.to_string(),
})),
};
let admission = match (local_admission, distributed_admission) {
(None, None) => None,
(local, distributed) => Some(serde_json::json!({
"local_streams": local,
"distributed_streams": distributed,
})),
};
let payload = serde_json::json!({
"node_id": node_id,
"heartbeat_session_id": heartbeat_session_id,
"heartbeat_id": heartbeat_id,
"active_connections": server.active_connections.load(Ordering::Acquire),
"total_requests": snapshot.requests,
"avg_latency_ms": avg_latency_ms,
"failed_requests": snapshot.failed,
"dns_failures": snapshot.dns_failures,
"stream_errors": snapshot.stream_errors,
"proxy_metadata": {
"version": CURRENT_VERSION,
"admission": admission,
},
});
Bytes::from(serde_json::to_vec(&payload).unwrap_or_default())
}
fn handle_ack(server: &ServerContext, payload: &[u8]) -> AckDecision {
if payload.is_empty() {
return AckDecision::Accept {
heartbeat_id: None,
upgrade_to: None,
};
}
#[derive(serde::Deserialize)]
struct AckPayload {
#[serde(default)]
remote_config: Option<RemoteConfig>,
#[serde(default)]
config_version: u64,
#[serde(default)]
heartbeat_id: Option<u64>,
#[serde(default)]
upgrade_to: Option<String>,
}
match serde_json::from_slice::<AckPayload>(payload) {
Ok(ack) => {
if let Some(ref rc) = ack.remote_config {
runtime::apply_remote_config(&server.dynamic, rc, ack.config_version);
}
AckDecision::Accept {
heartbeat_id: ack.heartbeat_id,
upgrade_to: ack.upgrade_to.and_then(normalize_upgrade_target),
}
}
Err(e) => {
warn!(error = %e, "failed to parse heartbeat ACK");
AckDecision::Ignore
}
}
}
fn normalize_upgrade_target(raw: String) -> Option<String> {
let trimmed = raw.trim();
if trimmed.is_empty() {
return None;
}
let normalized = trimmed.strip_prefix("proxy-v").unwrap_or(trimmed);
if normalized == CURRENT_VERSION {
return None;
}
Some(normalized.to_string())
}
fn maybe_trigger_upgrade(version: Option<String>) {
let Some(target_version) = version else {
return;
};
if !crate::setup::service::is_root() {
if NON_ROOT_UPGRADE_WARNED
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
{
warn!(
target_version = %target_version,
"remote upgrade skipped: root privileges are required"
);
}
return;
}
if UPGRADE_IN_PROGRESS
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
.is_err()
{
debug!(target_version = %target_version, "upgrade already in progress, ignoring");
return;
}
tokio::spawn(async move {
info!(target_version = %target_version, "received remote upgrade instruction");
match crate::setup::upgrade::perform_upgrade(&target_version).await {
Ok(()) => {
info!(target_version = %target_version, "remote upgrade finished");
}
Err(e) => {
warn!(
target_version = %target_version,
error = %e,
"remote upgrade failed"
);
UPGRADE_IN_PROGRESS.store(false, Ordering::Release);
}
}
});
}

View File

@@ -0,0 +1,496 @@
pub mod client;
pub mod dispatcher;
pub mod heartbeat;
pub mod protocol;
pub mod stream_handler;
pub mod writer;
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use tokio::sync::watch;
use tracing::{error, info};
use crate::state::{AppState, ServerContext};
/// If a tunnel stays connected at least this long, treat the next disconnect
/// as a non-failure and reset reconnect backoff.
const STABLE_SESSION_RESET_AFTER: Duration = Duration::from_secs(30);
/// Startup staggering step per secondary connection, used to avoid
/// simultaneous bursts when a pool of tunnels starts together.
const STARTUP_STAGGER_STEP_MS: u64 = 150;
/// Upper bound for startup staggering.
const MAX_STARTUP_STAGGER_MS: u64 = 1_500;
/// Keep a tiny floor for repeated reconnects; first retry is still immediate.
const MIN_RECONNECT_DELAY_MS: u64 = 50;
/// Even under sustained failures, keep probing frequently so recovery is fast
/// once cross-border network quality improves.
const RECONNECT_PROBE_MAX_DELAY_MS: u64 = 3_000;
/// Run the tunnel mode main loop (connect, dispatch, reconnect).
///
/// `conn_idx` identifies which connection in the pool this is (0-based).
/// Only connection 0 sends heartbeats to avoid resetting shared metrics.
pub async fn run(
state: &Arc<AppState>,
server: &Arc<ServerContext>,
conn_idx: usize,
mut shutdown: watch::Receiver<bool>,
) {
info!(server = %server.server_label, conn = conn_idx, "starting tunnel");
let reconnect_salt = compute_connection_salt(server, conn_idx);
let startup_delay = compute_startup_stagger(conn_idx, reconnect_salt);
if !startup_delay.is_zero() {
info!(
server = %server.server_label,
conn = conn_idx,
delay_ms = startup_delay.as_millis(),
"startup stagger before first connect"
);
tokio::select! {
_ = tokio::time::sleep(startup_delay) => {}
_ = shutdown.changed() => {
info!(server = %server.server_label, conn = conn_idx, "shutdown requested during startup stagger");
return;
}
}
}
let mut consecutive_failures: u32 = 0;
loop {
let started_at = Instant::now();
match client::connect_and_run(state, server, conn_idx, &mut shutdown).await {
Ok(client::TunnelOutcome::Shutdown) => {
info!(server = %server.server_label, conn = conn_idx, "tunnel shut down gracefully");
return;
}
Ok(client::TunnelOutcome::Disconnected) => {
info!(server = %server.server_label, conn = conn_idx, "tunnel disconnected, reconnecting");
}
Err(e) => {
error!(server = %server.server_label, conn = conn_idx, error = %e, "tunnel connection error, reconnecting");
}
}
if *shutdown.borrow() {
info!(server = %server.server_label, conn = conn_idx, "shutdown requested, not reconnecting");
return;
}
// Reset backoff after a stable session to keep recovery snappy when
// failures are only occasional.
let connected_for = started_at.elapsed();
if connected_for >= STABLE_SESSION_RESET_AFTER {
consecutive_failures = 0;
} else {
consecutive_failures = consecutive_failures.saturating_add(1);
}
let reconnect_delay = compute_reconnect_delay(
state.config.tunnel_reconnect_base_ms,
state.config.tunnel_reconnect_max_ms,
consecutive_failures,
reconnect_salt,
);
info!(
server = %server.server_label,
conn = conn_idx,
failures = consecutive_failures,
delay_ms = reconnect_delay.as_millis(),
"waiting before reconnect"
);
tokio::select! {
_ = tokio::time::sleep(reconnect_delay) => {}
_ = shutdown.changed() => {
info!(server = %server.server_label, conn = conn_idx, "shutdown requested during reconnect wait");
return;
}
}
}
}
fn compute_connection_salt(server: &ServerContext, conn_idx: usize) -> u64 {
// FNV-1a style hash over server label + connection index.
let mut h: u64 = 0xcbf29ce484222325;
for &b in server.server_label.as_bytes() {
h ^= b as u64;
h = h.wrapping_mul(0x100000001b3);
}
h ^= conn_idx as u64;
mix_u64(h)
}
fn compute_startup_stagger(conn_idx: usize, salt: u64) -> Duration {
if conn_idx == 0 {
return Duration::ZERO;
}
let base = (conn_idx as u64).saturating_mul(STARTUP_STAGGER_STEP_MS);
let jitter = mix_u64(salt) % 301; // 0..=300ms
Duration::from_millis((base + jitter).min(MAX_STARTUP_STAGGER_MS))
}
fn compute_reconnect_delay(
base_ms: u64,
max_ms: u64,
consecutive_failures: u32,
salt: u64,
) -> Duration {
// First retry should be immediate to maximize recovery speed on transient
// blips (the user's primary expectation in poor networks).
if consecutive_failures <= 1 {
return Duration::ZERO;
}
// Keep a sane minimum for repeated failures.
let base_ms = base_ms.max(MIN_RECONNECT_DELAY_MS);
let max_ms = max_ms.max(base_ms);
let cap_ms = compute_reconnect_cap_ms(base_ms, max_ms, consecutive_failures)
.min(RECONNECT_PROBE_MAX_DELAY_MS.max(base_ms));
// Equal-jitter: randomize in [cap/2, cap], preventing synchronized reconnect
// storms while keeping reconnect latency bounded.
if cap_ms <= 1 {
return Duration::from_millis(cap_ms);
}
let half = cap_ms / 2;
let span = cap_ms - half;
let now_nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.subsec_nanos() as u64)
.unwrap_or(0);
let mixed = mix_u64(now_nanos ^ salt);
let jitter = if span == 0 { 0 } else { mixed % (span + 1) };
Duration::from_millis(half + jitter)
}
fn compute_reconnect_cap_ms(base_ms: u64, max_ms: u64, consecutive_failures: u32) -> u64 {
if consecutive_failures <= 1 {
return base_ms.min(max_ms);
}
let shift = (consecutive_failures - 1).min(31);
let factor = 1u64 << shift;
base_ms.saturating_mul(factor).min(max_ms)
}
fn mix_u64(mut x: u64) -> u64 {
// SplitMix64 finalizer - cheap bit mixing for pseudo-random jitter.
x ^= x >> 30;
x = x.wrapping_mul(0xbf58476d1ce4e5b9);
x ^= x >> 27;
x = x.wrapping_mul(0x94d049bb133111eb);
x ^ (x >> 31)
}
#[cfg(test)]
mod tests {
use std::sync::atomic::AtomicU64;
use std::sync::{Arc, Once};
use std::time::Duration;
use aether_gateway::{build_router_with_state, AppState as GatewayAppState};
use arc_swap::ArcSwap;
use axum::Router;
use reqwest::StatusCode;
use tokio::sync::watch;
use crate::config::Config;
use crate::registration::client::AetherClient;
use crate::runtime::DynamicConfig;
use crate::state::{AppState as ProxyAppState, ProxyMetrics, ServerContext};
use crate::target_filter::DnsCache;
use crate::tunnel::protocol;
use crate::upstream_client;
use super::{
compute_reconnect_cap_ms, compute_reconnect_delay, compute_startup_stagger, run,
MAX_STARTUP_STAGGER_MS, RECONNECT_PROBE_MAX_DELAY_MS, STARTUP_STAGGER_STEP_MS,
};
#[test]
fn reconnect_cap_grows_exponentially_and_caps() {
let base = 500;
let max = 30_000;
assert_eq!(compute_reconnect_cap_ms(base, max, 0), 500);
assert_eq!(compute_reconnect_cap_ms(base, max, 1), 500);
assert_eq!(compute_reconnect_cap_ms(base, max, 2), 1_000);
assert_eq!(compute_reconnect_cap_ms(base, max, 3), 2_000);
assert_eq!(compute_reconnect_cap_ms(base, max, 4), 4_000);
assert_eq!(compute_reconnect_cap_ms(base, max, 5), 8_000);
assert_eq!(compute_reconnect_cap_ms(base, max, 6), 16_000);
assert_eq!(compute_reconnect_cap_ms(base, max, 7), 30_000);
assert_eq!(compute_reconnect_cap_ms(base, max, 20), 30_000);
}
#[test]
fn startup_stagger_is_zero_for_primary_and_bounded_for_secondary() {
assert_eq!(compute_startup_stagger(0, 42), Duration::ZERO);
let d1 = compute_startup_stagger(1, 42);
let d2 = compute_startup_stagger(2, 42);
assert!(d1 >= Duration::from_millis(STARTUP_STAGGER_STEP_MS));
assert!(d1 <= Duration::from_millis(MAX_STARTUP_STAGGER_MS));
assert!(d2 >= Duration::from_millis(STARTUP_STAGGER_STEP_MS * 2));
assert!(d2 <= Duration::from_millis(MAX_STARTUP_STAGGER_MS));
}
#[test]
fn reconnect_delay_is_immediate_on_first_failure() {
assert_eq!(compute_reconnect_delay(700, 45_000, 1, 123), Duration::ZERO);
}
#[test]
fn reconnect_delay_stays_within_probe_ceiling_after_many_failures() {
let d = compute_reconnect_delay(500, 45_000, 100, 12345);
assert!(d <= Duration::from_millis(RECONNECT_PROBE_MAX_DELAY_MS));
}
#[tokio::test]
async fn proxy_reconnects_after_gateway_restart() {
ensure_rustls_provider();
let gateway_port = reserve_local_port().expect("gateway port should reserve");
let gateway_base_url = format!("http://127.0.0.1:{gateway_port}");
let (gateway_state, mut gateway_handle) = start_gateway_on_port(gateway_port)
.await
.expect("gateway should start");
let state = sample_state(sample_config(&gateway_base_url));
let server = sample_server(&state, "node-recovery");
let (shutdown_tx, shutdown_rx) = watch::channel(false);
let proxy_task = tokio::spawn({
let state = Arc::clone(&state);
let server = Arc::clone(&server);
async move {
run(&state, &server, 0, shutdown_rx).await;
}
});
wait_until_relay_status(
&gateway_base_url,
"node-recovery",
StatusCode::GATEWAY_TIMEOUT,
)
.await;
assert_eq!(gateway_state.force_close_all_tunnel_proxies(), 1);
tokio::time::sleep(Duration::from_millis(200)).await;
gateway_handle.abort();
let (_restarted_gateway_state, restarted_gateway_handle) =
start_gateway_on_port_retry(gateway_port)
.await
.expect("gateway should restart on fixed port");
gateway_handle = restarted_gateway_handle;
wait_until_relay_status(
&gateway_base_url,
"node-recovery",
StatusCode::GATEWAY_TIMEOUT,
)
.await;
let _ = shutdown_tx.send(true);
tokio::time::timeout(Duration::from_secs(5), proxy_task)
.await
.expect("proxy task should stop")
.expect("proxy task should join");
gateway_handle.abort();
}
async fn wait_until_relay_status(gateway_base_url: &str, node_id: &str, expected: StatusCode) {
let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
let mut last_observed = None::<String>;
loop {
if let Some((status, body)) = probe_relay_status(gateway_base_url, node_id).await {
last_observed = Some(format!("{status} body={body}"));
if status == expected {
return;
}
}
assert!(
tokio::time::Instant::now() < deadline,
"relay status did not become {expected} within timeout; last={:?}",
last_observed
);
tokio::time::sleep(Duration::from_millis(25)).await;
}
}
async fn probe_relay_status(
gateway_base_url: &str,
node_id: &str,
) -> Option<(StatusCode, String)> {
let response = reqwest::Client::new()
.post(format!(
"{gateway_base_url}/api/internal/tunnel/relay/{node_id}"
))
.header("content-type", "application/octet-stream")
.body(relay_probe_envelope())
.send()
.await
.ok()?;
let status = response.status();
let body = response.text().await.unwrap_or_default();
Some((status, body))
}
fn relay_probe_envelope() -> Vec<u8> {
let meta = protocol::RequestMeta {
method: "GET".to_string(),
url: "http://127.0.0.1:80/blocked".to_string(),
headers: std::collections::HashMap::new(),
timeout: 5,
};
let meta_json =
serde_json::to_vec(&meta).expect("tunnel relay probe metadata should serialize");
let mut envelope = Vec::with_capacity(4 + meta_json.len());
envelope.extend_from_slice(&(meta_json.len() as u32).to_be_bytes());
envelope.extend_from_slice(&meta_json);
envelope
}
async fn start_gateway_on_port(
port: u16,
) -> Result<(GatewayAppState, tokio::task::JoinHandle<()>), std::io::Error> {
let state =
GatewayAppState::new("http://127.0.0.1:9").expect("gateway test state should build");
let router = build_router_with_state(state.clone());
let handle = spawn_router_on_port(port, router).await?;
Ok((state, handle))
}
async fn start_gateway_on_port_retry(
port: u16,
) -> Result<(GatewayAppState, tokio::task::JoinHandle<()>), std::io::Error> {
let mut attempts = 0usize;
loop {
match start_gateway_on_port(port).await {
Ok(server) => return Ok(server),
Err(err) => {
attempts += 1;
if attempts >= 20 {
return Err(err);
}
tokio::time::sleep(Duration::from_millis(50)).await;
}
}
}
}
async fn spawn_router_on_port(
port: u16,
app: Router,
) -> Result<tokio::task::JoinHandle<()>, std::io::Error> {
let listener = tokio::net::TcpListener::bind(("127.0.0.1", port)).await?;
Ok(tokio::spawn(async move {
axum::serve(
listener,
app.into_make_service_with_connect_info::<std::net::SocketAddr>(),
)
.await
.expect("gateway test server should run");
}))
}
fn reserve_local_port() -> Result<u16, std::io::Error> {
let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
let port = listener.local_addr()?.port();
drop(listener);
Ok(port)
}
fn sample_state(config: Config) -> Arc<ProxyAppState> {
let config = Arc::new(config);
let dns_cache = Arc::new(DnsCache::new(Duration::from_secs(60), 128));
let upstream_client =
upstream_client::build_upstream_client(&config, Arc::clone(&dns_cache));
Arc::new(ProxyAppState {
config,
dns_cache,
upstream_client,
tunnel_tls_config: Arc::new(crate::tunnel::client::build_tls_config()),
stream_gate: None,
distributed_stream_gate: None,
})
}
fn sample_server(state: &Arc<ProxyAppState>, node_id: &str) -> Arc<ServerContext> {
let config = Arc::clone(&state.config);
Arc::new(ServerContext {
server_label: "gateway-owned-tunnel".to_string(),
aether_url: config.aether_url.clone(),
management_token: config.management_token.clone(),
node_name: config.node_name.clone(),
node_id: Arc::new(std::sync::RwLock::new(node_id.to_string())),
aether_client: Arc::new(AetherClient::new(
&config,
&config.aether_url,
&config.management_token,
)),
dynamic: Arc::new(ArcSwap::from_pointee(DynamicConfig::from_config(&config))),
active_connections: Arc::new(AtomicU64::new(0)),
metrics: Arc::new(ProxyMetrics::new()),
})
}
fn sample_config(aether_url: &str) -> Config {
Config {
aether_url: aether_url.to_string(),
management_token: "token".to_string(),
public_ip: None,
node_name: "proxy-test".to_string(),
node_region: None,
heartbeat_interval: 1,
allowed_ports: vec![80, 443],
aether_request_timeout_secs: 10,
aether_connect_timeout_secs: 2,
aether_pool_max_idle_per_host: 8,
aether_pool_idle_timeout_secs: 90,
aether_tcp_keepalive_secs: 60,
aether_tcp_nodelay: true,
aether_http2: true,
aether_retry_max_attempts: 1,
aether_retry_base_delay_ms: 50,
aether_retry_max_delay_ms: 100,
max_concurrent_connections: None,
max_in_flight_streams: None,
distributed_stream_limit: None,
distributed_stream_redis_url: None,
distributed_stream_redis_key_prefix: None,
distributed_stream_lease_ttl_ms: 30_000,
distributed_stream_renew_interval_ms: 10_000,
distributed_stream_command_timeout_ms: 1_000,
dns_cache_ttl_secs: 60,
dns_cache_capacity: 128,
upstream_connect_timeout_secs: 30,
upstream_pool_max_idle_per_host: 4,
upstream_pool_idle_timeout_secs: 60,
upstream_tcp_keepalive_secs: 60,
upstream_tcp_nodelay: true,
log_level: "info".to_string(),
log_json: false,
tunnel_reconnect_base_ms: 50,
tunnel_reconnect_max_ms: 250,
tunnel_ping_interval_secs: 1,
tunnel_max_streams: Some(8),
tunnel_connect_timeout_secs: 2,
tunnel_tcp_keepalive_secs: 30,
tunnel_tcp_nodelay: true,
tunnel_stale_timeout_secs: 5,
tunnel_connections: 1,
}
}
fn ensure_rustls_provider() {
static INIT: Once = Once::new();
INIT.call_once(|| {
let _ = rustls::crypto::ring::default_provider().install_default();
});
}
}

View File

@@ -0,0 +1 @@
pub use aether_contracts::tunnel::*;

View File

@@ -0,0 +1,711 @@
//! Per-stream request handler.
//!
//! Receives request frames, executes the upstream HTTP request,
//! and sends response frames back through the writer channel.
use std::io;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::{Duration, Instant};
use aether_runtime::hold_admission_permit_until;
use bytes::Bytes;
use futures_util::stream;
use futures_util::StreamExt;
use http_body_util::BodyExt;
use hyper::body::Frame as BodyFrame;
use tokio::sync::mpsc;
use tracing::{debug, warn};
use crate::state::{AppState, ServerContext};
use crate::target_filter;
use crate::upstream_client;
use super::protocol::{
compress_payload, decompress_if_gzip, flags, Frame as TunnelFrame, MsgType, RequestMeta,
ResponseMeta,
};
use super::writer::FrameSender;
/// Maximum response body chunk size per frame (32 KB).
const MAX_CHUNK_SIZE: usize = 32 * 1024;
/// Timeout for sending a single frame to the writer channel.
/// If the writer is congested (TCP backpressure), we abandon the stream
/// rather than blocking indefinitely and exhausting the stream pool.
const FRAME_SEND_TIMEOUT: Duration = Duration::from_secs(30);
/// Minimum allowed upstream request timeout (seconds).
const MIN_TIMEOUT_SECS: u64 = 5;
/// Maximum allowed upstream request timeout (seconds).
const MAX_TIMEOUT_SECS: u64 = 300;
/// Headers that must not be forwarded to upstream (hop-by-hop or security-sensitive).
///
/// `host` and `content-length` are managed by the HTTP client (reqwest/hyper):
/// - `host` → translated to `:authority` pseudo-header in HTTP/2; forwarding
/// the original `host` alongside `:authority` triggers PROTOCOL_ERROR on
/// strict H2 implementations (e.g. Google APIs).
/// - `content-length` → recalculated by hyper from the actual body; a stale
/// value from the tunnel (body may have been re-compressed) causes H2
/// PROTOCOL_ERROR when it mismatches the real frame length.
const BLOCKED_HEADERS: &[&str] = &[
"connection",
"content-length",
"host",
"keep-alive",
"proxy-authenticate",
"proxy-authorization",
"proxy-connection",
"te",
"trailer",
"transfer-encoding",
"upgrade",
];
/// Handle a single stream: receive body, execute upstream, send response.
pub async fn handle_stream(
state: Arc<AppState>,
server: Arc<ServerContext>,
stream_id: u32,
meta: RequestMeta,
body_rx: mpsc::Receiver<TunnelFrame>,
frame_tx: FrameSender,
) {
let permit = match state.try_acquire_stream_permit().await {
Ok(permit) => permit,
Err(err) => {
let message = match err {
crate::state::ProxyAdmissionError::Saturated { .. } => "proxy overloaded",
crate::state::ProxyAdmissionError::Unavailable { .. } => {
"proxy admission unavailable"
}
};
send_error(&frame_tx, stream_id, message).await;
return;
}
};
server.active_connections.fetch_add(1, Ordering::Release);
let connect_elapsed = hold_admission_permit_until(permit, async {
handle_stream_inner(&state, &server, stream_id, meta, body_rx, &frame_tx).await
})
.await;
server.active_connections.fetch_sub(1, Ordering::Release);
if let Some(d) = connect_elapsed {
server.metrics.record_request(d);
}
}
/// Send a frame to the writer with a timeout. Returns false if send failed.
async fn send_frame(tx: &FrameSender, frame: TunnelFrame) -> bool {
match tokio::time::timeout(FRAME_SEND_TIMEOUT, tx.send(frame)).await {
Ok(Ok(())) => true,
Ok(Err(_)) => {
// Channel closed (writer exited)
false
}
Err(_) => {
// Timeout — writer is congested
warn!("frame send timeout (writer congested), abandoning stream");
false
}
}
}
/// Returns the connection-establishment duration (DNS + TCP/TLS + TTFB) if the
/// upstream request succeeded, or `None` if the request never reached the
/// response-headers stage.
async fn handle_stream_inner(
state: &AppState,
server: &ServerContext,
stream_id: u32,
meta: RequestMeta,
body_rx: mpsc::Receiver<TunnelFrame>,
frame_tx: &FrameSender,
) -> Option<Duration> {
// Validate target
let target_url = match url::Url::parse(&meta.url) {
Ok(u) => u,
Err(e) => {
send_error(frame_tx, stream_id, &format!("invalid URL: {e}")).await;
return None;
}
};
// Only allow http/https schemes (block file://, data://, etc.)
match target_url.scheme() {
"http" | "https" => {}
other => {
send_error(
frame_tx,
stream_id,
&format!("unsupported URL scheme: {other}"),
)
.await;
return None;
}
}
let host = match target_url.host_str() {
Some(h) => h.to_string(),
None => {
send_error(frame_tx, stream_id, "missing host in URL").await;
return None;
}
};
let port = target_url.port_or_known_default().unwrap_or(443);
// DNS + target validation (populates dns_cache for SafeDnsResolver)
let connect_start = Instant::now();
{
let allowed_ports = Arc::clone(&server.dynamic.load().allowed_ports);
if let Err(e) =
target_filter::validate_target(&host, port, &allowed_ports, &state.dns_cache).await
{
server.metrics.dns_failures.fetch_add(1, Ordering::Release);
send_error(frame_tx, stream_id, &format!("target blocked: {e}")).await;
return None;
}
}
let dns_ms = connect_start.elapsed().as_millis() as u64;
// Execute upstream request
let client = &state.upstream_client;
let timeout = Duration::from_secs(meta.timeout.clamp(MIN_TIMEOUT_SECS, MAX_TIMEOUT_SECS));
let request_body_size = Arc::new(AtomicUsize::new(0));
let request_body = build_streaming_request_body(body_rx, Arc::clone(&request_body_size));
let method: hyper::Method = meta.method.parse().unwrap_or(hyper::Method::GET);
let mut request = match hyper::Request::builder()
.method(method)
.uri(meta.url.as_str())
.body(request_body)
{
Ok(request) => request,
Err(e) => {
send_error(
frame_tx,
stream_id,
&format!("invalid upstream request: {e}"),
)
.await;
return None;
}
};
let headers = request.headers_mut();
for (k, v) in &meta.headers {
let k_lower = k.to_ascii_lowercase();
if BLOCKED_HEADERS.contains(&k_lower.as_str()) {
continue;
}
if let (Ok(name), Ok(value)) = (
hyper::header::HeaderName::from_bytes(k.as_bytes()),
hyper::header::HeaderValue::from_str(v),
) {
headers.insert(name, value);
}
}
let mut captured_connection = upstream_client::capture_connection(&mut request);
let connection_start = Instant::now();
let connection_capture = tokio::spawn(async move {
let connected = captured_connection.wait_for_connection_metadata().await;
connected
.as_ref()
.map(|_| connection_start.elapsed().as_millis() as u64)
});
let upstream_start = Instant::now();
let response = match tokio::time::timeout(timeout, client.request(request)).await {
Ok(Ok(response)) => response,
Ok(Err(e)) => {
connection_capture.abort();
server
.metrics
.failed_requests
.fetch_add(1, Ordering::Release);
let msg = if e.is_connect() {
format!("upstream connect error: {e}")
} else {
format!("upstream error: {e}")
};
send_error(frame_tx, stream_id, &msg).await;
return None;
}
Err(_) => {
connection_capture.abort();
server
.metrics
.failed_requests
.fetch_add(1, Ordering::Release);
send_error(frame_tx, stream_id, "upstream timeout").await;
return None;
}
};
// Capture connection-establishment duration (DNS + TCP/TLS + TTFB)
// before proceeding to stream the response body.
let connect_elapsed = connect_start.elapsed();
// Send RESPONSE_HEADERS
let status = response.status().as_u16();
let ttfb_ms = upstream_start.elapsed().as_millis() as u64;
// Short timeout: on connection reuse hyper may never fire the connect
// callback, so avoid blocking indefinitely.
let connection_acquire_ms =
match tokio::time::timeout(Duration::from_millis(100), connection_capture).await {
Ok(Ok(ms)) => ms,
Ok(Err(_)) => None, // JoinError (task panicked / cancelled)
Err(_) => None, // timeout -- task is detached but lightweight
};
let request_timing =
upstream_client::resolve_request_timing(&response, connection_acquire_ms, ttfb_ms);
let mut resp_headers: Vec<(String, String)> = Vec::with_capacity(response.headers().len() + 1);
for (k, v) in response.headers() {
if let Ok(vs) = v.to_str() {
resp_headers.push((k.as_str().to_string(), vs.to_string()));
}
}
let timing = serde_json::json!({
"dns_ms": dns_ms,
"connection_acquire_ms": request_timing.connection_acquire_ms,
"connection_reused": request_timing.connection_reused,
"connect_ms": request_timing.connect_ms,
"tls_ms": request_timing.tls_ms,
"ttfb_ms": ttfb_ms,
"upstream_ms": ttfb_ms,
"response_wait_ms": request_timing.response_wait_ms,
"upstream_processing_ms": request_timing.response_wait_ms,
"timing_source": "instrumented_connector",
"total_ms": connect_elapsed.as_millis() as u64,
"body_size": request_body_size.load(Ordering::Relaxed),
"mode": "tunnel",
});
resp_headers.push(("x-proxy-timing".to_string(), timing.to_string()));
let resp_meta = ResponseMeta {
status,
headers: resp_headers,
};
let meta_json: Bytes = serde_json::to_vec(&resp_meta).unwrap_or_default().into();
let (meta_payload, meta_flags) = compress_payload(meta_json);
if !send_frame(
frame_tx,
TunnelFrame::new(
stream_id,
MsgType::ResponseHeaders,
meta_flags,
meta_payload,
),
)
.await
{
return Some(connect_elapsed);
}
// Stream response body — relay upstream bytes through the tunnel.
// Apply tunnel-level frame compression for chunks that benefit from it
// (e.g. uncompressed SSE text). Already-compressed data (gzip/br from
// upstream Content-Encoding) won't shrink further and will be sent as-is
// thanks to the size check in compress_payload().
let mut stream = response.into_body().into_data_stream();
while let Some(chunk_result) = stream.next().await {
match chunk_result {
Ok(chunk) => {
if chunk.len() <= MAX_CHUNK_SIZE {
let (payload, extra_flags) = compress_payload(chunk);
if !send_frame(
frame_tx,
TunnelFrame::new(stream_id, MsgType::ResponseBody, extra_flags, payload),
)
.await
{
return Some(connect_elapsed);
}
} else {
// Split oversized chunks, compress each slice
let mut offset = 0;
while offset < chunk.len() {
let end = (offset + MAX_CHUNK_SIZE).min(chunk.len());
let slice = chunk.slice(offset..end);
let (payload, extra_flags) = compress_payload(slice);
if !send_frame(
frame_tx,
TunnelFrame::new(
stream_id,
MsgType::ResponseBody,
extra_flags,
payload,
),
)
.await
{
return Some(connect_elapsed);
}
offset = end;
}
}
}
Err(e) => {
server.metrics.stream_errors.fetch_add(1, Ordering::Release);
warn!(stream_id, error = %e, "upstream body read error");
send_error(frame_tx, stream_id, &format!("body read error: {e}")).await;
return Some(connect_elapsed);
}
}
}
// Send STREAM_END
let _ = send_frame(
frame_tx,
TunnelFrame::new(
stream_id,
MsgType::StreamEnd,
flags::END_STREAM,
Bytes::new(),
),
)
.await;
debug!(stream_id, status, "stream completed");
Some(connect_elapsed)
}
async fn send_error(tx: &FrameSender, stream_id: u32, msg: &str) {
// Error frames use best-effort delivery — don't block if writer is congested
let _ = send_frame(
tx,
TunnelFrame::new(
stream_id,
MsgType::StreamError,
0,
Bytes::from(msg.to_string()),
),
)
.await;
}
fn build_streaming_request_body(
body_rx: mpsc::Receiver<TunnelFrame>,
body_size: Arc<AtomicUsize>,
) -> upstream_client::UpstreamRequestBody {
let body_stream = stream::unfold(
(body_rx, body_size, false),
|(mut body_rx, body_size, finished)| async move {
if finished {
return None;
}
loop {
let frame = match body_rx.recv().await {
Some(frame) => frame,
None => return None,
};
match frame.msg_type {
MsgType::RequestBody => {
let end_stream = frame.is_end_stream();
let payload = match decompress_if_gzip(&frame) {
Ok(payload) => payload,
Err(error) => {
let err =
io::Error::other(format!("gzip decompress failed: {error}"));
return Some((Err(err), (body_rx, body_size, true)));
}
};
if payload.is_empty() {
if end_stream {
return None;
}
continue;
}
body_size.fetch_add(payload.len(), Ordering::Relaxed);
return Some((
Ok(BodyFrame::data(payload)),
(body_rx, body_size, end_stream),
));
}
MsgType::StreamError => {
let message = String::from_utf8(frame.payload.to_vec())
.unwrap_or_else(|_| "client cancelled request body".to_string());
return Some((Err(io::Error::other(message)), (body_rx, body_size, true)));
}
MsgType::StreamEnd => return None,
_ => continue,
}
}
},
);
upstream_client::stream_request_body(body_stream)
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use std::sync::atomic::AtomicU64;
use std::sync::Once;
use aether_runtime::{bounded_queue, ConcurrencyGate, DistributedConcurrencyGate};
use arc_swap::ArcSwap;
use super::*;
use crate::config::Config;
use crate::registration::client::AetherClient;
use crate::runtime::DynamicConfig;
use crate::state::ProxyMetrics;
use crate::target_filter::DnsCache;
use crate::tunnel::client::build_tls_config;
#[tokio::test]
async fn streaming_request_body_yields_chunks_and_tracks_size() {
let (tx, rx) = mpsc::channel(4);
let body_size = Arc::new(AtomicUsize::new(0));
let mut body = build_streaming_request_body(rx, Arc::clone(&body_size));
tx.send(TunnelFrame::new(
1,
MsgType::RequestBody,
0,
Bytes::from_static(b"abc"),
))
.await
.expect("send first chunk");
tx.send(TunnelFrame::new(
1,
MsgType::RequestBody,
flags::END_STREAM,
Bytes::from_static(b"def"),
))
.await
.expect("send final chunk");
drop(tx);
let first = body
.frame()
.await
.expect("first frame")
.expect("first frame ok")
.into_data()
.expect("first data frame");
let second = body
.frame()
.await
.expect("second frame")
.expect("second frame ok")
.into_data()
.expect("second data frame");
assert_eq!(first, Bytes::from_static(b"abc"));
assert_eq!(second, Bytes::from_static(b"def"));
assert!(body.frame().await.is_none());
assert_eq!(body_size.load(Ordering::Relaxed), 6);
}
#[tokio::test]
async fn streaming_request_body_surfaces_client_cancel_as_error() {
let (tx, rx) = mpsc::channel(4);
let body_size = Arc::new(AtomicUsize::new(0));
let mut body = build_streaming_request_body(rx, Arc::clone(&body_size));
tx.send(TunnelFrame::new(
1,
MsgType::StreamError,
0,
Bytes::from_static(b"client cancelled"),
))
.await
.expect("send cancel frame");
drop(tx);
let err = body
.frame()
.await
.expect("error frame present")
.expect_err("body should surface cancellation error");
assert!(err.to_string().contains("client cancelled"));
assert!(body.frame().await.is_none());
assert_eq!(body_size.load(Ordering::Relaxed), 0);
}
#[tokio::test]
async fn rejects_stream_when_local_admission_gate_is_saturated() {
let gate = Arc::new(ConcurrencyGate::new("proxy_streams", 1));
let _permit = gate.try_acquire().expect("first permit");
let state = sample_state(Some(gate), None);
let server = sample_server(&state);
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(4);
let (_body_tx, body_rx) = mpsc::channel(1);
handle_stream(
Arc::clone(&state),
server,
7,
sample_request_meta(),
body_rx,
frame_tx,
)
.await;
let frame = frame_rx.recv().await.expect("overload frame");
assert_eq!(frame.stream_id, 7);
assert_eq!(frame.msg_type, MsgType::StreamError);
assert_eq!(frame.payload, Bytes::from_static(b"proxy overloaded"));
assert_eq!(
state
.stream_gate
.as_ref()
.expect("stream gate")
.snapshot()
.rejected,
1
);
}
#[tokio::test]
async fn rejects_stream_when_distributed_admission_gate_is_saturated() {
let gate = Arc::new(DistributedConcurrencyGate::new_in_memory(
"proxy_streams_distributed",
1,
));
let _permit = gate.try_acquire().await.expect("first permit");
let state = sample_state(None, Some(gate));
let server = sample_server(&state);
let (frame_tx, mut frame_rx) = bounded_queue::<TunnelFrame>(4);
let (_body_tx, body_rx) = mpsc::channel(1);
handle_stream(
Arc::clone(&state),
server,
9,
sample_request_meta(),
body_rx,
frame_tx,
)
.await;
let frame = frame_rx.recv().await.expect("overload frame");
assert_eq!(frame.stream_id, 9);
assert_eq!(frame.msg_type, MsgType::StreamError);
assert_eq!(frame.payload, Bytes::from_static(b"proxy overloaded"));
assert_eq!(
state
.distributed_stream_gate
.as_ref()
.expect("distributed gate")
.snapshot()
.await
.expect("distributed snapshot")
.rejected,
1
);
}
fn sample_request_meta() -> RequestMeta {
RequestMeta {
method: "GET".to_string(),
url: "https://example.com/ok".to_string(),
headers: HashMap::new(),
timeout: 30,
}
}
fn sample_state(
stream_gate: Option<Arc<ConcurrencyGate>>,
distributed_stream_gate: Option<Arc<DistributedConcurrencyGate>>,
) -> Arc<AppState> {
ensure_rustls_provider();
let config = Arc::new(sample_config());
let dns_cache = Arc::new(DnsCache::new(Duration::from_secs(60), 128));
let upstream_client =
upstream_client::build_upstream_client(&config, Arc::clone(&dns_cache));
Arc::new(AppState {
config,
dns_cache,
upstream_client,
tunnel_tls_config: Arc::new(build_tls_config()),
stream_gate,
distributed_stream_gate,
})
}
fn sample_server(state: &Arc<AppState>) -> Arc<ServerContext> {
let config = Arc::clone(&state.config);
Arc::new(ServerContext {
server_label: "server".to_string(),
aether_url: config.aether_url.clone(),
management_token: config.management_token.clone(),
node_name: config.node_name.clone(),
node_id: Arc::new(std::sync::RwLock::new("node-1".to_string())),
aether_client: Arc::new(AetherClient::new(
&config,
&config.aether_url,
&config.management_token,
)),
dynamic: Arc::new(ArcSwap::from_pointee(DynamicConfig::from_config(&config))),
active_connections: Arc::new(AtomicU64::new(0)),
metrics: Arc::new(ProxyMetrics::new()),
})
}
fn sample_config() -> Config {
Config {
aether_url: "https://aether.example.com".to_string(),
management_token: "token".to_string(),
public_ip: None,
node_name: "proxy-test".to_string(),
node_region: None,
heartbeat_interval: 30,
allowed_ports: vec![80, 443],
aether_request_timeout_secs: 10,
aether_connect_timeout_secs: 10,
aether_pool_max_idle_per_host: 8,
aether_pool_idle_timeout_secs: 90,
aether_tcp_keepalive_secs: 60,
aether_tcp_nodelay: true,
aether_http2: true,
aether_retry_max_attempts: 3,
aether_retry_base_delay_ms: 200,
aether_retry_max_delay_ms: 2_000,
max_concurrent_connections: None,
max_in_flight_streams: None,
distributed_stream_limit: None,
distributed_stream_redis_url: None,
distributed_stream_redis_key_prefix: None,
distributed_stream_lease_ttl_ms: 30_000,
distributed_stream_renew_interval_ms: 10_000,
distributed_stream_command_timeout_ms: 1_000,
dns_cache_ttl_secs: 60,
dns_cache_capacity: 128,
upstream_connect_timeout_secs: 30,
upstream_pool_max_idle_per_host: 4,
upstream_pool_idle_timeout_secs: 60,
upstream_tcp_keepalive_secs: 60,
upstream_tcp_nodelay: true,
log_level: "info".to_string(),
log_json: false,
tunnel_reconnect_base_ms: 500,
tunnel_reconnect_max_ms: 30_000,
tunnel_ping_interval_secs: 15,
tunnel_max_streams: Some(8),
tunnel_connect_timeout_secs: 15,
tunnel_tcp_keepalive_secs: 30,
tunnel_tcp_nodelay: true,
tunnel_stale_timeout_secs: 45,
tunnel_connections: 1,
}
}
fn ensure_rustls_provider() {
static INIT: Once = Once::new();
INIT.call_once(|| {
let _ = rustls::crypto::ring::default_provider().install_default();
});
}
}

View File

@@ -0,0 +1,63 @@
//! Dedicated WebSocket writer task.
//!
//! All frame writes go through an mpsc channel to a single writer task,
//! avoiding contention on the WebSocket sink. The writer also sends
//! periodic WebSocket Ping frames to keep the connection alive through
//! intermediary proxies (Nginx, Cloudflare, etc.).
use std::time::Duration;
use aether_runtime::{bounded_queue, BoundedQueueSender};
use futures_util::SinkExt;
use tokio::task::JoinHandle;
use tokio_tungstenite::tungstenite::Message;
use tracing::{debug, error, trace};
use super::protocol::Frame;
/// Sender half — cloned by stream handlers and heartbeat.
pub type FrameSender = BoundedQueueSender<Frame>;
/// Spawn the writer task. Returns the sender and a JoinHandle for cleanup.
///
/// `ping_interval` controls WebSocket-level Ping frequency (typically 15s).
/// This keeps the connection alive through intermediary proxies/load-balancers.
pub fn spawn_writer<S>(mut sink: S, ping_interval: Duration) -> (FrameSender, JoinHandle<()>)
where
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
{
let (tx, mut rx) = bounded_queue::<Frame>(256);
let handle = tokio::spawn(async move {
let mut ping_ticker = tokio::time::interval(ping_interval);
ping_ticker.tick().await; // skip first immediate tick
loop {
tokio::select! {
frame = rx.recv() => {
match frame {
Some(frame) => {
let data = frame.encode();
if let Err(e) = sink.send(Message::Binary(data.into())).await {
error!(error = %e, "failed to write frame to WebSocket");
break;
}
}
None => break, // all senders dropped
}
}
_ = ping_ticker.tick() => {
if let Err(e) = sink.send(Message::Ping(vec![])).await {
error!(error = %e, "failed to send WebSocket ping");
break;
}
trace!("sent WebSocket ping");
}
}
}
debug!("writer task exiting");
let _ = sink.close().await;
});
(tx, handle)
}