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 # 注意：更换此密钥后需要在管理面板重新配置所有 Provider API Key
ENCRYPTION_KEY=change-this-to-another-secure-random-string 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_EMAIL=admin@example.com
ADMIN_USERNAME=admin 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) # Management Token (ae_xxx, must belong to an ADMIN user)
AETHER_PROXY_MANAGEMENT_TOKEN=ae_xxxxx 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 # Node identification
AETHER_PROXY_NODE_NAME=proxy-01 AETHER_PROXY_NODE_NAME=proxy-01

167
aether-proxy/Cargo.lock generated
View File

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

16
aether-proxy/Cargo.toml
View File

@@ -2,19 +2,13 @@
name = "aether-proxy" name = "aether-proxy"
version = "0.1.6" version = "0.1.6"
edition = "2021" edition = "2021"
description = "Forward proxy for Aether with HMAC authentication" description = "Tunnel proxy for Aether"
[dependencies] [dependencies]
tokio = { version = "1", features = ["full"] } 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"] } 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" futures-util = "0.3"
hmac = "0.12"
sha2 = "0.10"
subtle = "2"
base64 = "0.22" base64 = "0.22"
clap = { version = "4", features = ["derive", "env"] } clap = { version = "4", features = ["derive", "env"] }
tracing = "0.1" tracing = "0.1"
@@ -23,15 +17,11 @@ serde = { version = "1", features = ["derive"] }
serde_json = "1" serde_json = "1"
thiserror = "2" thiserror = "2"
bytes = "1" bytes = "1"
sha2 = "0.10"
hex = "0.4" hex = "0.4"
anyhow = "1" anyhow = "1"
toml = "0.8" toml = "0.8"
tokio-rustls = "0.26"
webpki-roots = "1"
rustls = { version = "0.23", features = ["ring"] } rustls = { version = "0.23", features = ["ring"] }
rustls-pki-types = "1"
rustls-pemfile = "2"
rcgen = "0.13"
ratatui = "0.30" ratatui = "0.30"
crossterm = "0.28" crossterm = "0.28"
url = "2" url = "2"

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

@@ -1,40 +1,38 @@
//! Application lifecycle: initialization, task orchestration, and shutdown. //! 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::sync::{Arc, RwLock};
use std::time::Duration; use std::time::Duration;
use tokio::signal; use tokio::signal;
use tokio::sync::{watch, Semaphore}; use tokio::sync::watch;
use tracing::{error, info}; use tracing::{error, info, warn};
use crate::config::Config; use crate::config::{Config, ServerEntry};
use crate::net; use crate::net;
use crate::registration::client::AetherClient; use crate::registration::client::AetherClient;
use crate::runtime::{self, DynamicConfig}; use crate::runtime::{self, DynamicConfig};
use crate::state::{AppState, ProxyMetrics}; use crate::state::{AppState, ProxyMetrics, ServerContext};
use crate::{hardware, proxy}; use crate::{hardware, target_filter, tunnel};
/// Run the full application lifecycle after config has been parsed. /// 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); init_tracing(&config);
info!( info!(
version = env!("CARGO_PKG_VERSION"), version = env!("CARGO_PKG_VERSION"),
port = config.listen_port,
node_name = %config.node_name, 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 { let public_ip = match &config.public_ip {
Some(ip) => ip.clone(), 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 // Auto-detect region if not configured
if config.node_region.is_none() { if config.node_region.is_none() {
@@ -43,118 +41,136 @@ pub async fn run(mut config: Config) -> anyhow::Result<()> {
} }
} }
// Initialize TLS if enabled // Collect hardware info (once at startup, sent during registration)
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)
let hw_info = hardware::collect(); let hw_info = hardware::collect();
let max_connections_raw = config // Auto-detect tunnel_max_streams from hardware if not explicitly set
.max_concurrent_connections if config.tunnel_max_streams.is_none() {
.unwrap_or(hw_info.estimated_max_concurrency) let auto = (hw_info.estimated_max_concurrency / 10).clamp(64, 1024) as u32;
.max(1); config.tunnel_max_streams = Some(auto);
let max_connections = usize::try_from(max_connections_raw).unwrap_or(usize::MAX); info!(
tunnel_max_streams = auto,
"auto-detected tunnel_max_streams from hardware"
);
}
info!( info!(
max_connections = max_connections_raw, max_concurrency = hw_info.estimated_max_concurrency,
"connection limit configured" "hardware info collected"
); );
let connection_semaphore = Arc::new(Semaphore::new(max_connections)); let dns_cache = Arc::new(target_filter::DnsCache::new(
let metrics = Arc::new(ProxyMetrics::new());
let dns_cache = Arc::new(proxy::target_filter::DnsCache::new(
Duration::from_secs(config.dns_cache_ttl_secs), Duration::from_secs(config.dns_cache_ttl_secs),
config.dns_cache_capacity, config.dns_cache_capacity,
)); ));
// Register with Aether // Build reqwest client for tunnel upstream requests (shared).
let aether_client = Arc::new(AetherClient::new(&config)); let reqwest_client = reqwest::Client::builder()
let node_id = aether_client .pool_max_idle_per_host(config.upstream_pool_max_idle_per_host)
.register( .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, &config,
&public_ip, &entry.aether_url,
config.enable_tls, &entry.management_token,
tls_fingerprint.as_deref(), ));
Some(&hw_info), match client
) .register(&config, &node_name, &public_ip, Some(&hw_info))
.await?; .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"); if server_contexts.is_empty() {
anyhow::bail!("no servers registered successfully");
// 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);
// Build shared application state // Build shared application state
let state = Arc::new(AppState { let state = Arc::new(AppState {
config: Arc::new(config), 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, dns_cache,
metrics, reqwest_client,
}); });
// Shutdown signal channel // Shutdown signal channel
let (shutdown_tx, shutdown_rx) = watch::channel(false); let (shutdown_tx, shutdown_rx) = watch::channel(false);
// Start heartbeat task info!(
let heartbeat_handle = { active_servers = server_contexts.len(),
let state = Arc::clone(&state); "running in tunnel mode"
);
// 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(); let rx = shutdown_rx.clone();
tokio::spawn(async move { tunnel_handles.push(tokio::spawn(async move {
crate::registration::heartbeat::run(&state, rx).await; tunnel::run(&s, &srv, rx).await;
}) }));
};
// Start proxy server
let server_handle = {
let state = Arc::clone(&state);
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");
}
})
};
// Wait for shutdown signal (SIGTERM or SIGINT)
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");
} }
// Wait for tasks to finish // Wait for shutdown signal
let _ = tokio::join!(heartbeat_handle, server_handle); wait_for_shutdown().await;
info!("shutdown signal received, cleaning up...");
let _ = shutdown_tx.send(true);
// 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 all tunnel tasks
for h in tunnel_handles {
let _ = h.await;
}
info!("aether-proxy stopped"); info!("aether-proxy stopped");
Ok(()) Ok(())
@@ -168,7 +184,6 @@ fn init_tracing(config: &Config) {
let (filter_layer, reload_handle) = reload::Layer::new(filter); let (filter_layer, reload_handle) = reload::Layer::new(filter);
// Register log-level hot-reloader
runtime::set_log_reloader(Box::new(move |level: &str| { runtime::set_log_reloader(Box::new(move |level: &str| {
if let Ok(new_filter) = EnvFilter::try_new(level) { if let Ok(new_filter) = EnvFilter::try_new(level) {
let _ = reload_handle.modify(|f| *f = new_filter); 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 clap::Parser;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
/// Aether forward proxy with HMAC authentication. /// Aether tunnel proxy.
/// ///
/// Deployed on overseas VPS to relay API traffic for Aether instances /// Deployed on overseas VPS to relay API traffic for Aether instances
/// behind the GFW. Registers with Aether, sends heartbeats, and validates /// behind the GFW. Connects to Aether via WebSocket tunnel, registers
/// incoming proxy requests via HMAC-SHA256 signatures in Basic Auth. /// with Aether, and relays upstream requests.
#[derive(Parser, Debug, Clone)] #[derive(Parser, Debug, Clone)]
#[command(version, about)] #[command(version, about)]
pub struct Config { pub struct Config {
@@ -19,14 +19,6 @@ pub struct Config {
#[arg(long, env = "AETHER_PROXY_MANAGEMENT_TOKEN")] #[arg(long, env = "AETHER_PROXY_MANAGEMENT_TOKEN")]
pub management_token: String, 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) /// Public IP address of this node (auto-detected if omitted)
#[arg(long, env = "AETHER_PROXY_PUBLIC_IP")] #[arg(long, env = "AETHER_PROXY_PUBLIC_IP")]
pub public_ip: Option<String>, pub public_ip: Option<String>,
@@ -52,10 +44,6 @@ pub struct Config {
)] )]
pub allowed_ports: Vec<u16>, 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 /// Aether API request timeout in seconds
#[arg( #[arg(
long, long,
@@ -128,14 +116,6 @@ pub struct Config {
#[arg(long, env = "AETHER_PROXY_MAX_CONCURRENT_CONNECTIONS")] #[arg(long, env = "AETHER_PROXY_MAX_CONCURRENT_CONNECTIONS")]
pub max_concurrent_connections: Option<u64>, 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 /// DNS cache TTL in seconds
#[arg(long, env = "AETHER_PROXY_DNS_CACHE_TTL", default_value_t = 60)] #[arg(long, env = "AETHER_PROXY_DNS_CACHE_TTL", default_value_t = 60)]
pub dns_cache_ttl_secs: u64, 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)] #[arg(long, env = "AETHER_PROXY_DNS_CACHE_CAPACITY", default_value_t = 1024)]
pub dns_cache_capacity: usize, pub dns_cache_capacity: usize,
/// Delegate HTTP client connect timeout in seconds /// Upstream HTTP client connect timeout in seconds
#[arg( #[arg(
long, long,
env = "AETHER_PROXY_DELEGATE_CONNECT_TIMEOUT", env = "AETHER_PROXY_UPSTREAM_CONNECT_TIMEOUT",
default_value_t = 30 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( #[arg(
long, long,
env = "AETHER_PROXY_DELEGATE_POOL_MAX_IDLE_PER_HOST", env = "AETHER_PROXY_UPSTREAM_POOL_MAX_IDLE_PER_HOST",
default_value_t = 64 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( #[arg(
long, long,
env = "AETHER_PROXY_DELEGATE_POOL_IDLE_TIMEOUT", env = "AETHER_PROXY_UPSTREAM_POOL_IDLE_TIMEOUT",
default_value_t = 300 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( #[arg(
long, long,
env = "AETHER_PROXY_DELEGATE_TCP_KEEPALIVE", env = "AETHER_PROXY_UPSTREAM_TCP_KEEPALIVE",
default_value_t = 60 default_value_t = 60
)] )]
pub delegate_tcp_keepalive_secs: u64, pub upstream_tcp_keepalive_secs: u64,
/// Delegate TCP_NODELAY /// Upstream TCP_NODELAY
#[arg( #[arg(
long, long,
env = "AETHER_PROXY_DELEGATE_TCP_NODELAY", env = "AETHER_PROXY_UPSTREAM_TCP_NODELAY",
default_value_t = true default_value_t = true
)] )]
pub delegate_tcp_nodelay: bool, pub upstream_tcp_nodelay: bool,
/// Log level (trace, debug, info, warn, error) /// Log level (trace, debug, info, warn, error)
#[arg(long, env = "AETHER_PROXY_LOG_LEVEL", default_value = "info")] #[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)] #[arg(long, env = "AETHER_PROXY_LOG_JSON", default_value_t = false)]
pub log_json: bool, pub log_json: bool,
/// Enable TLS encryption (dual-stack: accepts both HTTP and TLS on same port) /// WebSocket reconnect base delay in milliseconds
#[arg(long, env = "AETHER_PROXY_ENABLE_TLS", default_value_t = true)]
pub enable_tls: bool,
/// Path to TLS certificate PEM file
#[arg( #[arg(
long, long,
env = "AETHER_PROXY_TLS_CERT", env = "AETHER_PROXY_TUNNEL_RECONNECT_BASE_MS",
default_value = "aether-proxy-cert.pem" 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( #[arg(
long, long,
env = "AETHER_PROXY_TLS_KEY", env = "AETHER_PROXY_TUNNEL_RECONNECT_MAX_MS",
default_value = "aether-proxy-key.pem" 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. /// 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)] #[derive(Debug, Default, Serialize, Deserialize)]
pub struct ConfigFile { pub struct ConfigFile {
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
@@ -226,10 +219,6 @@ pub struct ConfigFile {
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
pub management_token: Option<String>, pub management_token: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")] #[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>, pub public_ip: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
pub node_name: Option<String>, pub node_name: Option<String>,
@@ -240,8 +229,6 @@ pub struct ConfigFile {
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
pub allowed_ports: Option<Vec<u16>>, pub allowed_ports: Option<Vec<u16>>,
#[serde(skip_serializing_if = "Option::is_none")] #[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>, pub aether_request_timeout_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
pub aether_connect_timeout_secs: Option<u64>, pub aether_connect_timeout_secs: Option<u64>,
@@ -264,33 +251,37 @@ pub struct ConfigFile {
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
pub max_concurrent_connections: Option<u64>, pub max_concurrent_connections: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")] #[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>, pub dns_cache_ttl_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
pub dns_cache_capacity: Option<usize>, pub dns_cache_capacity: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")] #[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")] #[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")] #[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")] #[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")] #[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")] #[serde(skip_serializing_if = "Option::is_none")]
pub log_level: Option<String>, pub log_level: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
pub log_json: Option<bool>, pub log_json: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")] #[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")] #[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")] #[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 { impl ConfigFile {
@@ -307,6 +298,24 @@ impl ConfigFile {
Ok(()) 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. /// Inject values as environment variables so clap picks them up.
/// ///
/// Only sets variables that are **not** already present in the /// 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); // When top-level fields are absent, fall back to the first [[servers]]
set!("AETHER_PROXY_HMAC_KEY", self.hmac_key); // entry so that clap's required `aether_url` / `management_token` are
set!("AETHER_PROXY_LISTEN_PORT", self.listen_port); // 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_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_NODE_REGION", self.node_region);
set!("AETHER_PROXY_HEARTBEAT_INTERVAL", self.heartbeat_interval); set!("AETHER_PROXY_HEARTBEAT_INTERVAL", self.heartbeat_interval);
set!("AETHER_PROXY_TIMESTAMP_TOLERANCE", self.timestamp_tolerance);
set!( set!(
"AETHER_PROXY_AETHER_REQUEST_TIMEOUT", "AETHER_PROXY_AETHER_REQUEST_TIMEOUT",
self.aether_request_timeout_secs self.aether_request_timeout_secs
@@ -379,38 +403,43 @@ impl ConfigFile {
"AETHER_PROXY_MAX_CONCURRENT_CONNECTIONS", "AETHER_PROXY_MAX_CONCURRENT_CONNECTIONS",
self.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_TTL", self.dns_cache_ttl_secs);
set!("AETHER_PROXY_DNS_CACHE_CAPACITY", self.dns_cache_capacity); set!("AETHER_PROXY_DNS_CACHE_CAPACITY", self.dns_cache_capacity);
set!( set!(
"AETHER_PROXY_DELEGATE_CONNECT_TIMEOUT", "AETHER_PROXY_UPSTREAM_CONNECT_TIMEOUT",
self.delegate_connect_timeout_secs self.upstream_connect_timeout_secs
); );
set!( set!(
"AETHER_PROXY_DELEGATE_POOL_MAX_IDLE_PER_HOST", "AETHER_PROXY_UPSTREAM_POOL_MAX_IDLE_PER_HOST",
self.delegate_pool_max_idle_per_host self.upstream_pool_max_idle_per_host
); );
set!( set!(
"AETHER_PROXY_DELEGATE_POOL_IDLE_TIMEOUT", "AETHER_PROXY_UPSTREAM_POOL_IDLE_TIMEOUT",
self.delegate_pool_idle_timeout_secs self.upstream_pool_idle_timeout_secs
); );
set!( set!(
"AETHER_PROXY_DELEGATE_TCP_KEEPALIVE", "AETHER_PROXY_UPSTREAM_TCP_KEEPALIVE",
self.delegate_tcp_keepalive_secs self.upstream_tcp_keepalive_secs
); );
set!( set!(
"AETHER_PROXY_DELEGATE_TCP_NODELAY", "AETHER_PROXY_UPSTREAM_TCP_NODELAY",
self.delegate_tcp_nodelay self.upstream_tcp_nodelay
); );
set!("AETHER_PROXY_LOG_LEVEL", self.log_level); set!("AETHER_PROXY_LOG_LEVEL", self.log_level);
set!("AETHER_PROXY_LOG_JSON", self.log_json); set!("AETHER_PROXY_LOG_JSON", self.log_json);
set!("AETHER_PROXY_ENABLE_TLS", self.enable_tls); set!(
set!("AETHER_PROXY_TLS_CERT", self.tls_cert); "AETHER_PROXY_TUNNEL_RECONNECT_BASE_MS",
set!("AETHER_PROXY_TLS_KEY", self.tls_key); 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) // allowed_ports needs special handling (comma-separated)
if let Some(ref ports) = self.allowed_ports { if let Some(ref ports) = self.allowed_ports {

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

@@ -1,13 +1,13 @@
mod app; mod app;
mod auth;
mod config; mod config;
mod hardware; mod hardware;
mod net; mod net;
mod proxy;
mod registration; mod registration;
mod runtime; mod runtime;
mod setup; mod setup;
mod state; mod state;
mod target_filter;
mod tunnel;
use std::path::PathBuf; use std::path::PathBuf;
@@ -135,5 +135,28 @@ async fn run_proxy(config: Config) -> anyhow::Result<()> {
std::process::exit(1); 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 reqwest::{Client, StatusCode};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use tokio::time::sleep; use tokio::time::sleep;
use tracing::{debug, error, info, warn}; use tracing::{debug, error, info};
use crate::config::Config; use crate::config::Config;
use crate::hardware::HardwareInfo; 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)] #[derive(Debug, Serialize)]
struct RegisterRequest { struct RegisterRequest {
name: String, name: String,
@@ -35,14 +16,11 @@ struct RegisterRequest {
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
region: Option<String>, region: Option<String>,
heartbeat_interval: u64, 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")] #[serde(skip_serializing_if = "Option::is_none")]
hardware_info: Option<serde_json::Value>, hardware_info: Option<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")] #[serde(skip_serializing_if = "Option::is_none")]
estimated_max_concurrency: Option<u64>, estimated_max_concurrency: Option<u64>,
tunnel_mode: bool,
} }
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]
@@ -50,17 +28,6 @@ pub struct RegisterResponse {
pub node_id: String, 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. /// Remote configuration pushed by the Aether management backend.
#[derive(Debug, Clone, Deserialize)] #[derive(Debug, Clone, Deserialize)]
pub struct RemoteConfig { pub struct RemoteConfig {
@@ -68,29 +35,6 @@ pub struct RemoteConfig {
pub allowed_ports: Option<Vec<u16>>, pub allowed_ports: Option<Vec<u16>>,
pub log_level: Option<String>, pub log_level: Option<String>,
pub heartbeat_interval: Option<u64>, 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)] #[derive(Debug, Serialize)]
@@ -109,7 +53,7 @@ pub struct AetherClient {
} }
impl 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() let mut builder = Client::builder()
.timeout(Duration::from_secs(config.aether_request_timeout_secs)) .timeout(Duration::from_secs(config.aether_request_timeout_secs))
.connect_timeout(Duration::from_secs(config.aether_connect_timeout_secs)) .connect_timeout(Duration::from_secs(config.aether_connect_timeout_secs))
@@ -136,8 +80,8 @@ impl AetherClient {
Self { Self {
http, http,
base_url: config.aether_url.trim_end_matches('/').to_string(), base_url: aether_url.trim_end_matches('/').to_string(),
token: config.management_token.clone(), token: management_token.to_string(),
retry_max_attempts: config.aether_retry_max_attempts.max(1), retry_max_attempts: config.aether_retry_max_attempts.max(1),
retry_base_delay, retry_base_delay,
retry_max_delay, retry_max_delay,
@@ -150,29 +94,26 @@ impl AetherClient {
pub async fn register( pub async fn register(
&self, &self,
config: &Config, config: &Config,
node_name: &str,
public_ip: &str, public_ip: &str,
tls_enabled: bool,
tls_cert_fingerprint: Option<&str>,
hw: Option<&HardwareInfo>, hw: Option<&HardwareInfo>,
) -> anyhow::Result<String> { ) -> anyhow::Result<String> {
let url = format!("{}/api/admin/proxy-nodes/register", self.base_url); let url = format!("{}/api/admin/proxy-nodes/register", self.base_url);
let body = RegisterRequest { let body = RegisterRequest {
name: config.node_name.clone(), name: node_name.to_string(),
ip: public_ip.to_string(), ip: public_ip.to_string(),
port: config.listen_port, port: 0,
region: config.node_region.clone(), region: config.node_region.clone(),
heartbeat_interval: config.heartbeat_interval, 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()), hardware_info: hw.and_then(|h| serde_json::to_value(h).ok()),
estimated_max_concurrency: hw.map(|h| h.estimated_max_concurrency), estimated_max_concurrency: hw.map(|h| h.estimated_max_concurrency),
tunnel_mode: true,
}; };
info!( info!(
url = %url, url = %url,
name = %body.name, name = %body.name,
ip = %body.ip, ip = %body.ip,
port = body.port,
"registering with Aether" "registering with Aether"
); );
@@ -199,80 +140,6 @@ impl AetherClient {
Ok(data.node_id) 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). /// Unregister this node from Aether (graceful shutdown).
pub async fn unregister(&self, node_id: &str) -> anyhow::Result<()> { pub async fn unregister(&self, node_id: &str) -> anyhow::Result<()> {
let url = format!("{}/api/admin/proxy-nodes/unregister", self.base_url); 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 client;
pub mod heartbeat;

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

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

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

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

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

@@ -1,42 +1,45 @@
//! Shared application state passed to all subsystems. //! 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::sync::{Arc, RwLock};
use std::time::Duration; use std::time::Duration;
use tokio::sync::Semaphore;
use tokio_rustls::TlsAcceptor;
use crate::config::Config; 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::registration::client::AetherClient;
use crate::runtime::SharedDynamicConfig; 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 struct AppState {
pub config: Arc<Config>, pub config: Arc<Config>,
pub node_id: Arc<RwLock<String>>, /// DNS cache for upstream target resolution (shared).
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.
pub dns_cache: Arc<DnsCache>, 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>, pub metrics: Arc<ProxyMetrics>,
/// Reconnect attempt counter (reset on successful connection).
pub reconnect_attempts: AtomicU32,
} }
/// Aggregate metrics for reporting to Aether. /// 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[] allowed_ports?: number[]
log_level?: string log_level?: string
heartbeat_interval?: number heartbeat_interval?: number
timestamp_tolerance?: number
} }
export interface ProxyNode { export interface ProxyNode {
@@ -16,6 +15,9 @@ export interface ProxyNode {
region: string | null region: string | null
status: 'online' | 'unhealthy' | 'offline' status: 'online' | 'unhealthy' | 'offline'
is_manual: boolean is_manual: boolean
tunnel_mode: boolean
tunnel_connected: boolean
tunnel_connected_at: string | null
// 手动节点专用字段 // 手动节点专用字段
proxy_url?: string proxy_url?: string
proxy_username?: string proxy_username?: string
@@ -101,9 +103,4 @@ export const proxyNodesApi = {
const response = await apiClient.post<ProxyNodeTestResult>('/api/admin/proxy-nodes/test-url', data) const response = await apiClient.post<ProxyNodeTestResult>('/api/admin/proxy-nodes/test-url', data)
return response.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> </SelectContent>
</Select> </Select>
<div class="h-4 w-px bg-border" /> <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 <Button
variant="ghost" variant="ghost"
size="icon" size="icon"
@@ -172,11 +162,18 @@
> >
手动 手动
</Badge> </Badge>
<Badge
v-if="node.tunnel_mode"
variant="outline"
class="text-[10px] px-1.5 py-0"
>
Tunnel
</Badge>
<HardwareTooltip :node="node" /> <HardwareTooltip :node="node" />
</div> </div>
</TableCell> </TableCell>
<TableCell class="py-4"> <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>
<TableCell class="py-4"> <TableCell class="py-4">
<span class="text-sm text-muted-foreground">{{ formatRegion(node.region) }}</span> <span class="text-sm text-muted-foreground">{{ formatRegion(node.region) }}</span>
@@ -282,9 +279,16 @@
> >
手动 手动
</Badge> </Badge>
<Badge
v-if="node.tunnel_mode"
variant="outline"
class="text-[10px] px-1.5 py-0"
>
Tunnel
</Badge>
<HardwareTooltip :node="node" /> <HardwareTooltip :node="node" />
</div> </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> </div>
<Badge <Badge
:variant="statusVariant(node.status)" :variant="statusVariant(node.status)"
@@ -521,15 +525,6 @@
max="600" max="600"
/> />
</div> </div>
<div class="space-y-1.5">
<Label>时间戳容差 (秒)</Label>
<Input
v-model="configForm.timestamp_tolerance"
type="number"
min="10"
max="3600"
/>
</div>
</div> </div>
<div <div
v-if="configNode" v-if="configNode"
@@ -560,7 +555,6 @@
import { ref, computed, onMounted, watch } from 'vue' import { ref, computed, onMounted, watch } from 'vue'
import { useProxyNodesStore } from '@/stores/proxy-nodes' import { useProxyNodesStore } from '@/stores/proxy-nodes'
import { useToast } from '@/composables/useToast' import { useToast } from '@/composables/useToast'
import { useClipboard } from '@/composables/useClipboard'
import { useConfirm } from '@/composables/useConfirm' import { useConfirm } from '@/composables/useConfirm'
import { proxyNodesApi, type ProxyNode, type ProxyNodeRemoteConfig } from '@/api/proxy-nodes' import { proxyNodesApi, type ProxyNode, type ProxyNodeRemoteConfig } from '@/api/proxy-nodes'
@@ -586,13 +580,12 @@ import {
Dialog, Dialog,
} from '@/components/ui' } 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 { parseApiError } from '@/utils/errorParser'
import { formatRegion } from '@/utils/region' import { formatRegion } from '@/utils/region'
import HardwareTooltip from './components/HardwareTooltip.vue' import HardwareTooltip from './components/HardwareTooltip.vue'
const { success, error: toastError } = useToast() const { success, error: toastError } = useToast()
const { copyToClipboard } = useClipboard()
const { confirmDanger } = useConfirm() const { confirmDanger } = useConfirm()
const store = useProxyNodesStore() const store = useProxyNodesStore()
@@ -621,7 +614,6 @@ const configForm = ref({
allowed_ports: '', allowed_ports: '',
log_level: 'info', log_level: 'info',
heartbeat_interval: '30', 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) { function handleEdit(node: ProxyNode) {
editingNode.value = node editingNode.value = node
addForm.value = { addForm.value = {
@@ -766,7 +749,6 @@ function handleConfig(node: ProxyNode) {
allowed_ports: rc.allowed_ports?.join(', ') || '', allowed_ports: rc.allowed_ports?.join(', ') || '',
log_level: rc.log_level || 'info', log_level: rc.log_level || 'info',
heartbeat_interval: String(rc.heartbeat_interval || node.heartbeat_interval || 30), heartbeat_interval: String(rc.heartbeat_interval || node.heartbeat_interval || 30),
timestamp_tolerance: String(rc.timestamp_tolerance || 300),
} }
showConfigDialog.value = true showConfigDialog.value = true
} }
@@ -800,10 +782,6 @@ async function handleSaveConfig() {
if (!isNaN(hb) && hb >= 5) { if (!isNaN(hb) && hb >= 5) {
data.heartbeat_interval = hb 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) await proxyNodesApi.updateNodeConfig(configNode.value.id, data)
success('远程配置已保存，将在下次心跳时生效') success('远程配置已保存，将在下次心跳时生效')
handleConfigDialogClose(false) handleConfigDialogClose(false)
@@ -817,7 +795,7 @@ async function handleSaveConfig() {
async function handleDelete(node: ProxyNode) { async function handleDelete(node: ProxyNode) {
const confirmed = await confirmDanger( const confirmed = await confirmDanger(
`确定要删除代理节点 "${node.name}" (${node.ip}:${node.port}) 吗？`, `确定要删除代理节点 "${node.name}" (${node.tunnel_mode ? node.ip : `${node.ip}:${node.port}`}) 吗？`,
'删除节点' '删除节点'
) )
if (!confirmed) return if (!confirmed) return
@@ -890,4 +868,10 @@ function formatTime(iso: string | null) {
if (diff < 86400) return `${Math.floor(diff / 3600)}小时前` 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' }) 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> </script>

8
generate_keys.py
View File

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

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

@@ -32,7 +32,7 @@ pipeline = ApiRequestPipeline()
class ProxyNodeRegisterRequest(BaseModel): class ProxyNodeRegisterRequest(BaseModel):
name: str = Field(..., min_length=1, max_length=100, description="节点名") name: str = Field(..., min_length=1, max_length=100, description="节点名")
ip: str = Field(..., description="公网 IP（IPv4/IPv6）") 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="区域标签") region: str | None = Field(None, max_length=100, description="区域标签")
heartbeat_interval: int = Field(30, ge=5, le=600, 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="累计请求数") total_requests: int | None = Field(None, ge=0, description="累计请求数")
avg_latency_ms: float | 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") hardware_info: dict | None = Field(None, description="硬件信息 JSON")
estimated_max_concurrency: int | None = Field(None, ge=0, description="估算最大并发连接数") estimated_max_concurrency: int | None = Field(None, ge=0, description="估算最大并发连接数")
# Tunnel 模式
tunnel_mode: bool = Field(False, description="是否使用 tunnel 模式连接")
@field_validator("ip") @field_validator("ip")
@classmethod @classmethod
def validate_ip(cls, v: str) -> str: def validate_ip(cls, v: str) -> str:
@@ -82,9 +79,6 @@ class ProxyNodeRemoteConfigRequest(BaseModel):
allowed_ports: list[int] | None = Field(None, description="允许代理的目标端口") allowed_ports: list[int] | None = Field(None, description="允许代理的目标端口")
log_level: str | None = Field(None, description="日志级别 (trace/debug/info/warn/error)") log_level: str | None = Field(None, description="日志级别 (trace/debug/info/warn/error)")
heartbeat_interval: int | None = Field(None, ge=5, le=600, description="心跳间隔（秒）") 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") @field_validator("allowed_ports")
@classmethod @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) 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") @router.post("/test-url")
async def test_proxy_url(request: Request, db: Session = Depends(get_db)) -> Any: async def test_proxy_url(request: Request, db: Session = Depends(get_db)) -> Any:
adapter = AdminTestProxyUrlAdapter() adapter = AdminTestProxyUrlAdapter()
@@ -273,14 +261,13 @@ class AdminRegisterProxyNodeAdapter(AdminApiAdapter):
port=req.port, port=req.port,
region=req.region, region=req.region,
heartbeat_interval=req.heartbeat_interval, heartbeat_interval=req.heartbeat_interval,
tls_enabled=req.tls_enabled,
tls_cert_fingerprint=req.tls_cert_fingerprint,
hardware_info=req.hardware_info, hardware_info=req.hardware_info,
estimated_max_concurrency=req.estimated_max_concurrency, estimated_max_concurrency=req.estimated_max_concurrency,
active_connections=req.active_connections, active_connections=req.active_connections,
total_requests=req.total_requests, total_requests=req.total_requests,
avg_latency_ms=req.avg_latency_ms, avg_latency_ms=req.avg_latency_ms,
registered_by=context.user.id if context.user else None, registered_by=context.user.id if context.user else None,
tunnel_mode=req.tunnel_mode,
) )
context.add_audit_metadata( context.add_audit_metadata(
@@ -376,11 +363,24 @@ class AdminDeleteProxyNodeAdapter(AdminApiAdapter):
) )
was_system_proxy = result["cleared_system_proxy"] 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 { return {
"message": msg, "message": ", ".join(parts),
"node_id": self.node_id, "node_id": self.node_id,
"cleared_system_proxy": was_system_proxy, "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 config_updates["log_level"] = req.log_level
if req.heartbeat_interval is not None: if req.heartbeat_interval is not None:
config_updates["heartbeat_interval"] = req.heartbeat_interval 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( node = ProxyNodeService.update_node_config(
context.db, node_id=self.node_id, config_updates=config_updates 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): class TestProxyUrlRequest(BaseModel):
proxy_url: str = Field(..., min_length=1, max_length=500) proxy_url: str = Field(..., min_length=1, max_length=500)
username: str | None = Field(None, max_length=255) 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.api.handlers.base.chat_handler_base import ChatHandlerBase
from src.core.api_format import ApiFamily, get_auth_handler from src.core.api_format import ApiFamily, get_auth_handler
from src.core.api_format.enums import AuthMethod 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.core.logger import logger
from src.models.gemini import GeminiRequest from src.models.gemini import GeminiRequest
from src.services.provider.transport import redact_url_for_log from src.services.provider.transport import redact_url_for_log
@@ -197,7 +198,7 @@ class GeminiChatAdapter(ChatAdapterBase):
else: else:
models_url = f"{base_url_clean}/v1beta/models?key={api_key}" 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: if extra_headers:
headers.update(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]] = {} _proxy_clients: dict[str, tuple[httpx.AsyncClient, float]] = {}
# 代理客户端缓存上限（避免内存泄漏） # 代理客户端缓存上限（避免内存泄漏）
_max_proxy_clients: int = 50 _max_proxy_clients: int = 50
# 代发客户端缓存：{tls: client, plain: client} # Tunnel 客户端缓存：{node_id: client}
_delegate_clients: dict[str, httpx.AsyncClient] = {} _tunnel_clients: dict[str, httpx.AsyncClient] = {}
def __new__(cls) -> "HTTPClientPool": def __new__(cls) -> "HTTPClientPool":
if cls._instance is None: if cls._instance is None:
@@ -289,15 +289,15 @@ class HTTPClientPool:
cls._proxy_clients.clear() cls._proxy_clients.clear()
# 关闭代发客户端缓存 # 关闭 tunnel 客户端缓存
for cache_key, client in cls._delegate_clients.items(): for nid, client in cls._tunnel_clients.items():
try: try:
await client.aclose() await client.aclose()
logger.debug("代发客户端已关闭: {}", cache_key) logger.debug("tunnel 客户端已关闭: {}", nid)
except Exception as e: except Exception as e:
logger.warning("关闭代发客户端失败: {}", e) logger.warning("关闭 tunnel 客户端失败: {}", e)
cls._delegate_clients.clear() cls._tunnel_clients.clear()
logger.info("所有HTTP客户端已关闭") logger.info("所有HTTP客户端已关闭")
@classmethod @classmethod
@@ -380,88 +380,6 @@ class HTTPClientPool:
client_config.update(kwargs) client_config.update(kwargs)
return httpx.AsyncClient(**client_config) # type: ignore[arg-type] 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 @classmethod
async def get_upstream_client( async def get_upstream_client(
cls, cls,
@@ -469,31 +387,70 @@ class HTTPClientPool:
proxy_config: dict[str, Any] | None = None, proxy_config: dict[str, Any] | None = None,
) -> httpx.AsyncClient: ) -> httpx.AsyncClient:
""" """
获取可复用的上游请求客户端（自动选择代发或代理模式） 获取可复用的上游请求客户端（自动选择 tunnel/代理模式）
代发模式(delegate_cfg非空)：返回代发客户端 tunnel 模式(delegate_cfg.tunnel=True)：返回 TunnelTransport 客户端
直连/代理模式：返回代理客户端（含系统默认代理回退） 直连/代理模式：返回代理客户端（含系统默认代理回退）
""" """
if delegate_cfg: if delegate_cfg and delegate_cfg.get("tunnel"):
return await cls.get_delegate_client(delegate_cfg) return await cls._get_tunnel_client(delegate_cfg["node_id"])
return await cls.get_proxy_client(proxy_config=proxy_config) return await cls.get_proxy_client(proxy_config=proxy_config)
@classmethod @classmethod
def create_upstream_stream_client( async def create_upstream_stream_client(
cls, cls,
delegate_cfg: dict[str, Any] | None, delegate_cfg: dict[str, Any] | None,
proxy_config: dict[str, Any] | None = None, proxy_config: dict[str, Any] | None = None,
timeout: httpx.Timeout | None = None, timeout: httpx.Timeout | None = None,
) -> httpx.AsyncClient: ) -> httpx.AsyncClient:
""" """
创建上游流式请求客户端（自动选择代发或代理模式） 创建上游流式请求客户端（自动选择 tunnel/代理模式）
调用者需负责关闭返回的客户端。 调用者需负责关闭返回的客户端。
""" """
if delegate_cfg: if delegate_cfg and delegate_cfg.get("tunnel"):
return cls.create_delegate_stream_client(delegate_cfg, timeout=timeout) return await cls._get_tunnel_client(delegate_cfg["node_id"], timeout=timeout)
return cls.create_client_with_proxy(proxy_config=proxy_config, 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 @classmethod
def get_pool_stats(cls) -> dict[str, Any]: def get_pool_stats(cls) -> dict[str, Any]:
"""获取连接池统计信息""" """获取连接池统计信息"""
@@ -502,7 +459,7 @@ class HTTPClientPool:
"named_clients_count": len(cls._clients), "named_clients_count": len(cls._clients),
"proxy_clients_count": len(cls._proxy_clients), "proxy_clients_count": len(cls._proxy_clients),
"max_proxy_clients": cls._max_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 文件加载配置 从环境变量或 .env 文件加载配置
""" """
import hashlib
import hmac
import os import os
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
@@ -48,13 +46,6 @@ class Config:
# 加密密钥配置（独立于JWT密钥，用于敏感数据加密） # 加密密钥配置（独立于JWT密钥，用于敏感数据加密）
self.encryption_key = os.getenv("ENCRYPTION_KEY", None) 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 部署默认为生产环境，本地开发默认为开发环境 # Docker 部署默认为生产环境，本地开发默认为开发环境
is_docker = ( is_docker = (
@@ -327,20 +318,6 @@ class Config:
# 验证连接池配置 # 验证连接池配置
self._validate_pool_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: def _auto_pool_size(self) -> int:
""" """
智能计算连接池大小 - 根据 Worker 数量和 PostgreSQL 限制计算 智能计算连接池大小 - 根据 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 collections.abc import Set as AbstractSet
from typing import Any from typing import Any
from src.core.api_format.enums import ApiFamily
from src.core.api_format.metadata import ( from src.core.api_format.metadata import (
get_auth_config_for_endpoint, get_auth_config_for_endpoint,
get_extra_headers_for_endpoint, get_extra_headers_for_endpoint,
get_protected_keys_for_endpoint, get_protected_keys_for_endpoint,
resolve_endpoint_definition,
) )
from src.core.api_format.signature import EndpointSignature, parse_signature_key from src.core.api_format.signature import EndpointSignature, parse_signature_key
from src.core.logger import logger 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( UPSTREAM_DROP_HEADERS: frozenset[str] = frozenset(
{ {
@@ -503,14 +536,23 @@ def build_adapter_base_headers_for_endpoint(
) -> dict[str, str]: ) -> dict[str, str]:
""" """
新模式：根据 endpoint signature 构建基础请求头。 新模式：根据 endpoint signature 构建基础请求头。
浏览器指纹 headers 作为底层默认值注入，Claude API family 额外注入 Anthropic 专属 header。
认证头和 extra_headers 会覆盖它们。
""" """
auth_header, auth_type = get_auth_config_for_endpoint(endpoint) auth_header, auth_type = get_auth_config_for_endpoint(endpoint)
auth_value = f"Bearer {api_key}" if auth_type == "bearer" else api_key auth_value = f"Bearer {api_key}" if auth_type == "bearer" else api_key
headers: dict[str, str] = { # 以浏览器指纹为底层默认值，绕过 Cloudflare 等反爬防护
auth_header: auth_value, headers: dict[str, str] = {**BROWSER_FINGERPRINT_HEADERS}
"Content-Type": "application/json",
} # 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: if include_extra:
extra = get_extra_headers_for_endpoint(endpoint) 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(public_router) # 公开API端点（用户可查看提供商和模型）
app.include_router(monitoring_router) # 监控端点 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: 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) total_requests = Column(BigInteger, default=0, nullable=False)
avg_latency_ms = Column(Float, nullable=True) 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 可扩展） # 硬件信息（注册时上报，JSON 可扩展）
hardware_info = Column( hardware_info = Column(
JSON, JSON,
@@ -936,6 +930,15 @@ class ProxyNode(Base):
Integer, nullable=True, comment="基于硬件估算的最大并发连接数" 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） # 管理端远程配置（通过心跳下发给 aether-proxy）
remote_config = Column( remote_config = Column(
JSON, JSON,

17
src/models/provider.py
View File

@@ -270,12 +270,6 @@ class ProxyNode(Base):
total_requests = Column(BigInteger, default=0, nullable=False) total_requests = Column(BigInteger, default=0, nullable=False)
avg_latency_ms = Column(Float, nullable=True) 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 可扩展） # 硬件信息（注册时上报，JSON 可扩展）
hardware_info = Column( hardware_info = Column(
JSON, JSON,
@@ -286,11 +280,20 @@ class ProxyNode(Base):
Integer, nullable=True, comment="基于硬件估算的最大并发连接数" 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） # 管理端远程配置（通过心跳下发给 aether-proxy）
remote_config = Column( remote_config = Column(
JSON, JSON,
nullable=True, nullable=True,
comment="管理端下发的远程配置 (allowed_ports, log_level, heartbeat_interval, timestamp_tolerance)", comment="管理端下发的远程配置 (allowed_ports, log_level, heartbeat_interval)",
) )
config_version = Column( config_version = Column(
Integer, default=0, nullable=False, comment="远程配置版本号，每次更新 +1" Integer, default=0, nullable=False, comment="远程配置版本号，每次更新 +1"

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

@@ -2,7 +2,7 @@
代理节点模块 代理节点模块
提供海外 VPS 代理节点的注册、心跳、管理功能。 提供海外 VPS 代理节点的注册、心跳、管理功能。
aether-proxy 部署在海外 VPS 上自动注册节点，Aether 通过 HMAC 签名认证转发请求。 aether-proxy 部署在海外 VPS 上，通过 WebSocket 隧道连接 Aether 转发 API 请求。
""" """
from __future__ import annotations from __future__ import annotations
@@ -77,15 +77,7 @@ async def _health_check() -> ModuleHealth:
def _validate_config(db: Session) -> tuple[bool, str]: def _validate_config(db: Session) -> tuple[bool, str]:
""" """验证配置（tunnel 模式无需额外密钥配置）"""
验证配置
代理节点模块需要 PROXY_HMAC_KEY 配置
"""
from src.config.settings import config
if not config.proxy_hmac_key:
return False, "PROXY_HMAC_KEY 未配置（也未设置 ENCRYPTION_KEY 用于自动派生）"
return True, "" return True, ""
@@ -93,7 +85,7 @@ proxy_nodes_module = ModuleDefinition(
metadata=ModuleMetadata( metadata=ModuleMetadata(
name="proxy_nodes", name="proxy_nodes",
display_name="代理节点", display_name="代理节点",
description="海外 VPS 代理节点管理，通过 HMAC 签名认证转发 API 请求", description="海外 VPS 代理节点管理，通过 WebSocket 隧道转发 API 请求",
category=ModuleCategory.INTEGRATION, category=ModuleCategory.INTEGRATION,
env_key="PROXY_NODES_AVAILABLE", env_key="PROXY_NODES_AVAILABLE",
default_available=True, 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 self._last_error: str | None = None
# 代理配置（支持 proxy_node_id 和旧的 proxy URL） # 代理配置（支持 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._proxy: str | httpx.Proxy | None = resolve_ops_proxy(self.config)
self._tunnel_node_id: str | None = resolve_ops_tunnel_node_id(self.config)
# HTTP 客户端配置 # HTTP 客户端配置
self._timeout = self.config.get("timeout", 30) self._timeout = self.config.get("timeout", 30)
@@ -117,13 +118,18 @@ class ProviderConnector(ABC):
""" """
获取已认证的 HTTP 客户端 获取已认证的 HTTP 客户端
使用 context manager 确保资源正确释放 使用 context manager 确保资源正确释放。
tunnel 模式下使用 TunnelTransport 替代 proxy transport。
Yields: Yields:
已配置认证信息的 AsyncClient 已配置认证信息的 AsyncClient
""" """
transport = None 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) transport = httpx.AsyncHTTPTransport(proxy=self._proxy)
async with httpx.AsyncClient( 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 import httpx
from src.core.api_format.headers import BROWSER_FINGERPRINT_HEADERS
from src.services.provider_ops.actions import ( from src.services.provider_ops.actions import (
NewApiBalanceAction, NewApiBalanceAction,
ProviderAction, ProviderAction,
@@ -75,6 +76,10 @@ class NewApiConnector(ProviderConnector):
def _apply_auth(self, request: httpx.Request) -> httpx.Request: 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: if self._api_key:
request.headers["Authorization"] = f"Bearer {self._api_key}" request.headers["Authorization"] = f"Bearer {self._api_key}"
if self._user_id: if self._user_id:
@@ -206,7 +211,8 @@ class NewApiArchitecture(ProviderArchitecture):
New API 特有：需要 New-Api-User Header 传递用户 ID New API 特有：需要 New-Api-User Header 传递用户 ID
""" """
headers: dict[str, str] = {} # 以浏览器指纹 Headers 为基础，绕过 Cloudflare 等防护
headers: dict[str, str] = {**BROWSER_FINGERPRINT_HEADERS}
# Bearer Token 认证 # Bearer Token 认证
api_key = credentials.get("api_key", "") api_key = credentials.get("api_key", "")

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

@@ -1033,10 +1033,11 @@ class ProviderOpsService:
list(headers.keys()), list(headers.keys()),
) )
# 获取代理配置（支持 proxy_node_id 和旧的 proxy URL） # 获取代理配置（支持 proxy_node_id、tunnel 模式和旧的 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
proxy = resolve_ops_proxy(config) proxy = resolve_ops_proxy(config)
tunnel_node_id = resolve_ops_tunnel_node_id(config)
try: try:
# 构建 httpx client 参数 # 构建 httpx client 参数
@@ -1044,7 +1045,12 @@ class ProviderOpsService:
"timeout": 30.0, "timeout": 30.0,
"verify": get_ssl_context(), "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 client_kwargs["proxy"] = proxy
logger.debug("使用代理: {}", 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 .health_scheduler import ProxyNodeHealthScheduler, get_proxy_node_health_scheduler
from .resolver import ( from .resolver import (
build_delegate_post_kwargs,
build_delegate_stream_kwargs,
build_hmac_proxy_url,
build_post_kwargs, build_post_kwargs,
build_proxy_url, build_proxy_url,
build_stream_kwargs, build_stream_kwargs,
@@ -12,10 +9,12 @@ from .resolver import (
get_proxy_label, get_proxy_label,
get_system_proxy_config, get_system_proxy_config,
inject_auth_into_proxy_url, inject_auth_into_proxy_url,
invalidate_proxy_node_cache,
invalidate_system_proxy_cache, invalidate_system_proxy_cache,
make_proxy_param, make_proxy_param,
resolve_delegate_config, resolve_delegate_config,
resolve_ops_proxy, resolve_ops_proxy,
resolve_ops_tunnel_node_id,
resolve_proxy_info, resolve_proxy_info,
) )
from .service import ProxyNodeService, node_to_dict from .service import ProxyNodeService, node_to_dict
@@ -25,9 +24,6 @@ __all__ = [
"get_proxy_node_health_scheduler", "get_proxy_node_health_scheduler",
"ProxyNodeService", "ProxyNodeService",
"node_to_dict", "node_to_dict",
"build_delegate_post_kwargs",
"build_delegate_stream_kwargs",
"build_hmac_proxy_url",
"build_post_kwargs", "build_post_kwargs",
"build_proxy_url", "build_proxy_url",
"build_stream_kwargs", "build_stream_kwargs",
@@ -36,8 +32,10 @@ __all__ = [
"make_proxy_param", "make_proxy_param",
"get_proxy_label", "get_proxy_label",
"get_system_proxy_config", "get_system_proxy_config",
"invalidate_proxy_node_cache",
"invalidate_system_proxy_cache", "invalidate_system_proxy_cache",
"resolve_delegate_config", "resolve_delegate_config",
"resolve_ops_proxy", "resolve_ops_proxy",
"resolve_ops_tunnel_node_id",
"resolve_proxy_info", "resolve_proxy_info",
] ]

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

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

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

@@ -7,18 +7,13 @@
from __future__ import annotations from __future__ import annotations
import base64
import gzip as _gzip
import hashlib import hashlib
import hmac as _hmac
import json as _json
import time import time
from typing import Any from typing import Any
from urllib.parse import quote, urlparse from urllib.parse import quote, urlparse
import httpx import httpx
from src.config import config
from src.core.exceptions import ProxyNodeUnavailableError from src.core.exceptions import ProxyNodeUnavailableError
from src.core.logger import logger 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, "name": node.name,
"ip": node.ip, "ip": node.ip,
"port": node.port, "port": node.port,
"tls_enabled": bool(node.tls_enabled), "tunnel_mode": bool(node.tunnel_mode),
"tls_cert_fingerprint": node.tls_cert_fingerprint, "tunnel_connected": bool(node.tunnel_connected),
} }
_proxy_node_cache[node_id] = (value, now + _PROXY_NODE_CACHE_TTL_SECONDS) _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() 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 _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: def invalidate_system_proxy_cache() -> None:
"""手动失效系统代理缓存（在删除节点等操作后调用）""" """手动失效系统代理缓存（在删除节点等操作后调用）"""
global _system_proxy_cache 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( def resolve_ops_proxy(
connector_config: dict[str, Any] | None, connector_config: dict[str, Any] | None,
) -> str | httpx.Proxy | None: ) -> str | httpx.Proxy | None:
@@ -235,37 +227,50 @@ def resolve_ops_proxy(
2. connector_config.proxy（旧格式 URL 字符串） 2. connector_config.proxy（旧格式 URL 字符串）
3. 系统默认代理节点 3. 系统默认代理节点
tunnel 模式节点不返回代理 URL（由 resolve_ops_tunnel_node_id 处理）。
Args: Args:
connector_config: connector 的 config 字典 connector_config: connector 的 config 字典
Returns: Returns:
httpx 可接受的代理参数（str 或 httpx.Proxy），或 None httpx 可接受的代理参数（str 或 httpx.Proxy），或 None
""" """
if connector_config: from .tunnel_transport import is_tunnel_node
# 新格式：proxy_node_id -> 通过 build_proxy_url 解析
node_id = connector_config.get("proxy_node_id")
if isinstance(node_id, str) and node_id.strip():
try:
url = build_proxy_url({"node_id": node_id.strip(), "enabled": True})
return make_proxy_param(url)
except Exception as exc:
logger.warning("解析 proxy_node_id={} 失败，回退到直连: {}", node_id, exc)
return None
# 旧格式：直接返回 proxy URL 字符串 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, "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") proxy = connector_config.get("proxy")
if isinstance(proxy, str) and proxy.strip(): if isinstance(proxy, str) and proxy.strip():
return proxy return proxy
# 回退：系统默认代理 return None
system_proxy = get_system_proxy_config()
if system_proxy:
try: def resolve_ops_tunnel_node_id(
url = build_proxy_url(system_proxy) connector_config: dict[str, Any] | None,
return make_proxy_param(url) ) -> str | None:
except Exception as exc: """
logger.warning("构建系统默认代理 URL 失败: {}", exc) 解析 ops connector 的 tunnel 节点 ID
return None
如果配置的代理节点是 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 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 inject_auth_into_proxy_url(manual_url, username, password)
return manual_url return manual_url
# aether-proxy 节点：使用 HMAC 认证 # tunnel 模式节点：不构建 proxy URL（通过 TunnelTransport 处理）
return build_hmac_proxy_url( from .tunnel_transport import is_tunnel_node
node_info["ip"],
node_info["port"], if is_tunnel_node(node_info):
tls_enabled=node_info.get("tls_enabled", False), return None
)
# aether-proxy 节点均为 tunnel 模式，不应走到这里
logger.warning("非 tunnel 模式的 aether-proxy 节点不再支持: node_id={}", node_id)
return None
proxy_url: str | None = proxy_config.get("url") proxy_url: str | None = proxy_config.get("url")
if not proxy_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): if not proxy_config.get("enabled", True):
return "__no_proxy__" return "__no_proxy__"
# ProxyNode 模式：基于 node_id + 时间桶缓存，避免签名随时间变化导致 cache key 爆炸 # ProxyNode 模式：基于 node_id 缓存
node_id = proxy_config.get("node_id") node_id = proxy_config.get("node_id")
if isinstance(node_id, str) and node_id.strip(): 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()}"
return f"proxy_node:{node_id.strip()}:{time_bucket}"
# 构建代理 URL 作为缓存键的基础 # 构建代理 URL 作为缓存键的基础
proxy_url = build_proxy_url(proxy_config) 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: 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()`` 闭包在每次请求 / 重试时 tunnel 模式节点返回 {"tunnel": True, "node_id": str}，
获取新鲜的 HMAC 签名，避免长生命周期内时间戳过期。 调用方应使用 TunnelTransport。
Returns: Returns:
{"delegate_url": str, "node_id": str, "tls_enabled": bool, {"tunnel": True, "node_id": str} 或 None
"auth_header": str, # 首次生成的签名（兼容旧调用）
"fresh_auth_header": Callable} # 延迟生成签名的闭包
或 None
""" """
effective_config = proxy_config 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") node_id = effective_config.get("node_id")
if not isinstance(node_id, str) or not node_id.strip(): if not isinstance(node_id, str) or not node_id.strip():
return None # 旧格式 URL 模式不支持代发 return None
node_id = node_id.strip() node_id = node_id.strip()
node_info = _get_proxy_node_info(node_id) node_info = _get_proxy_node_info(node_id)
if not node_info or node_info.get("is_manual"): if not node_info or node_info.get("is_manual"):
return None # 手动节点不支持代发 return None
tls_enabled = node_info.get("tls_enabled", False) from .tunnel_transport import is_tunnel_node
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"
# 每次调用生成新鲜签名（避免长连接内时间戳过期） if is_tunnel_node(node_info):
def _fresh() -> str: return {"tunnel": True, "node_id": node_id}
return _build_hmac_auth_header()
return { return None
"delegate_url": delegate_url,
"auth_header": _fresh(), # 立即生成一份，兼容旧调用方
"fresh_auth_header": _fresh,
"node_id": node_id,
"tls_enabled": tls_enabled,
}
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# 代发请求参数构建（消除 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( def build_post_kwargs(
delegate_cfg: dict[str, Any] | None, _delegate_cfg: dict[str, Any] | None = None,
*, *,
url: str, url: str,
headers: dict[str, str], headers: dict[str, str],
@@ -721,19 +582,13 @@ def build_post_kwargs(
refresh_auth: bool = False, refresh_auth: bool = False,
) -> dict[str, Any]: ) -> dict[str, Any]:
""" """
构建上游 POST 请求的 httpx kwargs（自动选择代发或直连模式） 构建上游 POST 请求的 httpx kwargs
返回的 dict 可直接传给 ``http_client.post(**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 { return {
"url": url, "url": url,
"json": payload, "json": payload,
@@ -743,7 +598,7 @@ def build_post_kwargs(
def build_stream_kwargs( def build_stream_kwargs(
delegate_cfg: dict[str, Any] | None, _delegate_cfg: dict[str, Any] | None = None,
*, *,
url: str, url: str,
headers: dict[str, str], headers: dict[str, str],
@@ -751,21 +606,13 @@ def build_stream_kwargs(
timeout: float | None = None, timeout: float | None = None,
) -> dict[str, Any]: ) -> dict[str, Any]:
""" """
构建上游 stream 请求的 httpx kwargs（自动选择代发或直连模式） 构建上游 stream 请求的 httpx kwargs
返回的 dict 可直接传给 ``http_client.stream(**kwargs)``。 返回的 dict 可直接传给 ``http_client.stream(**kwargs)``。
当 ``timeout`` 为 None 时由外层 asyncio.wait_for 控制超时。
当 ``timeout`` 为 None（直连模式下由外层 asyncio.wait_for 控制超时）， ``_delegate_cfg`` 已废弃，保留仅为兼容现有调用方签名。
直连分支不设置 timeout；代发分支始终携带 timeout（proxy 协议需要）。
""" """
if delegate_cfg:
return build_delegate_stream_kwargs(
delegate_cfg,
url=url,
headers=headers,
payload=payload,
timeout=timeout or 60,
)
kwargs: dict[str, Any] = { kwargs: dict[str, Any] = {
"method": "POST", "method": "POST",
"url": url, "url": url,

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

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