mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
feat(aether-proxy): 网络层全面优化与代理耗时追踪
aether-proxy: - 新增 20+ 可配置参数:Aether API/delegate/CONNECT 各阶段超时、连接池、TCP keepalive/nodelay - AetherClient 支持指数退避重试(可配置次数/延迟)和 HTTP/2 - 新增 DNS 缓存(TTL + 容量限制)避免重复解析 - 新增并发连接数限制(Semaphore,默认基于硬件估算) - CONNECT 隧道增加建连超时和升级超时 - TLS 增加握手超时、会话缓存(session ticket + memory cache)、ALPN 协商 - 新增 ProxyMetrics 收集请求计数和延迟,通过心跳上报 - delegate 响应注入 X-Proxy-Timing 头(dns_ms/upstream_ms/total_ms) 后端: - StreamContext 和 handler 提取 X-Proxy-Timing 写入 proxy_info 追踪数据 前端: - 请求时间线展示代理分阶段耗时(DNS/上游) - Endpoint 添加按钮改为文字按钮样式
This commit is contained in:
@@ -5,16 +5,17 @@
|
||||
|
||||
use std::sync::atomic::AtomicU64;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::signal;
|
||||
use tokio::sync::watch;
|
||||
use tokio::sync::{watch, Semaphore};
|
||||
use tracing::{error, info};
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::net;
|
||||
use crate::registration::client::AetherClient;
|
||||
use crate::runtime::{self, DynamicConfig};
|
||||
use crate::state::AppState;
|
||||
use crate::state::{AppState, ProxyMetrics};
|
||||
use crate::{hardware, proxy};
|
||||
|
||||
/// Run the full application lifecycle after config has been parsed.
|
||||
@@ -61,6 +62,23 @@ pub async fn run(mut config: Config) -> anyhow::Result<()> {
|
||||
// Collect hardware info (once at startup)
|
||||
let hw_info = hardware::collect();
|
||||
|
||||
let max_connections_raw = config
|
||||
.max_concurrent_connections
|
||||
.unwrap_or(hw_info.estimated_max_concurrency)
|
||||
.max(1);
|
||||
let max_connections = usize::try_from(max_connections_raw).unwrap_or(usize::MAX);
|
||||
info!(
|
||||
max_connections = max_connections_raw,
|
||||
"connection limit configured"
|
||||
);
|
||||
|
||||
let connection_semaphore = Arc::new(Semaphore::new(max_connections));
|
||||
let metrics = Arc::new(ProxyMetrics::new());
|
||||
let dns_cache = Arc::new(proxy::target_filter::DnsCache::new(
|
||||
Duration::from_secs(config.dns_cache_ttl_secs),
|
||||
config.dns_cache_capacity,
|
||||
));
|
||||
|
||||
// Register with Aether
|
||||
let aether_client = Arc::new(AetherClient::new(&config));
|
||||
let node_id = aether_client
|
||||
@@ -82,10 +100,21 @@ pub async fn run(mut config: Config) -> anyhow::Result<()> {
|
||||
// No overall timeout — SSE streams can last indefinitely.
|
||||
// Connect timeout limits connection establishment; Aether controls
|
||||
// first-byte / idle timeouts on its own side.
|
||||
let delegate_client = reqwest::Client::builder()
|
||||
.connect_timeout(std::time::Duration::from_secs(30))
|
||||
.pool_max_idle_per_host(20)
|
||||
.pool_idle_timeout(std::time::Duration::from_secs(90))
|
||||
let mut delegate_builder = reqwest::Client::builder()
|
||||
.connect_timeout(Duration::from_secs(config.delegate_connect_timeout_secs))
|
||||
.pool_max_idle_per_host(config.delegate_pool_max_idle_per_host)
|
||||
.pool_idle_timeout(Duration::from_secs(config.delegate_pool_idle_timeout_secs))
|
||||
.tcp_nodelay(config.delegate_tcp_nodelay);
|
||||
|
||||
if config.delegate_tcp_keepalive_secs > 0 {
|
||||
delegate_builder = delegate_builder.tcp_keepalive(Some(Duration::from_secs(
|
||||
config.delegate_tcp_keepalive_secs,
|
||||
)));
|
||||
} else {
|
||||
delegate_builder = delegate_builder.tcp_keepalive(None);
|
||||
}
|
||||
|
||||
let delegate_client = delegate_builder
|
||||
.build()
|
||||
.expect("failed to create delegate HTTP client");
|
||||
|
||||
@@ -101,6 +130,9 @@ pub async fn run(mut config: Config) -> anyhow::Result<()> {
|
||||
tls_acceptor,
|
||||
delegate_client,
|
||||
active_connections: Arc::new(AtomicU64::new(0)),
|
||||
connection_semaphore,
|
||||
dns_cache,
|
||||
metrics,
|
||||
});
|
||||
|
||||
// Shutdown signal channel
|
||||
|
||||
@@ -125,6 +125,26 @@ mod tests {
|
||||
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,
|
||||
|
||||
@@ -44,13 +44,146 @@ pub struct Config {
|
||||
pub heartbeat_interval: u64,
|
||||
|
||||
/// Allowed destination ports (default: 80,443,8080,8443)
|
||||
#[arg(long, env = "AETHER_PROXY_ALLOWED_PORTS", value_delimiter = ',', default_values_t = vec![80, 443, 8080, 8443])]
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_PROXY_ALLOWED_PORTS",
|
||||
value_delimiter = ',',
|
||||
default_values_t = vec![80, 443, 8080, 8443]
|
||||
)]
|
||||
pub allowed_ports: Vec<u16>,
|
||||
|
||||
/// Timestamp tolerance window in seconds for HMAC validation
|
||||
#[arg(long, env = "AETHER_PROXY_TIMESTAMP_TOLERANCE", default_value_t = 300)]
|
||||
pub timestamp_tolerance: u64,
|
||||
|
||||
/// Aether API request timeout in seconds
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_PROXY_AETHER_REQUEST_TIMEOUT",
|
||||
default_value_t = 10
|
||||
)]
|
||||
pub aether_request_timeout_secs: u64,
|
||||
|
||||
/// Aether API connect timeout in seconds
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_PROXY_AETHER_CONNECT_TIMEOUT",
|
||||
default_value_t = 10
|
||||
)]
|
||||
pub aether_connect_timeout_secs: u64,
|
||||
|
||||
/// Aether API max idle connections per host
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_PROXY_AETHER_POOL_MAX_IDLE_PER_HOST",
|
||||
default_value_t = 8
|
||||
)]
|
||||
pub aether_pool_max_idle_per_host: usize,
|
||||
|
||||
/// Aether API idle timeout in seconds
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_PROXY_AETHER_POOL_IDLE_TIMEOUT",
|
||||
default_value_t = 90
|
||||
)]
|
||||
pub aether_pool_idle_timeout_secs: u64,
|
||||
|
||||
/// Aether API TCP keepalive in seconds (0 disables)
|
||||
#[arg(long, env = "AETHER_PROXY_AETHER_TCP_KEEPALIVE", default_value_t = 60)]
|
||||
pub aether_tcp_keepalive_secs: u64,
|
||||
|
||||
/// Aether API TCP_NODELAY
|
||||
#[arg(long, env = "AETHER_PROXY_AETHER_TCP_NODELAY", default_value_t = true)]
|
||||
pub aether_tcp_nodelay: bool,
|
||||
|
||||
/// Enable HTTP/2 when talking to Aether API
|
||||
#[arg(long, env = "AETHER_PROXY_AETHER_HTTP2", default_value_t = true)]
|
||||
pub aether_http2: bool,
|
||||
|
||||
/// Aether API retry attempts (including initial)
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_PROXY_AETHER_RETRY_MAX_ATTEMPTS",
|
||||
default_value_t = 3
|
||||
)]
|
||||
pub aether_retry_max_attempts: u32,
|
||||
|
||||
/// Aether API retry base delay in milliseconds
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_PROXY_AETHER_RETRY_BASE_DELAY_MS",
|
||||
default_value_t = 200
|
||||
)]
|
||||
pub aether_retry_base_delay_ms: u64,
|
||||
|
||||
/// Aether API retry max delay in milliseconds
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_PROXY_AETHER_RETRY_MAX_DELAY_MS",
|
||||
default_value_t = 2000
|
||||
)]
|
||||
pub aether_retry_max_delay_ms: u64,
|
||||
|
||||
/// Maximum concurrent TCP connections (defaults to hardware estimate)
|
||||
#[arg(long, env = "AETHER_PROXY_MAX_CONCURRENT_CONNECTIONS")]
|
||||
pub max_concurrent_connections: Option<u64>,
|
||||
|
||||
/// Upstream TCP connect timeout in seconds for CONNECT tunnels
|
||||
#[arg(long, env = "AETHER_PROXY_CONNECT_TIMEOUT", default_value_t = 30)]
|
||||
pub connect_timeout_secs: u64,
|
||||
|
||||
/// TLS handshake timeout in seconds for incoming TLS connections
|
||||
#[arg(long, env = "AETHER_PROXY_TLS_HANDSHAKE_TIMEOUT", default_value_t = 10)]
|
||||
pub tls_handshake_timeout_secs: u64,
|
||||
|
||||
/// DNS cache TTL in seconds
|
||||
#[arg(long, env = "AETHER_PROXY_DNS_CACHE_TTL", default_value_t = 60)]
|
||||
pub dns_cache_ttl_secs: u64,
|
||||
|
||||
/// DNS cache capacity (entries)
|
||||
#[arg(long, env = "AETHER_PROXY_DNS_CACHE_CAPACITY", default_value_t = 1024)]
|
||||
pub dns_cache_capacity: usize,
|
||||
|
||||
/// Delegate HTTP client connect timeout in seconds
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_PROXY_DELEGATE_CONNECT_TIMEOUT",
|
||||
default_value_t = 30
|
||||
)]
|
||||
pub delegate_connect_timeout_secs: u64,
|
||||
|
||||
/// Delegate HTTP client max idle connections per host
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_PROXY_DELEGATE_POOL_MAX_IDLE_PER_HOST",
|
||||
default_value_t = 64
|
||||
)]
|
||||
pub delegate_pool_max_idle_per_host: usize,
|
||||
|
||||
/// Delegate HTTP client idle timeout in seconds
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_PROXY_DELEGATE_POOL_IDLE_TIMEOUT",
|
||||
default_value_t = 300
|
||||
)]
|
||||
pub delegate_pool_idle_timeout_secs: u64,
|
||||
|
||||
/// Delegate TCP keepalive in seconds (0 disables)
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_PROXY_DELEGATE_TCP_KEEPALIVE",
|
||||
default_value_t = 60
|
||||
)]
|
||||
pub delegate_tcp_keepalive_secs: u64,
|
||||
|
||||
/// Delegate TCP_NODELAY
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_PROXY_DELEGATE_TCP_NODELAY",
|
||||
default_value_t = true
|
||||
)]
|
||||
pub delegate_tcp_nodelay: bool,
|
||||
|
||||
/// Log level (trace, debug, info, warn, error)
|
||||
#[arg(long, env = "AETHER_PROXY_LOG_LEVEL", default_value = "info")]
|
||||
pub log_level: String,
|
||||
@@ -109,6 +242,46 @@ pub struct ConfigFile {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub timestamp_tolerance: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub aether_request_timeout_secs: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub aether_connect_timeout_secs: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub aether_pool_max_idle_per_host: Option<usize>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub aether_pool_idle_timeout_secs: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub aether_tcp_keepalive_secs: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub aether_tcp_nodelay: Option<bool>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub aether_http2: Option<bool>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub aether_retry_max_attempts: Option<u32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub aether_retry_base_delay_ms: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub aether_retry_max_delay_ms: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub max_concurrent_connections: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub connect_timeout_secs: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tls_handshake_timeout_secs: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub dns_cache_ttl_secs: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub dns_cache_capacity: Option<usize>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub delegate_connect_timeout_secs: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub delegate_pool_max_idle_per_host: Option<usize>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub delegate_pool_idle_timeout_secs: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub delegate_tcp_keepalive_secs: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub delegate_tcp_nodelay: Option<bool>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub log_level: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub log_json: Option<bool>,
|
||||
@@ -168,6 +341,71 @@ impl ConfigFile {
|
||||
set!("AETHER_PROXY_NODE_REGION", self.node_region);
|
||||
set!("AETHER_PROXY_HEARTBEAT_INTERVAL", self.heartbeat_interval);
|
||||
set!("AETHER_PROXY_TIMESTAMP_TOLERANCE", self.timestamp_tolerance);
|
||||
set!(
|
||||
"AETHER_PROXY_AETHER_REQUEST_TIMEOUT",
|
||||
self.aether_request_timeout_secs
|
||||
);
|
||||
set!(
|
||||
"AETHER_PROXY_AETHER_CONNECT_TIMEOUT",
|
||||
self.aether_connect_timeout_secs
|
||||
);
|
||||
set!(
|
||||
"AETHER_PROXY_AETHER_POOL_MAX_IDLE_PER_HOST",
|
||||
self.aether_pool_max_idle_per_host
|
||||
);
|
||||
set!(
|
||||
"AETHER_PROXY_AETHER_POOL_IDLE_TIMEOUT",
|
||||
self.aether_pool_idle_timeout_secs
|
||||
);
|
||||
set!(
|
||||
"AETHER_PROXY_AETHER_TCP_KEEPALIVE",
|
||||
self.aether_tcp_keepalive_secs
|
||||
);
|
||||
set!("AETHER_PROXY_AETHER_TCP_NODELAY", self.aether_tcp_nodelay);
|
||||
set!("AETHER_PROXY_AETHER_HTTP2", self.aether_http2);
|
||||
set!(
|
||||
"AETHER_PROXY_AETHER_RETRY_MAX_ATTEMPTS",
|
||||
self.aether_retry_max_attempts
|
||||
);
|
||||
set!(
|
||||
"AETHER_PROXY_AETHER_RETRY_BASE_DELAY_MS",
|
||||
self.aether_retry_base_delay_ms
|
||||
);
|
||||
set!(
|
||||
"AETHER_PROXY_AETHER_RETRY_MAX_DELAY_MS",
|
||||
self.aether_retry_max_delay_ms
|
||||
);
|
||||
set!(
|
||||
"AETHER_PROXY_MAX_CONCURRENT_CONNECTIONS",
|
||||
self.max_concurrent_connections
|
||||
);
|
||||
set!("AETHER_PROXY_CONNECT_TIMEOUT", self.connect_timeout_secs);
|
||||
set!(
|
||||
"AETHER_PROXY_TLS_HANDSHAKE_TIMEOUT",
|
||||
self.tls_handshake_timeout_secs
|
||||
);
|
||||
set!("AETHER_PROXY_DNS_CACHE_TTL", self.dns_cache_ttl_secs);
|
||||
set!("AETHER_PROXY_DNS_CACHE_CAPACITY", self.dns_cache_capacity);
|
||||
set!(
|
||||
"AETHER_PROXY_DELEGATE_CONNECT_TIMEOUT",
|
||||
self.delegate_connect_timeout_secs
|
||||
);
|
||||
set!(
|
||||
"AETHER_PROXY_DELEGATE_POOL_MAX_IDLE_PER_HOST",
|
||||
self.delegate_pool_max_idle_per_host
|
||||
);
|
||||
set!(
|
||||
"AETHER_PROXY_DELEGATE_POOL_IDLE_TIMEOUT",
|
||||
self.delegate_pool_idle_timeout_secs
|
||||
);
|
||||
set!(
|
||||
"AETHER_PROXY_DELEGATE_TCP_KEEPALIVE",
|
||||
self.delegate_tcp_keepalive_secs
|
||||
);
|
||||
set!(
|
||||
"AETHER_PROXY_DELEGATE_TCP_NODELAY",
|
||||
self.delegate_tcp_nodelay
|
||||
);
|
||||
set!("AETHER_PROXY_LOG_LEVEL", self.log_level);
|
||||
set!("AETHER_PROXY_LOG_JSON", self.log_json);
|
||||
set!("AETHER_PROXY_ENABLE_TLS", self.enable_tls);
|
||||
|
||||
@@ -1,14 +1,16 @@
|
||||
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;
|
||||
use crate::proxy::target_filter::{self, DnsCache};
|
||||
|
||||
/// Handle HTTP CONNECT tunnel requests.
|
||||
///
|
||||
@@ -18,6 +20,7 @@ pub async fn handle_connect(
|
||||
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
|
||||
@@ -44,30 +47,43 @@ pub async fn handle_connect(
|
||||
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).await {
|
||||
Ok(addr) => addr,
|
||||
Err(e) => {
|
||||
warn!(host = %host, port, error = %e, "CONNECT target rejected");
|
||||
return forbidden(&e.to_string());
|
||||
}
|
||||
};
|
||||
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 target_stream = match TcpStream::connect(target_addr).await {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
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 hyper::upgrade::on(req).await {
|
||||
Ok(upgraded) => {
|
||||
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;
|
||||
|
||||
@@ -85,9 +101,12 @@ pub async fn handle_connect(
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
Ok(Err(e)) => {
|
||||
warn!(target = %target_display, error = %e, "CONNECT upgrade failed");
|
||||
}
|
||||
Err(_) => {
|
||||
warn!(target = %target_display, "CONNECT upgrade timeout");
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
@@ -133,3 +152,12 @@ fn bad_gateway(msg: &str) -> Response<http_body_util::Empty<bytes::Bytes>> {
|
||||
.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()
|
||||
}
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use std::collections::HashMap;
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use futures_util::TryStreamExt;
|
||||
use http_body_util::{BodyExt, Full, Limited, StreamBody};
|
||||
@@ -13,7 +14,7 @@ use url::Url;
|
||||
use super::BoxBody;
|
||||
use crate::auth;
|
||||
use crate::config::Config;
|
||||
use crate::proxy::target_filter;
|
||||
use crate::proxy::target_filter::{self, DnsCache};
|
||||
|
||||
/// Delegation request payload sent by Aether.
|
||||
#[derive(Debug, Deserialize)]
|
||||
@@ -36,8 +37,11 @@ pub async fn handle_delegate(
|
||||
config: Arc<Config>,
|
||||
allowed_ports: &HashSet<u16>,
|
||||
timestamp_tolerance: u64,
|
||||
dns_cache: &DnsCache,
|
||||
http_client: &reqwest::Client,
|
||||
) -> Response<BoxBody> {
|
||||
let total_start = Instant::now();
|
||||
|
||||
// Authenticate via Authorization header (same HMAC scheme as Proxy-Authorization)
|
||||
let auth_header = req
|
||||
.headers()
|
||||
@@ -86,10 +90,12 @@ pub async fn handle_delegate(
|
||||
|
||||
let port = parsed_url.port_or_known_default().unwrap_or(443);
|
||||
|
||||
if let Err(e) = target_filter::validate_target(&host, port, allowed_ports).await {
|
||||
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 = %delegate_req.method,
|
||||
@@ -110,8 +116,8 @@ pub async fn handle_delegate(
|
||||
|
||||
// NOTE: We intentionally do NOT set a per-request timeout here.
|
||||
// reqwest's `.timeout()` caps the *entire* request including body streaming,
|
||||
// which would truncate long-lived SSE streams. The delegate_client already
|
||||
// has a 30s connect_timeout for connection establishment, and Aether controls
|
||||
// which would truncate long-lived SSE streams. The delegate_client already
|
||||
// has a configured connect_timeout for connection establishment, and Aether controls
|
||||
// first-byte / idle timeouts on its own side via asyncio.
|
||||
|
||||
// Set headers (skip `host` — reqwest sets it from the URL automatically,
|
||||
@@ -129,6 +135,7 @@ pub async fn handle_delegate(
|
||||
}
|
||||
|
||||
// Send upstream request
|
||||
let upstream_start = Instant::now();
|
||||
let upstream_resp = match upstream_req.send().await {
|
||||
Ok(resp) => resp,
|
||||
Err(e) => {
|
||||
@@ -142,12 +149,29 @@ pub async fn handle_delegate(
|
||||
return error_response(502, "upstream_connection_failed", &safe_detail);
|
||||
}
|
||||
};
|
||||
let upstream_ms = upstream_start.elapsed().as_millis() as u64;
|
||||
|
||||
// Build response: pass through upstream status + headers, stream body back
|
||||
let status = upstream_resp.status().as_u16();
|
||||
let upstream_headers = upstream_resp.headers().clone();
|
||||
|
||||
debug!(url = %delegate_req.url, status, "delegate upstream response");
|
||||
let total_ms = total_start.elapsed().as_millis() as u64;
|
||||
|
||||
debug!(
|
||||
url = %delegate_req.url,
|
||||
status,
|
||||
dns_ms,
|
||||
upstream_ms,
|
||||
total_ms,
|
||||
"delegate upstream response"
|
||||
);
|
||||
|
||||
// Inject proxy timing header for Aether to parse
|
||||
let timing = serde_json::json!({
|
||||
"dns_ms": dns_ms,
|
||||
"upstream_ms": upstream_ms,
|
||||
"total_ms": total_ms,
|
||||
});
|
||||
|
||||
// Stream the response body
|
||||
let body_stream = upstream_resp
|
||||
@@ -161,10 +185,14 @@ pub async fn handle_delegate(
|
||||
for (name, value) in upstream_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())
|
||||
builder.body(stream_body).unwrap_or_else(|_| {
|
||||
Response::builder()
|
||||
.status(500)
|
||||
.body(super::empty_box_body())
|
||||
.unwrap()
|
||||
})
|
||||
}
|
||||
|
||||
// ── Sanitisation ─────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
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;
|
||||
@@ -11,6 +12,7 @@ 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};
|
||||
@@ -39,6 +41,8 @@ pub async fn run(
|
||||
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() => {
|
||||
@@ -52,15 +56,38 @@ pub async fn run(
|
||||
|
||||
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 {
|
||||
if tls::is_tls_client_hello(&stream).await {
|
||||
match acceptor.clone().accept(stream).await {
|
||||
Ok(tls_stream) => {
|
||||
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),
|
||||
@@ -69,9 +96,12 @@ pub async fn run(
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Err(e) => {
|
||||
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;
|
||||
@@ -107,13 +137,18 @@ where
|
||||
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();
|
||||
@@ -121,13 +156,20 @@ where
|
||||
};
|
||||
|
||||
if req.method() == Method::CONNECT {
|
||||
let resp =
|
||||
connect::handle_connect(req, config, &allowed_ports, timestamp_tolerance).await;
|
||||
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
|
||||
{
|
||||
@@ -136,19 +178,23 @@ where
|
||||
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).
|
||||
Ok(Response::builder()
|
||||
let resp = Response::builder()
|
||||
.status(405)
|
||||
.header("Allow", "CONNECT")
|
||||
.header("Content-Length", "0")
|
||||
.body(crate::proxy::empty_box_body())
|
||||
.unwrap())
|
||||
.unwrap();
|
||||
metrics.record_request(start.elapsed());
|
||||
Ok(resp)
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
@@ -1,5 +1,8 @@
|
||||
use std::collections::HashSet;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
|
||||
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 {
|
||||
@@ -80,6 +83,85 @@ impl std::fmt::Display for FilterError {
|
||||
}
|
||||
}
|
||||
|
||||
struct DnsCacheEntry {
|
||||
addr: SocketAddr,
|
||||
expires_at: Instant,
|
||||
inserted_at: Instant,
|
||||
}
|
||||
|
||||
/// Lightweight DNS cache with TTL + capacity bounds.
|
||||
pub struct DnsCache {
|
||||
ttl: Duration,
|
||||
capacity: usize,
|
||||
entries: RwLock<HashMap<String, DnsCacheEntry>>,
|
||||
}
|
||||
|
||||
impl DnsCache {
|
||||
pub fn new(ttl: Duration, capacity: usize) -> Self {
|
||||
Self {
|
||||
ttl,
|
||||
capacity,
|
||||
entries: RwLock::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn get(&self, host: &str, port: u16) -> Option<SocketAddr> {
|
||||
if self.capacity == 0 || self.ttl.is_zero() {
|
||||
return None;
|
||||
}
|
||||
let key = Self::key(host, port);
|
||||
let now = Instant::now();
|
||||
|
||||
// Fast path: read lock for cache hit
|
||||
{
|
||||
let entries = self.entries.read().await;
|
||||
match entries.get(&key) {
|
||||
Some(entry) if entry.expires_at > now => return Some(entry.addr),
|
||||
None => return None,
|
||||
Some(_) => {} // expired, fall through to evict
|
||||
}
|
||||
}
|
||||
|
||||
// Slow path: write lock to remove expired entry
|
||||
let mut entries = self.entries.write().await;
|
||||
entries.remove(&key);
|
||||
None
|
||||
}
|
||||
|
||||
pub async fn insert(&self, host: &str, port: u16, addr: SocketAddr) {
|
||||
if self.capacity == 0 || self.ttl.is_zero() {
|
||||
return;
|
||||
}
|
||||
let key = Self::key(host, port);
|
||||
let now = Instant::now();
|
||||
let mut entries = self.entries.write().await;
|
||||
entries.retain(|_, entry| entry.expires_at > now);
|
||||
while entries.len() >= self.capacity {
|
||||
let oldest_key = entries
|
||||
.iter()
|
||||
.min_by_key(|(_, entry)| entry.inserted_at)
|
||||
.map(|(key, _)| key.clone());
|
||||
if let Some(key) = oldest_key {
|
||||
entries.remove(&key);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
entries.insert(
|
||||
key,
|
||||
DnsCacheEntry {
|
||||
addr,
|
||||
expires_at: now + self.ttl,
|
||||
inserted_at: now,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
fn key(host: &str, port: u16) -> String {
|
||||
format!("{}:{}", host, port)
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate that the target host:port is allowed.
|
||||
///
|
||||
/// Uses async DNS resolution (via `tokio::net::lookup_host`) to avoid
|
||||
@@ -90,6 +172,7 @@ pub async fn validate_target(
|
||||
host: &str,
|
||||
port: u16,
|
||||
allowed_ports: &HashSet<u16>,
|
||||
dns_cache: &DnsCache,
|
||||
) -> Result<SocketAddr, FilterError> {
|
||||
// Port whitelist check
|
||||
if !allowed_ports.contains(&port) {
|
||||
@@ -104,6 +187,10 @@ pub async fn validate_target(
|
||||
return Ok(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)
|
||||
@@ -123,7 +210,9 @@ pub async fn validate_target(
|
||||
}
|
||||
|
||||
// Return the first valid address
|
||||
Ok(addrs[0])
|
||||
let selected = addrs[0];
|
||||
dns_cache.insert(host, port, selected).await;
|
||||
Ok(selected)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -134,6 +223,10 @@ mod tests {
|
||||
[80, 443, 8080, 8443].into_iter().collect()
|
||||
}
|
||||
|
||||
fn cache() -> DnsCache {
|
||||
DnsCache::new(Duration::from_secs(60), 128)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_private_ipv4() {
|
||||
assert!(is_private_ip(&IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1))));
|
||||
@@ -162,19 +255,22 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_port_not_allowed() {
|
||||
let result = validate_target("8.8.8.8", 22, &ports()).await;
|
||||
let cache = cache();
|
||||
let result = validate_target("8.8.8.8", 22, &ports(), &cache).await;
|
||||
assert!(matches!(result, Err(FilterError::PortNotAllowed(22))));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_private_ip_blocked() {
|
||||
let result = validate_target("127.0.0.1", 80, &ports()).await;
|
||||
let cache = cache();
|
||||
let result = validate_target("127.0.0.1", 80, &ports(), &cache).await;
|
||||
assert!(matches!(result, Err(FilterError::PrivateIp(_))));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_public_ip_allowed() {
|
||||
let result = validate_target("8.8.8.8", 443, &ports()).await;
|
||||
let cache = cache();
|
||||
let result = validate_target("8.8.8.8", 443, &ports(), &cache).await;
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -9,6 +9,8 @@ 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`.
|
||||
@@ -73,10 +75,21 @@ pub fn build_tls_acceptor(cert_path: &Path, key_path: &Path) -> anyhow::Result<T
|
||||
rustls_pemfile::private_key(&mut BufReader::new(key_file))?
|
||||
.ok_or_else(|| anyhow::anyhow!("no private key found in {}", key_path.display()))?;
|
||||
|
||||
let config = rustls::ServerConfig::builder()
|
||||
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)))
|
||||
}
|
||||
|
||||
|
||||
@@ -1,5 +1,8 @@
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use reqwest::{Client, StatusCode};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::time::sleep;
|
||||
use tracing::{debug, error, info, warn};
|
||||
|
||||
use crate::config::Config;
|
||||
@@ -100,19 +103,44 @@ pub struct AetherClient {
|
||||
http: Client,
|
||||
base_url: String,
|
||||
token: String,
|
||||
retry_max_attempts: u32,
|
||||
retry_base_delay: Duration,
|
||||
retry_max_delay: Duration,
|
||||
}
|
||||
|
||||
impl AetherClient {
|
||||
pub fn new(config: &Config) -> Self {
|
||||
let http = Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(10))
|
||||
.build()
|
||||
.expect("failed to create HTTP client");
|
||||
let mut builder = Client::builder()
|
||||
.timeout(Duration::from_secs(config.aether_request_timeout_secs))
|
||||
.connect_timeout(Duration::from_secs(config.aether_connect_timeout_secs))
|
||||
.pool_max_idle_per_host(config.aether_pool_max_idle_per_host)
|
||||
.pool_idle_timeout(Duration::from_secs(config.aether_pool_idle_timeout_secs))
|
||||
.tcp_nodelay(config.aether_tcp_nodelay);
|
||||
|
||||
if config.aether_tcp_keepalive_secs > 0 {
|
||||
builder =
|
||||
builder.tcp_keepalive(Some(Duration::from_secs(config.aether_tcp_keepalive_secs)));
|
||||
} else {
|
||||
builder = builder.tcp_keepalive(None);
|
||||
}
|
||||
|
||||
if config.aether_http2 {
|
||||
builder = builder.http2_adaptive_window(true);
|
||||
}
|
||||
|
||||
let http = builder.build().expect("failed to create HTTP client");
|
||||
|
||||
let retry_base_delay = Duration::from_millis(config.aether_retry_base_delay_ms);
|
||||
let retry_max_delay =
|
||||
Duration::from_millis(config.aether_retry_max_delay_ms).max(retry_base_delay);
|
||||
|
||||
Self {
|
||||
http,
|
||||
base_url: config.aether_url.trim_end_matches('/').to_string(),
|
||||
token: config.management_token.clone(),
|
||||
retry_max_attempts: config.aether_retry_max_attempts.max(1),
|
||||
retry_base_delay,
|
||||
retry_max_delay,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -149,11 +177,15 @@ impl AetherClient {
|
||||
);
|
||||
|
||||
let resp = self
|
||||
.http
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.token))
|
||||
.json(&body)
|
||||
.send()
|
||||
.send_with_retry(
|
||||
|| {
|
||||
self.http
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.token))
|
||||
.json(&body)
|
||||
},
|
||||
"register",
|
||||
)
|
||||
.await?;
|
||||
|
||||
let status = resp.status();
|
||||
@@ -190,11 +222,15 @@ impl AetherClient {
|
||||
debug!(node_id = %node_id, "sending heartbeat");
|
||||
|
||||
let resp = self
|
||||
.http
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.token))
|
||||
.json(&body)
|
||||
.send()
|
||||
.send_with_retry(
|
||||
|| {
|
||||
self.http
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.token))
|
||||
.json(&body)
|
||||
},
|
||||
"heartbeat",
|
||||
)
|
||||
.await
|
||||
.map_err(|e| HeartbeatError::Other(e.into()))?;
|
||||
|
||||
@@ -247,11 +283,15 @@ impl AetherClient {
|
||||
info!(node_id = %node_id, "unregistering from Aether");
|
||||
|
||||
let resp = self
|
||||
.http
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.token))
|
||||
.json(&body)
|
||||
.send()
|
||||
.send_with_retry(
|
||||
|| {
|
||||
self.http
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.token))
|
||||
.json(&body)
|
||||
},
|
||||
"unregister",
|
||||
)
|
||||
.await;
|
||||
|
||||
match resp {
|
||||
@@ -271,4 +311,75 @@ impl AetherClient {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_with_retry<F>(
|
||||
&self,
|
||||
mut make_req: F,
|
||||
label: &str,
|
||||
) -> Result<reqwest::Response, reqwest::Error>
|
||||
where
|
||||
F: FnMut() -> reqwest::RequestBuilder,
|
||||
{
|
||||
let mut attempt: u32 = 0;
|
||||
let mut delay = self.retry_base_delay;
|
||||
|
||||
loop {
|
||||
attempt = attempt.saturating_add(1);
|
||||
let resp = make_req().send().await;
|
||||
match resp {
|
||||
Ok(resp) => {
|
||||
if should_retry_status(resp.status()) && attempt < self.retry_max_attempts {
|
||||
let sleep_for = jitter_delay(delay);
|
||||
debug!(
|
||||
attempt,
|
||||
status = %resp.status(),
|
||||
sleep_ms = sleep_for.as_millis(),
|
||||
label,
|
||||
"Aether request retrying"
|
||||
);
|
||||
sleep(sleep_for).await;
|
||||
let next_delay = delay.checked_mul(2).unwrap_or(self.retry_max_delay);
|
||||
delay = std::cmp::min(next_delay, self.retry_max_delay);
|
||||
continue;
|
||||
}
|
||||
return Ok(resp);
|
||||
}
|
||||
Err(e) => {
|
||||
if attempt < self.retry_max_attempts {
|
||||
let sleep_for = jitter_delay(delay);
|
||||
debug!(
|
||||
attempt,
|
||||
error = %e,
|
||||
sleep_ms = sleep_for.as_millis(),
|
||||
label,
|
||||
"Aether request retrying"
|
||||
);
|
||||
sleep(sleep_for).await;
|
||||
let next_delay = delay.checked_mul(2).unwrap_or(self.retry_max_delay);
|
||||
delay = std::cmp::min(next_delay, self.retry_max_delay);
|
||||
continue;
|
||||
}
|
||||
return Err(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn should_retry_status(status: StatusCode) -> bool {
|
||||
status.is_server_error()
|
||||
|| status == StatusCode::TOO_MANY_REQUESTS
|
||||
|| status == StatusCode::REQUEST_TIMEOUT
|
||||
}
|
||||
|
||||
fn jitter_delay(base: Duration) -> Duration {
|
||||
if base.is_zero() {
|
||||
return base;
|
||||
}
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.subsec_nanos() as u64)
|
||||
.unwrap_or(0);
|
||||
let jitter_ms = nanos % 100;
|
||||
base + Duration::from_millis(jitter_ms)
|
||||
}
|
||||
|
||||
@@ -34,9 +34,24 @@ pub async fn run(state: &Arc<AppState>, mut shutdown_rx: watch::Receiver<bool>)
|
||||
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(¤t_node_id, Some(active_conns), None, None)
|
||||
.heartbeat(
|
||||
¤t_node_id,
|
||||
Some(active_conns),
|
||||
Some(interval_requests_i64),
|
||||
avg_latency_ms,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(result) => {
|
||||
|
||||
@@ -209,6 +209,26 @@ impl App {
|
||||
heartbeat_interval: None,
|
||||
allowed_ports: None,
|
||||
timestamp_tolerance: None,
|
||||
aether_request_timeout_secs: None,
|
||||
aether_connect_timeout_secs: None,
|
||||
aether_pool_max_idle_per_host: None,
|
||||
aether_pool_idle_timeout_secs: None,
|
||||
aether_tcp_keepalive_secs: None,
|
||||
aether_tcp_nodelay: None,
|
||||
aether_http2: None,
|
||||
aether_retry_max_attempts: None,
|
||||
aether_retry_base_delay_ms: None,
|
||||
aether_retry_max_delay_ms: None,
|
||||
max_concurrent_connections: None,
|
||||
connect_timeout_secs: None,
|
||||
tls_handshake_timeout_secs: None,
|
||||
dns_cache_ttl_secs: None,
|
||||
dns_cache_capacity: None,
|
||||
delegate_connect_timeout_secs: None,
|
||||
delegate_pool_max_idle_per_host: None,
|
||||
delegate_pool_idle_timeout_secs: None,
|
||||
delegate_tcp_keepalive_secs: None,
|
||||
delegate_tcp_nodelay: None,
|
||||
log_level: get("log_level"),
|
||||
log_json: get("log_json").and_then(|v| v.parse().ok()),
|
||||
enable_tls: None,
|
||||
|
||||
@@ -3,13 +3,16 @@
|
||||
//! Consolidates the multiple `Arc<...>` parameters that were previously
|
||||
//! threaded individually through proxy server, heartbeat, and handlers.
|
||||
|
||||
use std::sync::atomic::AtomicU64;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::sync::Semaphore;
|
||||
use tokio_rustls::TlsAcceptor;
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::hardware::HardwareInfo;
|
||||
use crate::proxy::target_filter::DnsCache;
|
||||
use crate::registration::client::AetherClient;
|
||||
use crate::runtime::SharedDynamicConfig;
|
||||
|
||||
@@ -27,4 +30,31 @@ pub struct AppState {
|
||||
pub delegate_client: reqwest::Client,
|
||||
/// 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>,
|
||||
/// Request/latency metrics for heartbeat.
|
||||
pub metrics: Arc<ProxyMetrics>,
|
||||
}
|
||||
|
||||
/// Aggregate metrics for reporting to Aether.
|
||||
pub struct ProxyMetrics {
|
||||
pub total_requests: AtomicU64,
|
||||
pub total_latency_ns: AtomicU64,
|
||||
}
|
||||
|
||||
impl ProxyMetrics {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
total_requests: AtomicU64::new(0),
|
||||
total_latency_ns: 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);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user