feat(proxy): 安全加固与架构优化

- 引入 SafeDnsResolver 消除 DNS rebinding TOCTTOU 漏洞,DNS 缓存改为多地址存储
- 扩展私有 IP 检测范围(CGNAT 100.64/10、基准测试 198.18/15、保留 240/4)
- 请求处理增加 hop-by-hop 头过滤、URL scheme 校验、超时范围限制
- 动态配置从 RwLock 切换到 ArcSwap 实现无锁读取
- 启动注册失败的服务器支持后台自动重试
- WebSocket 帧大小上限提升至 64MiB 匹配 Python 端
- 心跳支持动态间隔更新,新增 failed_requests/dns_failures/stream_errors 指标
- 配置启动校验、systemd UMask=0077、配置文件权限 600
- Python 端支持 per-connection max_streams(X-Tunnel-Max-Streams)
This commit is contained in:
fawney19
2026-02-28 01:32:28 +08:00
parent 2a0c684e88
commit e748277902
22 changed files with 746 additions and 149 deletions

View File

@@ -13,6 +13,7 @@ name = "aether-proxy"
version = "0.2.0"
dependencies = [
"anyhow",
"arc-swap",
"base64",
"bytes",
"clap",
@@ -111,6 +112,15 @@ version = "1.0.101"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5f0e0fee31ef5ed1ba1316088939cea399010ed7731dba877ed44aeb407a75ea"
[[package]]
name = "arc-swap"
version = "1.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f9f3647c145568cec02c42054e07bdf9a5a698e15b466fb2341bfc393cd24aa5"
dependencies = [
"rustversion",
]
[[package]]
name = "atomic"
version = "0.6.1"

View File

@@ -20,6 +20,7 @@ bytes = "1"
sha2 = "0.10"
hex = "0.4"
anyhow = "1"
arc-swap = "1"
toml = "0.8"
rustls = { version = "0.23", features = ["ring"] }
ratatui = "0.30"

View File

@@ -7,6 +7,4 @@ RUN apt-get update && apt-get install -y --no-install-recommends ca-certificates
COPY build/linux-${TARGETARCH}/aether-proxy /usr/local/bin/aether-proxy
EXPOSE 18080
ENTRYPOINT ["aether-proxy"]

View File

@@ -1,18 +1,16 @@
# aether-proxy
Aether 正向代理节点,部署在海外 VPS 上,为墙内的 Aether 实例中转 API 流量。
Aether Tunnel 代理节点,部署在海外 VPS 上,通过 WebSocket 隧道为 Aether 实例中转 API 流量。
Tunnel 模式下代理节点**无需对外监听端口**,仅需出站连接到 Aether 服务器。
## 安装
### Docker Compose 部署
```bash
# 拉取镜像
docker pull ghcr.io/fawney19/aether-proxy:latest
# 或使用 docker compose
cp .env.example .env
# 编辑 .env 填入 AETHER_PROXY_AETHER_URL, MANAGEMENT_TOKEN, HMAC_KEY
# 编辑 .env 填入 AETHER_PROXY_AETHER_URL 和 AETHER_PROXY_MANAGEMENT_TOKEN
docker compose up -d
```
@@ -69,43 +67,79 @@ sudo aether-proxy uninstall
### 参数一览
#### 基础配置
| 参数 | 环境变量 | 默认值 | 说明 |
|------|----------|--------|------|
| `--aether-url` | `AETHER_PROXY_AETHER_URL` | **必填** | Aether 服务器地址 |
| `--management-token` | `AETHER_PROXY_MANAGEMENT_TOKEN` | **必填** | 管理员 Token`ae_xxx` 格式) |
| `--hmac-key` | `AETHER_PROXY_HMAC_KEY` | **必填** | HMAC 密钥,需与 Aether 端一致 |
| `--listen-port` | `AETHER_PROXY_LISTEN_PORT` | `18080` | 监听端口 |
| `--public-ip` | `AETHER_PROXY_PUBLIC_IP` | 自动检测 | 公网 IP |
| `--node-name` | `AETHER_PROXY_NODE_NAME` | `proxy-01` | 节点名称标识 |
| `--node-region` | `AETHER_PROXY_NODE_REGION` | 自动检测 | 地区标识 |
| `--heartbeat-interval` | `AETHER_PROXY_HEARTBEAT_INTERVAL` | `30` | 心跳间隔(秒) |
| `--allowed-ports` | `AETHER_PROXY_ALLOWED_PORTS` | `80,443,8080,8443` | 允许代理的目标端口 |
| `--timestamp-tolerance` | `AETHER_PROXY_TIMESTAMP_TOLERANCE` | `300` | HMAC 时间戳容差(秒) |
| `--aether-request-timeout` | `AETHER_PROXY_AETHER_REQUEST_TIMEOUT` | `10` | Aether API 请求总超时(秒) |
| `--aether-connect-timeout` | `AETHER_PROXY_AETHER_CONNECT_TIMEOUT` | `10` | Aether API 建连超时(秒) |
| `--aether-pool-max-idle-per-host` | `AETHER_PROXY_AETHER_POOL_MAX_IDLE_PER_HOST` | `8` | Aether API 每 Host 最大空闲连接数 |
| `--aether-pool-idle-timeout` | `AETHER_PROXY_AETHER_POOL_IDLE_TIMEOUT` | `90` | Aether API 连接池空闲超时(秒) |
| `--aether-tcp-keepalive` | `AETHER_PROXY_AETHER_TCP_KEEPALIVE` | `60` | Aether API TCP keepalive0 关闭) |
| `--aether-tcp-nodelay` | `AETHER_PROXY_AETHER_TCP_NODELAY` | `true` | Aether API 启用 TCP_NODELAY |
| `--aether-http2` | `AETHER_PROXY_AETHER_HTTP2` | `true` | Aether API 启用 HTTP/2 |
| `--aether-retry-max-attempts` | `AETHER_PROXY_AETHER_RETRY_MAX_ATTEMPTS` | `3` | Aether API 最大重试次数(含首次 |
| `--aether-retry-base-delay-ms` | `AETHER_PROXY_AETHER_RETRY_BASE_DELAY_MS` | `200` | Aether API 重试基础延迟(毫秒) |
| `--aether-retry-max-delay-ms` | `AETHER_PROXY_AETHER_RETRY_MAX_DELAY_MS` | `2000` | Aether API 重试最大延迟(毫秒) |
| `--max-concurrent-connections` | `AETHER_PROXY_MAX_CONCURRENT_CONNECTIONS` | 自动估算 | 最大并发连接数(默认硬件估算 |
| `--connect-timeout` | `AETHER_PROXY_CONNECT_TIMEOUT` | `30` | CONNECT 上游建连超时(秒) |
| `--tls-handshake-timeout` | `AETHER_PROXY_TLS_HANDSHAKE_TIMEOUT` | `10` | TLS 握手超时(秒) |
| `--dns-cache-ttl` | `AETHER_PROXY_DNS_CACHE_TTL` | `60` | DNS 缓存 TTL |
#### Tunnel 连接
| 参数 | 环境变量 | 默认值 | 说明 |
|------|----------|--------|------|
| `--tunnel-connections` | `AETHER_PROXY_TUNNEL_CONNECTIONS` | `3` | Aether 的连接池大小 |
| `--tunnel-max-streams` | `AETHER_PROXY_TUNNEL_MAX_STREAMS` | 自动(硬件估算) | 单连接最大并发 stream 数 |
| `--tunnel-connect-timeout-secs` | `AETHER_PROXY_TUNNEL_CONNECT_TIMEOUT_SECS` | `15` | TCP + TLS 握手超时(秒) |
| `--tunnel-tcp-keepalive-secs` | `AETHER_PROXY_TUNNEL_TCP_KEEPALIVE_SECS` | `30` | TCP keepalive 初始延迟(秒 |
| `--tunnel-tcp-nodelay` | `AETHER_PROXY_TUNNEL_TCP_NODELAY` | `true` | 禁用 Nagle 算法 |
| `--tunnel-ping-interval-secs` | `AETHER_PROXY_TUNNEL_PING_INTERVAL_SECS` | `15` | WebSocket Ping 频率(秒) |
| `--tunnel-stale-timeout-secs` | `AETHER_PROXY_TUNNEL_STALE_TIMEOUT_SECS` | `45` | 无数据断连阈值(秒 |
| `--tunnel-reconnect-base-ms` | `AETHER_PROXY_TUNNEL_RECONNECT_BASE_MS` | `500` | 指数退避基础延迟(毫秒) |
| `--tunnel-reconnect-max-ms` | `AETHER_PROXY_TUNNEL_RECONNECT_MAX_MS` | `30000` | 指数退避上限(毫秒) |
#### 上游 HTTP 请求
| 参数 | 环境变量 | 默认值 | 说明 |
|------|----------|--------|------|
| `--upstream-connect-timeout-secs` | `AETHER_PROXY_UPSTREAM_CONNECT_TIMEOUT_SECS` | `30` | 上游建连超时(秒) |
| `--upstream-pool-max-idle-per-host` | `AETHER_PROXY_UPSTREAM_POOL_MAX_IDLE_PER_HOST` | `64` | 每 Host 最大空闲连接数 |
| `--upstream-pool-idle-timeout-secs` | `AETHER_PROXY_UPSTREAM_POOL_IDLE_TIMEOUT_SECS` | `300` | 连接池空闲超时(秒) |
| `--upstream-tcp-keepalive-secs` | `AETHER_PROXY_UPSTREAM_TCP_KEEPALIVE_SECS` | `60` | TCP keepalive0 关闭) |
| `--upstream-tcp-nodelay` | `AETHER_PROXY_UPSTREAM_TCP_NODELAY` | `true` | 启用 TCP_NODELAY |
#### Aether API 客户端
| 参数 | 环境变量 | 默认值 | 说明 |
|------|----------|--------|------|
| `--aether-request-timeout-secs` | `AETHER_PROXY_AETHER_REQUEST_TIMEOUT_SECS` | `10` | 请求总超时(秒) |
| `--aether-connect-timeout-secs` | `AETHER_PROXY_AETHER_CONNECT_TIMEOUT_SECS` | `10` | 建连超时(秒) |
| `--aether-retry-max-attempts` | `AETHER_PROXY_AETHER_RETRY_MAX_ATTEMPTS` | `3` | 最大重试次数 |
#### DNS 与安全
| 参数 | 环境变量 | 默认值 | 说明 |
|------|----------|--------|------|
| `--dns-cache-ttl-secs` | `AETHER_PROXY_DNS_CACHE_TTL_SECS` | `60` | DNS 缓存 TTL |
| `--dns-cache-capacity` | `AETHER_PROXY_DNS_CACHE_CAPACITY` | `1024` | DNS 缓存容量(条目数) |
| `--delegate-connect-timeout` | `AETHER_PROXY_DELEGATE_CONNECT_TIMEOUT` | `30` | delegate 上游建连超时(秒) |
| `--delegate-pool-max-idle-per-host` | `AETHER_PROXY_DELEGATE_POOL_MAX_IDLE_PER_HOST` | `64` | delegate 每 Host 最大空闲连接数 |
| `--delegate-pool-idle-timeout` | `AETHER_PROXY_DELEGATE_POOL_IDLE_TIMEOUT` | `300` | delegate 连接池空闲超时(秒) |
| `--delegate-tcp-keepalive` | `AETHER_PROXY_DELEGATE_TCP_KEEPALIVE` | `60` | delegate TCP keepalive0 关闭) |
| `--delegate-tcp-nodelay` | `AETHER_PROXY_DELEGATE_TCP_NODELAY` | `true` | delegate 启用 TCP_NODELAY |
#### 日志
| 参数 | 环境变量 | 默认值 | 说明 |
|------|----------|--------|------|
| `--log-level` | `AETHER_PROXY_LOG_LEVEL` | `info` | 日志级别 |
| `--log-json` | `AETHER_PROXY_LOG_JSON` | `false` | JSON 格式日志 |
| `--enable-tls` | `AETHER_PROXY_ENABLE_TLS` | `true` | 启用 TLS |
| `--tls-cert` | `AETHER_PROXY_TLS_CERT` | `aether-proxy-cert.pem` | TLS 证书路径 |
| `--tls-key` | `AETHER_PROXY_TLS_KEY` | `aether-proxy-key.pem` | TLS 私钥路径 |
### 多服务器配置
`aether-proxy.toml` 中使用 `[[servers]]` 配置多个 Aether 服务器:
```toml
[[servers]]
aether_url = "https://aether-1.example.com"
management_token = "ae_xxx"
node_name = "jp-proxy-01"
[[servers]]
aether_url = "https://aether-2.example.com"
management_token = "ae_yyy"
node_name = "jp-proxy-02"
```
## 发布新版本
@@ -115,6 +149,6 @@ sudo aether-proxy uninstall
- 更新 README 中的下载链接表格
```bash
git tag proxy-v0.1.0
git push origin proxy-v0.1.0
git tag proxy-v0.2.0
git push origin proxy-v0.2.0
```

View File

@@ -3,12 +3,9 @@ services:
image: ghcr.io/fawney19/aether-proxy:latest
container_name: aether-proxy
restart: unless-stopped
ports:
- "${AETHER_PROXY_LISTEN_PORT:-18080}:18080"
env_file:
- .env
environment:
AETHER_PROXY_LISTEN_PORT: 18080
AETHER_PROXY_LOG_JSON: "true"
logging:
driver: json-file

View File

@@ -4,19 +4,22 @@ use std::sync::atomic::AtomicU64;
use std::sync::{Arc, RwLock};
use std::time::Duration;
use arc_swap::ArcSwap;
use tokio::signal;
use tokio::sync::watch;
use tokio::sync::{watch, Mutex};
use tracing::{error, info, warn};
use crate::config::{Config, ServerEntry};
use crate::net;
use crate::registration::client::AetherClient;
use crate::runtime::{self, DynamicConfig};
use crate::safe_dns::SafeDnsResolver;
use crate::state::{AppState, ProxyMetrics, ServerContext};
use crate::{hardware, target_filter, tunnel};
/// Run the full application lifecycle after config has been parsed.
pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Result<()> {
config.validate()?;
init_tracing(&config);
info!(
@@ -65,7 +68,12 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
));
// Build reqwest client for tunnel upstream requests (shared).
// Inject SafeDnsResolver so reqwest only connects to addresses that were
// validated by validate_target() — this eliminates the DNS rebinding
// TOCTTOU gap where a second DNS lookup could return a private IP.
let safe_resolver = SafeDnsResolver::new(Arc::clone(&dns_cache));
let mut reqwest_builder = reqwest::Client::builder()
.dns_resolver(Arc::new(safe_resolver))
.pool_max_idle_per_host(config.upstream_pool_max_idle_per_host)
.pool_idle_timeout(Duration::from_secs(config.upstream_pool_idle_timeout_secs))
.connect_timeout(Duration::from_secs(config.upstream_connect_timeout_secs))
@@ -81,8 +89,10 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
.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();
// Register with each Aether server and build per-server contexts.
// Wrapped in Arc<Mutex> so retry_failed_registrations can append later.
let server_contexts: Arc<Mutex<Vec<Arc<ServerContext>>>> = Arc::new(Mutex::new(Vec::new()));
let mut failed_entries: Vec<(String, ServerEntry)> = Vec::new();
for (i, entry) in servers.iter().enumerate() {
let label = if servers.len() == 1 {
"server".to_string()
@@ -104,14 +114,18 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
{
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 {
// Initialize dynamic config with per-server node_name (not global),
// so that the heartbeat and reconnect use the correct name.
let mut dynamic = DynamicConfig::from_config(&config);
dynamic.node_name = node_name.clone();
server_contexts.lock().await.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))),
dynamic: Arc::new(ArcSwap::from_pointee(dynamic)),
active_connections: Arc::new(AtomicU64::new(0)),
metrics: Arc::new(ProxyMetrics::new()),
}));
@@ -121,14 +135,24 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
server = %label,
url = %entry.aether_url,
error = %e,
"registration failed, skipping server"
"registration failed, will retry in background"
);
failed_entries.push((label, entry.clone()));
}
}
}
if server_contexts.is_empty() {
anyhow::bail!("no servers registered successfully");
{
let ctx_count = server_contexts.lock().await.len();
if ctx_count == 0 && failed_entries.is_empty() {
anyhow::bail!("no servers configured");
}
if ctx_count == 0 {
anyhow::bail!(
"no servers registered successfully (all {} failed)",
failed_entries.len()
);
}
}
// Build shared application state
@@ -144,14 +168,14 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
let (shutdown_tx, shutdown_rx) = watch::channel(false);
info!(
active_servers = server_contexts.len(),
active_servers = server_contexts.lock().await.len(),
"running in tunnel mode"
);
// Spawn tunnel connections per server (pool_size connections each)
let pool_size = state.config.tunnel_connections.max(1) as usize;
let mut tunnel_handles = Vec::new();
for server in &server_contexts {
for server in server_contexts.lock().await.iter() {
for conn_idx in 0..pool_size {
let s = Arc::clone(&state);
let srv = Arc::clone(server);
@@ -162,13 +186,35 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
}
}
// Spawn background retry for failed server registrations
if !failed_entries.is_empty() {
let retry_state = Arc::clone(&state);
let retry_contexts = Arc::clone(&server_contexts);
let retry_public_ip = public_ip.clone();
let retry_hw_info = hw_info.clone();
let retry_shutdown = shutdown_rx.clone();
let retry_pool_size = pool_size;
tokio::spawn(async move {
retry_failed_registrations(
retry_state,
retry_contexts,
failed_entries,
retry_public_ip,
retry_hw_info,
retry_pool_size,
retry_shutdown,
)
.await;
});
}
// Wait for shutdown signal
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 {
// Graceful unregister from all servers (including retry-registered ones)
for server in server_contexts.lock().await.iter() {
let node_id = server.node_id.read().unwrap().clone();
if let Err(e) = server.aether_client.unregister(&node_id).await {
error!(
@@ -188,6 +234,97 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
Ok(())
}
/// Retry interval for failed server registrations (5 minutes).
const REGISTRATION_RETRY_INTERVAL: Duration = Duration::from_secs(300);
/// Max registration retry attempts before giving up.
const REGISTRATION_RETRY_MAX: u32 = 12;
/// Background task that retries registration for servers that failed at startup.
async fn retry_failed_registrations(
state: Arc<AppState>,
server_contexts: Arc<Mutex<Vec<Arc<ServerContext>>>>,
failed: Vec<(String, ServerEntry)>,
public_ip: String,
hw_info: crate::hardware::HardwareInfo,
pool_size: usize,
mut shutdown: watch::Receiver<bool>,
) {
for (label, entry) in &failed {
let node_name = entry
.node_name
.clone()
.unwrap_or_else(|| state.config.node_name.clone());
let client = Arc::new(AetherClient::new(
&state.config,
&entry.aether_url,
&entry.management_token,
));
let mut attempt = 0u32;
let node_id = loop {
attempt += 1;
tokio::select! {
_ = tokio::time::sleep(REGISTRATION_RETRY_INTERVAL) => {}
_ = shutdown.changed() => {
info!(server = %label, "shutdown during registration retry");
return;
}
}
match client
.register(&state.config, &node_name, &public_ip, Some(&hw_info))
.await
{
Ok(id) => {
info!(server = %label, node_id = %id, attempt, "registration retry succeeded");
break id;
}
Err(e) => {
warn!(
server = %label,
attempt,
max = REGISTRATION_RETRY_MAX,
error = %e,
"registration retry failed"
);
if attempt >= REGISTRATION_RETRY_MAX {
error!(server = %label, "giving up registration after {} attempts", attempt);
return;
}
}
}
};
// Build server context and spawn tunnels
let mut dynamic = DynamicConfig::from_config(&state.config);
dynamic.node_name = node_name.clone();
let server = Arc::new(ServerContext {
server_label: label.clone(),
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(ArcSwap::from_pointee(dynamic)),
active_connections: Arc::new(AtomicU64::new(0)),
metrics: Arc::new(ProxyMetrics::new()),
});
// Add to shared list so shutdown can unregister this server
server_contexts.lock().await.push(Arc::clone(&server));
for conn_idx in 0..pool_size {
let s = Arc::clone(&state);
let srv = Arc::clone(&server);
let rx = shutdown.clone();
tokio::spawn(async move {
tunnel::run(&s, &srv, conn_idx, rx).await;
});
}
}
}
fn init_tracing(config: &Config) {
use tracing_subscriber::prelude::*;
use tracing_subscriber::{reload, EnvFilter};

View File

@@ -251,6 +251,57 @@ pub struct Config {
pub tunnel_connections: u32,
}
impl Config {
/// Validate configuration values are within sane ranges.
/// Called after parsing to catch misconfigurations early.
pub fn validate(&self) -> anyhow::Result<()> {
if self.heartbeat_interval == 0 {
anyhow::bail!("heartbeat_interval must be > 0");
}
if self.heartbeat_interval > 3600 {
anyhow::bail!("heartbeat_interval must be <= 3600");
}
if self.allowed_ports.is_empty() {
anyhow::bail!("allowed_ports must not be empty");
}
for &port in &self.allowed_ports {
if port == 0 {
anyhow::bail!("allowed_ports: port 0 is not valid");
}
}
if self.tunnel_connect_timeout_secs == 0 {
anyhow::bail!("tunnel_connect_timeout_secs must be > 0");
}
if self.tunnel_ping_interval_secs == 0 {
anyhow::bail!("tunnel_ping_interval_secs must be > 0");
}
if self.tunnel_stale_timeout_secs <= self.tunnel_ping_interval_secs {
anyhow::bail!(
"tunnel_stale_timeout_secs ({}) must be > tunnel_ping_interval_secs ({})",
self.tunnel_stale_timeout_secs,
self.tunnel_ping_interval_secs
);
}
if self.tunnel_connections == 0 {
anyhow::bail!("tunnel_connections must be > 0");
}
if self.aether_retry_max_attempts == 0 {
anyhow::bail!("aether_retry_max_attempts must be >= 1");
}
if self.tunnel_reconnect_base_ms > self.tunnel_reconnect_max_ms {
anyhow::bail!(
"tunnel_reconnect_base_ms ({}) must be <= tunnel_reconnect_max_ms ({})",
self.tunnel_reconnect_base_ms,
self.tunnel_reconnect_max_ms
);
}
if self.upstream_connect_timeout_secs == 0 {
anyhow::bail!("upstream_connect_timeout_secs must be > 0");
}
Ok(())
}
}
/// Per-server connection config (used in multi-server TOML `[[servers]]`).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ServerEntry {
@@ -490,16 +541,16 @@ impl ConfigFile {
let first_server = self.servers.first();
let aether_url = self
.aether_url
.clone()
.or_else(|| first_server.map(|s| s.aether_url.clone()));
.as_deref()
.or(first_server.map(|s| s.aether_url.as_str()));
let management_token = self
.management_token
.clone()
.or_else(|| first_server.map(|s| s.management_token.clone()));
.as_deref()
.or(first_server.map(|s| s.management_token.as_str()));
let node_name = self
.node_name
.clone()
.or_else(|| first_server.and_then(|s| s.node_name.clone()));
.as_deref()
.or(first_server.and_then(|s| s.node_name.as_deref()));
set!("AETHER_PROXY_AETHER_URL", aether_url);
set!("AETHER_PROXY_MANAGEMENT_TOKEN", management_token);

View File

@@ -4,6 +4,7 @@ mod hardware;
mod net;
mod registration;
mod runtime;
mod safe_dns;
mod setup;
mod state;
mod target_filter;

View File

@@ -5,17 +5,18 @@
//! management backend through the heartbeat response.
use std::collections::HashSet;
use std::sync::{Arc, OnceLock, RwLock};
use std::sync::{Arc, OnceLock};
use arc_swap::ArcSwap;
use tracing::info;
use crate::config::Config;
/// Configuration that can be changed at runtime without restart.
#[derive(Debug)]
#[derive(Debug, Clone)]
pub struct DynamicConfig {
pub node_name: String,
pub allowed_ports: HashSet<u16>,
pub allowed_ports: Arc<HashSet<u16>>,
pub log_level: String,
pub heartbeat_interval: u64,
/// Monotonically increasing version from the backend.
@@ -28,7 +29,7 @@ impl DynamicConfig {
pub fn from_config(config: &Config) -> Self {
Self {
node_name: config.node_name.clone(),
allowed_ports: config.allowed_ports.iter().copied().collect(),
allowed_ports: Arc::new(config.allowed_ports.iter().copied().collect()),
log_level: config.log_level.clone(),
heartbeat_interval: config.heartbeat_interval,
config_version: 0,
@@ -36,10 +37,10 @@ impl DynamicConfig {
}
}
/// Shared dynamic config handle.
pub type SharedDynamicConfig = Arc<RwLock<DynamicConfig>>;
/// Shared dynamic config handle (lock-free reads via ArcSwap).
pub type SharedDynamicConfig = Arc<ArcSwap<DynamicConfig>>;
// ── Log-level hot-reload ─────────────────────────────────────────────────────
// -- Log-level hot-reload -----
/// Global log-level reloader function, set during tracing init.
type LogReloader = Box<dyn Fn(&str) + Send + Sync>;
@@ -53,46 +54,50 @@ pub fn set_log_reloader(f: LogReloader) {
/// Apply a remote config update to the dynamic config.
///
/// Uses copy-on-write: loads the current snapshot, clones it, applies changes,
/// and stores the new Arc. Reads are always lock-free.
///
/// Returns `true` if the config was actually changed.
pub fn apply_remote_config(
dynamic: &SharedDynamicConfig,
remote: &crate::registration::client::RemoteConfig,
version: u64,
) -> bool {
let mut cfg = dynamic.write().unwrap();
let current = dynamic.load();
if version <= cfg.config_version {
if version <= current.config_version {
return false;
}
let mut new_cfg = (**current).clone();
let mut changed = Vec::new();
if let Some(ref name) = remote.node_name {
if *name != cfg.node_name {
changed.push(format!("node_name {}", name));
cfg.node_name = name.clone();
if *name != new_cfg.node_name {
changed.push(format!("node_name -> {}", name));
new_cfg.node_name = name.clone();
}
}
if let Some(ref ports) = remote.allowed_ports {
let new_set: HashSet<u16> = ports.iter().copied().collect();
if new_set != cfg.allowed_ports {
if new_set != *new_cfg.allowed_ports {
changed.push(format!("allowed_ports -> {:?}", ports));
cfg.allowed_ports = new_set;
new_cfg.allowed_ports = Arc::new(new_set);
}
}
if let Some(interval) = remote.heartbeat_interval {
if interval != cfg.heartbeat_interval {
changed.push(format!("heartbeat_interval {}s", interval));
cfg.heartbeat_interval = interval;
if interval != new_cfg.heartbeat_interval {
changed.push(format!("heartbeat_interval -> {}s", interval));
new_cfg.heartbeat_interval = interval;
}
}
if let Some(ref level) = remote.log_level {
if *level != cfg.log_level {
changed.push(format!("log_level {}", level));
cfg.log_level = level.clone();
if *level != new_cfg.log_level {
changed.push(format!("log_level -> {}", level));
new_cfg.log_level = level.clone();
// Hot-reload tracing filter
if let Some(reloader) = LOG_RELOADER.get() {
reloader(level);
@@ -100,15 +105,17 @@ pub fn apply_remote_config(
}
}
cfg.config_version = version;
let has_changes = !changed.is_empty();
if !changed.is_empty() {
if has_changes {
new_cfg.config_version = version;
info!(
version,
changes = %changed.join(", "),
"remote config applied"
);
dynamic.store(Arc::new(new_cfg));
}
!changed.is_empty()
has_changes
}

View File

@@ -0,0 +1,66 @@
//! Safe DNS resolver for reqwest that reuses validated addresses from DnsCache.
//!
//! This resolver ensures reqwest connects only to addresses that have been
//! previously validated by `target_filter::validate_target()`, eliminating
//! the TOCTTOU gap where DNS rebinding could redirect traffic to private IPs.
use std::net::SocketAddr;
use std::sync::Arc;
use reqwest::dns::{Addrs, Name, Resolve, Resolving};
use crate::target_filter::{self, DnsCache};
/// A DNS resolver that serves validated public addresses from the shared DnsCache.
///
/// When reqwest needs to resolve a hostname, this resolver returns addresses
/// from the cache (populated by `validate_target()` during request validation).
/// If the hostname is not in cache (shouldn't happen in normal flow), it
/// performs a fresh resolution with private-IP filtering.
pub struct SafeDnsResolver {
dns_cache: Arc<DnsCache>,
}
impl SafeDnsResolver {
pub fn new(dns_cache: Arc<DnsCache>) -> Self {
Self { dns_cache }
}
}
impl Resolve for SafeDnsResolver {
fn resolve(&self, name: Name) -> Resolving {
let dns_cache = Arc::clone(&self.dns_cache);
Box::pin(async move {
let host = name.as_str();
// Try cache first (should be populated by validate_target).
// reqwest resolves by hostname only (no port), so use host-only lookup.
if let Some(addrs) = dns_cache.get_by_host(host).await {
let socket_addrs: Vec<SocketAddr> = (*addrs).clone();
return Ok(Box::new(socket_addrs.into_iter()) as Addrs);
}
// Fallback: resolve with private-IP filtering (defensive).
// This path should rarely be hit since validate_target() runs first.
// We don't know the real port here (reqwest Resolve only gives hostname),
// so resolve directly without caching to avoid polluting the cache with
// an incorrect port-based key.
let addr_str = format!("{}:0", host);
let resolved: Vec<SocketAddr> = tokio::net::lookup_host(&addr_str)
.await
.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> { Box::new(e) })?
.filter(|addr| !target_filter::is_private_ip(&addr.ip()))
.collect();
if resolved.is_empty() {
return Err(Box::new(std::io::Error::other(format!(
"all resolved addresses for {} are private/reserved",
host
)))
as Box<dyn std::error::Error + Send + Sync>);
}
Ok(Box::new(resolved.into_iter()) as Addrs)
})
}
}

View File

@@ -67,6 +67,7 @@ pub fn install_service(config_path: &Path) -> anyhow::Result<()> {
Restart=on-failure\n\
RestartSec=5\n\
LimitNOFILE=65535\n\
UMask=0077\n\
\n\
[Install]\n\
WantedBy=multi-user.target\n",

View File

@@ -296,6 +296,13 @@ impl App {
fn save(&mut self) -> anyhow::Result<()> {
let cfg = self.to_config();
cfg.save(&self.config_path)?;
// Restrict config file permissions to owner-only (contains management token).
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let _ =
std::fs::set_permissions(&self.config_path, std::fs::Permissions::from_mode(0o600));
}
self.modified = false;
self.saved_once = true;
self.message = Some((

View File

@@ -28,7 +28,9 @@ pub struct ServerContext {
pub aether_url: String,
/// Management token for this server.
pub management_token: String,
/// Resolved node name (per-server override or global fallback).
/// Resolved node name at registration time (per-server override or global fallback).
/// After startup, the active node_name is read from `dynamic` (may be updated remotely).
#[allow(dead_code)]
pub node_name: String,
/// Node ID assigned by this Aether server.
pub node_id: Arc<RwLock<String>>,
@@ -46,6 +48,9 @@ pub struct ServerContext {
pub struct ProxyMetrics {
pub total_requests: AtomicU64,
pub total_latency_ns: AtomicU64,
pub failed_requests: AtomicU64,
pub dns_failures: AtomicU64,
pub stream_errors: AtomicU64,
}
impl ProxyMetrics {
@@ -53,12 +58,15 @@ impl ProxyMetrics {
Self {
total_requests: AtomicU64::new(0),
total_latency_ns: AtomicU64::new(0),
failed_requests: AtomicU64::new(0),
dns_failures: AtomicU64::new(0),
stream_errors: AtomicU64::new(0),
}
}
pub fn record_request(&self, elapsed: Duration) {
let nanos = u64::try_from(elapsed.as_nanos()).unwrap_or(u64::MAX);
self.total_requests.fetch_add(1, Ordering::Relaxed);
self.total_latency_ns.fetch_add(nanos, Ordering::Relaxed);
self.total_requests.fetch_add(1, Ordering::Release);
self.total_latency_ns.fetch_add(nanos, Ordering::Release);
}
}

View File

@@ -1,11 +1,12 @@
use std::collections::{HashMap, HashSet};
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
/// Check if an IP address belongs to a private/reserved network.
fn is_private_ip(ip: &IpAddr) -> bool {
pub fn is_private_ip(ip: &IpAddr) -> bool {
match ip {
IpAddr::V4(v4) => is_private_ipv4(v4),
IpAddr::V6(v6) => is_private_ipv6(v6),
@@ -38,6 +39,22 @@ fn is_private_ipv4(ip: &Ipv4Addr) -> bool {
if octets[0] == 0 {
return true;
}
// 100.64.0.0/10 (CGNAT / shared address space)
if octets[0] == 100 && (64..=127).contains(&octets[1]) {
return true;
}
// 192.0.0.0/24 (IETF protocol assignments)
if octets[0] == 192 && octets[1] == 0 && octets[2] == 0 {
return true;
}
// 198.18.0.0/15 (benchmark testing)
if octets[0] == 198 && (18..=19).contains(&octets[1]) {
return true;
}
// 240.0.0.0/4 (reserved for future use)
if octets[0] >= 240 {
return true;
}
false
}
@@ -71,6 +88,7 @@ pub enum FilterError {
PrivateIp(IpAddr),
PortNotAllowed(u16),
DnsResolutionFailed(String),
NoPublicAddrs(String),
}
impl std::fmt::Display for FilterError {
@@ -79,17 +97,26 @@ impl std::fmt::Display for FilterError {
Self::PrivateIp(ip) => write!(f, "target IP {} is in private/reserved range", ip),
Self::PortNotAllowed(port) => write!(f, "port {} not in allowed list", port),
Self::DnsResolutionFailed(host) => write!(f, "DNS resolution failed for {}", host),
Self::NoPublicAddrs(host) => {
write!(
f,
"all resolved addresses for {} are private/reserved",
host
)
}
}
}
}
struct DnsCacheEntry {
addr: SocketAddr,
addrs: Arc<Vec<SocketAddr>>,
expires_at: Instant,
inserted_at: Instant,
}
/// Lightweight DNS cache with TTL + capacity bounds.
/// Stores all public resolved addresses per host (used by SafeDnsResolver
/// to ensure reqwest connects to the same validated addresses).
pub struct DnsCache {
ttl: Duration,
capacity: usize,
@@ -105,7 +132,27 @@ impl DnsCache {
}
}
pub async fn get(&self, host: &str, port: u16) -> Option<SocketAddr> {
/// Look up cached public addresses for a host (any port).
///
/// Used by `SafeDnsResolver` which only knows the hostname — returns the
/// first unexpired entry whose key starts with `host:`.
pub async fn get_by_host(&self, host: &str) -> Option<Arc<Vec<SocketAddr>>> {
if self.capacity == 0 || self.ttl.is_zero() {
return None;
}
let prefix = format!("{}:", host.to_ascii_lowercase());
let now = Instant::now();
let entries = self.entries.read().await;
for (key, entry) in entries.iter() {
if key.starts_with(&prefix) && entry.expires_at > now {
return Some(Arc::clone(&entry.addrs));
}
}
None
}
/// Look up cached public addresses for a host + port.
pub async fn get(&self, host: &str, port: u16) -> Option<Arc<Vec<SocketAddr>>> {
if self.capacity == 0 || self.ttl.is_zero() {
return None;
}
@@ -116,7 +163,7 @@ impl DnsCache {
{
let entries = self.entries.read().await;
match entries.get(&key) {
Some(entry) if entry.expires_at > now => return Some(entry.addr),
Some(entry) if entry.expires_at > now => return Some(Arc::clone(&entry.addrs)),
None => return None,
Some(_) => {} // expired, fall through to evict
}
@@ -128,8 +175,9 @@ impl DnsCache {
None
}
pub async fn insert(&self, host: &str, port: u16, addr: SocketAddr) {
if self.capacity == 0 || self.ttl.is_zero() {
/// Insert resolved public addresses into cache.
pub async fn insert(&self, host: &str, port: u16, addrs: Arc<Vec<SocketAddr>>) {
if self.capacity == 0 || self.ttl.is_zero() || addrs.is_empty() {
return;
}
let key = Self::key(host, port);
@@ -150,7 +198,7 @@ impl DnsCache {
entries.insert(
key,
DnsCacheEntry {
addr,
addrs,
expires_at: now + self.ttl,
inserted_at: now,
},
@@ -158,22 +206,62 @@ impl DnsCache {
}
fn key(host: &str, port: u16) -> String {
format!("{}:{}", host, port)
format!("{}:{}", host.to_ascii_lowercase(), port)
}
}
/// Resolve a hostname to public (non-private) socket addresses.
///
/// Results are cached in `dns_cache`. Private/reserved IPs are filtered out.
/// Returns an error if no public addresses remain after filtering.
pub async fn resolve_public_addrs(
host: &str,
port: u16,
dns_cache: &DnsCache,
) -> Result<Vec<SocketAddr>, FilterError> {
// Cache hit
if let Some(addrs) = dns_cache.get(host, port).await {
return Ok((*addrs).clone());
}
// Async DNS resolution
let addr_str = format!("{}:{}", host, port);
let resolved: Vec<SocketAddr> = tokio::net::lookup_host(&addr_str)
.await
.map_err(|_| FilterError::DnsResolutionFailed(host.to_string()))?
.collect();
if resolved.is_empty() {
return Err(FilterError::DnsResolutionFailed(host.to_string()));
}
// Filter out private/reserved addresses
let public: Vec<SocketAddr> = resolved
.into_iter()
.filter(|addr| !is_private_ip(&addr.ip()))
.collect();
if public.is_empty() {
return Err(FilterError::NoPublicAddrs(host.to_string()));
}
// Cache the validated public addresses
let arc_addrs = Arc::new(public);
dns_cache.insert(host, port, Arc::clone(&arc_addrs)).await;
Ok((*arc_addrs).clone())
}
/// Validate that the target host:port is allowed.
///
/// Uses async DNS resolution (via `tokio::net::lookup_host`) to avoid
/// blocking the async runtime on potentially slow DNS lookups.
///
/// Returns the resolved socket address to connect to.
/// Performs port whitelist check, private IP filtering, and DNS resolution
/// with caching. The resolved addresses are stored in the shared DnsCache
/// so that the SafeDnsResolver can reuse them, eliminating the TOCTTOU gap.
pub async fn validate_target(
host: &str,
port: u16,
allowed_ports: &HashSet<u16>,
dns_cache: &DnsCache,
) -> Result<SocketAddr, FilterError> {
) -> Result<Vec<SocketAddr>, FilterError> {
// Port whitelist check
if !allowed_ports.contains(&port) {
return Err(FilterError::PortNotAllowed(port));
@@ -184,35 +272,11 @@ pub async fn validate_target(
if is_private_ip(&ip) {
return Err(FilterError::PrivateIp(ip));
}
return Ok(SocketAddr::new(ip, port));
return Ok(vec![SocketAddr::new(ip, port)]);
}
if let Some(addr) = dns_cache.get(host, port).await {
return Ok(addr);
}
// Async DNS resolution with private IP check (DNS rebinding protection)
let addr_str = format!("{}:{}", host, port);
let addrs: Vec<SocketAddr> = tokio::net::lookup_host(&addr_str)
.await
.map_err(|_| FilterError::DnsResolutionFailed(host.to_string()))?
.collect();
if addrs.is_empty() {
return Err(FilterError::DnsResolutionFailed(host.to_string()));
}
// All resolved addresses must be non-private
for addr in &addrs {
if is_private_ip(&addr.ip()) {
return Err(FilterError::PrivateIp(addr.ip()));
}
}
// Return the first valid address
let selected = addrs[0];
dns_cache.insert(host, port, selected).await;
Ok(selected)
// Resolve and validate DNS (populates cache for SafeDnsResolver)
resolve_public_addrs(host, port, dns_cache).await
}
#[cfg(test)]
@@ -235,6 +299,19 @@ mod tests {
assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1))));
assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(169, 254, 1, 1))));
assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0))));
// CGNAT
assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(100, 64, 0, 1))));
assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(
100, 127, 255, 254
))));
assert!(!is_private_ip(&IpAddr::V4(Ipv4Addr::new(
100, 63, 255, 254
))));
// Benchmark testing
assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(198, 18, 0, 1))));
// Reserved
assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(240, 0, 0, 1))));
// Public
assert!(!is_private_ip(&IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8))));
assert!(!is_private_ip(&IpAddr::V4(Ipv4Addr::new(203, 0, 113, 1))));
}
@@ -272,5 +349,33 @@ mod tests {
let cache = cache();
let result = validate_target("8.8.8.8", 443, &ports(), &cache).await;
assert!(result.is_ok());
let addrs = result.unwrap();
assert_eq!(addrs.len(), 1);
assert_eq!(addrs[0].ip(), IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)));
}
#[tokio::test]
async fn test_cache_stores_multiple_addrs() {
let cache = cache();
let addrs = vec![
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)), 443),
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 0, 0, 1)), 443),
];
cache
.insert("example.com", 443, Arc::new(addrs.clone()))
.await;
let cached = cache.get("example.com", 443).await.unwrap();
assert_eq!(*cached, addrs);
}
#[tokio::test]
async fn test_cache_key_case_insensitive() {
let cache = cache();
let addrs = vec![SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)), 443)];
cache
.insert("Example.COM", 443, Arc::new(addrs.clone()))
.await;
let cached = cache.get("example.com", 443).await.unwrap();
assert_eq!(*cached, addrs);
}
}

View File

@@ -8,6 +8,7 @@ use tokio::net::TcpStream;
use tokio::sync::watch;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http;
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
use tracing::{debug, info, warn};
use crate::state::{AppState, ServerContext};
@@ -44,10 +45,19 @@ pub async fn connect_and_run(
);
let node_id = server.node_id.read().unwrap().clone();
headers.insert("X-Node-Id", http::HeaderValue::from_str(&node_id)?);
// Use dynamic node_name (may be updated by remote config) instead of
// the static server.node_name, so that remote name changes take effect
// on the next reconnect.
let dynamic_node_name = server.dynamic.load().node_name.clone();
headers.insert(
"X-Node-Name",
http::HeaderValue::from_str(&server.node_name)?,
http::HeaderValue::from_str(&dynamic_node_name)?,
);
// Advertise per-connection max concurrent streams so the backend can
// respect the proxy's capacity limit (backward-compatible: old backends
// ignore this header).
let max_streams = state.config.tunnel_max_streams.unwrap_or(128);
headers.insert("X-Tunnel-Max-Streams", http::HeaderValue::from(max_streams));
// Parse host:port from URL
let uri: http::Uri = ws_url.parse()?;
@@ -79,10 +89,23 @@ pub async fn connect_and_run(
} else {
None
};
// Match Python-side _MAX_FRAME_SIZE (64 MiB) to prevent tungstenite's
// default 16 MiB limit from rejecting large AI API payloads (multi-image
// base64 requests can exceed 16 MiB).
let ws_config = WebSocketConfig {
max_frame_size: Some(64 << 20),
max_message_size: Some(64 << 20),
..Default::default()
};
let handshake_timeout = Duration::from_secs(state.config.tunnel_connect_timeout_secs);
let (ws_stream, _response) = tokio::time::timeout(
handshake_timeout,
tokio_tungstenite::client_async_tls_with_config(request, tcp_stream, None, connector),
tokio_tungstenite::client_async_tls_with_config(
request,
tcp_stream,
Some(ws_config),
connector,
),
)
.await
.map_err(|_| {
@@ -137,8 +160,17 @@ pub async fn connect_and_run(
Err(e) => return Err(e),
}
}
_ = &mut writer_handle => {
warn!("writer task exited, triggering reconnect");
writer_result = &mut writer_handle => {
match writer_result {
Ok(()) => warn!("writer task exited normally, triggering reconnect"),
Err(e) => {
if e.is_panic() {
tracing::error!(error = %e, "writer task panicked, triggering reconnect");
} else {
warn!(error = %e, "writer task cancelled, triggering reconnect");
}
}
}
TunnelOutcome::Disconnected
}
_ = shutdown.changed() => {

View File

@@ -139,7 +139,7 @@ where
}
// Create body channel and spawn handler
let (body_tx, body_rx) = mpsc::channel::<Frame>(16);
let (body_tx, body_rx) = mpsc::channel::<Frame>(64);
streams.insert(frame.stream_id, body_tx);
let state_clone = Arc::clone(&state);

View File

@@ -39,27 +39,44 @@ pub fn spawn_noop() -> HeartbeatHandle {
/// Spawn the heartbeat task. Returns a handle for forwarding ACKs.
pub fn spawn(
config: Arc<Config>,
_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
// Read initial interval from dynamic config (may be updated by remote config).
let initial_interval = Duration::from_secs(server.dynamic.load().heartbeat_interval);
let mut current_interval = initial_interval;
// Skip first immediate tick by sleeping first.
tokio::time::sleep(current_interval).await;
loop {
tokio::select! {
_ = ticker.tick() => {
_ = tokio::time::sleep(current_interval) => {
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");
// Re-read interval from dynamic config (remote config may have
// updated it since the last heartbeat).
let new_interval = Duration::from_secs(
server.dynamic.load().heartbeat_interval
);
if new_interval != current_interval {
debug!(
old_secs = current_interval.as_secs(),
new_secs = new_interval.as_secs(),
"heartbeat interval updated from dynamic config"
);
current_interval = new_interval;
}
}
Some(ack_payload) = ack_rx.recv() => {
handle_ack(&server, &ack_payload);
@@ -78,8 +95,11 @@ pub fn spawn(
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 interval_requests = server.metrics.total_requests.swap(0, Ordering::AcqRel);
let interval_latency_ns = server.metrics.total_latency_ns.swap(0, Ordering::AcqRel);
let interval_failed = server.metrics.failed_requests.swap(0, Ordering::AcqRel);
let interval_dns_failures = server.metrics.dns_failures.swap(0, Ordering::AcqRel);
let interval_stream_errors = server.metrics.stream_errors.swap(0, Ordering::AcqRel);
let avg_latency_ms = if interval_requests > 0 {
Some(interval_latency_ns as f64 / interval_requests as f64 / 1_000_000.0)
} else {
@@ -88,9 +108,12 @@ fn build_heartbeat_payload(server: &ServerContext) -> Bytes {
let payload = serde_json::json!({
"node_id": node_id,
"active_connections": server.active_connections.load(Ordering::Relaxed),
"active_connections": server.active_connections.load(Ordering::Acquire),
"total_requests": interval_requests,
"avg_latency_ms": avg_latency_ms,
"failed_requests": interval_failed,
"dns_failures": interval_dns_failures,
"stream_errors": interval_stream_errors,
});
Bytes::from(serde_json::to_vec(&payload).unwrap_or_default())

View File

@@ -47,7 +47,7 @@ pub async fn run(
let duration = connect_start.elapsed();
if duration >= MIN_STABLE_DURATION {
// Stable session -- reset backoff for quick reconnect
reconnect_attempts.store(0, Ordering::Relaxed);
reconnect_attempts.store(0, Ordering::Release);
info!(
server = %server.server_label,
conn = conn_idx,

View File

@@ -26,6 +26,24 @@ const MAX_CHUNK_SIZE: usize = 32 * 1024;
/// rather than blocking indefinitely and exhausting the stream pool.
const FRAME_SEND_TIMEOUT: Duration = Duration::from_secs(30);
/// Minimum allowed upstream request timeout (seconds).
const MIN_TIMEOUT_SECS: u64 = 5;
/// Maximum allowed upstream request timeout (seconds).
const MAX_TIMEOUT_SECS: u64 = 300;
/// Headers that must not be forwarded to upstream (hop-by-hop or security-sensitive).
const BLOCKED_HEADERS: &[&str] = &[
"connection",
"keep-alive",
"proxy-authenticate",
"proxy-authorization",
"proxy-connection",
"te",
"trailer",
"transfer-encoding",
"upgrade",
];
/// Handle a single stream: receive body, execute upstream, send response.
pub async fn handle_stream(
state: Arc<AppState>,
@@ -36,11 +54,11 @@ pub async fn handle_stream(
frame_tx: FrameSender,
) {
let start = Instant::now();
server.active_connections.fetch_add(1, Ordering::Relaxed);
server.active_connections.fetch_add(1, Ordering::Release);
handle_stream_inner(&state, &server, stream_id, meta, &mut body_rx, &frame_tx).await;
server.active_connections.fetch_sub(1, Ordering::Relaxed);
server.active_connections.fetch_sub(1, Ordering::Release);
server.metrics.record_request(start.elapsed());
}
@@ -134,6 +152,20 @@ async fn handle_stream_inner(
}
};
// Only allow http/https schemes (block file://, data://, etc.)
match target_url.scheme() {
"http" | "https" => {}
other => {
send_error(
frame_tx,
stream_id,
&format!("unsupported URL scheme: {other}"),
)
.await;
return;
}
}
let host = match target_url.host_str() {
Some(h) => h.to_string(),
None => {
@@ -143,13 +175,14 @@ async fn handle_stream_inner(
};
let port = target_url.port_or_known_default().unwrap_or(443);
// DNS + target validation (dns_cache is populated as a side effect)
// DNS + target validation (populates dns_cache for SafeDnsResolver)
let dns_start = Instant::now();
{
let allowed_ports = server.dynamic.read().unwrap().allowed_ports.clone();
let allowed_ports = Arc::clone(&server.dynamic.load().allowed_ports);
if let Err(e) =
target_filter::validate_target(&host, port, &allowed_ports, &state.dns_cache).await
{
server.metrics.dns_failures.fetch_add(1, Ordering::Release);
send_error(frame_tx, stream_id, &format!("target blocked: {e}")).await;
return;
}
@@ -158,12 +191,23 @@ async fn handle_stream_inner(
// Execute upstream request
let client = &state.reqwest_client;
let timeout = Duration::from_secs(meta.timeout);
let timeout = Duration::from_secs(meta.timeout.clamp(MIN_TIMEOUT_SECS, MAX_TIMEOUT_SECS));
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 k_lower = k.to_ascii_lowercase();
// Skip hop-by-hop and security-sensitive headers
if BLOCKED_HEADERS.contains(&k_lower.as_str()) {
continue;
}
// Validate header name/value are valid HTTP
if let (Ok(name), Ok(value)) = (
reqwest::header::HeaderName::from_bytes(k.as_bytes()),
reqwest::header::HeaderValue::from_str(v),
) {
req = req.header(name, value);
}
}
let body_size = body.len();
if !body.is_empty() {
@@ -175,6 +219,10 @@ async fn handle_stream_inner(
let response = match req.send().await {
Ok(r) => r,
Err(e) => {
server
.metrics
.failed_requests
.fetch_add(1, Ordering::Release);
let msg = if e.is_timeout() {
"upstream timeout".to_string()
} else if e.is_connect() {
@@ -190,7 +238,7 @@ async fn handle_stream_inner(
// 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();
let mut resp_headers: Vec<(String, String)> = Vec::with_capacity(response.headers().len() + 1);
for (k, v) in response.headers() {
if let Ok(vs) = v.to_str() {
resp_headers.push((k.as_str().to_string(), vs.to_string()));
@@ -253,6 +301,7 @@ async fn handle_stream_inner(
}
}
Err(e) => {
server.metrics.stream_errors.fetch_add(1, Ordering::Release);
warn!(stream_id, error = %e, "upstream body read error");
send_error(frame_tx, stream_id, &format!("body read error: {e}")).await;
return;