i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-01 17:00:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(proxy): 将 aether-proxy 从 HMAC 正向代理迁移到 WebSocket 隧道模式

Browse Source 
移除 HMAC 认证、TLS 自签名证书、HTTP CONNECT 代理和代发(delegate)模式，
改为 aether-proxy 主动通过 WebSocket 连接 Aether 服务端建立隧道。

Aether 服务端新增:
- WebSocket 隧道端点 (proxy_tunnel.py)
- TunnelManager 管理隧道连接和请求分发
- TunnelTransport 作为 httpx 自定义 transport 层
- 基于二进制帧的隧道协议 (tunnel_protocol.py)

aether-proxy (Rust) 重构:
- 新增 tunnel 模块 (client/dispatcher/stream_handler/protocol)
- 支持多 Aether 服务端连接 ([[servers]] 配置)
- 移除 proxy/auth/delegate 模块和 hyper 依赖
- 改用 tokio-tungstenite 实现 WebSocket 客户端

同时:
- 添加浏览器指纹 Headers 绕过 Cloudflare 防护
- 删除节点时自动清理 Provider/Endpoint 的代理引用
- 数据库迁移: 新增 tunnel_mode/tunnel_connected/tunnel_connected_at 字段
This commit is contained in:
fawney19
2026-02-25 21:59:29 +08:00
parent 39b036abd5
commit fd9040b9aa
53 changed files with 2938 additions and 2728 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
.env.example
View File

@@ -21,12 +21,6 @@ JWT_SECRET_KEY=change-this-to-a-secure-random-string
# 注意：更换此密钥后需要在管理面板重新配置所有 Provider API Key
ENCRYPTION_KEY=change-this-to-another-secure-random-string
# 代理节点 HMAC 密钥（用于 aether-proxy 认证）
# 可选：不设置时会从 ENCRYPTION_KEY 自动派生
# 显式设置时，aether-proxy.toml 的 hmac_key 配置相同值即可
# 可通过 python generate_keys.py 生成
# PROXY_HMAC_KEY=change-this-to-a-proxy-hmac-key
# 管理员账号（仅首次初始化时使用, 创建完成后可在系统内修改密码）
ADMIN_EMAIL=admin@example.com
ADMIN_USERNAME=admin

6
aether-proxy/.env.example
View File

@@ -4,11 +4,5 @@ AETHER_PROXY_AETHER_URL=https://aether.example.com
# Management Token (ae_xxx, must belong to an ADMIN user)
AETHER_PROXY_MANAGEMENT_TOKEN=ae_xxxxx
# HMAC key (must match Aether's PROXY_HMAC_KEY)
AETHER_PROXY_HMAC_KEY=
# Proxy listen port
AETHER_PROXY_LISTEN_PORT=18080
# Node identification
AETHER_PROXY_NODE_NAME=proxy-01

167
aether-proxy/Cargo.lock generated
View File

@@ -10,7 +10,7 @@ checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa"
[[package]]
name = "aether-proxy"
version = "0.1.4"
version = "0.1.6"
dependencies = [
"anyhow",
"base64",
@@ -20,32 +20,22 @@ dependencies = [
"flate2",
"futures-util",
"hex",
"hmac",
"http-body-util",
"hyper",
"hyper-util",
"libc",
"ratatui",
"rcgen",
"reqwest",
"rustls",
"rustls-pemfile",
"rustls-pki-types",
"serde",
"serde_json",
"sha2",
"subtle",
"sysinfo",
"tar",
"thiserror 2.0.18",
"tokio",
"tokio-rustls",
"tokio-tungstenite",
"toml",
"tower-service",
"tracing",
"tracing-subscriber",
"url",
"webpki-roots",
]
[[package]]
@@ -216,6 +206,12 @@ version = "1.25.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec"
[[package]]
name = "byteorder"
version = "1.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
[[package]]
name = "bytes"
version = "1.11.1"
@@ -479,6 +475,12 @@ dependencies = [
"syn 2.0.114",
]
[[package]]
name = "data-encoding"
version = "2.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d7a1e2f27636f116493b8b860f5546edb47c8d8f8ea73e1d2a20be88e28d1fea"
[[package]]
name = "deltae"
version = "0.3.2"
@@ -524,7 +526,6 @@ checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292"
dependencies = [
"block-buffer",
"crypto-common",
"subtle",
]
[[package]]
@@ -811,15 +812,6 @@ version = "0.4.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
[[package]]
name = "hmac"
version = "0.12.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e"
dependencies = [
"digest",
]
[[package]]
name = "http"
version = "1.4.0"
@@ -859,12 +851,6 @@ version = "1.10.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87"
[[package]]
name = "httpdate"
version = "1.0.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9"
[[package]]
name = "hyper"
version = "1.8.1"
@@ -879,7 +865,6 @@ dependencies = [
"http",
"http-body",
"httparse",
"httpdate",
"itoa",
"pin-project-lite",
"pin-utils",
@@ -902,7 +887,7 @@ dependencies = [
"tokio",
"tokio-rustls",
"tower-service",
"webpki-roots",
"webpki-roots 1.0.6",
]
[[package]]
@@ -1416,16 +1401,6 @@ dependencies = [
"windows-link",
]
[[package]]
name = "pem"
version = "3.0.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be"
dependencies = [
"base64",
"serde_core",
]
[[package]]
name = "percent-encoding"
version = "2.3.2"
@@ -1654,6 +1629,8 @@ version = "0.8.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
dependencies = [
"libc",
"rand_chacha 0.3.1",
"rand_core 0.6.4",
]
@@ -1663,10 +1640,20 @@ version = "0.9.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1"
dependencies = [
"rand_chacha",
"rand_chacha 0.9.0",
"rand_core 0.9.5",
]
[[package]]
name = "rand_chacha"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
dependencies = [
"ppv-lite86",
"rand_core 0.6.4",
]
[[package]]
name = "rand_chacha"
version = "0.9.0"
@@ -1682,6 +1669,9 @@ name = "rand_core"
version = "0.6.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
dependencies = [
"getrandom 0.2.17",
]
[[package]]
name = "rand_core"
@@ -1797,19 +1787,6 @@ dependencies = [
"crossbeam-utils",
]
[[package]]
name = "rcgen"
version = "0.13.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "75e669e5202259b5314d1ea5397316ad400819437857b90861765f24c4cf80a2"
dependencies = [
"pem",
"ring",
"rustls-pki-types",
"time",
"yasna",
]
[[package]]
name = "redox_syscall"
version = "0.5.18"
@@ -1896,7 +1873,7 @@ dependencies = [
"wasm-bindgen-futures",
"wasm-streams",
"web-sys",
"webpki-roots",
"webpki-roots 1.0.6",
]
[[package]]
@@ -1970,15 +1947,6 @@ dependencies = [
"zeroize",
]
[[package]]
name = "rustls-pemfile"
version = "2.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dce314e5fee3f39953d46bb63bb8a46d40c2f8fb7cc5a3b6cab2bde9721d6e50"
dependencies = [
"rustls-pki-types",
]
[[package]]
name = "rustls-pki-types"
version = "1.14.0"
@@ -2089,6 +2057,17 @@ dependencies = [
"serde",
]
[[package]]
name = "sha1"
version = "0.10.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba"
dependencies = [
"cfg-if",
"cpufeatures",
"digest",
]
[[package]]
name = "sha2"
version = "0.10.9"
@@ -2488,6 +2467,22 @@ dependencies = [
"tokio",
]
[[package]]
name = "tokio-tungstenite"
version = "0.24.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "edc5f74e248dc973e0dbb7b74c7e0d6fcc301c694ff50049504004ef4d0cdcd9"
dependencies = [
"futures-util",
"log",
"rustls",
"rustls-pki-types",
"tokio",
"tokio-rustls",
"tungstenite",
"webpki-roots 0.26.11",
]
[[package]]
name = "tokio-util"
version = "0.7.18"
@@ -2667,6 +2662,26 @@ version = "0.2.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b"
[[package]]
name = "tungstenite"
version = "0.24.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "18e5b8366ee7a95b16d32197d0b2604b43a0be89dc5fac9f8e96ccafbaedda8a"
dependencies = [
"byteorder",
"bytes",
"data-encoding",
"http",
"httparse",
"log",
"rand 0.8.5",
"rustls",
"rustls-pki-types",
"sha1",
"thiserror 1.0.69",
"utf-8",
]
[[package]]
name = "typenum"
version = "1.19.0"
@@ -2726,6 +2741,12 @@ dependencies = [
"serde",
]
[[package]]
name = "utf-8"
version = "0.7.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9"
[[package]]
name = "utf8_iter"
version = "1.0.4"
@@ -2887,6 +2908,15 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "webpki-roots"
version = "0.26.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9"
dependencies = [
"webpki-roots 1.0.6",
]
[[package]]
name = "webpki-roots"
version = "1.0.6"
@@ -3236,15 +3266,6 @@ dependencies = [
"rustix 1.1.3",
]
[[package]]
name = "yasna"
version = "0.5.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e17bb3549cc1321ae1296b9cdc2698e2b6cb1992adfa19a8c72e5b7a738f44cd"
dependencies = [
"time",
]
[[package]]
name = "yoke"
version = "0.8.1"

16
aether-proxy/Cargo.toml
View File

@@ -2,19 +2,13 @@
name = "aether-proxy"
version = "0.1.6"
edition = "2021"
description = "Forward proxy for Aether with HMAC authentication"
description = "Tunnel proxy for Aether"
[dependencies]
tokio = { version = "1", features = ["full"] }
hyper = { version = "1", features = ["http1", "server"] }
hyper-util = { version = "0.1", features = ["tokio", "http1", "http2", "server", "client-legacy"] }
tower-service = "0.3"
http-body-util = "0.1"
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "stream", "http2"] }
tokio-tungstenite = { version = "0.24", features = ["rustls-tls-webpki-roots"] }
futures-util = "0.3"
hmac = "0.12"
sha2 = "0.10"
subtle = "2"
base64 = "0.22"
clap = { version = "4", features = ["derive", "env"] }
tracing = "0.1"
@@ -23,15 +17,11 @@ serde = { version = "1", features = ["derive"] }
serde_json = "1"
thiserror = "2"
bytes = "1"
sha2 = "0.10"
hex = "0.4"
anyhow = "1"
toml = "0.8"
tokio-rustls = "0.26"
webpki-roots = "1"
rustls = { version = "0.23", features = ["ring"] }
rustls-pki-types = "1"
rustls-pemfile = "2"
rcgen = "0.13"
ratatui = "0.30"
crossterm = "0.28"
url = "2"

209
aether-proxy/src/app.rs
View File

@@ -1,40 +1,38 @@
//! Application lifecycle: initialization, task orchestration, and shutdown.
//!
//! Extracted from `main.rs` to keep the entry point minimal and consolidate
//! the startup sequence, tracing init, and graceful shutdown logic.
use std::sync::atomic::AtomicU64;
use std::sync::atomic::{AtomicU32, AtomicU64};
use std::sync::{Arc, RwLock};
use std::time::Duration;
use tokio::signal;
use tokio::sync::{watch, Semaphore};
use tracing::{error, info};
use tokio::sync::watch;
use tracing::{error, info, warn};
use crate::config::Config;
use crate::config::{Config, ServerEntry};
use crate::net;
use crate::registration::client::AetherClient;
use crate::runtime::{self, DynamicConfig};
use crate::state::{AppState, ProxyMetrics};
use crate::{hardware, proxy};
use crate::state::{AppState, ProxyMetrics, ServerContext};
use crate::{hardware, target_filter, tunnel};
/// Run the full application lifecycle after config has been parsed.
pub async fn run(mut config: Config) -> anyhow::Result<()> {
pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Result<()> {
init_tracing(&config);
info!(
version = env!("CARGO_PKG_VERSION"),
port = config.listen_port,
node_name = %config.node_name,
"aether-proxy starting"
server_count = servers.len(),
"aether-proxy starting (tunnel mode)"
);
// Resolve public IP
// Resolve public IP (best-effort for region info)
let public_ip = match &config.public_ip {
Some(ip) => ip.clone(),
None => net::detect_public_ip().await?,
None => net::detect_public_ip()
.await
.unwrap_or_else(|_| "0.0.0.0".to_string()),
};
info!(public_ip = %public_ip, "using public IP");
// Auto-detect region if not configured
if config.node_region.is_none() {
@@ -43,118 +41,136 @@ pub async fn run(mut config: Config) -> anyhow::Result<()> {
}
}
// Initialize TLS if enabled
let (tls_acceptor, tls_fingerprint) = if config.enable_tls {
let cert_path = std::path::PathBuf::from(&config.tls_cert);
let key_path = std::path::PathBuf::from(&config.tls_key);
proxy::tls::ensure_self_signed_cert(&cert_path, &key_path)?;
let acceptor = proxy::tls::build_tls_acceptor(&cert_path, &key_path)?;
let fingerprint = proxy::tls::cert_sha256_fingerprint(&cert_path)?;
info!(fingerprint = %fingerprint, "TLS enabled");
(Some(acceptor), Some(fingerprint))
} else {
info!("TLS disabled");
(None, None)
};
// Collect hardware info (once at startup)
// Collect hardware info (once at startup, sent during registration)
let hw_info = hardware::collect();
let max_connections_raw = config
.max_concurrent_connections
.unwrap_or(hw_info.estimated_max_concurrency)
.max(1);
let max_connections = usize::try_from(max_connections_raw).unwrap_or(usize::MAX);
// Auto-detect tunnel_max_streams from hardware if not explicitly set
if config.tunnel_max_streams.is_none() {
let auto = (hw_info.estimated_max_concurrency / 10).clamp(64, 1024) as u32;
config.tunnel_max_streams = Some(auto);
info!(
max_connections = max_connections_raw,
"connection limit configured"
tunnel_max_streams = auto,
"auto-detected tunnel_max_streams from hardware"
);
}
info!(
max_concurrency = hw_info.estimated_max_concurrency,
"hardware info collected"
);
let connection_semaphore = Arc::new(Semaphore::new(max_connections));
let metrics = Arc::new(ProxyMetrics::new());
let dns_cache = Arc::new(proxy::target_filter::DnsCache::new(
let dns_cache = Arc::new(target_filter::DnsCache::new(
Duration::from_secs(config.dns_cache_ttl_secs),
config.dns_cache_capacity,
));
// Register with Aether
let aether_client = Arc::new(AetherClient::new(&config));
let node_id = aether_client
.register(
// Build reqwest client for tunnel upstream requests (shared).
let reqwest_client = reqwest::Client::builder()
.pool_max_idle_per_host(config.upstream_pool_max_idle_per_host)
.pool_idle_timeout(Duration::from_secs(config.upstream_pool_idle_timeout_secs))
.connect_timeout(Duration::from_secs(config.upstream_connect_timeout_secs))
.tcp_nodelay(config.upstream_tcp_nodelay)
.build()
.expect("failed to build reqwest client");
// Register with each Aether server and build per-server contexts
let mut server_contexts: Vec<Arc<ServerContext>> = Vec::new();
for (i, entry) in servers.iter().enumerate() {
let label = if servers.len() == 1 {
"server".to_string()
} else {
format!("server-{}", i)
};
let node_name = entry
.node_name
.clone()
.unwrap_or_else(|| config.node_name.clone());
let client = Arc::new(AetherClient::new(
&config,
&public_ip,
config.enable_tls,
tls_fingerprint.as_deref(),
Some(&hw_info),
)
.await?;
&entry.aether_url,
&entry.management_token,
));
match client
.register(&config, &node_name, &public_ip, Some(&hw_info))
.await
{
Ok(node_id) => {
info!(server = %label, node_id = %node_id, url = %entry.aether_url, node_name = %node_name, "registered");
server_contexts.push(Arc::new(ServerContext {
server_label: label,
aether_url: entry.aether_url.clone(),
management_token: entry.management_token.clone(),
node_name,
node_id: Arc::new(RwLock::new(node_id)),
aether_client: client,
dynamic: Arc::new(RwLock::new(DynamicConfig::from_config(&config))),
active_connections: Arc::new(AtomicU64::new(0)),
metrics: Arc::new(ProxyMetrics::new()),
reconnect_attempts: AtomicU32::new(0),
}));
}
Err(e) => {
warn!(
server = %label,
url = %entry.aether_url,
error = %e,
"registration failed, skipping server"
);
}
}
}
info!(node_id = %node_id, "node registered");
// Build DynamicConfig before moving config into Arc
let dynamic = Arc::new(RwLock::new(DynamicConfig::from_config(&config)));
// Build delegate HTTP client (for proxy-initiated upstream requests).
let delegate_client = proxy::delegate_client::build_delegate_client(&config);
if server_contexts.is_empty() {
anyhow::bail!("no servers registered successfully");
}
// Build shared application state
let state = Arc::new(AppState {
config: Arc::new(config),
node_id: Arc::new(RwLock::new(node_id)),
dynamic,
aether_client,
hardware_info: Arc::new(hw_info),
public_ip,
tls_fingerprint,
tls_acceptor,
delegate_client,
active_connections: Arc::new(AtomicU64::new(0)),
connection_semaphore,
dns_cache,
metrics,
reqwest_client,
});
// Shutdown signal channel
let (shutdown_tx, shutdown_rx) = watch::channel(false);
// Start heartbeat task
let heartbeat_handle = {
let state = Arc::clone(&state);
let rx = shutdown_rx.clone();
tokio::spawn(async move {
crate::registration::heartbeat::run(&state, rx).await;
})
};
info!(
active_servers = server_contexts.len(),
"running in tunnel mode"
);
// Start proxy server
let server_handle = {
let state = Arc::clone(&state);
// Spawn one tunnel task per server
let mut tunnel_handles = Vec::new();
for server in &server_contexts {
let s = Arc::clone(&state);
let srv = Arc::clone(server);
let rx = shutdown_rx.clone();
tokio::spawn(async move {
if let Err(e) = proxy::server::run(&state, rx).await {
error!(error = %e, "proxy server error");
tunnel_handles.push(tokio::spawn(async move {
tunnel::run(&s, &srv, rx).await;
}));
}
})
};
// Wait for shutdown signal (SIGTERM or SIGINT)
// Wait for shutdown signal
wait_for_shutdown().await;
info!("shutdown signal received, cleaning up...");
// Signal all tasks to stop
let _ = shutdown_tx.send(true);
// Graceful unregister (best-effort)
let current_node_id = state.node_id.read().unwrap().clone();
if let Err(e) = state.aether_client.unregister(&current_node_id).await {
error!(error = %e, "unregister failed during shutdown");
// Graceful unregister from all servers
for server in &server_contexts {
let node_id = server.node_id.read().unwrap().clone();
if let Err(e) = server.aether_client.unregister(&node_id).await {
error!(
server = %server.server_label,
error = %e,
"unregister failed during shutdown"
);
}
}
// Wait for tasks to finish
let _ = tokio::join!(heartbeat_handle, server_handle);
// Wait for all tunnel tasks
for h in tunnel_handles {
let _ = h.await;
}
info!("aether-proxy stopped");
Ok(())
@@ -168,7 +184,6 @@ fn init_tracing(config: &Config) {
let (filter_layer, reload_handle) = reload::Layer::new(filter);
// Register log-level hot-reloader
runtime::set_log_reloader(Box::new(move |level: &str| {
if let Ok(new_filter) = EnvFilter::try_new(level) {
let _ = reload_handle.modify(|f| *f = new_filter);

196
aether-proxy/src/auth/hmac.rs
View File

@@ -1,196 +0,0 @@
use base64::Engine;
use hmac::{Hmac, Mac};
use sha2::Sha256;
use subtle::ConstantTimeEq;
use crate::config::Config;
type HmacSha256 = Hmac<Sha256>;
#[derive(Debug)]
pub enum AuthError {
MissingHeader,
InvalidBasicAuth,
InvalidUsername,
InvalidPasswordFormat,
TimestampParseError,
TimestampExpired,
SignatureMismatch,
}
impl std::fmt::Display for AuthError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::MissingHeader => write!(f, "missing Proxy-Authorization header"),
Self::InvalidBasicAuth => write!(f, "invalid Basic auth encoding"),
Self::InvalidUsername => write!(f, "username must be 'hmac'"),
Self::InvalidPasswordFormat => {
write!(f, "password format must be 'timestamp.signature'")
}
Self::TimestampParseError => write!(f, "invalid timestamp"),
Self::TimestampExpired => write!(f, "timestamp outside tolerance window"),
Self::SignatureMismatch => write!(f, "HMAC signature mismatch"),
}
}
}
/// Validate Proxy-Authorization header.
///
/// Expected format: `Basic base64(hmac:{timestamp}.{signature})`
/// where signature = hex(HMAC-SHA256(hmac_key, "{timestamp}"))
///
/// The signature no longer includes `node_id`, eliminating race conditions
/// during re-registration where the Aether server's cached `node_id` could
/// differ from the proxy's freshly assigned `node_id`.
///
/// `timestamp_tolerance` is accepted separately so the caller can supply
/// the value from [`DynamicConfig`](crate::runtime::DynamicConfig) (which
/// may be updated remotely).
pub fn validate_proxy_auth(
proxy_auth_header: Option<&str>,
config: &Config,
timestamp_tolerance: u64,
) -> Result<(), AuthError> {
let header = proxy_auth_header.ok_or(AuthError::MissingHeader)?;
let encoded = header
.strip_prefix("Basic ")
.or_else(|| header.strip_prefix("basic "))
.ok_or(AuthError::InvalidBasicAuth)?;
let decoded_bytes = base64::engine::general_purpose::STANDARD
.decode(encoded.trim())
.map_err(|_| AuthError::InvalidBasicAuth)?;
let decoded = String::from_utf8(decoded_bytes).map_err(|_| AuthError::InvalidBasicAuth)?;
// format: hmac:{timestamp}.{signature}
let (username, password) = decoded.split_once(':').ok_or(AuthError::InvalidBasicAuth)?;
if username != "hmac" {
return Err(AuthError::InvalidUsername);
}
let (timestamp_str, signature_hex) = password
.split_once('.')
.ok_or(AuthError::InvalidPasswordFormat)?;
// Validate timestamp window
let timestamp: u64 = timestamp_str
.parse()
.map_err(|_| AuthError::TimestampParseError)?;
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("system clock before epoch")
.as_secs();
let diff = now.abs_diff(timestamp);
if diff > timestamp_tolerance {
return Err(AuthError::TimestampExpired);
}
// Recompute signature: HMAC-SHA256(key, timestamp)
let mut mac =
HmacSha256::new_from_slice(config.hmac_key.as_bytes()).expect("HMAC accepts any key size");
mac.update(timestamp_str.as_bytes());
let expected = mac.finalize().into_bytes();
let expected_hex = hex::encode(expected);
// Constant-time comparison
let sig_bytes = signature_hex.as_bytes();
let exp_bytes = expected_hex.as_bytes();
if sig_bytes.len() != exp_bytes.len() || sig_bytes.ct_eq(exp_bytes).unwrap_u8() != 1 {
return Err(AuthError::SignatureMismatch);
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
fn make_config() -> Config {
Config {
aether_url: String::new(),
management_token: String::new(),
hmac_key: "test-hmac-key".to_string(),
listen_port: 18080,
public_ip: None,
node_name: "test".to_string(),
node_region: None,
heartbeat_interval: 30,
allowed_ports: vec![80, 443],
timestamp_tolerance: 300,
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: 2000,
max_concurrent_connections: None,
connect_timeout_secs: 30,
tls_handshake_timeout_secs: 10,
dns_cache_ttl_secs: 60,
dns_cache_capacity: 1024,
delegate_connect_timeout_secs: 30,
delegate_pool_max_idle_per_host: 64,
delegate_pool_idle_timeout_secs: 300,
delegate_tcp_keepalive_secs: 60,
delegate_tcp_nodelay: true,
log_level: "info".to_string(),
log_json: false,
enable_tls: false,
tls_cert: String::new(),
tls_key: String::new(),
}
}
fn make_valid_auth(config: &Config) -> String {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs();
let mut mac = HmacSha256::new_from_slice(config.hmac_key.as_bytes()).unwrap();
mac.update(now.to_string().as_bytes());
let sig = hex::encode(mac.finalize().into_bytes());
let cred = format!("hmac:{}.{}", now, sig);
let encoded = base64::engine::general_purpose::STANDARD.encode(cred);
format!("Basic {}", encoded)
}
#[test]
fn test_valid_auth() {
let config = make_config();
let header = make_valid_auth(&config);
assert!(validate_proxy_auth(Some(&header), &config, config.timestamp_tolerance).is_ok());
}
#[test]
fn test_missing_header() {
let config = make_config();
assert!(matches!(
validate_proxy_auth(None, &config, config.timestamp_tolerance),
Err(AuthError::MissingHeader)
));
}
#[test]
fn test_wrong_username() {
let cred = "user:12345.abc";
let encoded = base64::engine::general_purpose::STANDARD.encode(cred);
let header = format!("Basic {}", encoded);
let config = make_config();
assert!(matches!(
validate_proxy_auth(Some(&header), &config, config.timestamp_tolerance),
Err(AuthError::InvalidUsername)
));
}
}

3
aether-proxy/src/auth/mod.rs
View File

@@ -1,3 +0,0 @@
pub mod hmac;
pub use self::hmac::validate_proxy_auth;

215
aether-proxy/src/config.rs
View File

@@ -3,11 +3,11 @@ use std::path::Path;
use clap::Parser;
use serde::{Deserialize, Serialize};
/// Aether forward proxy with HMAC authentication.
/// Aether tunnel proxy.
///
/// Deployed on overseas VPS to relay API traffic for Aether instances
/// behind the GFW. Registers with Aether, sends heartbeats, and validates
/// incoming proxy requests via HMAC-SHA256 signatures in Basic Auth.
/// behind the GFW. Connects to Aether via WebSocket tunnel, registers
/// with Aether, and relays upstream requests.
#[derive(Parser, Debug, Clone)]
#[command(version, about)]
pub struct Config {
@@ -19,14 +19,6 @@ pub struct Config {
#[arg(long, env = "AETHER_PROXY_MANAGEMENT_TOKEN")]
pub management_token: String,
/// HMAC-SHA256 key for proxy authentication
#[arg(long, env = "AETHER_PROXY_HMAC_KEY")]
pub hmac_key: String,
/// Port to listen on for proxy connections
#[arg(long, env = "AETHER_PROXY_LISTEN_PORT", default_value_t = 18080)]
pub listen_port: u16,
/// Public IP address of this node (auto-detected if omitted)
#[arg(long, env = "AETHER_PROXY_PUBLIC_IP")]
pub public_ip: Option<String>,
@@ -52,10 +44,6 @@ pub struct Config {
)]
pub allowed_ports: Vec<u16>,
/// Timestamp tolerance window in seconds for HMAC validation
#[arg(long, env = "AETHER_PROXY_TIMESTAMP_TOLERANCE", default_value_t = 300)]
pub timestamp_tolerance: u64,
/// Aether API request timeout in seconds
#[arg(
long,
@@ -128,14 +116,6 @@ pub struct Config {
#[arg(long, env = "AETHER_PROXY_MAX_CONCURRENT_CONNECTIONS")]
pub max_concurrent_connections: Option<u64>,
/// Upstream TCP connect timeout in seconds for CONNECT tunnels
#[arg(long, env = "AETHER_PROXY_CONNECT_TIMEOUT", default_value_t = 30)]
pub connect_timeout_secs: u64,
/// TLS handshake timeout in seconds for incoming TLS connections
#[arg(long, env = "AETHER_PROXY_TLS_HANDSHAKE_TIMEOUT", default_value_t = 10)]
pub tls_handshake_timeout_secs: u64,
/// DNS cache TTL in seconds
#[arg(long, env = "AETHER_PROXY_DNS_CACHE_TTL", default_value_t = 60)]
pub dns_cache_ttl_secs: u64,
@@ -144,45 +124,45 @@ pub struct Config {
#[arg(long, env = "AETHER_PROXY_DNS_CACHE_CAPACITY", default_value_t = 1024)]
pub dns_cache_capacity: usize,
/// Delegate HTTP client connect timeout in seconds
/// Upstream HTTP client connect timeout in seconds
#[arg(
long,
env = "AETHER_PROXY_DELEGATE_CONNECT_TIMEOUT",
env = "AETHER_PROXY_UPSTREAM_CONNECT_TIMEOUT",
default_value_t = 30
)]
pub delegate_connect_timeout_secs: u64,
pub upstream_connect_timeout_secs: u64,
/// Delegate HTTP client max idle connections per host
/// Upstream HTTP client max idle connections per host
#[arg(
long,
env = "AETHER_PROXY_DELEGATE_POOL_MAX_IDLE_PER_HOST",
env = "AETHER_PROXY_UPSTREAM_POOL_MAX_IDLE_PER_HOST",
default_value_t = 64
)]
pub delegate_pool_max_idle_per_host: usize,
pub upstream_pool_max_idle_per_host: usize,
/// Delegate HTTP client idle timeout in seconds
/// Upstream HTTP client idle timeout in seconds
#[arg(
long,
env = "AETHER_PROXY_DELEGATE_POOL_IDLE_TIMEOUT",
env = "AETHER_PROXY_UPSTREAM_POOL_IDLE_TIMEOUT",
default_value_t = 300
)]
pub delegate_pool_idle_timeout_secs: u64,
pub upstream_pool_idle_timeout_secs: u64,
/// Delegate TCP keepalive in seconds (0 disables)
/// Upstream TCP keepalive in seconds (0 disables)
#[arg(
long,
env = "AETHER_PROXY_DELEGATE_TCP_KEEPALIVE",
env = "AETHER_PROXY_UPSTREAM_TCP_KEEPALIVE",
default_value_t = 60
)]
pub delegate_tcp_keepalive_secs: u64,
pub upstream_tcp_keepalive_secs: u64,
/// Delegate TCP_NODELAY
/// Upstream TCP_NODELAY
#[arg(
long,
env = "AETHER_PROXY_DELEGATE_TCP_NODELAY",
env = "AETHER_PROXY_UPSTREAM_TCP_NODELAY",
default_value_t = true
)]
pub delegate_tcp_nodelay: bool,
pub upstream_tcp_nodelay: bool,
/// Log level (trace, debug, info, warn, error)
#[arg(long, env = "AETHER_PROXY_LOG_LEVEL", default_value = "info")]
@@ -192,25 +172,38 @@ pub struct Config {
#[arg(long, env = "AETHER_PROXY_LOG_JSON", default_value_t = false)]
pub log_json: bool,
/// Enable TLS encryption (dual-stack: accepts both HTTP and TLS on same port)
#[arg(long, env = "AETHER_PROXY_ENABLE_TLS", default_value_t = true)]
pub enable_tls: bool,
/// Path to TLS certificate PEM file
/// WebSocket reconnect base delay in milliseconds
#[arg(
long,
env = "AETHER_PROXY_TLS_CERT",
default_value = "aether-proxy-cert.pem"
env = "AETHER_PROXY_TUNNEL_RECONNECT_BASE_MS",
default_value_t = 1000
)]
pub tls_cert: String,
pub tunnel_reconnect_base_ms: u64,
/// Path to TLS private key PEM file
/// WebSocket reconnect max delay in milliseconds
#[arg(
long,
env = "AETHER_PROXY_TLS_KEY",
default_value = "aether-proxy-key.pem"
env = "AETHER_PROXY_TUNNEL_RECONNECT_MAX_MS",
default_value_t = 30000
)]
pub tls_key: String,
pub tunnel_reconnect_max_ms: u64,
/// WebSocket tunnel ping interval in seconds
#[arg(long, env = "AETHER_PROXY_TUNNEL_PING_INTERVAL", default_value_t = 15)]
pub tunnel_ping_interval_secs: u64,
/// Maximum concurrent streams over tunnel (auto-detected from hardware if omitted)
#[arg(long, env = "AETHER_PROXY_TUNNEL_MAX_STREAMS")]
pub tunnel_max_streams: Option<u32>,
}
/// Per-server connection config (used in multi-server TOML `[[servers]]`).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ServerEntry {
pub aether_url: String,
pub management_token: String,
/// Per-server node name override. Falls back to the global `node_name`.
pub node_name: Option<String>,
}
// ---------------------------------------------------------------------------
@@ -218,7 +211,7 @@ pub struct Config {
// ---------------------------------------------------------------------------
/// Serializable config for TOML file persistence.
/// All fields are optional — only populated values are written.
/// All fields are optional -- only populated values are written.
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct ConfigFile {
#[serde(skip_serializing_if = "Option::is_none")]
@@ -226,10 +219,6 @@ pub struct ConfigFile {
#[serde(skip_serializing_if = "Option::is_none")]
pub management_token: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub hmac_key: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub listen_port: Option<u16>,
#[serde(skip_serializing_if = "Option::is_none")]
pub public_ip: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub node_name: Option<String>,
@@ -240,8 +229,6 @@ pub struct ConfigFile {
#[serde(skip_serializing_if = "Option::is_none")]
pub allowed_ports: Option<Vec<u16>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub timestamp_tolerance: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub aether_request_timeout_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub aether_connect_timeout_secs: Option<u64>,
@@ -264,33 +251,37 @@ pub struct ConfigFile {
#[serde(skip_serializing_if = "Option::is_none")]
pub max_concurrent_connections: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub connect_timeout_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tls_handshake_timeout_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub dns_cache_ttl_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub dns_cache_capacity: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub delegate_connect_timeout_secs: Option<u64>,
pub upstream_connect_timeout_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub delegate_pool_max_idle_per_host: Option<usize>,
pub upstream_pool_max_idle_per_host: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub delegate_pool_idle_timeout_secs: Option<u64>,
pub upstream_pool_idle_timeout_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub delegate_tcp_keepalive_secs: Option<u64>,
pub upstream_tcp_keepalive_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub delegate_tcp_nodelay: Option<bool>,
pub upstream_tcp_nodelay: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub log_level: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub log_json: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub enable_tls: Option<bool>,
pub tunnel_reconnect_base_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tls_cert: Option<String>,
pub tunnel_reconnect_max_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tls_key: Option<String>,
pub tunnel_ping_interval_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tunnel_max_streams: Option<u32>,
/// Multi-server config: each entry connects to a separate Aether instance.
/// When present, top-level aether_url/management_token are ignored for
/// tunnel connections (but still injected as env for clap compatibility).
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub servers: Vec<ServerEntry>,
}
impl ConfigFile {
@@ -307,6 +298,24 @@ impl ConfigFile {
Ok(())
}
/// Resolve the effective server list.
///
/// If `[[servers]]` is present, use it. Otherwise fall back to the
/// top-level `aether_url` + `management_token` as a single server.
pub fn effective_servers(&self) -> Vec<ServerEntry> {
if !self.servers.is_empty() {
return self.servers.clone();
}
match (&self.aether_url, &self.management_token) {
(Some(url), Some(token)) => vec![ServerEntry {
aether_url: url.clone(),
management_token: token.clone(),
node_name: None,
}],
_ => vec![],
}
}
/// Inject values as environment variables so clap picks them up.
///
/// Only sets variables that are **not** already present in the
@@ -332,15 +341,30 @@ impl ConfigFile {
}
};
}
set!("AETHER_PROXY_AETHER_URL", self.aether_url);
set!("AETHER_PROXY_MANAGEMENT_TOKEN", self.management_token);
set!("AETHER_PROXY_HMAC_KEY", self.hmac_key);
set!("AETHER_PROXY_LISTEN_PORT", self.listen_port);
// When top-level fields are absent, fall back to the first [[servers]]
// entry so that clap's required `aether_url` / `management_token` are
// satisfied even with the new config format.
let first_server = self.servers.first();
let aether_url = self
.aether_url
.clone()
.or_else(|| first_server.map(|s| s.aether_url.clone()));
let management_token = self
.management_token
.clone()
.or_else(|| first_server.map(|s| s.management_token.clone()));
let node_name = self
.node_name
.clone()
.or_else(|| first_server.and_then(|s| s.node_name.clone()));
set!("AETHER_PROXY_AETHER_URL", aether_url);
set!("AETHER_PROXY_MANAGEMENT_TOKEN", management_token);
set!("AETHER_PROXY_PUBLIC_IP", self.public_ip);
set!("AETHER_PROXY_NODE_NAME", self.node_name);
set!("AETHER_PROXY_NODE_NAME", node_name);
set!("AETHER_PROXY_NODE_REGION", self.node_region);
set!("AETHER_PROXY_HEARTBEAT_INTERVAL", self.heartbeat_interval);
set!("AETHER_PROXY_TIMESTAMP_TOLERANCE", self.timestamp_tolerance);
set!(
"AETHER_PROXY_AETHER_REQUEST_TIMEOUT",
self.aether_request_timeout_secs
@@ -379,38 +403,43 @@ impl ConfigFile {
"AETHER_PROXY_MAX_CONCURRENT_CONNECTIONS",
self.max_concurrent_connections
);
set!("AETHER_PROXY_CONNECT_TIMEOUT", self.connect_timeout_secs);
set!(
"AETHER_PROXY_TLS_HANDSHAKE_TIMEOUT",
self.tls_handshake_timeout_secs
);
set!("AETHER_PROXY_DNS_CACHE_TTL", self.dns_cache_ttl_secs);
set!("AETHER_PROXY_DNS_CACHE_CAPACITY", self.dns_cache_capacity);
set!(
"AETHER_PROXY_DELEGATE_CONNECT_TIMEOUT",
self.delegate_connect_timeout_secs
"AETHER_PROXY_UPSTREAM_CONNECT_TIMEOUT",
self.upstream_connect_timeout_secs
);
set!(
"AETHER_PROXY_DELEGATE_POOL_MAX_IDLE_PER_HOST",
self.delegate_pool_max_idle_per_host
"AETHER_PROXY_UPSTREAM_POOL_MAX_IDLE_PER_HOST",
self.upstream_pool_max_idle_per_host
);
set!(
"AETHER_PROXY_DELEGATE_POOL_IDLE_TIMEOUT",
self.delegate_pool_idle_timeout_secs
"AETHER_PROXY_UPSTREAM_POOL_IDLE_TIMEOUT",
self.upstream_pool_idle_timeout_secs
);
set!(
"AETHER_PROXY_DELEGATE_TCP_KEEPALIVE",
self.delegate_tcp_keepalive_secs
"AETHER_PROXY_UPSTREAM_TCP_KEEPALIVE",
self.upstream_tcp_keepalive_secs
);
set!(
"AETHER_PROXY_DELEGATE_TCP_NODELAY",
self.delegate_tcp_nodelay
"AETHER_PROXY_UPSTREAM_TCP_NODELAY",
self.upstream_tcp_nodelay
);
set!("AETHER_PROXY_LOG_LEVEL", self.log_level);
set!("AETHER_PROXY_LOG_JSON", self.log_json);
set!("AETHER_PROXY_ENABLE_TLS", self.enable_tls);
set!("AETHER_PROXY_TLS_CERT", self.tls_cert);
set!("AETHER_PROXY_TLS_KEY", self.tls_key);
set!(
"AETHER_PROXY_TUNNEL_RECONNECT_BASE_MS",
self.tunnel_reconnect_base_ms
);
set!(
"AETHER_PROXY_TUNNEL_RECONNECT_MAX_MS",
self.tunnel_reconnect_max_ms
);
set!(
"AETHER_PROXY_TUNNEL_PING_INTERVAL",
self.tunnel_ping_interval_secs
);
set!("AETHER_PROXY_TUNNEL_MAX_STREAMS", self.tunnel_max_streams);
// allowed_ports needs special handling (comma-separated)
if let Some(ref ports) = self.allowed_ports {

29
aether-proxy/src/main.rs
View File

@@ -1,13 +1,13 @@
mod app;
mod auth;
mod config;
mod hardware;
mod net;
mod proxy;
mod registration;
mod runtime;
mod setup;
mod state;
mod target_filter;
mod tunnel;
use std::path::PathBuf;
@@ -135,5 +135,28 @@ async fn run_proxy(config: Config) -> anyhow::Result<()> {
std::process::exit(1);
}
app::run(config).await
// Resolve server list: prefer [[servers]] from TOML, fall back to CLI/env single server.
let config_path =
std::env::var("AETHER_PROXY_CONFIG").unwrap_or_else(|_| DEFAULT_CONFIG.to_string());
let servers = if std::path::Path::new(&config_path).exists() {
config::ConfigFile::load(std::path::Path::new(&config_path))
.ok()
.map(|f| f.effective_servers())
.filter(|s| !s.is_empty())
.unwrap_or_else(|| {
vec![config::ServerEntry {
aether_url: config.aether_url.clone(),
management_token: config.management_token.clone(),
node_name: None,
}]
})
} else {
vec![config::ServerEntry {
aether_url: config.aether_url.clone(),
management_token: config.management_token.clone(),
node_name: None,
}]
};
app::run(config, servers).await
}

163
aether-proxy/src/proxy/connect.rs
View File

@@ -1,163 +0,0 @@
use std::collections::HashSet;
use std::sync::Arc;
use std::time::Duration;
use hyper::body::Incoming;
use hyper::{Request, Response};
use tokio::net::TcpStream;
use tokio::time::timeout;
use tracing::{debug, warn};
use crate::auth;
use crate::config::Config;
use crate::proxy::target_filter::{self, DnsCache};
/// Handle HTTP CONNECT tunnel requests.
///
/// Flow: validate auth -> check target filter -> TCP connect -> 200 -> bidirectional copy
pub async fn handle_connect(
req: Request<Incoming>,
config: Arc<Config>,
allowed_ports: &HashSet<u16>,
timestamp_tolerance: u64,
dns_cache: &DnsCache,
) -> Response<http_body_util::Empty<bytes::Bytes>> {
// Extract Proxy-Authorization header
let proxy_auth = req
.headers()
.get("proxy-authorization")
.and_then(|v| v.to_str().ok());
// HMAC authentication
if let Err(e) = auth::validate_proxy_auth(proxy_auth, &config, timestamp_tolerance) {
warn!(error = %e, "CONNECT auth failed");
return proxy_auth_required(&e.to_string());
}
// Parse target host:port from CONNECT URI
let authority = match req.uri().authority() {
Some(auth) => auth.clone(),
None => {
warn!("CONNECT request missing authority");
return bad_request("missing target authority");
}
};
let host = authority.host().to_string();
let port = authority.port_u16().unwrap_or(443);
// Target filter: private IP + port whitelist
let target_addr =
match target_filter::validate_target(&host, port, allowed_ports, dns_cache).await {
Ok(addr) => addr,
Err(e) => {
warn!(host = %host, port, error = %e, "CONNECT target rejected");
return forbidden(&e.to_string());
}
};
debug!(target = %target_addr, "CONNECT tunnel establishing");
// Connect to target
let connect_timeout = Duration::from_secs(config.connect_timeout_secs);
let target_stream = match timeout(connect_timeout, TcpStream::connect(target_addr)).await {
Ok(Ok(s)) => s,
Ok(Err(e)) => {
warn!(target = %target_addr, error = %e, "CONNECT target connection failed");
return bad_gateway(&e.to_string());
}
Err(_) => {
warn!(target = %target_addr, "CONNECT target connection timeout");
return gateway_timeout("connect timeout");
}
};
if let Err(e) = target_stream.set_nodelay(true) {
debug!(target = %target_addr, error = %e, "failed to set TCP_NODELAY");
}
// Respond 200 and upgrade connection to raw TCP tunnel
let target_display = target_addr.to_string();
// Reuse connect_timeout for upgrade: both are connection-phase operations
// and should complete within the same order of magnitude.
let upgrade_timeout = Duration::from_secs(config.connect_timeout_secs);
tokio::task::spawn(async move {
match timeout(upgrade_timeout, hyper::upgrade::on(req)).await {
Ok(Ok(upgraded)) => {
let mut upgraded = hyper_util::rt::TokioIo::new(upgraded);
let mut target = target_stream;
match tokio::io::copy_bidirectional(&mut upgraded, &mut target).await {
Ok((from_client, from_target)) => {
debug!(
target = %target_display,
from_client,
from_target,
"CONNECT tunnel closed"
);
}
Err(e) => {
debug!(target = %target_display, error = %e, "CONNECT tunnel error");
}
}
}
Ok(Err(e)) => {
warn!(target = %target_display, error = %e, "CONNECT upgrade failed");
}
Err(_) => {
warn!(target = %target_display, "CONNECT upgrade timeout");
}
}
});
Response::builder()
.status(200)
.body(http_body_util::Empty::new())
.unwrap()
}
fn proxy_auth_required(msg: &str) -> Response<http_body_util::Empty<bytes::Bytes>> {
Response::builder()
.status(407)
.header("Proxy-Authenticate", "HMAC-SHA256")
.header("Content-Length", "0")
.header("X-Error", msg)
.body(http_body_util::Empty::new())
.unwrap()
}
fn forbidden(msg: &str) -> Response<http_body_util::Empty<bytes::Bytes>> {
Response::builder()
.status(403)
.header("Content-Length", "0")
.header("X-Error", msg)
.body(http_body_util::Empty::new())
.unwrap()
}
fn bad_request(msg: &str) -> Response<http_body_util::Empty<bytes::Bytes>> {
Response::builder()
.status(400)
.header("Content-Length", "0")
.header("X-Error", msg)
.body(http_body_util::Empty::new())
.unwrap()
}
fn bad_gateway(msg: &str) -> Response<http_body_util::Empty<bytes::Bytes>> {
Response::builder()
.status(502)
.header("Content-Length", "0")
.header("X-Error", msg)
.body(http_body_util::Empty::new())
.unwrap()
}
fn gateway_timeout(msg: &str) -> Response<http_body_util::Empty<bytes::Bytes>> {
Response::builder()
.status(504)
.header("Content-Length", "0")
.header("X-Error", msg)
.body(http_body_util::Empty::new())
.unwrap()
}

389
aether-proxy/src/proxy/delegate.rs
View File

@@ -1,389 +0,0 @@
use std::collections::HashMap;
use std::collections::HashSet;
use std::error::Error as StdError;
use std::sync::Arc;
use std::time::Instant;
use futures_util::StreamExt;
use http_body_util::{BodyExt, Full, Limited, StreamBody};
use hyper::body::{Frame, Incoming};
use hyper::header::{HeaderName, HeaderValue};
use hyper::{Method, Request, Response, Uri};
use tracing::{debug, warn};
use url::Url;
use super::BoxBody;
use crate::auth;
use crate::config::Config;
use crate::proxy::delegate_client::{ConnectTiming, DelegateClient};
use crate::proxy::target_filter::{self, DnsCache};
/// Handle delegation requests: Aether sends a full request description,
/// and the proxy issues the actual upstream HTTP call using its own TLS stack.
///
/// Endpoint: POST /_aether/delegate
///
/// Wire format: metadata in HTTP headers, upstream body sent directly
/// as HTTP body (optionally gzip-compressed via `Content-Encoding: gzip`).
///
/// Headers:
/// X-Delegate-Method: POST
/// X-Delegate-Url: https://api.anthropic.com/v1/messages
/// X-Delegate-Headers: base64-encoded JSON {"Authorization": "Bearer ...", ...}
/// X-Delegate-Timeout: 30 (accepted but not used — Aether controls timeouts)
/// Content-Encoding: gzip (optional, indicates body is gzip-compressed)
pub async fn handle_delegate(
req: Request<Incoming>,
config: Arc<Config>,
allowed_ports: &HashSet<u16>,
timestamp_tolerance: u64,
dns_cache: &DnsCache,
http_client: &DelegateClient,
) -> Response<BoxBody> {
let total_start = Instant::now();
// ── Auth ──
let auth_header = req
.headers()
.get("authorization")
.and_then(|v| v.to_str().ok());
if let Err(e) = auth::validate_proxy_auth(auth_header, &config, timestamp_tolerance) {
warn!(error = %e, "delegate auth failed");
return error_response(401, "authentication_failed", &e.to_string());
}
let auth_ms = total_start.elapsed().as_millis() as u64;
// ── Parse metadata from headers ──
let meta_start = Instant::now();
let method_str = match req
.headers()
.get("x-delegate-method")
.and_then(|v| v.to_str().ok())
{
Some(m) => m.to_string(),
None => {
warn!("delegate missing X-Delegate-Method");
return error_response(400, "bad_request", "missing X-Delegate-Method header");
}
};
let target_url = match req
.headers()
.get("x-delegate-url")
.and_then(|v| v.to_str().ok())
{
Some(u) => u.to_string(),
None => {
warn!("delegate missing X-Delegate-Url");
return error_response(400, "bad_request", "missing X-Delegate-Url header");
}
};
let upstream_headers: HashMap<String, String> = match req
.headers()
.get("x-delegate-headers")
.and_then(|v| v.to_str().ok())
{
Some(b64) => {
match base64::Engine::decode(&base64::engine::general_purpose::STANDARD, b64) {
Ok(decoded) => match serde_json::from_slice(&decoded) {
Ok(h) => h,
Err(e) => {
warn!(error = %e, "delegate invalid X-Delegate-Headers JSON");
return error_response(
400,
"bad_request",
"invalid X-Delegate-Headers JSON",
);
}
},
Err(e) => {
warn!(error = %e, "delegate invalid X-Delegate-Headers base64");
return error_response(400, "bad_request", "invalid X-Delegate-Headers base64");
}
}
}
None => HashMap::new(),
};
let is_gzip = req
.headers()
.get("content-encoding")
.and_then(|v| v.to_str().ok())
.map(|v| v.eq_ignore_ascii_case("gzip"))
.unwrap_or(false);
let req_content_length: u64 = req
.headers()
.get("content-length")
.and_then(|v| v.to_str().ok())
.and_then(|v| v.parse().ok())
.unwrap_or(0);
let meta_ms = meta_start.elapsed().as_millis() as u64;
// ── Target validation ──
let parsed_url = match Url::parse(&target_url) {
Ok(u) => u,
Err(e) => {
warn!(url = %target_url, error = %e, "delegate invalid target URL");
return error_response(400, "bad_request", &format!("invalid URL: {}", e));
}
};
let host = match parsed_url.host_str() {
Some(h) => h.to_string(),
None => {
warn!(url = %target_url, "delegate target URL missing host");
return error_response(400, "bad_request", "URL missing host");
}
};
let port = parsed_url.port_or_known_default().unwrap_or(443);
let dns_start = Instant::now();
if let Err(e) = target_filter::validate_target(&host, port, allowed_ports, dns_cache).await {
warn!(host = %host, port, error = %e, "delegate target rejected");
return error_response(403, "target_not_allowed", &e.to_string());
}
let dns_ms = dns_start.elapsed().as_millis() as u64;
debug!(method = %method_str, url = %target_url, is_gzip, "delegate request");
// ── Build upstream request ──
let method = match method_str.parse::<Method>() {
Ok(m) => m,
Err(e) => {
warn!(error = %e, method = %method_str, "delegate invalid HTTP method");
return error_response(400, "bad_request", &format!("invalid method: {}", e));
}
};
let uri = match target_url.parse::<Uri>() {
Ok(u) => u,
Err(e) => {
warn!(error = %e, url = %target_url, "delegate invalid target URI");
return error_response(400, "bad_request", &format!("invalid URL: {}", e));
}
};
// ── Stream body passthrough ──
// When body is gzip-compressed, forward it directly to upstream with
// Content-Encoding: gzip header — no collect/decompress needed.
// All major AI API providers (Anthropic, OpenAI, Google) accept gzip request bodies.
let wire_size: u64;
let upstream_body: BoxBody;
if is_gzip {
let body_stream =
http_body_util::BodyStream::new(req.into_body()).filter_map(|result| async {
match result {
Ok(frame) => frame.into_data().ok().map(|data| {
Ok::<_, Box<dyn std::error::Error + Send + Sync>>(Frame::data(data))
}),
Err(e) => Some(Err(Box::new(e) as Box<dyn std::error::Error + Send + Sync>)),
}
});
let stream_body = StreamBody::new(body_stream);
upstream_body = BodyExt::boxed(stream_body);
// wire_size will be reported from Content-Length if available, otherwise 0
wire_size = req_content_length;
} else {
// Non-gzip: read body into memory (legacy path)
const MAX_BODY: usize = 10 * 1024 * 1024;
let body_bytes = match Limited::new(req.into_body(), MAX_BODY).collect().await {
Ok(collected) => collected.to_bytes(),
Err(e) => {
warn!(error = %e, "delegate failed to read request body");
return error_response(413, "payload_too_large", "request body exceeds 10MB limit");
}
};
wire_size = body_bytes.len() as u64;
let body = Full::new(body_bytes)
.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { match e {} })
.boxed();
upstream_body = body;
}
let mut upstream_req = Request::new(upstream_body);
*upstream_req.method_mut() = method;
*upstream_req.uri_mut() = uri;
{
let headers = upstream_req.headers_mut();
// Set headers (skip `host` — hyper sets it from the URI automatically,
// and a duplicate Host header can confuse certain upstreams)
for (name, value) in &upstream_headers {
if name.eq_ignore_ascii_case("host") {
continue;
}
let header_name = match HeaderName::from_bytes(name.as_bytes()) {
Ok(n) => n,
Err(_) => {
warn!(header = %name, "delegate invalid header name");
return error_response(400, "bad_request", "invalid header name");
}
};
let header_value = match HeaderValue::from_str(value) {
Ok(v) => v,
Err(_) => {
warn!(header = %name, "delegate invalid header value");
return error_response(400, "bad_request", "invalid header value");
}
};
headers.insert(header_name, header_value);
}
if is_gzip {
headers.insert(
hyper::header::CONTENT_ENCODING,
HeaderValue::from_static("gzip"),
);
}
}
// ── Send upstream request ──
// NOTE: We intentionally do NOT set a per-request timeout here.
// Connect timeout limits connection establishment; Aether controls
// first-byte / idle timeouts on its own side via asyncio.
let upstream_start = Instant::now();
let upstream_resp = match http_client.request(upstream_req).await {
Ok(resp) => resp,
Err(e) => {
warn!(url = %target_url, error = %e, "delegate upstream request failed");
let safe_detail = sanitize_upstream_error(&root_error_message(&e));
if is_timeout_error(&e) {
return error_response(504, "upstream_timeout", &safe_detail);
}
return error_response(502, "upstream_connection_failed", &safe_detail);
}
};
let ttfb_ms = upstream_start.elapsed().as_millis() as u64;
// ── Build response ──
let status = upstream_resp.status().as_u16();
let resp_headers = upstream_resp.headers().clone();
let (connect_ms, tls_ms) = upstream_resp
.extensions()
.get::<ConnectTiming>()
.map(|t| (t.connect_ms, t.tls_ms))
.unwrap_or((0, 0));
let upstream_processing_ms = ttfb_ms.saturating_sub(connect_ms.saturating_add(tls_ms));
let total_ms = total_start.elapsed().as_millis() as u64;
debug!(
url = %target_url,
status,
dns_ms,
connect_ms,
tls_ms,
ttfb_ms,
upstream_processing_ms,
total_ms,
wire_size,
is_gzip,
"delegate upstream response"
);
let timing = serde_json::json!({
"auth_ms": auth_ms,
"meta_ms": meta_ms,
"wire_size": wire_size,
"passthrough": is_gzip,
"dns_ms": dns_ms,
"connect_ms": connect_ms,
"tls_ms": tls_ms,
"ttfb_ms": ttfb_ms,
"upstream_ms": ttfb_ms,
"upstream_processing_ms": upstream_processing_ms,
"total_ms": total_ms,
});
let stream_body: BoxBody = upstream_resp
.into_body()
.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { Box::new(e) })
.boxed();
let mut builder = Response::builder().status(status);
for (name, value) in resp_headers.iter() {
builder = builder.header(name, value);
}
builder = builder.header("X-Proxy-Timing", timing.to_string());
builder.body(stream_body).unwrap_or_else(|_| {
Response::builder()
.status(500)
.body(super::empty_box_body())
.unwrap()
})
}
fn root_error_message(err: &dyn StdError) -> String {
let mut current = err;
while let Some(source) = current.source() {
current = source;
}
current.to_string()
}
fn is_timeout_error(err: &(dyn StdError + 'static)) -> bool {
if err.is::<tokio::time::error::Elapsed>() {
return true;
}
if let Some(io_err) = err.downcast_ref::<std::io::Error>() {
if io_err.kind() == std::io::ErrorKind::TimedOut {
return true;
}
}
if let Some(source) = err.source() {
// source() returns &(dyn Error + 'static), so this is safe
return is_timeout_error(source);
}
false
}
// ── Sanitisation ─────────────────────────────────────────────────────────────
/// Strip full URLs from error messages to prevent leaking upstream API keys,
/// paths, or query parameters in the delegate error response.
///
/// Replaces `https://api.example.com/v1/chat?key=xxx` with `api.example.com`.
fn sanitize_upstream_error(msg: &str) -> String {
// Simple regex-free approach: find "https://..." or "http://..." spans and
// replace them with just the host portion.
let mut result = msg.to_string();
for scheme in &["https://", "http://"] {
while let Some(start) = result.find(scheme) {
let after_scheme = start + scheme.len();
// Host ends at '/', '?', '#', ' ', or end of string
let host_end = result[after_scheme..]
.find(['/', '?', '#', ' '])
.map(|i| after_scheme + i)
.unwrap_or(result.len());
let host = &result[after_scheme..host_end];
result = format!("{}{}{}", &result[..start], host, &result[host_end..]);
}
}
result
}
// ── Error response helpers ───────────────────────────────────────────────────
fn error_response(status: u16, error: &str, detail: &str) -> Response<BoxBody> {
let body = serde_json::json!({
"error": error,
"detail": detail,
});
let body_bytes = bytes::Bytes::from(body.to_string());
Response::builder()
.status(status)
.header("Content-Type", "application/json")
.header("X-Delegate-Error", "true")
.body(
Full::new(body_bytes)
.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { match e {} })
.boxed(),
)
.unwrap()
}

282
aether-proxy/src/proxy/delegate_client.rs
View File

@@ -1,282 +0,0 @@
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::{Duration, Instant};
use hyper::rt;
use hyper::Uri;
use hyper_util::client::legacy::connect::{Connected, Connection, HttpConnector};
use hyper_util::client::legacy::Client;
use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
use rustls::ClientConfig;
use rustls_pki_types::ServerName;
use tokio_rustls::TlsConnector;
use tower_service::Service;
use crate::config::Config;
use crate::proxy::BoxBody;
type BoxError = Box<dyn std::error::Error + Send + Sync>;
type DelegateStream = MaybeHttpsStream<TokioIo<tokio::net::TcpStream>>;
type DelegateConn = TimedConn<DelegateStream>;
pub(crate) type DelegateClient = Client<InstrumentedConnector, BoxBody>;
#[derive(Clone, Copy, Debug, Default)]
pub(crate) struct ConnectTiming {
pub connect_ms: u64,
pub tls_ms: u64,
}
pub(crate) fn build_delegate_client(config: &Config) -> DelegateClient {
let mut http = HttpConnector::new();
http.enforce_http(false);
http.set_connect_timeout(Some(Duration::from_secs(
config.delegate_connect_timeout_secs,
)));
http.set_nodelay(config.delegate_tcp_nodelay);
if config.delegate_tcp_keepalive_secs > 0 {
http.set_keepalive(Some(Duration::from_secs(
config.delegate_tcp_keepalive_secs,
)));
} else {
http.set_keepalive(None);
}
let connector = InstrumentedConnector {
http,
tls_config: build_tls_config(),
};
let mut builder = Client::builder(TokioExecutor::new());
builder.pool_max_idle_per_host(config.delegate_pool_max_idle_per_host);
builder.pool_idle_timeout(Duration::from_secs(config.delegate_pool_idle_timeout_secs));
builder.pool_timer(TokioTimer::new());
builder.build::<_, BoxBody>(connector)
}
fn build_tls_config() -> Arc<ClientConfig> {
let root_store =
rustls::RootCertStore::from_iter(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
let mut config = ClientConfig::builder()
.with_root_certificates(root_store)
.with_no_client_auth();
config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
Arc::new(config)
}
#[derive(Clone)]
pub(crate) struct InstrumentedConnector {
http: HttpConnector,
tls_config: Arc<ClientConfig>,
}
impl Service<Uri> for InstrumentedConnector {
type Response = DelegateConn;
type Error = BoxError;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, BoxError>> + Send>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.http.poll_ready(cx).map_err(Into::into)
}
fn call(&mut self, dst: Uri) -> Self::Future {
let scheme = dst.scheme_str().map(|s| s.to_ascii_lowercase());
let tls_config = self.tls_config.clone();
let connecting = self.http.call(dst.clone());
let connect_start = Instant::now();
Box::pin(async move {
match scheme.as_deref() {
Some("http") => {
let tcp = connecting.await.map_err(|e| Box::new(e) as BoxError)?;
let connect_ms = connect_start.elapsed().as_millis() as u64;
Ok(TimedConn::new(
MaybeHttpsStream::Http(tcp),
ConnectTiming {
connect_ms,
tls_ms: 0,
},
))
}
Some("https") => {
let server_name = resolve_server_name(&dst)?;
let tcp = connecting.await.map_err(|e| Box::new(e) as BoxError)?;
let connect_ms = connect_start.elapsed().as_millis() as u64;
let tls_start = Instant::now();
let tls_stream = TlsConnector::from(tls_config)
.connect(server_name, TokioIo::new(tcp))
.await
.map_err(std::io::Error::other)?;
let tls_ms = tls_start.elapsed().as_millis() as u64;
Ok(TimedConn::new(
MaybeHttpsStream::Https(TokioIo::new(tls_stream)),
ConnectTiming { connect_ms, tls_ms },
))
}
Some(other) => {
Err(std::io::Error::other(format!("unsupported scheme {other}")).into())
}
None => Err(std::io::Error::other("missing scheme").into()),
}
})
}
}
fn resolve_server_name(uri: &Uri) -> Result<ServerName<'static>, BoxError> {
let host = uri.host().ok_or("missing host")?;
let host = host.trim_start_matches('[').trim_end_matches(']');
Ok(ServerName::try_from(host.to_string())?)
}
pub(crate) struct TimedConn<T> {
inner: T,
timing: ConnectTiming,
}
impl<T> TimedConn<T> {
fn new(inner: T, timing: ConnectTiming) -> Self {
Self { inner, timing }
}
}
impl<T: Connection> Connection for TimedConn<T> {
fn connected(&self) -> Connected {
self.inner.connected().extra(self.timing)
}
}
impl<T: rt::Read + Unpin> rt::Read for TimedConn<T> {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: rt::ReadBufCursor<'_>,
) -> Poll<Result<(), std::io::Error>> {
Pin::new(&mut self.inner).poll_read(cx, buf)
}
}
impl<T: rt::Write + Unpin> rt::Write for TimedConn<T> {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, std::io::Error>> {
Pin::new(&mut self.inner).poll_write(cx, buf)
}
fn poll_flush(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Result<(), std::io::Error>> {
Pin::new(&mut self.inner).poll_flush(cx)
}
fn poll_shutdown(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Result<(), std::io::Error>> {
Pin::new(&mut self.inner).poll_shutdown(cx)
}
fn is_write_vectored(&self) -> bool {
self.inner.is_write_vectored()
}
fn poll_write_vectored(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
bufs: &[std::io::IoSlice<'_>],
) -> Poll<Result<usize, std::io::Error>> {
Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
}
}
#[allow(clippy::large_enum_variant)]
pub(crate) enum MaybeHttpsStream<T> {
Http(T),
Https(TokioIo<tokio_rustls::client::TlsStream<TokioIo<T>>>),
}
impl<T: rt::Read + rt::Write + Connection + Unpin> Connection for MaybeHttpsStream<T> {
fn connected(&self) -> Connected {
match self {
Self::Http(stream) => stream.connected(),
Self::Https(stream) => {
let (tcp, tls) = stream.inner().get_ref();
if tls.alpn_protocol() == Some(b"h2") {
tcp.inner().connected().negotiated_h2()
} else {
tcp.inner().connected()
}
}
}
}
}
impl<T: rt::Read + rt::Write + Unpin> rt::Read for MaybeHttpsStream<T> {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: rt::ReadBufCursor<'_>,
) -> Poll<Result<(), std::io::Error>> {
match Pin::get_mut(self) {
Self::Http(stream) => Pin::new(stream).poll_read(cx, buf),
Self::Https(stream) => Pin::new(stream).poll_read(cx, buf),
}
}
}
impl<T: rt::Write + rt::Read + Unpin> rt::Write for MaybeHttpsStream<T> {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, std::io::Error>> {
match Pin::get_mut(self) {
Self::Http(stream) => Pin::new(stream).poll_write(cx, buf),
Self::Https(stream) => Pin::new(stream).poll_write(cx, buf),
}
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), std::io::Error>> {
match Pin::get_mut(self) {
Self::Http(stream) => Pin::new(stream).poll_flush(cx),
Self::Https(stream) => Pin::new(stream).poll_flush(cx),
}
}
fn poll_shutdown(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Result<(), std::io::Error>> {
match Pin::get_mut(self) {
Self::Http(stream) => Pin::new(stream).poll_shutdown(cx),
Self::Https(stream) => Pin::new(stream).poll_shutdown(cx),
}
}
fn is_write_vectored(&self) -> bool {
match self {
Self::Http(stream) => stream.is_write_vectored(),
Self::Https(stream) => stream.is_write_vectored(),
}
}
fn poll_write_vectored(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
bufs: &[std::io::IoSlice<'_>],
) -> Poll<Result<usize, std::io::Error>> {
match Pin::get_mut(self) {
Self::Http(stream) => Pin::new(stream).poll_write_vectored(cx, bufs),
Self::Https(stream) => Pin::new(stream).poll_write_vectored(cx, bufs),
}
}
}

19
aether-proxy/src/proxy/mod.rs
View File

@@ -1,19 +0,0 @@
pub mod connect;
pub mod delegate;
pub mod delegate_client;
pub mod server;
pub mod target_filter;
pub mod tls;
use http_body_util::BodyExt;
/// Boxed body type used across proxy handlers.
pub type BoxBody =
http_body_util::combinators::BoxBody<bytes::Bytes, Box<dyn std::error::Error + Send + Sync>>;
/// Create an empty [`BoxBody`] (for error responses, 405, etc.).
pub fn empty_box_body() -> BoxBody {
http_body_util::Full::new(bytes::Bytes::new())
.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { match e {} })
.boxed()
}

213
aether-proxy/src/proxy/server.rs
View File

@@ -1,213 +0,0 @@
use std::net::SocketAddr;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::{Duration, Instant};
use http_body_util::BodyExt;
use hyper::body::Incoming;
use hyper::rt::{Read, Write};
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper::{Method, Request, Response};
use hyper_util::rt::TokioIo;
use tokio::net::TcpListener;
use tokio::sync::watch;
use tokio::time::timeout;
use tracing::{debug, info, warn};
use crate::proxy::{connect, delegate, tls, BoxBody};
use crate::state::AppState;
/// Start the proxy server.
///
/// Listens for incoming TCP connections and dispatches:
/// - CONNECT requests -> tunnel handler
/// - POST /_aether/delegate -> delegate handler
/// - Other requests -> 405 Method Not Allowed
///
/// When TLS is configured, the server operates in dual-stack mode:
/// it peeks at the first byte of each connection to distinguish TLS ClientHello
/// (0x16) from plain HTTP, and handles both on the same port.
pub async fn run(
state: &Arc<AppState>,
mut shutdown_rx: watch::Receiver<bool>,
) -> anyhow::Result<()> {
let addr = SocketAddr::from(([0, 0, 0, 0], state.config.listen_port));
let listener = TcpListener::bind(addr).await?;
if state.tls_acceptor.is_some() {
info!(addr = %addr, "proxy server listening (HTTP+TLS dual-stack)");
} else {
info!(addr = %addr, "proxy server listening (HTTP only)");
}
let handshake_timeout = Duration::from_secs(state.config.tls_handshake_timeout_secs);
loop {
tokio::select! {
result = listener.accept() => {
let (stream, peer_addr) = match result {
Ok(v) => v,
Err(e) => {
warn!(error = %e, "failed to accept connection");
continue;
}
};
debug!(peer = %peer_addr, "new connection");
if let Err(e) = stream.set_nodelay(true) {
debug!(peer = %peer_addr, error = %e, "failed to set TCP_NODELAY");
}
let permit = match state.connection_semaphore.clone().try_acquire_owned() {
Ok(permit) => permit,
Err(_) => {
warn!(peer = %peer_addr, "connection rejected: limit reached");
continue;
}
};
let state = Arc::clone(state);
state.active_connections.fetch_add(1, Ordering::Relaxed);
tokio::task::spawn(async move {
let _permit = permit;
// Dual-stack: peek first byte to decide TLS vs plain HTTP
if let Some(ref acceptor) = state.tls_acceptor {
let is_tls = match timeout(handshake_timeout, tls::is_tls_client_hello(&stream)).await {
Ok(v) => v,
Err(_) => {
debug!(peer = %peer_addr, "TLS detection timeout");
state.active_connections.fetch_sub(1, Ordering::Relaxed);
return;
}
};
if is_tls {
match timeout(handshake_timeout, acceptor.clone().accept(stream)).await {
Ok(Ok(tls_stream)) => {
debug!(peer = %peer_addr, "TLS handshake ok");
serve_connection(
TokioIo::new(tls_stream),
peer_addr,
&state,
)
.await;
}
Ok(Err(e)) => {
debug!(peer = %peer_addr, error = %e, "TLS handshake failed");
}
Err(_) => {
debug!(peer = %peer_addr, "TLS handshake timeout");
}
}
state.active_connections.fetch_sub(1, Ordering::Relaxed);
return;
}
}
// Plain HTTP
serve_connection(
TokioIo::new(stream),
peer_addr,
&state,
)
.await;
state.active_connections.fetch_sub(1, Ordering::Relaxed);
});
}
_ = shutdown_rx.changed() => {
info!("proxy server shutting down");
break;
}
}
}
Ok(())
}
/// Serve a single HTTP/1.1 connection (works over both plain TCP and TLS).
async fn serve_connection<I>(io: I, peer_addr: SocketAddr, state: &Arc<AppState>)
where
I: Read + Write + Unpin + Send + 'static,
{
let config = Arc::clone(&state.config);
let dynamic = Arc::clone(&state.dynamic);
let delegate_client = state.delegate_client.clone();
let dns_cache = Arc::clone(&state.dns_cache);
let metrics = Arc::clone(&state.metrics);
let service = service_fn(move |req: Request<Incoming>| {
let config = Arc::clone(&config);
let dynamic = Arc::clone(&dynamic);
let delegate_client = delegate_client.clone();
let dns_cache = Arc::clone(&dns_cache);
let metrics = Arc::clone(&metrics);
async move {
let start = Instant::now();
// Snapshot current dynamic values (may be updated by remote config)
let (allowed_ports, timestamp_tolerance) = {
let d = dynamic.read().unwrap();
(d.allowed_ports.clone(), d.timestamp_tolerance)
};
if req.method() == Method::CONNECT {
let resp = connect::handle_connect(
req,
config,
&allowed_ports,
timestamp_tolerance,
dns_cache.as_ref(),
)
.await;
let resp = resp.map(|_| -> BoxBody {
http_body_util::Empty::new()
.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { match e {} })
.boxed()
});
metrics.record_request(start.elapsed());
Ok::<_, hyper::Error>(resp)
} else if req.uri().path() == "/_aether/delegate" && req.method() == hyper::Method::POST
{
let resp = delegate::handle_delegate(
req,
config,
&allowed_ports,
timestamp_tolerance,
dns_cache.as_ref(),
&delegate_client,
)
.await;
metrics.record_request(start.elapsed());
Ok(resp)
} else {
// Only CONNECT tunnels and /_aether/delegate are supported;
// plain HTTP forward proxy was removed (all API traffic is HTTPS).
let resp = Response::builder()
.status(405)
.header("Allow", "CONNECT")
.header("Content-Length", "0")
.body(crate::proxy::empty_box_body())
.unwrap();
metrics.record_request(start.elapsed());
Ok(resp)
}
}
});
if let Err(e) = http1::Builder::new()
.preserve_header_case(true)
.title_case_headers(false)
.serve_connection(io, service)
.with_upgrades()
.await
{
if !e.to_string().contains("connection closed") {
debug!(peer = %peer_addr, error = %e, "connection error");
}
}
}

125
aether-proxy/src/proxy/tls.rs
View File

@@ -1,125 +0,0 @@
use std::fs;
use std::io::BufReader;
use std::path::Path;
use std::sync::Arc;
use rcgen::{CertificateParams, KeyPair};
use rustls_pki_types::{CertificateDer, PrivateKeyDer};
use sha2::{Digest, Sha256};
use tokio_rustls::TlsAcceptor;
use tracing::{info, warn};
const SESSION_CACHE_SIZE: usize = 1024;
/// Generate a self-signed certificate if the files do not already exist.
///
/// The certificate includes SANs: `localhost` and `aether-proxy`.
/// The private key file is set to mode 0600 on unix.
pub fn ensure_self_signed_cert(cert_path: &Path, key_path: &Path) -> anyhow::Result<()> {
if cert_path.exists() && key_path.exists() {
info!(
cert = %cert_path.display(),
key = %key_path.display(),
"using existing TLS certificate"
);
return Ok(());
}
info!("generating self-signed TLS certificate");
let mut params = CertificateParams::new(vec!["localhost".into(), "aether-proxy".into()])?;
params.distinguished_name = rcgen::DistinguishedName::new();
params
.distinguished_name
.push(rcgen::DnType::CommonName, "aether-proxy");
let key_pair = KeyPair::generate()?;
let cert = params.self_signed(&key_pair)?;
let cert_pem = cert.pem();
let key_pem = key_pair.serialize_pem();
fs::write(cert_path, &cert_pem)?;
fs::write(key_path, &key_pem)?;
// Set key file permissions to 0600 on unix
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let perms = fs::Permissions::from_mode(0o600);
fs::set_permissions(key_path, perms)?;
}
info!(
cert = %cert_path.display(),
key = %key_path.display(),
"self-signed TLS certificate generated"
);
Ok(())
}
/// Build a `TlsAcceptor` from PEM certificate and key files.
pub fn build_tls_acceptor(cert_path: &Path, key_path: &Path) -> anyhow::Result<TlsAcceptor> {
let cert_file = fs::File::open(cert_path)?;
let key_file = fs::File::open(key_path)?;
let certs: Vec<CertificateDer<'static>> =
rustls_pemfile::certs(&mut BufReader::new(cert_file)).collect::<Result<Vec<_>, _>>()?;
if certs.is_empty() {
anyhow::bail!("no certificates found in {}", cert_path.display());
}
let key: PrivateKeyDer<'static> =
rustls_pemfile::private_key(&mut BufReader::new(key_file))?
.ok_or_else(|| anyhow::anyhow!("no private key found in {}", key_path.display()))?;
let mut config = rustls::ServerConfig::builder()
.with_no_client_auth()
.with_single_cert(certs, key)?;
config.alpn_protocols = vec![b"http/1.1".to_vec()];
config.session_storage = rustls::server::ServerSessionMemoryCache::new(SESSION_CACHE_SIZE);
match rustls::crypto::ring::Ticketer::new() {
Ok(ticketer) => {
config.ticketer = ticketer;
}
Err(e) => {
warn!(error = %e, "failed to init TLS ticketer; tickets disabled");
}
}
Ok(TlsAcceptor::from(Arc::new(config)))
}
/// Compute the SHA-256 fingerprint of the first certificate in a PEM file.
///
/// Returns the hex-encoded fingerprint (lowercase, no separators).
pub fn cert_sha256_fingerprint(cert_path: &Path) -> anyhow::Result<String> {
let cert_file = fs::File::open(cert_path)?;
let certs: Vec<CertificateDer<'static>> =
rustls_pemfile::certs(&mut BufReader::new(cert_file)).collect::<Result<Vec<_>, _>>()?;
let cert = certs
.first()
.ok_or_else(|| anyhow::anyhow!("no certificates found in {}", cert_path.display()))?;
let digest = Sha256::digest(cert.as_ref());
Ok(hex::encode(digest))
}
/// Peek at the first byte of a TCP stream to determine if it is a TLS ClientHello.
///
/// Returns `true` if the first byte is 0x16 (TLS record type: Handshake).
pub async fn is_tls_client_hello(stream: &tokio::net::TcpStream) -> bool {
let mut buf = [0u8; 1];
match stream.peek(&mut buf).await {
Ok(1) => buf[0] == 0x16,
Ok(_) => false,
Err(e) => {
warn!(error = %e, "failed to peek first byte");
false
}
}
}

151
aether-proxy/src/registration/client.rs
View File

@@ -3,30 +3,11 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH};
use reqwest::{Client, StatusCode};
use serde::{Deserialize, Serialize};
use tokio::time::sleep;
use tracing::{debug, error, info, warn};
use tracing::{debug, error, info};
use crate::config::Config;
use crate::hardware::HardwareInfo;
/// Heartbeat-specific error that distinguishes "node not found" (needs
/// re-registration) from transient / other failures.
#[derive(Debug)]
pub enum HeartbeatError {
/// HTTP 404 – the node_id is no longer known to Aether.
NodeNotFound(String),
/// Any other failure (network, 5xx, etc.).
Other(anyhow::Error),
}
impl std::fmt::Display for HeartbeatError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::NodeNotFound(msg) => write!(f, "node not found: {}", msg),
Self::Other(e) => write!(f, "{}", e),
}
}
}
#[derive(Debug, Serialize)]
struct RegisterRequest {
name: String,
@@ -35,14 +16,11 @@ struct RegisterRequest {
#[serde(skip_serializing_if = "Option::is_none")]
region: Option<String>,
heartbeat_interval: u64,
#[serde(skip_serializing_if = "std::ops::Not::not")]
tls_enabled: bool,
#[serde(skip_serializing_if = "Option::is_none")]
tls_cert_fingerprint: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
hardware_info: Option<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
estimated_max_concurrency: Option<u64>,
tunnel_mode: bool,
}
#[derive(Debug, Deserialize)]
@@ -50,17 +28,6 @@ pub struct RegisterResponse {
pub node_id: String,
}
#[derive(Debug, Serialize)]
struct HeartbeatRequest {
node_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
active_connections: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
total_requests: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
avg_latency_ms: Option<f64>,
}
/// Remote configuration pushed by the Aether management backend.
#[derive(Debug, Clone, Deserialize)]
pub struct RemoteConfig {
@@ -68,29 +35,6 @@ pub struct RemoteConfig {
pub allowed_ports: Option<Vec<u16>>,
pub log_level: Option<String>,
pub heartbeat_interval: Option<u64>,
pub timestamp_tolerance: Option<u64>,
}
/// Parsed heartbeat response from Aether.
#[derive(Debug, Deserialize)]
struct HeartbeatResponseBody {
#[serde(default)]
node: Option<HeartbeatNodeInfo>,
}
#[derive(Debug, Deserialize)]
struct HeartbeatNodeInfo {
#[serde(default)]
remote_config: Option<RemoteConfig>,
#[serde(default)]
config_version: Option<u64>,
}
/// Heartbeat result returned to the caller.
#[derive(Debug)]
pub struct HeartbeatResult {
pub remote_config: Option<RemoteConfig>,
pub config_version: u64,
}
#[derive(Debug, Serialize)]
@@ -109,7 +53,7 @@ pub struct AetherClient {
}
impl AetherClient {
pub fn new(config: &Config) -> Self {
pub fn new(config: &Config, aether_url: &str, management_token: &str) -> Self {
let mut builder = Client::builder()
.timeout(Duration::from_secs(config.aether_request_timeout_secs))
.connect_timeout(Duration::from_secs(config.aether_connect_timeout_secs))
@@ -136,8 +80,8 @@ impl AetherClient {
Self {
http,
base_url: config.aether_url.trim_end_matches('/').to_string(),
token: config.management_token.clone(),
base_url: aether_url.trim_end_matches('/').to_string(),
token: management_token.to_string(),
retry_max_attempts: config.aether_retry_max_attempts.max(1),
retry_base_delay,
retry_max_delay,
@@ -150,29 +94,26 @@ impl AetherClient {
pub async fn register(
&self,
config: &Config,
node_name: &str,
public_ip: &str,
tls_enabled: bool,
tls_cert_fingerprint: Option<&str>,
hw: Option<&HardwareInfo>,
) -> anyhow::Result<String> {
let url = format!("{}/api/admin/proxy-nodes/register", self.base_url);
let body = RegisterRequest {
name: config.node_name.clone(),
name: node_name.to_string(),
ip: public_ip.to_string(),
port: config.listen_port,
port: 0,
region: config.node_region.clone(),
heartbeat_interval: config.heartbeat_interval,
tls_enabled,
tls_cert_fingerprint: tls_cert_fingerprint.map(|s| s.to_string()),
hardware_info: hw.and_then(|h| serde_json::to_value(h).ok()),
estimated_max_concurrency: hw.map(|h| h.estimated_max_concurrency),
tunnel_mode: true,
};
info!(
url = %url,
name = %body.name,
ip = %body.ip,
port = body.port,
"registering with Aether"
);
@@ -199,80 +140,6 @@ impl AetherClient {
Ok(data.node_id)
}
/// Send heartbeat to Aether.
///
/// On success, returns any remote config included in the response.
/// Returns [`HeartbeatError::NodeNotFound`] on HTTP 404 so the caller
/// can trigger re-registration.
pub async fn heartbeat(
&self,
node_id: &str,
active_connections: Option<i64>,
total_requests: Option<i64>,
avg_latency_ms: Option<f64>,
) -> Result<HeartbeatResult, HeartbeatError> {
let url = format!("{}/api/admin/proxy-nodes/heartbeat", self.base_url);
let body = HeartbeatRequest {
node_id: node_id.to_string(),
active_connections,
total_requests,
avg_latency_ms,
};
debug!(node_id = %node_id, "sending heartbeat");
let resp = self
.send_with_retry(
|| {
self.http
.post(&url)
.header("Authorization", format!("Bearer {}", self.token))
.json(&body)
},
"heartbeat",
)
.await
.map_err(|e| HeartbeatError::Other(e.into()))?;
let status = resp.status();
if !status.is_success() {
let text = resp.text().await.unwrap_or_default();
warn!(status = %status, body = %text, "heartbeat failed");
if status == StatusCode::NOT_FOUND {
return Err(HeartbeatError::NodeNotFound(text));
}
return Err(HeartbeatError::Other(anyhow::anyhow!(
"heartbeat failed (HTTP {}): {}",
status,
text
)));
}
// Parse remote config from response (best-effort)
let result = match resp.json::<HeartbeatResponseBody>().await {
Ok(body) => {
let (remote_config, config_version) = match body.node {
Some(node) => (node.remote_config, node.config_version.unwrap_or(0)),
None => (None, 0),
};
HeartbeatResult {
remote_config,
config_version,
}
}
Err(e) => {
debug!(error = %e, "failed to parse heartbeat response body");
HeartbeatResult {
remote_config: None,
config_version: 0,
}
}
};
debug!(node_id = %node_id, config_version = result.config_version, "heartbeat ok");
Ok(result)
}
/// Unregister this node from Aether (graceful shutdown).
pub async fn unregister(&self, node_id: &str) -> anyhow::Result<()> {
let url = format!("{}/api/admin/proxy-nodes/unregister", self.base_url);

127
aether-proxy/src/registration/heartbeat.rs
View File

@@ -1,127 +0,0 @@
use std::sync::atomic::Ordering;
use std::sync::Arc;
use tokio::sync::watch;
use tracing::{debug, error, info, warn};
use crate::registration::client::HeartbeatError;
use crate::runtime;
use crate::state::AppState;
/// Run periodic heartbeat task until shutdown signal.
///
/// When Aether responds with 404 (node not found), this task automatically
/// re-registers the node and updates the shared `node_id` so the proxy
/// server and future heartbeats use the new identity.
///
/// When the heartbeat response includes a `remote_config`, it is applied
/// to the [`DynamicConfig`](crate::runtime::DynamicConfig) so the proxy
/// picks up changes without a restart.
pub async fn run(state: &Arc<AppState>, mut shutdown_rx: watch::Receiver<bool>) {
let mut consecutive_failures: u32 = 0;
// Skip the first tick (registration already acts as initial heartbeat)
let initial_interval = state.dynamic.read().unwrap().heartbeat_interval;
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_secs(initial_interval)) => {}
_ = shutdown_rx.changed() => {
debug!("heartbeat task stopping (during initial wait)");
return;
}
}
loop {
let current_node_id = state.node_id.read().unwrap().clone();
let active_conns = state.active_connections.load(Ordering::Relaxed) as i64;
// Swap-and-reset: report incremental metrics since last heartbeat
let interval_requests = state.metrics.total_requests.swap(0, Ordering::Relaxed);
let interval_latency_ns = state.metrics.total_latency_ns.swap(0, Ordering::Relaxed);
let interval_requests_i64 = i64::try_from(interval_requests).unwrap_or(i64::MAX);
let avg_latency_ms = if interval_requests > 0 {
Some(interval_latency_ns as f64 / interval_requests as f64 / 1_000_000.0)
} else {
None
};
match state
.aether_client
.heartbeat(
&current_node_id,
Some(active_conns),
Some(interval_requests_i64),
avg_latency_ms,
)
.await
{
Ok(result) => {
if consecutive_failures > 0 {
info!(
previous_failures = consecutive_failures,
"heartbeat recovered"
);
}
consecutive_failures = 0;
// Apply remote config if present and version changed
if let Some(ref remote) = result.remote_config {
runtime::apply_remote_config(&state.dynamic, remote, result.config_version);
}
}
Err(HeartbeatError::NodeNotFound(_)) => {
warn!(
old_node_id = %current_node_id,
"node not found, re-registering"
);
match state
.aether_client
.register(
&state.config,
&state.public_ip,
state.config.enable_tls,
state.tls_fingerprint.as_deref(),
Some(&state.hardware_info),
)
.await
{
Ok(new_id) => {
info!(
old_node_id = %current_node_id,
new_node_id = %new_id,
"re-registered successfully"
);
*state.node_id.write().unwrap() = new_id;
consecutive_failures = 0;
}
Err(e) => {
consecutive_failures += 1;
error!(
error = %e,
consecutive_failures,
"re-registration failed"
);
}
}
}
Err(HeartbeatError::Other(e)) => {
consecutive_failures += 1;
warn!(
error = %e,
consecutive_failures,
"heartbeat failed"
);
}
}
// Read interval from dynamic config (may have been updated remotely)
let interval_secs = state.dynamic.read().unwrap().heartbeat_interval;
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_secs(interval_secs)) => {}
_ = shutdown_rx.changed() => {
debug!("heartbeat task stopping");
break;
}
}
}
}

1
aether-proxy/src/registration/mod.rs
View File

@@ -1,2 +1 @@
pub mod client;
pub mod heartbeat;

11
aether-proxy/src/runtime.rs
View File

@@ -16,7 +16,6 @@ use crate::config::Config;
pub struct DynamicConfig {
pub node_name: String,
pub allowed_ports: HashSet<u16>,
pub timestamp_tolerance: u64,
pub log_level: String,
pub heartbeat_interval: u64,
/// Monotonically increasing version from the backend.
@@ -30,7 +29,6 @@ impl DynamicConfig {
Self {
node_name: config.node_name.clone(),
allowed_ports: config.allowed_ports.iter().copied().collect(),
timestamp_tolerance: config.timestamp_tolerance,
log_level: config.log_level.clone(),
heartbeat_interval: config.heartbeat_interval,
config_version: 0,
@@ -79,18 +77,11 @@ pub fn apply_remote_config(
if let Some(ref ports) = remote.allowed_ports {
let new_set: HashSet<u16> = ports.iter().copied().collect();
if new_set != cfg.allowed_ports {
changed.push(format!("allowed_ports → {:?}", ports));
changed.push(format!("allowed_ports -> {:?}", ports));
cfg.allowed_ports = new_set;
}
}
if let Some(tol) = remote.timestamp_tolerance {
if tol != cfg.timestamp_tolerance {
changed.push(format!("timestamp_tolerance → {}", tol));
cfg.timestamp_tolerance = tol;
}
}
if let Some(interval) = remote.heartbeat_interval {
if interval != cfg.heartbeat_interval {
changed.push(format!("heartbeat_interval → {}s", interval));

548
aether-proxy/src/setup/tui.rs
View File

@@ -2,7 +2,8 @@
//!
//! Launched via `aether-proxy setup [path]`. Presents a full-screen form
//! backed by ratatui where the user can navigate fields, edit values, and
//! save to a TOML config file.
//! save to a TOML config file. Supports multi-server configuration via
//! a tabbed interface.
use std::io;
use std::path::PathBuf;
@@ -19,13 +20,13 @@ use ratatui::widgets::{Block, Borders, Paragraph};
use ratatui::Frame;
use ratatui::Terminal;
use crate::config::ConfigFile;
use crate::config::{ConfigFile, ServerEntry};
/// Outcome of the setup wizard, returned to the caller.
pub enum SetupOutcome {
/// Config saved; systemd service installed and started.
ServiceInstalled,
/// Config saved; no service — caller should start the proxy directly.
/// Config saved; no service -- caller should start the proxy directly.
ReadyToRun(PathBuf),
/// User quit without saving.
Cancelled,
@@ -34,13 +35,12 @@ pub enum SetupOutcome {
/// Column width reserved for the field label (chars).
const LABEL_WIDTH: usize = 22;
// ── Field types ──────────────────────────────────────────────────────────────
// -- Field types --------------------------------------------------------------
#[derive(Clone, Copy, PartialEq)]
enum FieldKind {
Text,
Secret,
Number,
Bool,
LogLevel,
}
@@ -53,31 +53,15 @@ struct Field {
required: bool,
help: &'static str,
}
// -- Server tab ---------------------------------------------------------------
// ── App state ────────────────────────────────────────────────────────────────
#[derive(PartialEq)]
enum Mode {
Normal,
Editing,
}
struct App {
/// A single server tab's editable fields.
struct ServerTab {
fields: Vec<Field>,
selected: usize,
mode: Mode,
edit_buffer: String,
edit_cursor: usize, // char index
config_path: PathBuf,
modified: bool,
message: Option<(String, Instant, bool)>, // (text, when, is_error)
scroll_offset: usize,
saved_once: bool,
pending_quit: bool, // true after first q/Esc with unsaved changes
}
impl App {
fn new(config_path: PathBuf) -> Self {
impl ServerTab {
fn new() -> Self {
Self {
fields: vec![
Field {
@@ -86,7 +70,7 @@ impl App {
value: String::new(),
kind: FieldKind::Text,
required: true,
help: "Aether 服务器 URL (如 https://aether.example.com)",
help: "Aether URL (e.g. https://aether.example.com)",
},
Field {
label: "Management Token",
@@ -94,23 +78,7 @@ impl App {
value: String::new(),
kind: FieldKind::Secret,
required: true,
help: "Aether 管理 API Token (ae_xxx)",
},
Field {
label: "HMAC Key",
key: "hmac_key",
value: String::new(),
kind: FieldKind::Secret,
required: true,
help: "HMAC-SHA256 签名密钥，用于代理请求认证",
},
Field {
label: "Listen Port",
key: "listen_port",
value: "18080".into(),
kind: FieldKind::Number,
required: true,
help: "代理服务监听端口",
help: "Aether Management Token (ae_xxx)",
},
Field {
label: "Node Name",
@@ -118,15 +86,60 @@ impl App {
value: "proxy-01".into(),
kind: FieldKind::Text,
required: true,
help: "节点名称，用于在 Aether 后台识别",
help: "Node name for identification in Aether dashboard",
},
],
}
}
fn from_entry(entry: &ServerEntry) -> Self {
let mut tab = Self::new();
tab.fields[0].value = entry.aether_url.clone();
tab.fields[1].value = entry.management_token.clone();
if let Some(ref name) = entry.node_name {
tab.fields[2].value = name.clone();
}
tab
}
}
// -- App state ----------------------------------------------------------------
#[derive(PartialEq)]
enum Mode {
Normal,
Editing,
}
struct App {
server_tabs: Vec<ServerTab>,
active_tab: usize,
global_fields: Vec<Field>,
selected: usize,
mode: Mode,
edit_buffer: String,
edit_cursor: usize,
config_path: PathBuf,
modified: bool,
message: Option<(String, Instant, bool)>,
scroll_offset: usize,
saved_once: bool,
pending_quit: bool,
confirm_delete: bool,
}
impl App {
fn new(config_path: PathBuf) -> Self {
Self {
server_tabs: vec![ServerTab::new()],
active_tab: 0,
global_fields: vec![
Field {
label: "Log Level",
key: "log_level",
value: "info".into(),
kind: FieldKind::LogLevel,
required: true,
help: "日志级别 -- Enter 切换: trace / debug / info / warn / error",
help: "Log level -- Enter to cycle: trace / debug / info / warn / error",
},
Field {
label: "Log JSON",
@@ -134,7 +147,7 @@ impl App {
value: "false".into(),
kind: FieldKind::Bool,
required: true,
help: "是否以 JSON 格式输出日志 -- Enter 切换",
help: "Output logs as JSON -- Enter to toggle",
},
Field {
label: "Install Service",
@@ -147,7 +160,7 @@ impl App {
.into(),
kind: FieldKind::Bool,
required: true,
help: "注册为 systemd 开机启动服务 (需要 root 权限) -- Enter 切换",
help: "Install as systemd service (requires root) -- Enter to toggle",
},
],
selected: 0,
@@ -160,10 +173,47 @@ impl App {
scroll_offset: 0,
saved_once: false,
pending_quit: false,
confirm_delete: false,
}
}
// ── Config ↔ fields ──────────────────────────────────────────────────
// -- Field accessors (unified index across server + global) ---------------
fn server_field_count(&self) -> usize {
self.server_tabs[self.active_tab].fields.len()
}
fn total_field_count(&self) -> usize {
self.server_field_count() + self.global_fields.len()
}
fn selected_field(&self) -> &Field {
let sc = self.server_field_count();
if self.selected < sc {
&self.server_tabs[self.active_tab].fields[self.selected]
} else {
&self.global_fields[self.selected - sc]
}
}
fn selected_field_mut(&mut self) -> &mut Field {
let sc = self.server_field_count();
if self.selected < sc {
&mut self.server_tabs[self.active_tab].fields[self.selected]
} else {
&mut self.global_fields[self.selected - sc]
}
}
fn clamp_selection(&mut self) {
let max = self.total_field_count();
if self.selected >= max {
self.selected = max.saturating_sub(1);
}
self.scroll_offset = 0;
self.confirm_delete = false;
}
// -- Config <-> fields -----------------------------------------------------
fn load_from_file(&mut self) {
if let Ok(cfg) = ConfigFile::load(&self.config_path) {
@@ -172,13 +222,9 @@ impl App {
}
fn apply_config(&mut self, cfg: &ConfigFile) {
for field in &mut self.fields {
// Global fields
for field in &mut self.global_fields {
let val: Option<String> = match field.key {
"aether_url" => cfg.aether_url.clone(),
"management_token" => cfg.management_token.clone(),
"hmac_key" => cfg.hmac_key.clone(),
"listen_port" => cfg.listen_port.map(|v| v.to_string()),
"node_name" => cfg.node_name.clone(),
"log_level" => cfg.log_level.clone(),
"log_json" => cfg.log_json.map(|v| v.to_string()),
_ => None,
@@ -187,54 +233,64 @@ impl App {
field.value = v;
}
}
// Server tabs
let servers = cfg.effective_servers();
if servers.is_empty() {
let mut tab = ServerTab::new();
// Single-server fallback: use top-level node_name
if let Some(ref name) = cfg.node_name {
tab.fields[2].value = name.clone();
}
self.server_tabs = vec![tab];
} else {
self.server_tabs = servers.iter().map(ServerTab::from_entry).collect();
// For single-server mode, node_name might be in top-level only
if self.server_tabs.len() == 1 && self.server_tabs[0].fields[2].value.is_empty() {
if let Some(ref name) = cfg.node_name {
self.server_tabs[0].fields[2].value = name.clone();
}
}
}
self.active_tab = 0;
self.selected = 0;
self.scroll_offset = 0;
}
fn to_config(&self) -> ConfigFile {
let get = |key: &str| -> Option<String> {
self.fields
let get_global = |key: &str| -> Option<String> {
self.global_fields
.iter()
.find(|f| f.key == key)
.map(|f| f.value.clone())
.filter(|v| !v.is_empty())
};
ConfigFile {
aether_url: get("aether_url"),
management_token: get("management_token"),
hmac_key: get("hmac_key"),
listen_port: get("listen_port").and_then(|v| v.parse().ok()),
public_ip: None,
node_name: get("node_name"),
node_region: None,
heartbeat_interval: None,
allowed_ports: None,
timestamp_tolerance: None,
aether_request_timeout_secs: None,
aether_connect_timeout_secs: None,
aether_pool_max_idle_per_host: None,
aether_pool_idle_timeout_secs: None,
aether_tcp_keepalive_secs: None,
aether_tcp_nodelay: None,
aether_http2: None,
aether_retry_max_attempts: None,
aether_retry_base_delay_ms: None,
aether_retry_max_delay_ms: None,
max_concurrent_connections: None,
connect_timeout_secs: None,
tls_handshake_timeout_secs: None,
dns_cache_ttl_secs: None,
dns_cache_capacity: None,
delegate_connect_timeout_secs: None,
delegate_pool_max_idle_per_host: None,
delegate_pool_idle_timeout_secs: None,
delegate_tcp_keepalive_secs: None,
delegate_tcp_nodelay: None,
log_level: get("log_level"),
log_json: get("log_json").and_then(|v| v.parse().ok()),
enable_tls: None,
tls_cert: None,
tls_key: None,
}
let get_tab = |tab: &ServerTab, key: &str| -> Option<String> {
tab.fields
.iter()
.find(|f| f.key == key)
.map(|f| f.value.clone())
.filter(|v| !v.is_empty())
};
let mut cfg = ConfigFile {
log_level: get_global("log_level"),
log_json: get_global("log_json").and_then(|v| v.parse().ok()),
..ConfigFile::default()
};
// Always write [[servers]] format; old top-level fields are read-only compat
cfg.servers = self
.server_tabs
.iter()
.map(|tab| ServerEntry {
aether_url: get_tab(tab, "aether_url").unwrap_or_default(),
management_token: get_tab(tab, "management_token").unwrap_or_default(),
node_name: get_tab(tab, "node_name"),
})
.collect();
cfg
}
fn save(&mut self) -> anyhow::Result<()> {
@@ -249,27 +305,33 @@ impl App {
));
Ok(())
}
// ── Scrolling ────────────────────────────────────────────────────────
// -- Scrolling ---------------------------------------------------------------
fn ensure_visible(&mut self, visible_rows: usize) {
if visible_rows == 0 {
return;
}
if self.selected < self.scroll_offset {
self.scroll_offset = self.selected;
} else if self.selected >= self.scroll_offset + visible_rows {
self.scroll_offset = self.selected - visible_rows + 1;
// Account for separator line between server and global fields
let display_row = if self.selected >= self.server_field_count() {
self.selected + 1
} else {
self.selected
};
if display_row < self.scroll_offset {
self.scroll_offset = display_row;
} else if display_row >= self.scroll_offset + visible_rows {
self.scroll_offset = display_row - visible_rows + 1;
}
}
// ── Key handling ─────────────────────────────────────────────────────
// -- Key handling -------------------------------------------------------------
/// Returns `true` when the app should exit.
fn handle_key(&mut self, key: KeyEvent) -> bool {
// Expire old messages (but keep quit-confirmation messages alive)
if let Some((_, when, _)) = &self.message {
if !self.pending_quit && when.elapsed() > Duration::from_secs(4) {
if !self.pending_quit && !self.confirm_delete && when.elapsed() > Duration::from_secs(4)
{
self.message = None;
}
}
@@ -284,7 +346,7 @@ impl App {
}
fn handle_normal(&mut self, key: KeyEvent) -> bool {
// ── Quit handling (with unsaved-changes confirmation) ─────────
// -- Quit handling (with unsaved-changes confirmation) -----------------
let is_quit_key = matches!(key.code, KeyCode::Char('q') | KeyCode::Esc);
if is_quit_key {
@@ -292,6 +354,7 @@ impl App {
return true;
}
self.pending_quit = true;
self.confirm_delete = false;
self.message = Some((
"unsaved changes! q again to discard, ^S to save".into(),
Instant::now(),
@@ -300,14 +363,21 @@ impl App {
return false;
}
// Any other key cancels the pending quit
// Any other key cancels pending quit / pending delete
if self.pending_quit {
self.pending_quit = false;
self.message = None;
}
if self.confirm_delete && !matches!(key.code, KeyCode::Delete | KeyCode::Char('x')) {
self.confirm_delete = false;
self.message = None;
}
match key.code {
KeyCode::Char('s') if key.modifiers.contains(KeyModifiers::CONTROL) => {
KeyCode::Char('s')
if key.modifiers.contains(KeyModifiers::CONTROL)
|| key.modifiers.contains(KeyModifiers::SUPER) =>
{
if let Err(e) = self.save() {
self.message = Some((format!("error: {}", e), Instant::now(), true));
}
@@ -316,23 +386,20 @@ impl App {
self.selected = self.selected.saturating_sub(1);
}
KeyCode::Down | KeyCode::Char('j') => {
if self.selected + 1 < self.fields.len() {
if self.selected + 1 < self.total_field_count() {
self.selected += 1;
}
}
KeyCode::Home => self.selected = 0,
KeyCode::End => self.selected = self.fields.len() - 1,
KeyCode::End => self.selected = self.total_field_count() - 1,
KeyCode::Enter | KeyCode::Char(' ') => {
let field = &self.fields[self.selected];
match field.kind {
let kind = self.selected_field().kind;
let key_str = self.selected_field().key;
let value = self.selected_field().value.clone();
match kind {
FieldKind::Bool => {
let toggled = if field.value == "true" {
"false"
} else {
"true"
};
// Block enabling service install without root/systemd
if field.key == "install_service"
let toggled = if value == "true" { "false" } else { "true" };
if key_str == "install_service"
&& toggled == "true"
&& !super::service::is_available()
{
@@ -342,27 +409,79 @@ impl App {
true,
));
} else {
self.fields[self.selected].value = toggled.into();
self.selected_field_mut().value = toggled.into();
self.modified = true;
}
}
FieldKind::LogLevel => {
const LEVELS: &[&str] = &["trace", "debug", "info", "warn", "error"];
let idx = LEVELS.iter().position(|l| *l == field.value).unwrap_or(2);
self.fields[self.selected].value = LEVELS[(idx + 1) % LEVELS.len()].into();
let idx = LEVELS.iter().position(|l| *l == value).unwrap_or(2);
self.selected_field_mut().value = LEVELS[(idx + 1) % LEVELS.len()].into();
self.modified = true;
}
_ => {
self.edit_buffer = field.value.clone();
self.edit_buffer = value;
self.edit_cursor = self.edit_buffer.chars().count();
self.mode = Mode::Editing;
}
}
}
// -- Tab navigation --
KeyCode::Tab => {
// Quick save shortcut
if let Err(e) = self.save() {
self.message = Some((format!("error: {}", e), Instant::now(), true));
if self.server_tabs.len() > 1 {
self.active_tab = (self.active_tab + 1) % self.server_tabs.len();
self.clamp_selection();
}
}
KeyCode::BackTab => {
if self.server_tabs.len() > 1 {
self.active_tab = if self.active_tab == 0 {
self.server_tabs.len() - 1
} else {
self.active_tab - 1
};
self.clamp_selection();
}
}
KeyCode::Char(c @ '1'..='9') if !key.modifiers.contains(KeyModifiers::CONTROL) => {
let idx = (c as usize) - ('1' as usize);
if idx < self.server_tabs.len() && idx != self.active_tab {
self.active_tab = idx;
self.clamp_selection();
}
}
// -- Add / remove server --
KeyCode::Char('+') | KeyCode::Char('a') => {
self.server_tabs.push(ServerTab::new());
self.active_tab = self.server_tabs.len() - 1;
self.selected = 0;
self.scroll_offset = 0;
self.modified = true;
self.message = Some((
format!("added server {}", self.server_tabs.len()),
Instant::now(),
false,
));
}
KeyCode::Delete | KeyCode::Char('x') => {
if self.server_tabs.len() <= 1 {
self.message =
Some(("cannot remove the last server".into(), Instant::now(), true));
} else if self.confirm_delete {
let removed = self.active_tab + 1;
self.server_tabs.remove(self.active_tab);
self.active_tab = self.active_tab.min(self.server_tabs.len() - 1);
self.clamp_selection();
self.modified = true;
self.message =
Some((format!("server {} removed", removed), Instant::now(), false));
} else {
self.confirm_delete = true;
self.message = Some((
"press Delete/x again to remove this server".into(),
Instant::now(),
true,
));
}
}
_ => {}
@@ -373,12 +492,11 @@ impl App {
fn handle_edit(&mut self, key: KeyEvent) {
match key.code {
KeyCode::Esc => {
// Cancel -- discard changes to this field
self.mode = Mode::Normal;
}
KeyCode::Enter => {
if self.validate_edit() {
self.fields[self.selected].value = self.edit_buffer.clone();
self.selected_field_mut().value = self.edit_buffer.clone();
self.modified = true;
self.mode = Mode::Normal;
} else {
@@ -419,12 +537,7 @@ impl App {
}
fn validate_edit(&self) -> bool {
let kind = self.fields[self.selected].kind;
let buf = &self.edit_buffer;
match kind {
FieldKind::Number => buf.is_empty() || buf.parse::<u64>().is_ok(),
_ => true,
}
true
}
/// Byte offset of the char at `char_idx`.
@@ -436,13 +549,11 @@ impl App {
.unwrap_or(self.edit_buffer.len())
}
}
// ── Rendering ────────────────────────────────────────────────────────────────
// -- Rendering ----------------------------------------------------------------
fn ui(f: &mut Frame, app: &mut App) {
let area = f.area();
// Outer block
let title = if app.modified {
" Aether Proxy Setup [*] "
} else {
@@ -458,28 +569,82 @@ fn ui(f: &mut Frame, app: &mut App) {
let inner = outer.inner(area);
f.render_widget(outer, area);
// Split: fields | footer
let chunks = Layout::vertical([Constraint::Min(1), Constraint::Length(4)]).split(inner);
// Split: fields | tab bar | footer
let chunks = Layout::vertical([
Constraint::Min(1),
Constraint::Length(1),
Constraint::Length(4),
])
.split(inner);
let fields_area = chunks[0];
let footer_area = chunks[1];
render_fields(f, app, fields_area);
render_footer(f, app, footer_area);
render_fields(f, app, chunks[0]);
render_tab_bar(f, app, chunks[1]);
render_footer(f, app, chunks[2]);
}
fn render_fields(f: &mut Frame, app: &mut App, area: Rect) {
let visible = area.height as usize;
app.ensure_visible(visible);
let server_count = app.server_field_count();
let mut lines: Vec<Line> = Vec::new();
// display_row tracks the actual row index (including separator)
let mut display_row: usize = 0;
for (i, field) in app.fields.iter().enumerate() {
if i < app.scroll_offset || i >= app.scroll_offset + visible {
continue;
// Server fields
for i in 0..server_count {
if display_row >= app.scroll_offset && display_row < app.scroll_offset + visible {
lines.push(build_field_line(app, i, display_row));
}
display_row += 1;
}
let selected = i == app.selected;
// Separator line
if display_row >= app.scroll_offset && display_row < app.scroll_offset + visible {
lines.push(Line::from(Span::styled(
" ----------------------------------------",
Style::default().fg(Color::DarkGray),
)));
}
display_row += 1;
// Global fields
for i in 0..app.global_fields.len() {
let field_idx = server_count + i;
if display_row >= app.scroll_offset && display_row < app.scroll_offset + visible {
lines.push(build_field_line(app, field_idx, display_row));
}
display_row += 1;
}
let paragraph = Paragraph::new(lines);
f.render_widget(paragraph, area);
// Cursor position while editing
if app.mode == Mode::Editing {
let sel_display_row = if app.selected >= server_count {
app.selected + 1
} else {
app.selected
};
let row_in_view = sel_display_row.saturating_sub(app.scroll_offset);
let prefix: u16 = 3 + LABEL_WIDTH as u16 + 2;
let cx = area.x + prefix + app.edit_cursor as u16;
let cy = area.y + row_in_view as u16;
if cx < area.x + area.width && cy < area.y + area.height {
f.set_cursor_position((cx, cy));
}
}
}
fn build_field_line(app: &App, field_idx: usize, _display_row: usize) -> Line<'static> {
let sc = app.server_field_count();
let field = if field_idx < sc {
&app.server_tabs[app.active_tab].fields[field_idx]
} else {
&app.global_fields[field_idx - sc]
};
let selected = field_idx == app.selected;
let indicator = if selected { " > " } else { " " };
let label_style = if selected {
@@ -492,35 +657,18 @@ fn render_fields(f: &mut Frame, app: &mut App, area: Rect) {
let padded_label = format!("{:<width$}", field.label, width = LABEL_WIDTH);
// Value display
let (value_text, value_style) = if app.mode == Mode::Editing && selected {
(app.edit_buffer.clone(), Style::default().fg(Color::Yellow))
} else {
field_display(field)
};
lines.push(Line::from(vec![
Span::styled(indicator, label_style),
Line::from(vec![
Span::styled(indicator.to_string(), label_style),
Span::styled(padded_label, label_style),
Span::raw(" "),
Span::styled(value_text, value_style),
]));
}
let paragraph = Paragraph::new(lines);
f.render_widget(paragraph, area);
// Cursor position while editing
if app.mode == Mode::Editing {
let row_in_view = app.selected - app.scroll_offset;
// prefix: 3 (indicator) + LABEL_WIDTH + 2 (gap) = 27
let prefix: u16 = 3 + LABEL_WIDTH as u16 + 2;
let cx = area.x + prefix + app.edit_cursor as u16;
let cy = area.y + row_in_view as u16;
if cx < area.x + area.width && cy < area.y + area.height {
f.set_cursor_position((cx, cy));
}
}
])
}
/// Returns (display_text, style) for a field in normal mode.
@@ -565,18 +713,54 @@ fn field_display(field: &Field) -> (String, Style) {
_ => (field.value.clone(), Style::default().fg(Color::White)),
}
}
fn render_tab_bar(f: &mut Frame, app: &App, area: Rect) {
let mut spans: Vec<Span> = Vec::new();
spans.push(Span::raw(" "));
for (i, tab) in app.server_tabs.iter().enumerate() {
let num = i + 1;
let name = tab
.fields
.iter()
.find(|f| f.key == "node_name")
.filter(|f| !f.value.is_empty())
.map(|f| f.value.clone())
.unwrap_or_else(|| format!("Server {}", num));
let label = format!(" {} {} ", num, name);
if i == app.active_tab {
spans.push(Span::styled(
label,
Style::default()
.fg(Color::Black)
.bg(Color::Cyan)
.add_modifier(Modifier::BOLD),
));
} else {
spans.push(Span::styled(label, Style::default().fg(Color::DarkGray)));
}
spans.push(Span::raw(" "));
}
spans.push(Span::styled(" + Add ", Style::default().fg(Color::Green)));
f.render_widget(Paragraph::new(Line::from(spans)), area);
}
fn render_footer(f: &mut Frame, app: &App, area: Rect) {
let help = app.fields[app.selected].help;
let help = app.selected_field().help;
let keybindings = if app.mode == Mode::Editing {
"Enter confirm Esc cancel"
} else if app.server_tabs.len() > 1 {
"j/k select Enter edit Tab switch + add x remove ^S save q quit"
} else {
"Up/Down select Enter edit ^S save q quit"
"j/k select Enter edit + add server ^S save q quit"
};
let mut status_spans: Vec<Span> = vec![Span::styled(
keybindings,
format!(" {}", keybindings),
Style::default().fg(Color::DarkGray),
)];
@@ -592,18 +776,7 @@ fn render_footer(f: &mut Frame, app: &App, area: Rect) {
format!(" {}", help),
Style::default().fg(Color::DarkGray),
)),
Line::from(
status_spans
.into_iter()
.map(|mut s| {
// add left padding to first span
if s.content.as_ref() == keybindings {
s.content = format!(" {}", s.content).into();
}
s
})
.collect::<Vec<_>>(),
),
Line::from(status_spans),
];
let footer = Paragraph::new(footer_text).block(
@@ -614,11 +787,9 @@ fn render_footer(f: &mut Frame, app: &App, area: Rect) {
f.render_widget(footer, area);
}
// ── Entry point ──────────────────────────────────────────────────────────────
// -- Entry point --------------------------------------------------------------
pub fn run(config_path: PathBuf) -> anyhow::Result<SetupOutcome> {
// Setup terminal
terminal::enable_raw_mode()?;
let mut stdout = io::stdout();
execute!(stdout, EnterAlternateScreen)?;
@@ -630,14 +801,13 @@ pub fn run(config_path: PathBuf) -> anyhow::Result<SetupOutcome> {
let result = event_loop(&mut terminal, &mut app);
// Restore terminal
terminal::disable_raw_mode()?;
execute!(terminal.backend_mut(), LeaveAlternateScreen)?;
terminal.show_cursor()?;
result?;
// ── Post-TUI: decide outcome ─────────────────────────────────────
// -- Post-TUI: decide outcome ---------------------------------------------
if !app.saved_once {
return Ok(SetupOutcome::Cancelled);
@@ -648,7 +818,7 @@ pub fn run(config_path: PathBuf) -> anyhow::Result<SetupOutcome> {
eprintln!();
let wants_service = app
.fields
.global_fields
.iter()
.find(|f| f.key == "install_service")
.map(|f| f.value == "true")
@@ -662,15 +832,12 @@ pub fn run(config_path: PathBuf) -> anyhow::Result<SetupOutcome> {
eprintln!(" Starting proxy directly instead.\n");
}
}
} else {
// Uninstall service if it was previously installed but toggled off
if super::service::is_installed() {
} else if super::service::is_installed() {
if let Err(e) = super::service::uninstall_service() {
eprintln!(" Service uninstall failed: {}", e);
eprintln!();
}
}
}
Ok(SetupOutcome::ReadyToRun(config_path))
}
@@ -684,7 +851,6 @@ fn event_loop(
if event::poll(Duration::from_millis(200))? {
if let Event::Key(key) = event::read()? {
// Only handle Press events (ignore Release on Windows)
if key.kind == KeyEventKind::Press && app.handle_key(key) {
break;
}

55
aether-proxy/src/state.rs
View File

@@ -1,42 +1,45 @@
//! Shared application state passed to all subsystems.
//!
//! Consolidates the multiple `Arc<...>` parameters that were previously
//! threaded individually through proxy server, heartbeat, and handlers.
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::{Arc, RwLock};
use std::time::Duration;
use tokio::sync::Semaphore;
use tokio_rustls::TlsAcceptor;
use crate::config::Config;
use crate::hardware::HardwareInfo;
use crate::proxy::delegate_client::DelegateClient;
use crate::proxy::target_filter::DnsCache;
use crate::registration::client::AetherClient;
use crate::runtime::SharedDynamicConfig;
use crate::target_filter::DnsCache;
/// Central application state shared across all tasks.
/// Central application state shared across all servers/tunnels.
pub struct AppState {
pub config: Arc<Config>,
pub node_id: Arc<RwLock<String>>,
pub dynamic: SharedDynamicConfig,
pub aether_client: Arc<AetherClient>,
pub hardware_info: Arc<HardwareInfo>,
pub public_ip: String,
pub tls_fingerprint: Option<String>,
pub tls_acceptor: Option<TlsAcceptor>,
/// Shared delegate client for proxy-initiated upstream requests.
pub delegate_client: DelegateClient,
/// Active connection count for metrics reporting.
pub active_connections: Arc<AtomicU64>,
/// Connection concurrency limiter.
pub connection_semaphore: Arc<Semaphore>,
/// DNS cache for upstream target resolution.
/// DNS cache for upstream target resolution (shared).
pub dns_cache: Arc<DnsCache>,
/// Request/latency metrics for heartbeat.
/// Reqwest client for tunnel upstream requests (shared).
pub reqwest_client: reqwest::Client,
}
/// Per-server state: one instance per Aether server connection.
pub struct ServerContext {
/// Human-readable label for logging (e.g. "server-0").
pub server_label: String,
/// Aether server URL for this connection.
pub aether_url: String,
/// Management token for this server.
pub management_token: String,
/// Resolved node name (per-server override or global fallback).
pub node_name: String,
/// Node ID assigned by this Aether server.
pub node_id: Arc<RwLock<String>>,
/// API client for this server.
pub aether_client: Arc<AetherClient>,
/// Dynamic config from this server's heartbeat ACKs.
pub dynamic: SharedDynamicConfig,
/// Per-server active connection count.
pub active_connections: Arc<AtomicU64>,
/// Per-server request/latency metrics.
pub metrics: Arc<ProxyMetrics>,
/// Reconnect attempt counter (reset on successful connection).
pub reconnect_attempts: AtomicU32,
}
/// Aggregate metrics for reporting to Aether.

0
aether-proxy/src/proxy/target_filter.rs → aether-proxy/src/target_filter.rs
View File

138
aether-proxy/src/tunnel/client.rs Normal file
View File

@@ -0,0 +1,138 @@
//! WebSocket tunnel client: connect, authenticate, and run the tunnel.
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::watch;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http;
use tracing::{debug, info};
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.
pub async fn connect_and_run(
state: &Arc<AppState>,
server: &Arc<ServerContext>,
shutdown: &mut watch::Receiver<bool>,
) -> Result<TunnelOutcome, anyhow::Error> {
let ws_url = build_tunnel_url(server);
info!(url = %ws_url, "connecting tunnel");
// Build WebSocket request with auth headers
let mut request = ws_url.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)?);
headers.insert(
"X-Node-Name",
http::HeaderValue::from_str(&server.node_name)?,
);
// Connect
let (ws_stream, _response) = tokio_tungstenite::connect_async(request).await?;
info!("tunnel connected");
// Reset reconnect counter on success
server.reconnect_attempts.store(0, Ordering::Relaxed);
// Split into read/write halves
let (ws_sink, ws_read) = futures_util::StreamExt::split(ws_stream);
// Spawn writer task
let (frame_tx, writer_handle) = writer::spawn_writer(ws_sink);
// Spawn heartbeat task
let hb_handle = heartbeat::spawn(
Arc::clone(&state.config),
Arc::clone(server),
frame_tx.clone(),
shutdown.clone(),
);
// Run dispatcher (blocks until disconnect or shutdown)
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),
}
}
_ = 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.
let _ = tokio::time::timeout(Duration::from_secs(35), writer_handle).await;
info!("tunnel disconnected");
Ok(outcome)
}
/// Calculate next reconnect delay with exponential backoff + jitter.
pub fn next_reconnect_delay(state: &Arc<AppState>, server: &Arc<ServerContext>) -> Duration {
let attempt = server.reconnect_attempts.fetch_add(1, Ordering::Relaxed);
let base_ms = state.config.tunnel_reconnect_base_ms;
let max_ms = state.config.tunnel_reconnect_max_ms;
let delay_ms = base_ms.saturating_mul(1u64 << attempt.min(10)).min(max_ms);
let jitter = (delay_ms / 4).max(1);
let jitter_ms = rand_u64() % jitter;
Duration::from_millis(delay_ms + jitter_ms)
}
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)
}
/// Simple pseudo-random u64 (no external crate needed).
fn rand_u64() -> u64 {
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
static COUNTER: AtomicU64 = AtomicU64::new(0);
let seed = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_nanos() as u64;
let cnt = COUNTER.fetch_add(1, Ordering::Relaxed);
let mut x = seed ^ cnt;
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
x
}

201
aether-proxy/src/tunnel/dispatcher.rs Normal file
View File

@@ -0,0 +1,201 @@
//! 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, warn};
use crate::state::{AppState, ServerContext};
use super::heartbeat::HeartbeatHandle;
use super::protocol::{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 read_err = loop {
let msg_result = match ws_stream.next().await {
Some(r) => r,
None => break None, // stream ended
};
let msg = match msg_result {
Ok(m) => m,
Err(e) => {
error!(error = %e, "WebSocket read error");
break Some(e);
}
};
let data = match msg {
Message::Binary(data) => Bytes::from(data),
Message::Ping(_) => continue,
Message::Pong(_) => continue,
Message::Close(_) => {
debug!("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 => {
// Parse request metadata
let meta: RequestMeta = match serde_json::from_slice(&frame.payload) {
Ok(m) => m,
Err(e) => {
warn!(stream_id = frame.stream_id, error = %e, "invalid request metadata");
let _ = frame_tx
.send(Frame::new(
frame.stream_id,
MsgType::StreamError,
0,
Bytes::from(format!("invalid request metadata: {e}")),
))
.await;
continue;
}
};
if streams.len() >= max_streams {
warn!(
stream_id = frame.stream_id,
"max concurrent streams reached"
);
let _ = frame_tx
.send(Frame::new(
frame.stream_id,
MsgType::StreamError,
0,
Bytes::from("max concurrent streams reached"),
))
.await;
continue;
}
// Create body channel and spawn handler
let (body_tx, body_rx) = mpsc::channel::<Frame>(16);
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
streams.remove(&frame.stream_id);
}
MsgType::Ping => {
let _ = frame_tx
.send(Frame::control(MsgType::Pong, frame.payload))
.await;
}
MsgType::HeartbeatAck => {
heartbeat.on_ack(frame.payload).await;
}
MsgType::GoAway => {
debug!("received GOAWAY");
break None;
}
_ => {
debug!(msg_type = ?frame.msg_type, "ignoring unexpected frame type");
}
}
// Periodically clean up finished handles to avoid unbounded growth
if handler_handles.len() > max_streams {
handler_handles.retain(|h| !h.is_finished());
}
};
// 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;
}

113
aether-proxy/src/tunnel/heartbeat.rs Normal file
View File

@@ -0,0 +1,113 @@
//! Tunnel heartbeat: sends metrics over the tunnel, processes ACKs.
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;
use bytes::Bytes;
use tokio::sync::watch;
use tracing::{debug, warn};
use crate::config::Config;
use crate::registration::client::RemoteConfig;
use crate::runtime;
use crate::state::ServerContext;
use super::protocol::{Frame, MsgType};
use super::writer::FrameSender;
/// 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;
}
}
/// Spawn the heartbeat task. Returns a handle for forwarding ACKs.
pub fn spawn(
config: Arc<Config>,
server: Arc<ServerContext>,
frame_tx: FrameSender,
mut shutdown: watch::Receiver<bool>,
) -> HeartbeatHandle {
let (ack_tx, mut ack_rx) = tokio::sync::mpsc::channel::<Bytes>(4);
let interval = Duration::from_secs(config.heartbeat_interval);
tokio::spawn(async move {
let mut ticker = tokio::time::interval(interval);
ticker.tick().await; // Skip first immediate tick
loop {
tokio::select! {
_ = ticker.tick() => {
let payload = build_heartbeat_payload(&server);
let frame = Frame::control(MsgType::HeartbeatData, payload);
if frame_tx.send(frame).await.is_err() {
break; // Writer closed
}
debug!("sent heartbeat data");
}
Some(ack_payload) = ack_rx.recv() => {
handle_ack(&server, &ack_payload);
}
_ = shutdown.changed() => {
debug!("heartbeat task shutting down");
break;
}
}
}
});
HeartbeatHandle { ack_tx }
}
fn build_heartbeat_payload(server: &ServerContext) -> Bytes {
let node_id = server.node_id.read().unwrap().clone();
let interval_requests = server.metrics.total_requests.swap(0, Ordering::Relaxed);
let interval_latency_ns = server.metrics.total_latency_ns.swap(0, Ordering::Relaxed);
let avg_latency_ms = if interval_requests > 0 {
Some(interval_latency_ns as f64 / interval_requests as f64 / 1_000_000.0)
} else {
None
};
let payload = serde_json::json!({
"node_id": node_id,
"active_connections": server.active_connections.load(Ordering::Relaxed),
"total_requests": interval_requests,
"avg_latency_ms": avg_latency_ms,
});
Bytes::from(serde_json::to_vec(&payload).unwrap_or_default())
}
fn handle_ack(server: &ServerContext, payload: &[u8]) {
if payload.is_empty() {
return;
}
#[derive(serde::Deserialize)]
struct AckPayload {
#[serde(default)]
remote_config: Option<RemoteConfig>,
#[serde(default)]
config_version: u64,
}
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);
}
}
Err(e) => {
warn!(error = %e, "failed to parse heartbeat ACK");
}
}
}

53
aether-proxy/src/tunnel/mod.rs Normal file
View File

@@ -0,0 +1,53 @@
pub mod client;
pub mod dispatcher;
pub mod heartbeat;
pub mod protocol;
pub mod stream_handler;
pub mod writer;
use std::sync::Arc;
use tokio::sync::watch;
use tracing::{error, info};
use crate::state::{AppState, ServerContext};
/// Run the tunnel mode main loop (connect, dispatch, reconnect).
pub async fn run(
state: &Arc<AppState>,
server: &Arc<ServerContext>,
mut shutdown: watch::Receiver<bool>,
) {
info!(server = %server.server_label, "starting tunnel");
loop {
match client::connect_and_run(state, server, &mut shutdown).await {
Ok(client::TunnelOutcome::Shutdown) => {
info!(server = %server.server_label, "tunnel shut down gracefully");
return;
}
Ok(client::TunnelOutcome::Disconnected) => {
info!(server = %server.server_label, "tunnel disconnected, will reconnect");
}
Err(e) => {
error!(server = %server.server_label, error = %e, "tunnel connection lost");
}
}
if *shutdown.borrow() {
info!(server = %server.server_label, "shutdown requested, not reconnecting");
return;
}
let delay = client::next_reconnect_delay(state, server);
info!(server = %server.server_label, delay_ms = delay.as_millis(), "reconnecting tunnel");
tokio::select! {
_ = tokio::time::sleep(delay) => {}
_ = shutdown.changed() => {
info!(server = %server.server_label, "shutdown requested during reconnect wait");
return;
}
}
}
}

161
aether-proxy/src/tunnel/protocol.rs Normal file
View File

@@ -0,0 +1,161 @@
//! Binary frame protocol for WebSocket tunnel multiplexing.
//!
//! Frame layout (10-byte header + variable payload):
//! ```text
//! | stream_id (4B) | msg_type (1B) | flags (1B) | payload_len (4B) | payload (NB) |
//! ```
use bytes::{Buf, BufMut, Bytes, BytesMut};
pub const HEADER_SIZE: usize = 10;
/// Frame flags.
pub mod flags {
pub const END_STREAM: u8 = 0x01;
pub const GZIP_COMPRESSED: u8 = 0x02;
}
/// Message types for the tunnel protocol.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum MsgType {
RequestHeaders = 0x01,
RequestBody = 0x02,
ResponseHeaders = 0x03,
ResponseBody = 0x04,
StreamEnd = 0x05,
StreamError = 0x06,
Ping = 0x10,
Pong = 0x11,
GoAway = 0x12,
HeartbeatData = 0x13,
HeartbeatAck = 0x14,
}
impl MsgType {
pub fn from_u8(v: u8) -> Option<Self> {
match v {
0x01 => Some(Self::RequestHeaders),
0x02 => Some(Self::RequestBody),
0x03 => Some(Self::ResponseHeaders),
0x04 => Some(Self::ResponseBody),
0x05 => Some(Self::StreamEnd),
0x06 => Some(Self::StreamError),
0x10 => Some(Self::Ping),
0x11 => Some(Self::Pong),
0x12 => Some(Self::GoAway),
0x13 => Some(Self::HeartbeatData),
0x14 => Some(Self::HeartbeatAck),
_ => None,
}
}
}
/// A single multiplexed frame.
#[derive(Debug, Clone)]
pub struct Frame {
pub stream_id: u32,
pub msg_type: MsgType,
pub flags: u8,
pub payload: Bytes,
}
impl Frame {
pub fn new(stream_id: u32, msg_type: MsgType, flags: u8, payload: impl Into<Bytes>) -> Self {
Self {
stream_id,
msg_type,
flags,
payload: payload.into(),
}
}
/// Control frame (stream_id = 0).
pub fn control(msg_type: MsgType, payload: impl Into<Bytes>) -> Self {
Self::new(0, msg_type, 0, payload)
}
pub fn is_end_stream(&self) -> bool {
self.flags & flags::END_STREAM != 0
}
pub fn is_gzip(&self) -> bool {
self.flags & flags::GZIP_COMPRESSED != 0
}
/// Encode into a binary buffer.
pub fn encode(&self) -> Bytes {
let mut buf = BytesMut::with_capacity(HEADER_SIZE + self.payload.len());
buf.put_u32(self.stream_id);
buf.put_u8(self.msg_type as u8);
buf.put_u8(self.flags);
buf.put_u32(self.payload.len() as u32);
buf.put(self.payload.clone());
buf.freeze()
}
/// Decode from a binary buffer.
pub fn decode(mut data: Bytes) -> Result<Self, ProtocolError> {
if data.len() < HEADER_SIZE {
return Err(ProtocolError::TooShort {
expected: HEADER_SIZE,
actual: data.len(),
});
}
let stream_id = data.get_u32();
let msg_type_raw = data.get_u8();
let frame_flags = data.get_u8();
let payload_len = data.get_u32() as usize;
if data.remaining() < payload_len {
return Err(ProtocolError::Incomplete {
expected: HEADER_SIZE + payload_len,
actual: HEADER_SIZE + data.remaining(),
});
}
let msg_type =
MsgType::from_u8(msg_type_raw).ok_or(ProtocolError::UnknownMsgType(msg_type_raw))?;
let payload = data.split_to(payload_len);
Ok(Self {
stream_id,
msg_type,
flags: frame_flags,
payload,
})
}
}
/// Protocol errors.
#[derive(Debug, thiserror::Error)]
pub enum ProtocolError {
#[error("frame too short: expected {expected} bytes, got {actual}")]
TooShort { expected: usize, actual: usize },
#[error("frame incomplete: expected {expected} bytes, got {actual}")]
Incomplete { expected: usize, actual: usize },
#[error("unknown message type: 0x{0:02x}")]
UnknownMsgType(u8),
}
/// JSON payload for REQUEST_HEADERS frames.
#[derive(Debug, serde::Deserialize)]
pub struct RequestMeta {
pub method: String,
pub url: String,
pub headers: std::collections::HashMap<String, String>,
#[serde(default = "default_timeout")]
pub timeout: u64,
}
fn default_timeout() -> u64 {
60
}
/// JSON payload for RESPONSE_HEADERS frames.
#[derive(Debug, serde::Serialize)]
pub struct ResponseMeta {
pub status: u16,
/// Header list preserving duplicates (e.g. multiple Set-Cookie).
pub headers: Vec<(String, String)>,
}

264
aether-proxy/src/tunnel/stream_handler.rs Normal file
View File

@@ -0,0 +1,264 @@
//! Per-stream request handler.
//!
//! Receives request frames, executes the upstream HTTP request,
//! and sends response frames back through the writer channel.
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::{Duration, Instant};
use bytes::Bytes;
use futures_util::StreamExt;
use tokio::sync::mpsc;
use tracing::{debug, warn};
use crate::state::{AppState, ServerContext};
use crate::target_filter;
use super::protocol::{flags, Frame, MsgType, RequestMeta, ResponseMeta};
use super::writer::FrameSender;
/// Maximum response body chunk size per frame (32 KB).
const MAX_CHUNK_SIZE: usize = 32 * 1024;
/// 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,
mut body_rx: mpsc::Receiver<Frame>,
frame_tx: FrameSender,
) {
let start = Instant::now();
server.active_connections.fetch_add(1, Ordering::Relaxed);
handle_stream_inner(&state, &server, stream_id, meta, &mut body_rx, &frame_tx).await;
server.active_connections.fetch_sub(1, Ordering::Relaxed);
server.metrics.record_request(start.elapsed());
}
async fn handle_stream_inner(
state: &AppState,
server: &ServerContext,
stream_id: u32,
meta: RequestMeta,
body_rx: &mut mpsc::Receiver<Frame>,
frame_tx: &FrameSender,
) {
// Collect request body
let mut body_parts: Vec<Bytes> = Vec::new();
let mut body_done = false;
// Drain body frames
while !body_done {
match body_rx.recv().await {
Some(frame) => {
if frame.msg_type == MsgType::RequestBody {
let payload = if frame.is_gzip() {
match decompress_gzip(&frame.payload) {
Ok(d) => d,
Err(e) => {
send_error(
frame_tx,
stream_id,
&format!("gzip decompress failed: {e}"),
)
.await;
return;
}
}
} else {
frame.payload.clone()
};
if !payload.is_empty() {
body_parts.push(payload);
}
if frame.is_end_stream() {
body_done = true;
}
} else if frame.msg_type == MsgType::StreamEnd
|| frame.msg_type == MsgType::StreamError
{
body_done = true;
if frame.msg_type == MsgType::StreamError {
return; // Client cancelled
}
}
}
None => return, // Channel closed
}
}
let body: Bytes = if body_parts.is_empty() {
Bytes::new()
} else if body_parts.len() == 1 {
body_parts.into_iter().next().unwrap()
} else {
let total: usize = body_parts.iter().map(|b| b.len()).sum();
let mut combined = Vec::with_capacity(total);
for part in &body_parts {
combined.extend_from_slice(part);
}
Bytes::from(combined)
};
// 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;
}
};
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;
}
};
let port = target_url.port_or_known_default().unwrap_or(443);
// DNS + target validation (dns_cache is populated as a side effect)
let dns_start = Instant::now();
{
let allowed_ports = server.dynamic.read().unwrap().allowed_ports.clone();
if let Err(e) =
target_filter::validate_target(&host, port, &allowed_ports, &state.dns_cache).await
{
send_error(frame_tx, stream_id, &format!("target blocked: {e}")).await;
return;
}
}
let dns_ms = dns_start.elapsed().as_millis() as u64;
// Execute upstream request
let client = &state.reqwest_client;
let timeout = Duration::from_secs(meta.timeout);
let method: reqwest::Method = meta.method.parse().unwrap_or(reqwest::Method::GET);
let mut req = client.request(method, &meta.url);
for (k, v) in &meta.headers {
req = req.header(k.as_str(), v.as_str());
}
let body_size = body.len();
if !body.is_empty() {
req = req.body(body);
}
req = req.timeout(timeout);
let upstream_start = Instant::now();
let response = match req.send().await {
Ok(r) => r,
Err(e) => {
let msg = if e.is_timeout() {
"upstream timeout".to_string()
} else if e.is_connect() {
format!("upstream connect error: {e}")
} else {
format!("upstream error: {e}")
};
send_error(frame_tx, stream_id, &msg).await;
return;
}
};
// Send RESPONSE_HEADERS
let status = response.status().as_u16();
let ttfb_ms = upstream_start.elapsed().as_millis() as u64;
let mut resp_headers: Vec<(String, String)> = Vec::new();
for (k, v) in response.headers() {
if let Ok(vs) = v.to_str() {
resp_headers.push((k.as_str().to_string(), vs.to_string()));
}
}
// Inject proxy timing (same format as delegate mode)
let timing = serde_json::json!({
"dns_ms": dns_ms,
"ttfb_ms": ttfb_ms,
"upstream_ms": ttfb_ms,
"upstream_processing_ms": ttfb_ms.saturating_sub(dns_ms),
"body_size": body_size,
"mode": "tunnel",
});
resp_headers.push(("x-proxy-timing".to_string(), timing.to_string()));
let resp_meta = ResponseMeta {
status,
headers: resp_headers,
};
let meta_json = serde_json::to_vec(&resp_meta).unwrap_or_default();
let _ = frame_tx
.send(Frame::new(
stream_id,
MsgType::ResponseHeaders,
0,
meta_json,
))
.await;
// Stream response body
let mut stream = response.bytes_stream();
while let Some(chunk_result) = stream.next().await {
match chunk_result {
Ok(chunk) => {
if chunk.len() <= MAX_CHUNK_SIZE {
// 大多数 chunk 无需分割，直接零拷贝发送
let _ = frame_tx
.send(Frame::new(stream_id, MsgType::ResponseBody, 0, chunk))
.await;
} else {
// 超大 chunk 按 MAX_CHUNK_SIZE 分割（使用 Bytes::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 _ = frame_tx
.send(Frame::new(stream_id, MsgType::ResponseBody, 0, slice))
.await;
offset = end;
}
}
}
Err(e) => {
warn!(stream_id, error = %e, "upstream body read error");
send_error(frame_tx, stream_id, &format!("body read error: {e}")).await;
return;
}
}
}
// Send STREAM_END
let _ = frame_tx
.send(Frame::new(
stream_id,
MsgType::StreamEnd,
flags::END_STREAM,
Bytes::new(),
))
.await;
debug!(stream_id, status, "stream completed");
}
async fn send_error(tx: &FrameSender, stream_id: u32, msg: &str) {
let _ = tx
.send(Frame::new(
stream_id,
MsgType::StreamError,
0,
Bytes::from(msg.to_string()),
))
.await;
}
fn decompress_gzip(data: &[u8]) -> Result<Bytes, std::io::Error> {
use flate2::read::GzDecoder;
use std::io::Read;
let mut decoder = GzDecoder::new(data);
let mut buf = Vec::new();
decoder.read_to_end(&mut buf)?;
Ok(Bytes::from(buf))
}

37
aether-proxy/src/tunnel/writer.rs Normal file
View File

@@ -0,0 +1,37 @@
//! Dedicated WebSocket writer task.
//!
//! All frame writes go through an mpsc channel to a single writer task,
//! avoiding contention on the WebSocket sink.
use futures_util::SinkExt;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use tokio_tungstenite::tungstenite::Message;
use tracing::{debug, error};
use super::protocol::Frame;
/// Sender half — cloned by stream handlers and heartbeat.
pub type FrameSender = mpsc::Sender<Frame>;
/// Spawn the writer task. Returns the sender and a JoinHandle for cleanup.
pub fn spawn_writer<S>(mut sink: S) -> (FrameSender, JoinHandle<()>)
where
S: SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin + Send + 'static,
{
let (tx, mut rx) = mpsc::channel::<Frame>(256);
let handle = tokio::spawn(async move {
while let Some(frame) = rx.recv().await {
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;
}
}
debug!("writer task exiting");
let _ = sink.close().await;
});
(tx, handle)
}

106
alembic/versions/20260224_1700_9a0b1c2d3e4f_add_tunnel_mode_fields.py Normal file
View File

@@ -0,0 +1,106 @@
"""Add tunnel mode fields and remove IP forwarding fields
Revision ID: 9a0b1c2d3e4f
Revises: 8f9a0b1c2d3e
Create Date: 2026-02-24 17:00:00.000000
"""
from __future__ import annotations
from collections.abc import Sequence
import sqlalchemy as sa
from sqlalchemy import inspect
from alembic import op
revision: str = "9a0b1c2d3e4f"
down_revision: str | None = "8f9a0b1c2d3e"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def column_exists(table_name: str, column_name: str) -> bool:
bind = op.get_bind()
inspector = inspect(bind)
columns = [c["name"] for c in inspector.get_columns(table_name)]
return column_name in columns
def upgrade() -> None:
# 添加 tunnel 模式字段
if not column_exists("proxy_nodes", "tunnel_mode"):
op.add_column(
"proxy_nodes",
sa.Column(
"tunnel_mode",
sa.Boolean(),
nullable=False,
server_default=sa.text("false"),
comment="是否使用 WebSocket 隧道模式",
),
)
if not column_exists("proxy_nodes", "tunnel_connected"):
op.add_column(
"proxy_nodes",
sa.Column(
"tunnel_connected",
sa.Boolean(),
nullable=False,
server_default=sa.text("false"),
comment="隧道是否已连接",
),
)
if not column_exists("proxy_nodes", "tunnel_connected_at"):
op.add_column(
"proxy_nodes",
sa.Column(
"tunnel_connected_at",
sa.DateTime(timezone=True),
nullable=True,
comment="隧道最近一次建立时间",
),
)
# tunnel 模式节点不需要 port，将其置零
op.execute("UPDATE proxy_nodes SET port = 0 WHERE tunnel_mode = true")
# 移除旧的 IP 转发字段
if column_exists("proxy_nodes", "tls_enabled"):
op.drop_column("proxy_nodes", "tls_enabled")
if column_exists("proxy_nodes", "tls_cert_fingerprint"):
op.drop_column("proxy_nodes", "tls_cert_fingerprint")
def downgrade() -> None:
# 恢复 IP 转发字段
if not column_exists("proxy_nodes", "tls_cert_fingerprint"):
op.add_column(
"proxy_nodes",
sa.Column(
"tls_cert_fingerprint",
sa.String(128),
nullable=True,
comment="TLS 证书 SHA-256 指纹（hex）",
),
)
if not column_exists("proxy_nodes", "tls_enabled"):
op.add_column(
"proxy_nodes",
sa.Column(
"tls_enabled",
sa.Boolean(),
nullable=False,
server_default=sa.text("false"),
comment="是否启用 TLS 加密",
),
)
# 移除 tunnel 模式字段
if column_exists("proxy_nodes", "tunnel_connected_at"):
op.drop_column("proxy_nodes", "tunnel_connected_at")
if column_exists("proxy_nodes", "tunnel_connected"):
op.drop_column("proxy_nodes", "tunnel_connected")
if column_exists("proxy_nodes", "tunnel_mode"):
op.drop_column("proxy_nodes", "tunnel_mode")

9
frontend/src/api/proxy-nodes.ts
View File

@@ -5,7 +5,6 @@ export interface ProxyNodeRemoteConfig {
allowed_ports?: number[]
log_level?: string
heartbeat_interval?: number
timestamp_tolerance?: number
}
export interface ProxyNode {
@@ -16,6 +15,9 @@ export interface ProxyNode {
region: string | null
status: 'online' | 'unhealthy' | 'offline'
is_manual: boolean
tunnel_mode: boolean
tunnel_connected: boolean
tunnel_connected_at: string | null
// 手动节点专用字段
proxy_url?: string
proxy_username?: string
@@ -101,9 +103,4 @@ export const proxyNodesApi = {
const response = await apiClient.post<ProxyNodeTestResult>('/api/admin/proxy-nodes/test-url', data)
return response.data
},
async getHmacKey(): Promise<{ proxy_hmac_key: string }> {
const response = await apiClient.get<{ proxy_hmac_key: string }>('/api/admin/proxy-nodes/hmac-key')
return response.data
},
}

64
frontend/src/views/admin/ProxyNodes.vue
View File

@@ -95,16 +95,6 @@
</SelectContent>
</Select>
<div class="h-4 w-px bg-border" />
<Button
variant="ghost"
size="icon"
class="h-8 w-8"
title="复制 HMAC Key"
@click="copyHmacKey"
>
<Copy class="w-3.5 h-3.5" />
</Button>
<div class="h-4 w-px bg-border" />
<Button
variant="ghost"
size="icon"
@@ -172,11 +162,18 @@
>
手动
</Badge>
<Badge
v-if="node.tunnel_mode"
variant="outline"
class="text-[10px] px-1.5 py-0"
>
Tunnel
</Badge>
<HardwareTooltip :node="node" />
</div>
</TableCell>
<TableCell class="py-4">
<code class="text-xs text-muted-foreground">{{ node.is_manual ? (node.proxy_url || `${node.ip}:${node.port}`) : `${node.ip}:${node.port}` }}</code>
<code class="text-xs text-muted-foreground">{{ nodeAddress(node) }}</code>
</TableCell>
<TableCell class="py-4">
<span class="text-sm text-muted-foreground">{{ formatRegion(node.region) }}</span>
@@ -282,9 +279,16 @@
>
手动
</Badge>
<Badge
v-if="node.tunnel_mode"
variant="outline"
class="text-[10px] px-1.5 py-0"
>
Tunnel
</Badge>
<HardwareTooltip :node="node" />
</div>
<code class="text-xs text-muted-foreground">{{ node.is_manual ? (node.proxy_url || `${node.ip}:${node.port}`) : `${node.ip}:${node.port}` }}</code>
<code class="text-xs text-muted-foreground">{{ nodeAddress(node) }}</code>
</div>
<Badge
:variant="statusVariant(node.status)"
@@ -521,15 +525,6 @@
max="600"
/>
</div>
<div class="space-y-1.5">
<Label>时间戳容差 (秒)</Label>
<Input
v-model="configForm.timestamp_tolerance"
type="number"
min="10"
max="3600"
/>
</div>
</div>
<div
v-if="configNode"
@@ -560,7 +555,6 @@
import { ref, computed, onMounted, watch } from 'vue'
import { useProxyNodesStore } from '@/stores/proxy-nodes'
import { useToast } from '@/composables/useToast'
import { useClipboard } from '@/composables/useClipboard'
import { useConfirm } from '@/composables/useConfirm'
import { proxyNodesApi, type ProxyNode, type ProxyNodeRemoteConfig } from '@/api/proxy-nodes'
@@ -586,13 +580,12 @@ import {
Dialog,
} from '@/components/ui'
import { Search, Trash2, Plus, SquarePen, Activity, Loader2, Settings, Copy } from 'lucide-vue-next'
import { Search, Trash2, Plus, SquarePen, Activity, Loader2, Settings } from 'lucide-vue-next'
import { parseApiError } from '@/utils/errorParser'
import { formatRegion } from '@/utils/region'
import HardwareTooltip from './components/HardwareTooltip.vue'
const { success, error: toastError } = useToast()
const { copyToClipboard } = useClipboard()
const { confirmDanger } = useConfirm()
const store = useProxyNodesStore()
@@ -621,7 +614,6 @@ const configForm = ref({
allowed_ports: '',
log_level: 'info',
heartbeat_interval: '30',
timestamp_tolerance: '300',
})
// 测试连通性
@@ -686,15 +678,6 @@ async function handleTestUrl() {
}
}
async function copyHmacKey() {
try {
const { proxy_hmac_key } = await proxyNodesApi.getHmacKey()
await copyToClipboard(proxy_hmac_key)
} catch (err: unknown) {
toastError(parseApiError(err, '获取 HMAC Key 失败'))
}
}
function handleEdit(node: ProxyNode) {
editingNode.value = node
addForm.value = {
@@ -766,7 +749,6 @@ function handleConfig(node: ProxyNode) {
allowed_ports: rc.allowed_ports?.join(', ') || '',
log_level: rc.log_level || 'info',
heartbeat_interval: String(rc.heartbeat_interval || node.heartbeat_interval || 30),
timestamp_tolerance: String(rc.timestamp_tolerance || 300),
}
showConfigDialog.value = true
}
@@ -800,10 +782,6 @@ async function handleSaveConfig() {
if (!isNaN(hb) && hb >= 5) {
data.heartbeat_interval = hb
}
const tt = parseInt(configForm.value.timestamp_tolerance)
if (!isNaN(tt) && tt >= 10) {
data.timestamp_tolerance = tt
}
await proxyNodesApi.updateNodeConfig(configNode.value.id, data)
success('远程配置已保存，将在下次心跳时生效')
handleConfigDialogClose(false)
@@ -817,7 +795,7 @@ async function handleSaveConfig() {
async function handleDelete(node: ProxyNode) {
const confirmed = await confirmDanger(
`确定要删除代理节点 "${node.name}" (${node.ip}:${node.port}) 吗？`,
`确定要删除代理节点 "${node.name}" (${node.tunnel_mode ? node.ip : `${node.ip}:${node.port}`}) 吗？`,
'删除节点'
)
if (!confirmed) return
@@ -890,4 +868,10 @@ function formatTime(iso: string | null) {
if (diff < 86400) return `${Math.floor(diff / 3600)}小时前`
return d.toLocaleDateString('zh-CN', { month: '2-digit', day: '2-digit', hour: '2-digit', minute: '2-digit' })
}
function nodeAddress(node: ProxyNode) {
if (node.is_manual) return node.proxy_url || `${node.ip}:${node.port}`
if (node.tunnel_mode) return node.ip || 'WebSocket Tunnel'
return `${node.ip}:${node.port}`
}
</script>

8
generate_keys.py
View File

@@ -6,7 +6,7 @@
import secrets
def main():
def main() -> None:
# 生成JWT密钥
jwt_key = secrets.token_urlsafe(32)
@@ -16,22 +16,18 @@ def main():
# 生成 Redis 密码
redis_password = secrets.token_urlsafe(32)
# 生成代理节点 HMAC 密钥（独立密钥，Aether 服务端和 aether-proxy 配置相同值）
proxy_hmac_key = secrets.token_urlsafe(32)
print("\n将以下内容添加到 .env 文件：\n")
print(f"JWT_SECRET_KEY={jwt_key}")
print(f"ENCRYPTION_KEY={encryption_key}")
print(f"REDIS_PASSWORD={redis_password}")
print(f"PROXY_HMAC_KEY={proxy_hmac_key}")
print()
print("注意:")
print(" - JWT_SECRET_KEY 用于用户登录 token 签名")
print(" - ENCRYPTION_KEY 用于敏感数据加密（如 Provider API Keys）")
print(" - REDIS_PASSWORD 用于 Redis 连接认证（并发控制）")
print(" - PROXY_HMAC_KEY 用于 aether-proxy 代理请求认证（两端配置相同值）")
print(" - 这些密钥应该独立设置，避免相互耦合")
print()
if __name__ == "__main__":
main()

59
src/api/admin/proxy_nodes/routes.py
View File

@@ -32,7 +32,7 @@ pipeline = ApiRequestPipeline()
class ProxyNodeRegisterRequest(BaseModel):
name: str = Field(..., min_length=1, max_length=100, description="节点名")
ip: str = Field(..., description="公网 IP（IPv4/IPv6）")
port: int = Field(..., ge=1, le=65535, description="代理端口")
port: int = Field(0, ge=0, le=65535, description="代理端口（tunnel 模式下为 0）")
region: str | None = Field(None, max_length=100, description="区域标签")
heartbeat_interval: int = Field(30, ge=5, le=600, description="心跳间隔（秒）")
@@ -41,16 +41,13 @@ class ProxyNodeRegisterRequest(BaseModel):
total_requests: int | None = Field(None, ge=0, description="累计请求数")
avg_latency_ms: float | None = Field(None, ge=0, description="平均延迟（毫秒）")
# TLS
tls_enabled: bool = Field(False, description="是否启用 TLS 加密")
tls_cert_fingerprint: str | None = Field(
None, max_length=128, description="TLS 证书 SHA-256 指纹"
)
# 硬件信息
hardware_info: dict | None = Field(None, description="硬件信息 JSON")
estimated_max_concurrency: int | None = Field(None, ge=0, description="估算最大并发连接数")
# Tunnel 模式
tunnel_mode: bool = Field(False, description="是否使用 tunnel 模式连接")
@field_validator("ip")
@classmethod
def validate_ip(cls, v: str) -> str:
@@ -82,9 +79,6 @@ class ProxyNodeRemoteConfigRequest(BaseModel):
allowed_ports: list[int] | None = Field(None, description="允许代理的目标端口")
log_level: str | None = Field(None, description="日志级别 (trace/debug/info/warn/error)")
heartbeat_interval: int | None = Field(None, ge=5, le=600, description="心跳间隔（秒）")
timestamp_tolerance: int | None = Field(
None, ge=10, le=3600, description="HMAC 时间戳容差（秒）"
)
@field_validator("allowed_ports")
@classmethod
@@ -216,12 +210,6 @@ async def test_proxy_node(node_id: str, request: Request, db: Session = Depends(
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
@router.get("/hmac-key")
async def get_proxy_hmac_key(request: Request, db: Session = Depends(get_db)) -> Any:
adapter = AdminGetProxyHmacKeyAdapter()
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
@router.post("/test-url")
async def test_proxy_url(request: Request, db: Session = Depends(get_db)) -> Any:
adapter = AdminTestProxyUrlAdapter()
@@ -273,14 +261,13 @@ class AdminRegisterProxyNodeAdapter(AdminApiAdapter):
port=req.port,
region=req.region,
heartbeat_interval=req.heartbeat_interval,
tls_enabled=req.tls_enabled,
tls_cert_fingerprint=req.tls_cert_fingerprint,
hardware_info=req.hardware_info,
estimated_max_concurrency=req.estimated_max_concurrency,
active_connections=req.active_connections,
total_requests=req.total_requests,
avg_latency_ms=req.avg_latency_ms,
registered_by=context.user.id if context.user else None,
tunnel_mode=req.tunnel_mode,
)
context.add_audit_metadata(
@@ -376,11 +363,24 @@ class AdminDeleteProxyNodeAdapter(AdminApiAdapter):
)
was_system_proxy = result["cleared_system_proxy"]
msg = "deleted, system default proxy cleared" if was_system_proxy else "deleted"
cleared_providers = result.get("cleared_providers", 0)
cleared_endpoints = result.get("cleared_endpoints", 0)
parts = ["deleted"]
if was_system_proxy:
parts.append("system default proxy cleared")
if cleared_providers or cleared_endpoints:
parts.append(
f"cleared proxy from {cleared_providers} provider(s) "
f"and {cleared_endpoints} endpoint(s)"
)
return {
"message": msg,
"message": ", ".join(parts),
"node_id": self.node_id,
"cleared_system_proxy": was_system_proxy,
"cleared_providers": cleared_providers,
"cleared_endpoints": cleared_endpoints,
}
@@ -478,8 +478,6 @@ class AdminUpdateProxyNodeConfigAdapter(AdminApiAdapter):
config_updates["log_level"] = req.log_level
if req.heartbeat_interval is not None:
config_updates["heartbeat_interval"] = req.heartbeat_interval
if req.timestamp_tolerance is not None:
config_updates["timestamp_tolerance"] = req.timestamp_tolerance
node = ProxyNodeService.update_node_config(
context.db, node_id=self.node_id, config_updates=config_updates
@@ -499,23 +497,6 @@ class AdminUpdateProxyNodeConfigAdapter(AdminApiAdapter):
}
@dataclass
class AdminGetProxyHmacKeyAdapter(AdminApiAdapter):
"""获取 proxy_hmac_key 供管理员复制到 aether-proxy 部署"""
name: str = "admin_get_proxy_hmac_key"
async def handle(self, context: ApiRequestContext) -> Any:
from src.config.settings import config
key = config.proxy_hmac_key
if not key:
raise InvalidRequestException(
"PROXY_HMAC_KEY 未配置（也未设置 ENCRYPTION_KEY 用于自动派生）"
)
return {"proxy_hmac_key": key}
class TestProxyUrlRequest(BaseModel):
proxy_url: str = Field(..., min_length=1, max_length=500)
username: str | None = Field(None, max_length=255)

171
src/api/admin/proxy_tunnel.py Normal file
View File

@@ -0,0 +1,171 @@
"""
WebSocket 隧道端点
aether-proxy 通过此端点建立 tunnel 连接。
路径: /api/internal/proxy-tunnel
"""
from __future__ import annotations
import asyncio
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from src.core.logger import logger
from src.services.proxy_node.tunnel_manager import (
TunnelConnection,
get_tunnel_manager,
)
from src.services.proxy_node.tunnel_protocol import Frame
router = APIRouter()
# 单帧最大 64 MB -- AI API 请求体可能包含多张 base64 图片，需要足够余量
_MAX_FRAME_SIZE = 64 * 1024 * 1024
# WebSocket 空闲超时（秒）-- proxy 端 ping 间隔默认 15s，3 倍余量
_IDLE_TIMEOUT = 90.0
async def _authenticate(ws: WebSocket) -> tuple[str, str] | None:
"""验证 WebSocket 连接的认证信息，返回 (node_id, node_name) 或 None
认证方式：Bearer <management_token>，通过 Management Token 系统验证。
authenticate_management_token 是 async 方法（内部有 Redis 速率限制），
因此直接 await 调用。节点存在性检查复用同一 session。
"""
auth = ws.headers.get("authorization", "")
if not auth.startswith("Bearer "):
return None
token = auth[7:]
if not token or not token.startswith("ae_"):
return None
client_ip = getattr(ws.client, "host", "unknown") if ws.client else "unknown"
node_id_header = ws.headers.get("x-node-id", "").strip()
node_name_header = ws.headers.get("x-node-name", "").strip()
if not node_id_header:
return None
from src.database import create_session
from src.models.database import ProxyNode
from src.services.auth.service import AuthService
db = create_session()
try:
result = await AuthService.authenticate_management_token(db, token, client_ip)
if not result:
return None
# 节点存在性检查（复用同一 session，避免额外连接开销）
exists = db.query(
db.query(ProxyNode).filter(ProxyNode.id == node_id_header).exists()
).scalar()
if not exists:
logger.warning("tunnel auth: node_id={} not found in DB", node_id_header)
return None
finally:
db.close()
return node_id_header, node_name_header or node_id_header
@router.websocket("/api/internal/proxy-tunnel")
async def proxy_tunnel_ws(ws: WebSocket) -> None:
"""aether-proxy tunnel WebSocket 端点"""
try:
auth = await _authenticate(ws)
except Exception as e:
logger.warning("tunnel auth error: {}", e)
await ws.accept()
await ws.close(code=4002, reason="authentication error")
return
if not auth:
await ws.accept()
await ws.close(code=4001, reason="unauthorized")
return
node_id: str = auth[0]
node_name: str = auth[1]
await ws.accept()
manager = get_tunnel_manager()
conn = TunnelConnection(node_id, node_name, ws)
manager.register(conn)
# 更新 DB: tunnel_connected = True
await _update_tunnel_status(node_id, connected=True)
try:
oversized_count = 0
while True:
try:
data = await asyncio.wait_for(ws.receive_bytes(), timeout=_IDLE_TIMEOUT)
except asyncio.TimeoutError:
logger.warning("tunnel idle timeout for node_id={}", node_id)
await ws.close(code=4004, reason="idle timeout")
break
if len(data) > _MAX_FRAME_SIZE:
oversized_count += 1
logger.warning("tunnel frame too large from {}: {} bytes", node_id, len(data))
if oversized_count >= 5:
logger.warning("too many oversized frames from {}, closing", node_id)
await ws.close(code=4003, reason="too many oversized frames")
break
continue
oversized_count = 0 # 正常帧重置计数
try:
frame = Frame.decode(data)
except ValueError as e:
logger.warning("tunnel frame decode error from {}: {}", node_id, e)
continue
await manager.handle_incoming_frame(node_id, frame)
except WebSocketDisconnect:
logger.info("tunnel WebSocket disconnected: node_id={}", node_id)
except Exception as e:
logger.error("tunnel WebSocket error for node_id={}: {}", node_id, e)
finally:
manager.unregister(node_id)
await _update_tunnel_status(node_id, connected=False)
async def _update_tunnel_status(node_id: str, *, connected: bool) -> None:
"""更新 ProxyNode 的 tunnel 连接状态（在线程池中执行，避免阻塞 event loop）"""
def _sync_update() -> None:
from datetime import datetime, timezone
from src.database import create_session
from src.models.database import ProxyNode, ProxyNodeStatus
db = create_session()
try:
node = db.query(ProxyNode).filter(ProxyNode.id == node_id).first()
if node:
node.tunnel_connected = connected
now = datetime.now(timezone.utc)
if connected:
node.tunnel_connected_at = now
node.status = ProxyNodeStatus.ONLINE
else:
# 记录断开时刻，供 health_scheduler 计算 UNHEALTHY 缓冲期
node.tunnel_connected_at = now
node.status = ProxyNodeStatus.UNHEALTHY
db.commit()
finally:
db.close()
try:
await asyncio.to_thread(_sync_update)
except Exception as e:
logger.warning("failed to update tunnel status for {}: {}", node_id, e)
# 清除节点信息缓存，确保后续请求能立即感知连接状态变化
from src.services.proxy_node.resolver import invalidate_proxy_node_cache
invalidate_proxy_node_cache(node_id)

3
src/api/handlers/gemini/adapter.py
View File

@@ -16,6 +16,7 @@ from src.api.handlers.base.chat_adapter_base import ChatAdapterBase, register_ad
from src.api.handlers.base.chat_handler_base import ChatHandlerBase
from src.core.api_format import ApiFamily, get_auth_handler
from src.core.api_format.enums import AuthMethod
from src.core.api_format.headers import BROWSER_FINGERPRINT_HEADERS
from src.core.logger import logger
from src.models.gemini import GeminiRequest
from src.services.provider.transport import redact_url_for_log
@@ -197,7 +198,7 @@ class GeminiChatAdapter(ChatAdapterBase):
else:
models_url = f"{base_url_clean}/v1beta/models?key={api_key}"
headers: dict[str, str] = {}
headers: dict[str, str] = {**BROWSER_FINGERPRINT_HEADERS}
if extra_headers:
headers.update(extra_headers)

153
src/clients/http_client.py
View File

@@ -51,8 +51,8 @@ class HTTPClientPool:
_proxy_clients: dict[str, tuple[httpx.AsyncClient, float]] = {}
# 代理客户端缓存上限（避免内存泄漏）
_max_proxy_clients: int = 50
# 代发客户端缓存：{tls: client, plain: client}
_delegate_clients: dict[str, httpx.AsyncClient] = {}
# Tunnel 客户端缓存：{node_id: client}
_tunnel_clients: dict[str, httpx.AsyncClient] = {}
def __new__(cls) -> "HTTPClientPool":
if cls._instance is None:
@@ -289,15 +289,15 @@ class HTTPClientPool:
cls._proxy_clients.clear()
# 关闭代发客户端缓存
for cache_key, client in cls._delegate_clients.items():
# 关闭 tunnel 客户端缓存
for nid, client in cls._tunnel_clients.items():
try:
await client.aclose()
logger.debug("代发客户端已关闭: {}", cache_key)
logger.debug("tunnel 客户端已关闭: {}", nid)
except Exception as e:
logger.warning("关闭代发客户端失败: {}", e)
logger.warning("关闭 tunnel 客户端失败: {}", e)
cls._delegate_clients.clear()
cls._tunnel_clients.clear()
logger.info("所有HTTP客户端已关闭")
@classmethod
@@ -380,88 +380,6 @@ class HTTPClientPool:
client_config.update(kwargs)
return httpx.AsyncClient(**client_config) # type: ignore[arg-type]
@classmethod
def create_delegate_stream_client(
cls,
delegate_config: dict[str, Any],
timeout: httpx.Timeout | None = None,
) -> httpx.AsyncClient:
"""
创建用于代发流式请求的 httpx 客户端
代发模式下不配置 proxy，直接 POST 到 proxy 的 /_aether/delegate 端点。
调用者需要负责关闭返回的客户端。
"""
client_config: dict[str, Any] = {
"http2": False,
"follow_redirects": False,
}
if timeout:
client_config["timeout"] = timeout
else:
client_config["timeout"] = httpx.Timeout(
connect=config.http_connect_timeout,
read=config.http_read_timeout,
write=config.http_write_timeout,
pool=config.http_pool_timeout,
)
if delegate_config.get("tls_enabled"):
from src.utils.ssl_utils import get_proxy_ssl_context
client_config["verify"] = get_proxy_ssl_context()
else:
client_config["verify"] = get_ssl_context()
return httpx.AsyncClient(**client_config)
@classmethod
async def get_delegate_client(
cls,
delegate_config: dict[str, Any],
) -> httpx.AsyncClient:
"""
获取可复用的代发客户端（非流式请求用）
根据 TLS 状态缓存两个客户端（tls / plain），避免每次请求创建新客户端。
当 tls_enabled=True 时使用 get_proxy_ssl_context()（信任自签名证书）。
"""
cache_key = "tls" if delegate_config.get("tls_enabled") else "plain"
lock = cls._get_proxy_clients_lock()
async with lock:
existing = cls._delegate_clients.get(cache_key)
if existing and not existing.is_closed:
return existing
if cache_key == "tls":
from src.utils.ssl_utils import get_proxy_ssl_context
verify: Any = get_proxy_ssl_context()
else:
verify = get_ssl_context()
client = httpx.AsyncClient(
http2=False,
verify=verify,
follow_redirects=False,
timeout=httpx.Timeout(
connect=config.http_connect_timeout,
read=config.http_read_timeout,
write=config.http_write_timeout,
pool=config.http_pool_timeout,
),
limits=httpx.Limits(
max_connections=config.http_max_connections,
max_keepalive_connections=config.http_keepalive_connections,
keepalive_expiry=config.http_keepalive_expiry,
),
)
cls._delegate_clients[cache_key] = client
logger.debug("创建代发客户端(缓存): {}", cache_key)
return client
@classmethod
async def get_upstream_client(
cls,
@@ -469,31 +387,70 @@ class HTTPClientPool:
proxy_config: dict[str, Any] | None = None,
) -> httpx.AsyncClient:
"""
获取可复用的上游请求客户端（自动选择代发或代理模式）
获取可复用的上游请求客户端（自动选择 tunnel/代理模式）
代发模式(delegate_cfg非空)：返回代发客户端
tunnel 模式(delegate_cfg.tunnel=True)：返回 TunnelTransport 客户端
直连/代理模式：返回代理客户端（含系统默认代理回退）
"""
if delegate_cfg:
return await cls.get_delegate_client(delegate_cfg)
if delegate_cfg and delegate_cfg.get("tunnel"):
return await cls._get_tunnel_client(delegate_cfg["node_id"])
return await cls.get_proxy_client(proxy_config=proxy_config)
@classmethod
def create_upstream_stream_client(
async def create_upstream_stream_client(
cls,
delegate_cfg: dict[str, Any] | None,
proxy_config: dict[str, Any] | None = None,
timeout: httpx.Timeout | None = None,
) -> httpx.AsyncClient:
"""
创建上游流式请求客户端（自动选择代发或代理模式）
创建上游流式请求客户端（自动选择 tunnel/代理模式）
调用者需负责关闭返回的客户端。
"""
if delegate_cfg:
return cls.create_delegate_stream_client(delegate_cfg, timeout=timeout)
if delegate_cfg and delegate_cfg.get("tunnel"):
return await cls._get_tunnel_client(delegate_cfg["node_id"], timeout=timeout)
return cls.create_client_with_proxy(proxy_config=proxy_config, timeout=timeout)
@classmethod
async def _get_tunnel_client(
cls,
node_id: str,
timeout: httpx.Timeout | None = None,
) -> httpx.AsyncClient:
"""获取使用 TunnelTransport 的 httpx 客户端
当 timeout 为 None 时（非流式请求），返回按 node_id 缓存的 client，
调用方不应关闭此 client，其生命周期由 HTTPClientPool 管理。
当 timeout 非 None 时（流式请求），每次创建新 client，由调用方负责关闭。
"""
from src.services.proxy_node.tunnel_transport import TunnelTransport
t = timeout or httpx.Timeout(
connect=config.http_connect_timeout,
read=config.http_read_timeout,
write=config.http_write_timeout,
pool=config.http_pool_timeout,
)
timeout_secs = t.read if isinstance(t, httpx.Timeout) else 60.0
# 流式请求：每次创建新 client（调用方负责关闭）
if timeout is not None:
transport = TunnelTransport(node_id, timeout=timeout_secs or 60.0)
return httpx.AsyncClient(transport=transport, timeout=t)
# 非流式请求：复用缓存的 client（加锁与 proxy_clients 保持一致）
lock = cls._get_proxy_clients_lock()
async with lock:
existing = cls._tunnel_clients.get(node_id)
if existing and not existing.is_closed:
return existing
transport = TunnelTransport(node_id, timeout=timeout_secs or 60.0)
client = httpx.AsyncClient(transport=transport, timeout=t)
cls._tunnel_clients[node_id] = client
return client
@classmethod
def get_pool_stats(cls) -> dict[str, Any]:
"""获取连接池统计信息"""
@@ -502,7 +459,7 @@ class HTTPClientPool:
"named_clients_count": len(cls._clients),
"proxy_clients_count": len(cls._proxy_clients),
"max_proxy_clients": cls._max_proxy_clients,
"delegate_clients_count": len(cls._delegate_clients),
"tunnel_clients_count": len(cls._tunnel_clients),
}

23
src/config/settings.py
View File

@@ -3,8 +3,6 @@
从环境变量或 .env 文件加载配置
"""
import hashlib
import hmac
import os
from pathlib import Path
from typing import Any
@@ -48,13 +46,6 @@ class Config:
# 加密密钥配置（独立于JWT密钥，用于敏感数据加密）
self.encryption_key = os.getenv("ENCRYPTION_KEY", None)
# 代理节点 HMAC 密钥（用于 aether-proxy 认证）
proxy_hmac_key_env = os.getenv("PROXY_HMAC_KEY")
if proxy_hmac_key_env and proxy_hmac_key_env.strip():
self.proxy_hmac_key = proxy_hmac_key_env.strip()
else:
self.proxy_hmac_key = self._derive_proxy_hmac_key()
# 环境配置 - 智能检测
# Docker 部署默认为生产环境，本地开发默认为开发环境
is_docker = (
@@ -327,20 +318,6 @@ class Config:
# 验证连接池配置
self._validate_pool_config()
def _derive_proxy_hmac_key(self) -> str:
"""
从 ENCRYPTION_KEY 派生 PROXY_HMAC_KEY
目的：避免把 ENCRYPTION_KEY 直接下发到 VPS（aether-proxy）。
"""
if not self.encryption_key:
return ""
return hmac.new(
self.encryption_key.encode("utf-8"),
b"aether-proxy-hmac-key-v1",
hashlib.sha256,
).hexdigest()
def _auto_pool_size(self) -> int:
"""
智能计算连接池大小 - 根据 Worker 数量和 PostgreSQL 限制计算

50
src/core/api_format/headers.py
View File

@@ -15,10 +15,12 @@ from __future__ import annotations
from collections.abc import Set as AbstractSet
from typing import Any
from src.core.api_format.enums import ApiFamily
from src.core.api_format.metadata import (
get_auth_config_for_endpoint,
get_extra_headers_for_endpoint,
get_protected_keys_for_endpoint,
resolve_endpoint_definition,
)
from src.core.api_format.signature import EndpointSignature, parse_signature_key
from src.core.logger import logger
@@ -27,6 +29,37 @@ from src.core.logger import logger
# 头部常量定义
# =============================================================================
# 通用浏览器指纹 Headers，用于绕过 Cloudflare 等反爬防护
# 基于 Electron 桌面客户端的真实请求头构建，作为所有 adapter 请求的底层默认值
BROWSER_FINGERPRINT_HEADERS: dict[str, str] = {
"User-Agent": (
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
"AppleWebKit/537.36 (KHTML, like Gecko) "
"Chrome/140.0.7339.249 Electron/38.7.0 Safari/537.36"
),
"Accept": "application/json",
"Accept-Language": "zh-CN",
"sec-ch-ua": '"Not=A?Brand";v="24", "Chromium";v="140"',
"sec-ch-ua-mobile": "?0",
"sec-ch-ua-platform": '"macOS"',
"Sec-Fetch-Site": "cross-site",
"Sec-Fetch-Mode": "cors",
"Sec-Fetch-Dest": "empty",
}
# Anthropic/Claude 专属 Headers（仅 Claude API family 使用）
# 包含 Stainless SDK 指纹和 direct-browser-access 标记
_ANTHROPIC_EXTRA_HEADERS: dict[str, str] = {
"anthropic-dangerous-direct-browser-access": "true",
"x-stainless-os": "Unknown",
"x-stainless-runtime": "browser:chrome",
"x-stainless-arch": "unknown",
"x-stainless-lang": "js",
"x-stainless-package-version": "0.41.0",
"x-stainless-runtime-version": "140.0.7339",
"x-stainless-retry-count": "0",
}
# 转发给上游时需要剔除的头部（系统管理 + 认证替换 + 客户端/代理元数据）
UPSTREAM_DROP_HEADERS: frozenset[str] = frozenset(
{
@@ -503,14 +536,23 @@ def build_adapter_base_headers_for_endpoint(
) -> dict[str, str]:
"""
新模式：根据 endpoint signature 构建基础请求头。
浏览器指纹 headers 作为底层默认值注入，Claude API family 额外注入 Anthropic 专属 header。
认证头和 extra_headers 会覆盖它们。
"""
auth_header, auth_type = get_auth_config_for_endpoint(endpoint)
auth_value = f"Bearer {api_key}" if auth_type == "bearer" else api_key
headers: dict[str, str] = {
auth_header: auth_value,
"Content-Type": "application/json",
}
# 以浏览器指纹为底层默认值，绕过 Cloudflare 等反爬防护
headers: dict[str, str] = {**BROWSER_FINGERPRINT_HEADERS}
# Claude API family 额外注入 Anthropic 专属 header
definition = resolve_endpoint_definition(endpoint)
if definition and definition.api_family == ApiFamily.CLAUDE:
headers.update(_ANTHROPIC_EXTRA_HEADERS)
headers[auth_header] = auth_value
headers["Content-Type"] = "application/json"
if include_extra:
extra = get_extra_headers_for_endpoint(endpoint)

5
src/main.py
View File

@@ -497,6 +497,11 @@ app.include_router(dashboard_router) # 仪表盘端点
app.include_router(public_router) # 公开API端点（用户可查看提供商和模型）
app.include_router(monitoring_router) # 监控端点
# WebSocket 隧道端点（aether-proxy tunnel 模式）
from src.api.admin.proxy_tunnel import router as proxy_tunnel_router
app.include_router(proxy_tunnel_router)
def main() -> Any:
# 初始化新日志系统

15
src/models/database.py
View File

@@ -920,12 +920,6 @@ class ProxyNode(Base):
total_requests = Column(BigInteger, default=0, nullable=False)
avg_latency_ms = Column(Float, nullable=True)
# TLS 加密
tls_enabled = Column(Boolean, default=False, nullable=False, comment="是否启用 TLS 加密")
tls_cert_fingerprint = Column(
String(128), nullable=True, comment="TLS 证书 SHA-256 指纹（hex）"
)
# 硬件信息（注册时上报，JSON 可扩展）
hardware_info = Column(
JSON,
@@ -936,6 +930,15 @@ class ProxyNode(Base):
Integer, nullable=True, comment="基于硬件估算的最大并发连接数"
)
# 隧道模式（proxy 主动连接 Aether 的 WebSocket 隧道）
tunnel_mode = Column(
Boolean, default=False, nullable=False, comment="是否使用 WebSocket 隧道模式"
)
tunnel_connected = Column(Boolean, default=False, nullable=False, comment="隧道是否已连接")
tunnel_connected_at = Column(
DateTime(timezone=True), nullable=True, comment="隧道最近一次建立时间"
)
# 管理端远程配置（通过心跳下发给 aether-proxy）
remote_config = Column(
JSON,

17
src/models/provider.py
View File

@@ -270,12 +270,6 @@ class ProxyNode(Base):
total_requests = Column(BigInteger, default=0, nullable=False)
avg_latency_ms = Column(Float, nullable=True)
# TLS 加密
tls_enabled = Column(Boolean, default=False, nullable=False, comment="是否启用 TLS 加密")
tls_cert_fingerprint = Column(
String(128), nullable=True, comment="TLS 证书 SHA-256 指纹（hex）"
)
# 硬件信息（注册时上报，JSON 可扩展）
hardware_info = Column(
JSON,
@@ -286,11 +280,20 @@ class ProxyNode(Base):
Integer, nullable=True, comment="基于硬件估算的最大并发连接数"
)
# 隧道模式（proxy 主动连接 Aether 的 WebSocket 隧道）
tunnel_mode = Column(
Boolean, default=False, nullable=False, comment="是否使用 WebSocket 隧道模式"
)
tunnel_connected = Column(Boolean, default=False, nullable=False, comment="隧道是否已连接")
tunnel_connected_at = Column(
DateTime(timezone=True), nullable=True, comment="隧道最近一次建立时间"
)
# 管理端远程配置（通过心跳下发给 aether-proxy）
remote_config = Column(
JSON,
nullable=True,
comment="管理端下发的远程配置 (allowed_ports, log_level, heartbeat_interval, timestamp_tolerance)",
comment="管理端下发的远程配置 (allowed_ports, log_level, heartbeat_interval)",
)
config_version = Column(
Integer, default=0, nullable=False, comment="远程配置版本号，每次更新 +1"

14
src/modules/proxy_nodes/__init__.py
View File

@@ -2,7 +2,7 @@
代理节点模块
提供海外 VPS 代理节点的注册、心跳、管理功能。
aether-proxy 部署在海外 VPS 上自动注册节点，Aether 通过 HMAC 签名认证转发请求。
aether-proxy 部署在海外 VPS 上，通过 WebSocket 隧道连接 Aether 转发 API 请求。
"""
from __future__ import annotations
@@ -77,15 +77,7 @@ async def _health_check() -> ModuleHealth:
def _validate_config(db: Session) -> tuple[bool, str]:
"""
验证配置
代理节点模块需要 PROXY_HMAC_KEY 配置
"""
from src.config.settings import config
if not config.proxy_hmac_key:
return False, "PROXY_HMAC_KEY 未配置（也未设置 ENCRYPTION_KEY 用于自动派生）"
"""验证配置（tunnel 模式无需额外密钥配置）"""
return True, ""
@@ -93,7 +85,7 @@ proxy_nodes_module = ModuleDefinition(
metadata=ModuleMetadata(
name="proxy_nodes",
display_name="代理节点",
description="海外 VPS 代理节点管理，通过 HMAC 签名认证转发 API 请求",
description="海外 VPS 代理节点管理，通过 WebSocket 隧道转发 API 请求",
category=ModuleCategory.INTEGRATION,
env_key="PROXY_NODES_AVAILABLE",
default_available=True,

12
src/services/provider_ops/architectures/base.py
View File

@@ -51,9 +51,10 @@ class ProviderConnector(ABC):
self._last_error: str | None = None
# 代理配置（支持 proxy_node_id 和旧的 proxy URL）
from src.services.proxy_node.resolver import resolve_ops_proxy
from src.services.proxy_node.resolver import resolve_ops_proxy, resolve_ops_tunnel_node_id
self._proxy: str | httpx.Proxy | None = resolve_ops_proxy(self.config)
self._tunnel_node_id: str | None = resolve_ops_tunnel_node_id(self.config)
# HTTP 客户端配置
self._timeout = self.config.get("timeout", 30)
@@ -117,13 +118,18 @@ class ProviderConnector(ABC):
"""
获取已认证的 HTTP 客户端
使用 context manager 确保资源正确释放
使用 context manager 确保资源正确释放。
tunnel 模式下使用 TunnelTransport 替代 proxy transport。
Yields:
已配置认证信息的 AsyncClient
"""
transport = None
if self._proxy:
if self._tunnel_node_id:
from src.services.proxy_node.tunnel_transport import TunnelTransport
transport = TunnelTransport(self._tunnel_node_id, timeout=self._timeout)
elif self._proxy:
transport = httpx.AsyncHTTPTransport(proxy=self._proxy)
async with httpx.AsyncClient(

8
src/services/provider_ops/architectures/new_api.py
View File

@@ -8,6 +8,7 @@ from typing import Any
import httpx
from src.core.api_format.headers import BROWSER_FINGERPRINT_HEADERS
from src.services.provider_ops.actions import (
NewApiBalanceAction,
ProviderAction,
@@ -75,6 +76,10 @@ class NewApiConnector(ProviderConnector):
def _apply_auth(self, request: httpx.Request) -> httpx.Request:
"""为请求应用认证信息"""
# 添加浏览器指纹 Headers 以绕过 Cloudflare 等防护
for key, value in BROWSER_FINGERPRINT_HEADERS.items():
request.headers.setdefault(key, value)
if self._api_key:
request.headers["Authorization"] = f"Bearer {self._api_key}"
if self._user_id:
@@ -206,7 +211,8 @@ class NewApiArchitecture(ProviderArchitecture):
New API 特有：需要 New-Api-User Header 传递用户 ID
"""
headers: dict[str, str] = {}
# 以浏览器指纹 Headers 为基础，绕过 Cloudflare 等防护
headers: dict[str, str] = {**BROWSER_FINGERPRINT_HEADERS}
# Bearer Token 认证
api_key = credentials.get("api_key", "")

12
src/services/provider_ops/service.py
View File

@@ -1033,10 +1033,11 @@ class ProviderOpsService:
list(headers.keys()),
)
# 获取代理配置（支持 proxy_node_id 和旧的 proxy URL）
from src.services.proxy_node.resolver import resolve_ops_proxy
# 获取代理配置（支持 proxy_node_id、tunnel 模式和旧的 proxy URL）
from src.services.proxy_node.resolver import resolve_ops_proxy, resolve_ops_tunnel_node_id
proxy = resolve_ops_proxy(config)
tunnel_node_id = resolve_ops_tunnel_node_id(config)
try:
# 构建 httpx client 参数
@@ -1044,7 +1045,12 @@ class ProviderOpsService:
"timeout": 30.0,
"verify": get_ssl_context(),
}
if proxy:
if tunnel_node_id:
from src.services.proxy_node.tunnel_transport import TunnelTransport
client_kwargs["transport"] = TunnelTransport(tunnel_node_id, timeout=30.0)
logger.debug("使用 tunnel 代理: node_id={}", tunnel_node_id)
elif proxy:
client_kwargs["proxy"] = proxy
logger.debug("使用代理: {}", proxy)

10
src/services/proxy_node/__init__.py
View File

@@ -2,9 +2,6 @@
from .health_scheduler import ProxyNodeHealthScheduler, get_proxy_node_health_scheduler
from .resolver import (
build_delegate_post_kwargs,
build_delegate_stream_kwargs,
build_hmac_proxy_url,
build_post_kwargs,
build_proxy_url,
build_stream_kwargs,
@@ -12,10 +9,12 @@ from .resolver import (
get_proxy_label,
get_system_proxy_config,
inject_auth_into_proxy_url,
invalidate_proxy_node_cache,
invalidate_system_proxy_cache,
make_proxy_param,
resolve_delegate_config,
resolve_ops_proxy,
resolve_ops_tunnel_node_id,
resolve_proxy_info,
)
from .service import ProxyNodeService, node_to_dict
@@ -25,9 +24,6 @@ __all__ = [
"get_proxy_node_health_scheduler",
"ProxyNodeService",
"node_to_dict",
"build_delegate_post_kwargs",
"build_delegate_stream_kwargs",
"build_hmac_proxy_url",
"build_post_kwargs",
"build_proxy_url",
"build_stream_kwargs",
@@ -36,8 +32,10 @@ __all__ = [
"make_proxy_param",
"get_proxy_label",
"get_system_proxy_config",
"invalidate_proxy_node_cache",
"invalidate_system_proxy_cache",
"resolve_delegate_config",
"resolve_ops_proxy",
"resolve_ops_tunnel_node_id",
"resolve_proxy_info",
]

29
src/services/proxy_node/health_scheduler.py
View File

@@ -1,9 +1,10 @@
"""
ProxyNode 心跳检测调度器
定期检查 proxy_nodes 的 last_heartbeat_at，更新节点状态：
- elapsed > interval * 3 -> unhealthy
- elapsed > interval * 10 -> offline
定期检查 proxy_nodes 的 tunnel 连接状态，更新节点状态：
- tunnel_connected=True -> ONLINE
- tunnel 刚断开 (<60s) -> UNHEALTHY（缓冲期，避免正在进行的请求被立即切走）
- tunnel 断开超过 60s -> OFFLINE
"""
from __future__ import annotations
@@ -56,6 +57,7 @@ class ProxyNodeHealthScheduler:
try:
now = datetime.now(timezone.utc)
# 仅检查非手动节点（手动节点无心跳，始终保持 ONLINE）
# 非手动节点均为 tunnel 模式，由 tunnel 连接状态决定
nodes = (
db.query(ProxyNode)
.filter(
@@ -69,19 +71,16 @@ class ProxyNodeHealthScheduler:
changed = 0
for node in nodes:
interval = int(node.heartbeat_interval or 30)
last = node.last_heartbeat_at
if last is None:
new_status = ProxyNodeStatus.OFFLINE
else:
elapsed = (now - last).total_seconds()
if elapsed > interval * 10:
new_status = ProxyNodeStatus.OFFLINE
elif elapsed > interval * 3:
new_status = ProxyNodeStatus.UNHEALTHY
else:
if node.tunnel_connected:
new_status = ProxyNodeStatus.ONLINE
elif node.tunnel_connected_at:
# tunnel 刚断开：给 60s 缓冲期标记为 UNHEALTHY
elapsed = (now - node.tunnel_connected_at).total_seconds()
new_status = (
ProxyNodeStatus.UNHEALTHY if elapsed < 60 else ProxyNodeStatus.OFFLINE
)
else:
new_status = ProxyNodeStatus.OFFLINE
if node.status != new_status:
node.status = new_status

335
src/services/proxy_node/resolver.py
View File

@@ -7,18 +7,13 @@
from __future__ import annotations
import base64
import gzip as _gzip
import hashlib
import hmac as _hmac
import json as _json
import time
from typing import Any
from urllib.parse import quote, urlparse
import httpx
from src.config import config
from src.core.exceptions import ProxyNodeUnavailableError
from src.core.logger import logger
@@ -81,8 +76,8 @@ def _get_proxy_node_info(node_id: str) -> dict[str, Any] | None:
"name": node.name,
"ip": node.ip,
"port": node.port,
"tls_enabled": bool(node.tls_enabled),
"tls_cert_fingerprint": node.tls_cert_fingerprint,
"tunnel_mode": bool(node.tunnel_mode),
"tunnel_connected": bool(node.tunnel_connected),
}
_proxy_node_cache[node_id] = (value, now + _PROXY_NODE_CACHE_TTL_SECONDS)
@@ -91,40 +86,6 @@ def _get_proxy_node_info(node_id: str) -> dict[str, Any] | None:
db.close()
# ---------------------------------------------------------------------------
# HMAC 签名
# ---------------------------------------------------------------------------
def build_hmac_proxy_url(ip: str, port: int, *, tls_enabled: bool = False) -> str:
"""
构建带 HMAC BasicAuth 的 httpx proxy URL
格式: http(s)://hmac:{timestamp}.{signature}@{ip}:{port}
signature = HMAC-SHA256(PROXY_HMAC_KEY, "{timestamp}") 的 hex
签名不再包含 node_id，避免 proxy 重新注册后 Aether 端缓存的旧 node_id
与 proxy 端新 node_id 不一致导致的认证失败窗口。
当 tls_enabled=True 时使用 https:// scheme。
"""
if not config.proxy_hmac_key:
logger.error("PROXY_HMAC_KEY 未配置，无法使用 ProxyNode 代理")
raise ProxyNodeUnavailableError("PROXY_HMAC_KEY 未配置，无法使用 ProxyNode 代理")
timestamp = str(int(time.time()))
payload = timestamp.encode("utf-8")
signature = _hmac.new(
config.proxy_hmac_key.encode("utf-8"),
payload,
hashlib.sha256,
).hexdigest()
host = f"[{ip}]" if ":" in ip else ip
scheme = "https" if tls_enabled else "http"
return f"{scheme}://hmac:{timestamp}.{signature}@{host}:{int(port)}"
# ---------------------------------------------------------------------------
# 系统默认代理
# ---------------------------------------------------------------------------
@@ -132,6 +93,11 @@ _system_proxy_cache: tuple[dict[str, Any] | None, float] | None = None
_SYSTEM_PROXY_CACHE_TTL = 60.0
def invalidate_proxy_node_cache(node_id: str) -> None:
"""主动清除指定节点的信息缓存（tunnel 断开时调用，避免使用过期的连接状态）"""
_proxy_node_cache.pop(node_id, None)
def invalidate_system_proxy_cache() -> None:
"""手动失效系统代理缓存（在删除节点等操作后调用）"""
global _system_proxy_cache
@@ -224,6 +190,32 @@ def make_proxy_param(proxy_url: str | None) -> str | httpx.Proxy | None:
# ---------------------------------------------------------------------------
def _resolve_effective_node(
connector_config: dict[str, Any] | None,
) -> tuple[str | None, dict[str, Any] | None]:
"""
从 connector_config 或系统默认代理中解析有效的 proxy_node_id 及其信息。
Returns:
(node_id, node_info) 或 (None, None)
"""
if connector_config:
node_id = connector_config.get("proxy_node_id")
if isinstance(node_id, str) and node_id.strip():
nid = node_id.strip()
return nid, _get_proxy_node_info(nid)
# 回退：系统默认代理
system_proxy = get_system_proxy_config()
if system_proxy:
node_id_sys = system_proxy.get("node_id")
if isinstance(node_id_sys, str) and node_id_sys.strip():
nid = node_id_sys.strip()
return nid, _get_proxy_node_info(nid)
return None, None
def resolve_ops_proxy(
connector_config: dict[str, Any] | None,
) -> str | httpx.Proxy | None:
@@ -235,38 +227,51 @@ def resolve_ops_proxy(
2. connector_config.proxy（旧格式 URL 字符串）
3. 系统默认代理节点
tunnel 模式节点不返回代理 URL（由 resolve_ops_tunnel_node_id 处理）。
Args:
connector_config: connector 的 config 字典
Returns:
httpx 可接受的代理参数（str 或 httpx.Proxy），或 None
"""
if connector_config:
# 新格式：proxy_node_id -> 通过 build_proxy_url 解析
node_id = connector_config.get("proxy_node_id")
if isinstance(node_id, str) and node_id.strip():
from .tunnel_transport import is_tunnel_node
node_id, node_info = _resolve_effective_node(connector_config)
if node_id and node_info:
if is_tunnel_node(node_info):
return None # tunnel 模式不使用 proxy URL
try:
url = build_proxy_url({"node_id": node_id.strip(), "enabled": True})
url = build_proxy_url({"node_id": node_id, "enabled": True})
return make_proxy_param(url)
except Exception as exc:
logger.warning("解析 proxy_node_id={} 失败，回退到直连: {}", node_id, exc)
return None
# 旧格式：直接返回 proxy URL 字符串
if connector_config:
proxy = connector_config.get("proxy")
if isinstance(proxy, str) and proxy.strip():
return proxy
# 回退：系统默认代理
system_proxy = get_system_proxy_config()
if system_proxy:
try:
url = build_proxy_url(system_proxy)
return make_proxy_param(url)
except Exception as exc:
logger.warning("构建系统默认代理 URL 失败: {}", exc)
return None
def resolve_ops_tunnel_node_id(
connector_config: dict[str, Any] | None,
) -> str | None:
"""
解析 ops connector 的 tunnel 节点 ID
如果配置的代理节点是 tunnel 模式且已连接，返回 node_id。
否则返回 None（含系统默认代理回退）。
"""
from .tunnel_transport import is_tunnel_node
node_id, node_info = _resolve_effective_node(connector_config)
if node_id and node_info and is_tunnel_node(node_info):
return node_id
return None
@@ -401,12 +406,15 @@ def build_proxy_url(proxy_config: dict[str, Any]) -> str | None:
return inject_auth_into_proxy_url(manual_url, username, password)
return manual_url
# aether-proxy 节点：使用 HMAC 认证
return build_hmac_proxy_url(
node_info["ip"],
node_info["port"],
tls_enabled=node_info.get("tls_enabled", False),
)
# tunnel 模式节点：不构建 proxy URL（通过 TunnelTransport 处理）
from .tunnel_transport import is_tunnel_node
if is_tunnel_node(node_info):
return None
# aether-proxy 节点均为 tunnel 模式，不应走到这里
logger.warning("非 tunnel 模式的 aether-proxy 节点不再支持: node_id={}", node_id)
return None
proxy_url: str | None = proxy_config.get("url")
if not proxy_url:
@@ -504,11 +512,10 @@ def compute_proxy_cache_key(proxy_config: dict[str, Any] | None) -> str:
if not proxy_config.get("enabled", True):
return "__no_proxy__"
# ProxyNode 模式：基于 node_id + 时间桶缓存，避免签名随时间变化导致 cache key 爆炸
# ProxyNode 模式：基于 node_id 缓存
node_id = proxy_config.get("node_id")
if isinstance(node_id, str) and node_id.strip():
time_bucket = int(time.time() / 240) # 240s bucket, within 300s HMAC tolerance
return f"proxy_node:{node_id.strip()}:{time_bucket}"
return f"proxy_node:{node_id.strip()}"
# 构建代理 URL 作为缓存键的基础
proxy_url = build_proxy_url(proxy_config)
@@ -520,46 +527,20 @@ def compute_proxy_cache_key(proxy_config: dict[str, Any] | None) -> str:
# ---------------------------------------------------------------------------
# 代发模式 (Delegate API)
# Tunnel 代理配置解析
# ---------------------------------------------------------------------------
def _build_hmac_auth_header() -> str:
"""
构建代发请求的 Authorization 头
格式: Basic base64(hmac:{timestamp}.{signature})
签名算法与 build_hmac_proxy_url 相同（仅使用 timestamp，不含 node_id）。
"""
if not config.proxy_hmac_key:
raise ProxyNodeUnavailableError("PROXY_HMAC_KEY 未配置，无法使用代发模式")
timestamp = str(int(time.time()))
payload = timestamp.encode("utf-8")
signature = _hmac.new(
config.proxy_hmac_key.encode("utf-8"),
payload,
hashlib.sha256,
).hexdigest()
cred = f"hmac:{timestamp}.{signature}"
encoded = base64.b64encode(cred.encode()).decode()
return f"Basic {encoded}"
def resolve_delegate_config(proxy_config: dict[str, Any] | None) -> dict[str, Any] | None:
"""
解析代发配置（仅 aether-proxy 节点支持，手动节点/旧格式 URL 不支持）
解析 tunnel 代理配置（仅 aether-proxy tunnel 节点支持）
无特定代理时自动回退到系统默认代理。
auth_header 延迟生成：通过 ``fresh_auth_header()`` 闭包在每次请求 / 重试时
获取新鲜的 HMAC 签名，避免长生命周期内时间戳过期。
tunnel 模式节点返回 {"tunnel": True, "node_id": str}，
调用方应使用 TunnelTransport。
Returns:
{"delegate_url": str, "node_id": str, "tls_enabled": bool,
"auth_header": str, # 首次生成的签名（兼容旧调用）
"fresh_auth_header": Callable} # 延迟生成签名的闭包
或 None
{"tunnel": True, "node_id": str} 或 None
"""
effective_config = proxy_config
@@ -571,148 +552,28 @@ def resolve_delegate_config(proxy_config: dict[str, Any] | None) -> dict[str, An
node_id = effective_config.get("node_id")
if not isinstance(node_id, str) or not node_id.strip():
return None # 旧格式 URL 模式不支持代发
return None
node_id = node_id.strip()
node_info = _get_proxy_node_info(node_id)
if not node_info or node_info.get("is_manual"):
return None # 手动节点不支持代发
return None
tls_enabled = node_info.get("tls_enabled", False)
host = f"[{node_info['ip']}]" if ":" in node_info["ip"] else node_info["ip"]
scheme = "https" if tls_enabled else "http"
delegate_url = f"{scheme}://{host}:{int(node_info['port'])}/_aether/delegate"
from .tunnel_transport import is_tunnel_node
# 每次调用生成新鲜签名（避免长连接内时间戳过期）
def _fresh() -> str:
return _build_hmac_auth_header()
if is_tunnel_node(node_info):
return {"tunnel": True, "node_id": node_id}
return {
"delegate_url": delegate_url,
"auth_header": _fresh(), # 立即生成一份，兼容旧调用方
"fresh_auth_header": _fresh,
"node_id": node_id,
"tls_enabled": tls_enabled,
}
return None
# ---------------------------------------------------------------------------
# 代发请求参数构建（消除 handler 层重复代码）
# ---------------------------------------------------------------------------
_JSON_CT = "application/json"
def _build_delegate_kwargs_core(
delegate_cfg: dict[str, Any],
*,
url: str,
headers: dict[str, str],
payload: Any,
timeout: float,
refresh_auth: bool = False,
) -> dict[str, Any]:
"""
构建代发请求的核心参数（post/stream 共用）
元数据通过 HTTP headers 传递（X-Delegate-Method/Url/Headers），
上游请求体 gzip 压缩后直接作为 HTTP body 发送，大幅减少跨国传输耗时。
Args:
delegate_cfg: resolve_delegate_config 返回的配置
url: 上游实际 URL
headers: 上游请求头
payload: 上游 JSON body（可以为 None）
timeout: 上游超时秒数
refresh_auth: 为 True 时重新生成 HMAC 签名（用于 retry）
"""
auth = (
delegate_cfg["fresh_auth_header"]()
if refresh_auth
else delegate_cfg.get("auth_header") or delegate_cfg["fresh_auth_header"]()
)
# 上游 headers base64 编码
headers_b64 = base64.b64encode(_json.dumps(headers, ensure_ascii=False).encode("utf-8")).decode(
"ascii"
)
# 构建代发请求 headers（元数据）
delegate_headers: dict[str, str] = {
"Authorization": auth,
"X-Delegate-Method": "POST",
"X-Delegate-Url": url,
"X-Delegate-Headers": headers_b64,
"X-Delegate-Timeout": str(int(timeout)),
}
kwargs: dict[str, Any] = {
"url": delegate_cfg["delegate_url"],
"headers": delegate_headers,
"timeout": httpx.Timeout(timeout + 10),
}
# body gzip 压缩后直接作为 HTTP content
if payload is not None:
body_bytes = _json.dumps(payload, ensure_ascii=False).encode("utf-8")
compressed = _gzip.compress(body_bytes)
kwargs["content"] = compressed
kwargs["headers"]["Content-Encoding"] = "gzip"
kwargs["headers"]["Content-Type"] = _JSON_CT
return kwargs
def build_delegate_post_kwargs(
delegate_cfg: dict[str, Any],
*,
url: str,
headers: dict[str, str],
payload: Any,
timeout: float,
refresh_auth: bool = False,
) -> dict[str, Any]:
"""构建代发 POST 请求的 httpx kwargs（非流式，传给 client.post）"""
return _build_delegate_kwargs_core(
delegate_cfg,
url=url,
headers=headers,
payload=payload,
timeout=timeout,
refresh_auth=refresh_auth,
)
def build_delegate_stream_kwargs(
delegate_cfg: dict[str, Any],
*,
url: str,
headers: dict[str, str],
payload: Any,
timeout: float,
refresh_auth: bool = False,
) -> dict[str, Any]:
"""构建代发 stream 请求的 httpx kwargs（传给 client.stream）"""
kwargs = _build_delegate_kwargs_core(
delegate_cfg,
url=url,
headers=headers,
payload=payload,
timeout=timeout,
refresh_auth=refresh_auth,
)
# stream() 需要显式 method 参数
kwargs["method"] = "POST"
return kwargs
# ---------------------------------------------------------------------------
# 统一上游请求参数构建（消除 handler 层 delegate/直连 分支重复）
# 统一上游请求参数构建
# ---------------------------------------------------------------------------
def build_post_kwargs(
delegate_cfg: dict[str, Any] | None,
_delegate_cfg: dict[str, Any] | None = None,
*,
url: str,
headers: dict[str, str],
@@ -721,19 +582,13 @@ def build_post_kwargs(
refresh_auth: bool = False,
) -> dict[str, Any]:
"""
构建上游 POST 请求的 httpx kwargs（自动选择代发或直连模式）
构建上游 POST 请求的 httpx kwargs
返回的 dict 可直接传给 ``http_client.post(**kwargs)``。
``_delegate_cfg`` 和 ``refresh_auth`` 已废弃（tunnel 模式下认证由 transport 层处理），
保留仅为兼容现有调用方签名。
"""
if delegate_cfg:
return build_delegate_post_kwargs(
delegate_cfg,
url=url,
headers=headers,
payload=payload,
timeout=timeout,
refresh_auth=refresh_auth,
)
return {
"url": url,
"json": payload,
@@ -743,7 +598,7 @@ def build_post_kwargs(
def build_stream_kwargs(
delegate_cfg: dict[str, Any] | None,
_delegate_cfg: dict[str, Any] | None = None,
*,
url: str,
headers: dict[str, str],
@@ -751,21 +606,13 @@ def build_stream_kwargs(
timeout: float | None = None,
) -> dict[str, Any]:
"""
构建上游 stream 请求的 httpx kwargs（自动选择代发或直连模式）
构建上游 stream 请求的 httpx kwargs
返回的 dict 可直接传给 ``http_client.stream(**kwargs)``。
当 ``timeout`` 为 None 时由外层 asyncio.wait_for 控制超时。
当 ``timeout`` 为 None（直连模式下由外层 asyncio.wait_for 控制超时），
直连分支不设置 timeout；代发分支始终携带 timeout（proxy 协议需要）。
``_delegate_cfg`` 已废弃，保留仅为兼容现有调用方签名。
"""
if delegate_cfg:
return build_delegate_stream_kwargs(
delegate_cfg,
url=url,
headers=headers,
payload=payload,
timeout=timeout or 60,
)
kwargs: dict[str, Any] = {
"method": "POST",
"url": url,

51
src/services/proxy_node/service.py
View File

@@ -17,10 +17,9 @@ import httpx
from sqlalchemy.orm import Session
from src.core.exceptions import InvalidRequestException, NotFoundException
from src.models.database import ProxyNode, ProxyNodeStatus, SystemConfig
from src.models.database import Provider, ProviderEndpoint, ProxyNode, ProxyNodeStatus, SystemConfig
from .resolver import (
build_hmac_proxy_url,
inject_auth_into_proxy_url,
invalidate_system_proxy_cache,
make_proxy_param,
@@ -50,14 +49,15 @@ def node_to_dict(node: ProxyNode) -> dict[str, Any]:
"region": node.region,
"status": node.status.value if node.status else None,
"is_manual": bool(node.is_manual),
"tunnel_mode": bool(node.tunnel_mode),
"tunnel_connected": bool(node.tunnel_connected),
"tunnel_connected_at": node.tunnel_connected_at,
"registered_by": node.registered_by,
"last_heartbeat_at": node.last_heartbeat_at,
"heartbeat_interval": node.heartbeat_interval,
"active_connections": node.active_connections,
"total_requests": node.total_requests,
"avg_latency_ms": node.avg_latency_ms,
"tls_enabled": bool(node.tls_enabled),
"tls_cert_fingerprint": node.tls_cert_fingerprint,
"hardware_info": node.hardware_info,
"estimated_max_concurrency": node.estimated_max_concurrency,
"remote_config": node.remote_config,
@@ -166,8 +166,8 @@ def _build_test_proxy_url(node: ProxyNode) -> str:
)
return proxy_url
else:
# aether-proxy: 使用 HMAC 认证构建代理 URL
return build_hmac_proxy_url(node.ip, node.port, tls_enabled=bool(node.tls_enabled))
# aether-proxy 节点均为 tunnel 模式，不支持通过代理 URL 测试
raise InvalidRequestException("aether-proxy tunnel 节点不支持代理 URL 连通性测试")
# ---------------------------------------------------------------------------
@@ -187,18 +187,28 @@ class ProxyNodeService:
port: int,
region: str | None = None,
heartbeat_interval: int = 30,
tls_enabled: bool = False,
tls_cert_fingerprint: str | None = None,
hardware_info: dict[str, Any] | None = None,
estimated_max_concurrency: int | None = None,
active_connections: int | None = None,
total_requests: int | None = None,
avg_latency_ms: float | None = None,
registered_by: str | None = None,
tunnel_mode: bool = False,
) -> ProxyNode:
"""注册或更新 aether-proxy 节点（按 ip+port upsert）"""
"""注册或更新 aether-proxy 节点
tunnel 模式按 name upsert（port 固定为 0，同 IP 可能有多个实例）；
旧模式按 ip+port upsert（向后兼容）。
"""
now = datetime.now(timezone.utc)
if tunnel_mode:
node = (
db.query(ProxyNode)
.filter(ProxyNode.name == name, ProxyNode.is_manual == False) # noqa: E712
.first()
)
else:
node = db.query(ProxyNode).filter(ProxyNode.ip == ip, ProxyNode.port == port).first()
if node:
node.name = name
@@ -206,8 +216,7 @@ class ProxyNodeService:
node.status = ProxyNodeStatus.ONLINE
node.last_heartbeat_at = now
node.heartbeat_interval = heartbeat_interval
node.tls_enabled = tls_enabled
node.tls_cert_fingerprint = tls_cert_fingerprint
node.tunnel_mode = tunnel_mode
if hardware_info is not None:
node.hardware_info = hardware_info
if estimated_max_concurrency is not None:
@@ -232,10 +241,9 @@ class ProxyNodeService:
active_connections=active_connections or 0,
total_requests=total_requests or 0,
avg_latency_ms=avg_latency_ms,
tls_enabled=tls_enabled,
tls_cert_fingerprint=tls_cert_fingerprint,
hardware_info=hardware_info,
estimated_max_concurrency=estimated_max_concurrency,
tunnel_mode=tunnel_mode,
created_at=now,
updated_at=now,
)
@@ -424,6 +432,21 @@ class ProxyNodeService:
sys_cfg.value = None
was_system_proxy = True
# 清理引用该节点的 Provider / ProviderEndpoint 的 proxy 字段（批量 SQL 更新）
cleared_providers = (
db.query(Provider)
.filter(Provider.proxy.isnot(None), Provider.proxy["node_id"].as_string() == node_id)
.update({"proxy": None}, synchronize_session="fetch")
)
cleared_endpoints = (
db.query(ProviderEndpoint)
.filter(
ProviderEndpoint.proxy.isnot(None),
ProviderEndpoint.proxy["node_id"].as_string() == node_id,
)
.update({"proxy": None}, synchronize_session="fetch")
)
node_info = {"proxy_node_ip": node.ip, "proxy_node_port": node.port}
db.delete(node)
db.commit()
@@ -435,6 +458,8 @@ class ProxyNodeService:
"node_id": node_id,
"node_info": node_info,
"cleared_system_proxy": was_system_proxy,
"cleared_providers": cleared_providers,
"cleared_endpoints": cleared_endpoints,
}
@staticmethod

334
src/services/proxy_node/tunnel_manager.py Normal file
View File

@@ -0,0 +1,334 @@
"""
WebSocket 隧道管理器
管理所有活跃的 aether-proxy tunnel 连接，提供通过隧道发送 HTTP 请求的能力。
每个 proxy node 最多一条 tunnel 连接。
"""
from __future__ import annotations
import asyncio
import json
import time
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import AsyncGenerator
from starlette.websockets import WebSocket, WebSocketState
from src.core.logger import logger
from .tunnel_protocol import Frame, FrameFlags, MsgType
class TunnelConnection:
"""单条 tunnel 连接"""
__slots__ = (
"node_id",
"node_name",
"ws",
"connected_at",
"_pending_streams",
"_write_lock",
"_next_stream_id",
)
def __init__(self, node_id: str, node_name: str, ws: WebSocket) -> None:
self.node_id = node_id
self.node_name = node_name
self.ws = ws
self.connected_at = time.time()
self._pending_streams: dict[int, _StreamState] = {}
self._write_lock = asyncio.Lock()
# Per-connection stream ID 分配器（Aether 端使用偶数，从 2 开始）
self._next_stream_id: int = 2
@property
def is_alive(self) -> bool:
return self.ws.client_state == WebSocketState.CONNECTED
async def send_frame(self, frame: Frame) -> None:
async with self._write_lock:
await self.ws.send_bytes(frame.encode())
def create_stream(self, stream_id: int) -> _StreamState:
state = _StreamState(stream_id)
self._pending_streams[stream_id] = state
return state
def get_stream(self, stream_id: int) -> _StreamState | None:
return self._pending_streams.get(stream_id)
def remove_stream(self, stream_id: int) -> None:
self._pending_streams.pop(stream_id, None)
@property
def stream_count(self) -> int:
return len(self._pending_streams)
def has_stream(self, stream_id: int) -> bool:
return stream_id in self._pending_streams
def alloc_stream_id(self, max_streams: int) -> int:
"""分配一个未被占用的偶数 stream_id，回绕时跳过飞行中的 ID"""
# 最多尝试 max_streams + 16 次（飞行中的 stream 数量不超过 max_streams）
for _ in range(max_streams + 16):
sid = self._next_stream_id
self._next_stream_id += 2
if self._next_stream_id > 0xFFFF_FFFE:
self._next_stream_id = 2
if sid not in self._pending_streams:
return sid
raise TunnelStreamError("stream ID space exhausted")
def cancel_all_streams(self) -> None:
for state in self._pending_streams.values():
state.set_error("tunnel disconnected")
self._pending_streams.clear()
class _StreamState:
"""跟踪单个 stream 的响应状态"""
__slots__ = (
"stream_id",
"status",
"headers",
"_header_event",
"_body_chunks",
"_done_event",
"_error",
)
def __init__(self, stream_id: int) -> None:
self.stream_id = stream_id
self.status: int = 0
self.headers: list[list[str]] = []
self._header_event = asyncio.Event()
self._body_chunks: asyncio.Queue[bytes | None] = asyncio.Queue()
self._done_event = asyncio.Event()
self._error: str | None = None
def set_response_headers(self, status: int, headers: list[list[str]] | dict[str, str]) -> None:
self.status = status
# headers 可能是 [[k, v], ...] (多值) 或 {k: v} (旧格式兼容)
if isinstance(headers, list):
self.headers = headers # type: ignore[assignment]
else:
self.headers = list(headers.items()) # type: ignore[assignment]
self._header_event.set()
def push_body_chunk(self, data: bytes) -> None:
self._body_chunks.put_nowait(data)
def set_done(self) -> None:
self._body_chunks.put_nowait(None) # sentinel
self._done_event.set()
def set_error(self, msg: str) -> None:
self._error = msg
self._header_event.set()
self._body_chunks.put_nowait(None)
self._done_event.set()
async def wait_headers(self, timeout: float = 60.0) -> None:
await asyncio.wait_for(self._header_event.wait(), timeout=timeout)
if self._error:
raise TunnelStreamError(self._error)
async def iter_body(self, chunk_timeout: float = 60.0) -> AsyncGenerator[bytes, None]:
while True:
try:
chunk = await asyncio.wait_for(self._body_chunks.get(), timeout=chunk_timeout)
except asyncio.TimeoutError:
self._error = "body chunk timeout"
self._done_event.set()
raise TunnelStreamError("body chunk timeout")
if chunk is None:
if self._error:
raise TunnelStreamError(self._error)
return
yield chunk
class TunnelStreamError(Exception):
pass
# ---------------------------------------------------------------------------
# 全局 TunnelManager 单例
# ---------------------------------------------------------------------------
class TunnelManager:
"""管理所有活跃的 tunnel 连接"""
# 单条 tunnel 上允许的最大并发 stream 数（超出时拒绝新请求）
MAX_STREAMS_PER_CONN = 2048
def __init__(self) -> None:
self._connections: dict[str, TunnelConnection] = {} # node_id -> conn
@property
def active_count(self) -> int:
return len(self._connections)
def get_connection(self, node_id: str) -> TunnelConnection | None:
conn = self._connections.get(node_id)
if conn and not conn.is_alive:
self._connections.pop(node_id, None)
conn.cancel_all_streams()
return None
return conn
def register(self, conn: TunnelConnection) -> None:
old = self._connections.get(conn.node_id)
if old:
old.cancel_all_streams()
self._connections[conn.node_id] = conn
logger.info("tunnel connected: node_id={}, name={}", conn.node_id, conn.node_name)
def unregister(self, node_id: str) -> None:
conn = self._connections.pop(node_id, None)
if conn:
conn.cancel_all_streams()
logger.info("tunnel disconnected: node_id={}, name={}", node_id, conn.node_name)
def has_tunnel(self, node_id: str) -> bool:
conn = self.get_connection(node_id)
return conn is not None
async def send_request(
self,
node_id: str,
*,
method: str,
url: str,
headers: dict[str, str],
body: bytes | None = None,
timeout: float = 60.0,
) -> _StreamState:
"""
通过 tunnel 发送 HTTP 请求，返回 StreamState 用于读取响应。
"""
conn = self.get_connection(node_id)
if not conn:
raise TunnelStreamError(f"tunnel not connected for node {node_id}")
if conn.stream_count >= self.MAX_STREAMS_PER_CONN:
raise TunnelStreamError(
f"tunnel stream limit reached ({self.MAX_STREAMS_PER_CONN}) for node {node_id}"
)
stream_id = conn.alloc_stream_id(self.MAX_STREAMS_PER_CONN)
stream_state = conn.create_stream(stream_id)
try:
# 发送 REQUEST_HEADERS
meta = json.dumps(
{
"method": method,
"url": url,
"headers": headers,
"timeout": int(timeout),
}
).encode()
await conn.send_frame(Frame(stream_id, MsgType.REQUEST_HEADERS, 0, meta))
# 发送 REQUEST_BODY + END_STREAM
body_data = body or b""
await conn.send_frame(
Frame(stream_id, MsgType.REQUEST_BODY, FrameFlags.END_STREAM, body_data)
)
except Exception:
conn.remove_stream(stream_id)
raise
return stream_state
async def handle_incoming_frame(self, node_id: str, frame: Frame) -> None:
"""处理从 proxy 收到的响应帧"""
conn = self.get_connection(node_id)
if not conn:
return
stream = conn.get_stream(frame.stream_id)
if frame.msg_type == MsgType.RESPONSE_HEADERS:
if not stream:
return
try:
meta = json.loads(frame.payload)
stream.set_response_headers(meta["status"], meta.get("headers", []))
except Exception as e:
stream.set_error(f"invalid response headers: {e}")
elif frame.msg_type == MsgType.RESPONSE_BODY:
if stream:
stream.push_body_chunk(frame.payload)
elif frame.msg_type == MsgType.STREAM_END:
if stream:
stream.set_done()
conn.remove_stream(frame.stream_id)
elif frame.msg_type == MsgType.STREAM_ERROR:
if stream:
msg = frame.payload.decode(errors="replace") if frame.payload else "stream error"
stream.set_error(msg)
conn.remove_stream(frame.stream_id)
elif frame.msg_type == MsgType.HEARTBEAT_DATA:
await self._handle_heartbeat(conn, frame)
elif frame.msg_type == MsgType.PING:
await conn.send_frame(Frame(0, MsgType.PONG, 0, frame.payload))
async def _handle_heartbeat(self, conn: TunnelConnection, frame: Frame) -> None:
"""处理 proxy 上报的心跳数据，更新 DB，返回 ACK"""
try:
data = json.loads(frame.payload) if frame.payload else {}
except Exception:
data = {}
def _sync_heartbeat() -> dict[str, Any]:
from src.database import create_session
from src.services.proxy_node.service import ProxyNodeService
db = create_session()
try:
node = ProxyNodeService.heartbeat(
db,
node_id=conn.node_id,
active_connections=data.get("active_connections"),
total_requests=data.get("total_requests"),
avg_latency_ms=data.get("avg_latency_ms"),
)
result: dict[str, Any] = {}
if node.remote_config:
result["remote_config"] = node.remote_config
result["config_version"] = node.config_version or 0
return result
finally:
db.close()
try:
ack = await asyncio.to_thread(_sync_heartbeat)
except Exception as e:
logger.warning("tunnel heartbeat DB update failed: {}", e)
ack = {}
await conn.send_frame(Frame(0, MsgType.HEARTBEAT_ACK, 0, json.dumps(ack).encode()))
# 全局单例
_tunnel_manager: TunnelManager | None = None
def get_tunnel_manager() -> TunnelManager:
global _tunnel_manager
if _tunnel_manager is None:
_tunnel_manager = TunnelManager()
return _tunnel_manager

99
src/services/proxy_node/tunnel_protocol.py Normal file
View File

@@ -0,0 +1,99 @@
"""
WebSocket \u96a7\u9053\u4e8c\u8fdb\u5236\u5e27\u534f\u8bae
\u5e27\u683c\u5f0f:
| stream_id (4B) | msg_type (1B) | flags (1B) | payload_len (4B) | payload (NB) |
\u7528\u4e8e Aether \u4e0e aether-proxy \u4e4b\u95f4\u7684 WebSocket \u96a7\u9053\u591a\u8def\u590d\u7528\u901a\u4fe1\u3002
"""
import struct
from enum import IntEnum
from typing import Self
HEADER_SIZE = 10 # 4 + 1 + 1 + 4 bytes
class MsgType(IntEnum):
"""\u6d88\u606f\u7c7b\u578b"""
REQUEST_HEADERS = 0x01 # Aether -> Proxy: \u8bf7\u6c42\u5143\u6570\u636e (JSON)
REQUEST_BODY = 0x02 # Aether -> Proxy: \u8bf7\u6c42\u4f53
RESPONSE_HEADERS = 0x03 # Proxy -> Aether: \u54cd\u5e94\u72b6\u6001\u7801 + headers (JSON)
RESPONSE_BODY = 0x04 # Proxy -> Aether: \u54cd\u5e94\u4f53\uff08\u6d41\u5f0f\u5206\u5757\uff09
STREAM_END = 0x05 # \u53cc\u5411: \u6d41\u7ed3\u675f
STREAM_ERROR = 0x06 # \u53cc\u5411: \u6d41\u9519\u8bef
PING = 0x10 # \u53cc\u5411: \u5fc3\u8df3 (stream_id=0)
PONG = 0x11 # \u53cc\u5411: \u5fc3\u8df3\u54cd\u5e94 (stream_id=0)
GOAWAY = 0x12 # \u53cc\u5411: \u4f18\u96c5\u5173\u95ed (stream_id=0)
HEARTBEAT_DATA = 0x13 # Proxy -> Aether: \u6307\u6807\u4e0a\u62a5
HEARTBEAT_ACK = 0x14 # Aether -> Proxy: \u5fc3\u8df3\u786e\u8ba4 + \u8fdc\u7a0b\u914d\u7f6e
class FrameFlags:
"""\u5e27\u6807\u5fd7\u4f4d"""
END_STREAM = 0x01
GZIP_COMPRESSED = 0x02
class Frame:
"""WebSocket \u96a7\u9053\u5e27"""
__slots__ = ("stream_id", "msg_type", "flags", "payload")
def __init__(
self,
stream_id: int,
msg_type: MsgType,
flags: int = 0,
payload: bytes = b"",
) -> None:
self.stream_id = stream_id
self.msg_type = msg_type
self.flags = flags
self.payload = payload
def encode(self) -> bytes:
header = struct.pack(
"!IBBI",
self.stream_id,
self.msg_type,
self.flags,
len(self.payload),
)
return header + self.payload
@classmethod
def decode(cls, data: bytes) -> Self:
if len(data) < HEADER_SIZE:
raise ValueError(
f"\u5e27\u6570\u636e\u592a\u77ed: \u9700\u8981 {HEADER_SIZE} \u5b57\u8282, \u5b9e\u9645 {len(data)}"
)
stream_id, msg_type_raw, flags, payload_len = struct.unpack("!IBBI", data[:HEADER_SIZE])
expected_total = HEADER_SIZE + payload_len
if len(data) < expected_total:
raise ValueError(
f"\u5e27\u6570\u636e\u4e0d\u5b8c\u6574: \u9700\u8981 {expected_total} \u5b57\u8282, \u5b9e\u9645 {len(data)}"
)
try:
msg_type = MsgType(msg_type_raw)
except ValueError:
raise ValueError(f"\u672a\u77e5\u6d88\u606f\u7c7b\u578b: 0x{msg_type_raw:02x}")
payload = data[HEADER_SIZE:expected_total]
return cls(stream_id, msg_type, flags, payload)
@property
def is_end_stream(self) -> bool:
return bool(self.flags & FrameFlags.END_STREAM)
@property
def is_gzip(self) -> bool:
return bool(self.flags & FrameFlags.GZIP_COMPRESSED)
def __repr__(self) -> str:
return (
f"Frame(stream={self.stream_id}, type={self.msg_type.name}, "
f"flags=0x{self.flags:02x}, payload_len={len(self.payload)})"
)

133
src/services/proxy_node/tunnel_transport.py Normal file
View File

@@ -0,0 +1,133 @@
"""
Tunnel httpx Transport
自定义 httpx AsyncBaseTransport，将 HTTP 请求通过 WebSocket tunnel 发送到 aether-proxy。
对 handler 层完全透明 -- 只需在创建 httpx.AsyncClient 时使用此 transport。
"""
from __future__ import annotations
import asyncio
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import AsyncGenerator
import httpx
from .tunnel_manager import TunnelManager, TunnelStreamError, _StreamState, get_tunnel_manager
_HOP_BY_HOP_HEADERS = frozenset(
{
"host",
"transfer-encoding",
"content-length",
"connection",
"upgrade",
"keep-alive",
"proxy-authorization",
"proxy-connection",
"te",
"trailer",
}
)
# bytes 版本，用于直接比较 httpx raw headers（key 已经是小写 bytes）
_HOP_BY_HOP_HEADERS_BYTES = frozenset(h.encode("ascii") for h in _HOP_BY_HOP_HEADERS)
class TunnelTransport(httpx.AsyncBaseTransport):
"""通过 WebSocket tunnel 发送请求的 httpx transport"""
def __init__(self, node_id: str, timeout: float = 60.0) -> None:
self._node_id = node_id
self._timeout = timeout
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
manager = get_tunnel_manager()
# 构建 headers dict（跳过 hop-by-hop 和 httpx 内部 headers）
# request.headers.raw 返回 (bytes, bytes) 元组，key 已经是小写
headers: dict[str, str] = {}
for key, value in request.headers.raw:
if key not in _HOP_BY_HOP_HEADERS_BYTES:
headers[key.decode("latin-1")] = value.decode("latin-1")
# 读取 body -- request.content 在 json= 传参时已由 httpx 序列化好；
# 对 stream 类型的 request 需要先 read() 才能拿到完整 content。
body = request.content or await request.aread() or None
stream_state: _StreamState | None = None
try:
stream_state = await manager.send_request(
self._node_id,
method=request.method,
url=str(request.url),
headers=headers,
body=body,
timeout=self._timeout,
)
# 等待响应头
await stream_state.wait_headers(timeout=self._timeout)
# 构建 httpx.Response（流式 body）
resp_headers = httpx.Headers(stream_state.headers)
return httpx.Response(
status_code=stream_state.status,
headers=resp_headers,
stream=TunnelResponseStream(
manager, self._node_id, stream_state, timeout=self._timeout
),
)
except TunnelStreamError as e:
self._cleanup_stream(manager, stream_state)
# 区分连接阶段和响应阶段的错误
if stream_state and stream_state.status > 0:
raise httpx.ReadError(str(e)) from e
raise httpx.ConnectError(str(e)) from e
except asyncio.TimeoutError:
self._cleanup_stream(manager, stream_state)
raise httpx.ReadTimeout("tunnel request timeout") from None
def _cleanup_stream(self, manager: TunnelManager, stream_state: _StreamState | None) -> None:
if stream_state is None:
return
conn = manager.get_connection(self._node_id)
if conn:
conn.remove_stream(stream_state.stream_id)
class TunnelResponseStream(httpx.AsyncByteStream):
"""将 tunnel stream 的 body chunks 包装为 httpx AsyncByteStream"""
def __init__(
self,
manager: TunnelManager,
node_id: str,
stream_state: _StreamState,
timeout: float = 60.0,
) -> None:
self._manager = manager
self._node_id = node_id
self._stream_state = stream_state
self._timeout = timeout
async def __aiter__(self) -> AsyncGenerator[bytes, None]:
async for chunk in self._stream_state.iter_body(chunk_timeout=self._timeout):
yield chunk
async def aclose(self) -> None:
# 确保 stream 从 connection 的 pending 列表中移除，防止内存泄漏
conn = self._manager.get_connection(self._node_id)
if conn:
conn.remove_stream(self._stream_state.stream_id)
def is_tunnel_node(node_info: dict[str, Any] | None) -> bool:
"""检查节点是否为 tunnel 模式且已连接"""
if not node_info:
return False
return bool(node_info.get("tunnel_mode")) and bool(node_info.get("tunnel_connected"))