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

Merge branch 'pr-599'

Browse Source
This commit is contained in:
fawney19
2026-05-29 02:46:07 +08:00
parent 8bd82c8c95 eaabfb83ed
commit 3e4dfd2bac
12 changed files with 273 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/internal/gateway.rs
+4 -4
View File
@@ -810,11 +810,11 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
node_id: node_id.clone(),
heartbeat_interval: payload.heartbeat_interval,
active_connections: payload.active_connections,
total_requests_delta: payload.total_requests,
total_requests_delta: payload.window_total_requests.or(payload.total_requests),
avg_latency_ms: payload.avg_latency_ms,
failed_requests_delta: payload.failed_requests,
dns_failures_delta: payload.dns_failures,
stream_errors_delta: payload.stream_errors,
failed_requests_delta: payload.window_failed_requests.or(payload.failed_requests),
dns_failures_delta: payload.window_dns_failures.or(payload.dns_failures),
stream_errors_delta: payload.window_stream_errors.or(payload.stream_errors),
proxy_metadata: payload.proxy_metadata,
proxy_version: payload.proxy_version,
};
apps/aether-gateway/src/handlers/internal/gateway_helpers.rs
+6
View File
@@ -402,10 +402,16 @@ pub(crate) fn parse_internal_tunnel_heartbeat_request(
.is_some_and(|value| !(5..=600).contains(&value))
|| payload.active_connections.is_some_and(|value| value < 0)
|| payload.total_requests.is_some_and(|value| value < 0)
|| payload.window_total_requests.is_some_and(|value| value < 0)
|| payload.avg_latency_ms.is_some_and(|value| value < 0.0)
|| payload.failed_requests.is_some_and(|value| value < 0)
|| payload
.window_failed_requests
.is_some_and(|value| value < 0)
|| payload.dns_failures.is_some_and(|value| value < 0)
|| payload.window_dns_failures.is_some_and(|value| value < 0)
|| payload.stream_errors.is_some_and(|value| value < 0)
|| payload.window_stream_errors.is_some_and(|value| value < 0)
|| payload
.proxy_version
.as_deref()
apps/aether-gateway/src/handlers/shared/payloads.rs
+8
View File
@@ -12,14 +12,22 @@ pub(crate) struct InternalTunnelHeartbeatRequest {
#[serde(default)]
pub(crate) total_requests: Option<i64>,
#[serde(default)]
pub(crate) window_total_requests: Option<i64>,
#[serde(default)]
pub(crate) avg_latency_ms: Option<f64>,
#[serde(default)]
pub(crate) failed_requests: Option<i64>,
#[serde(default)]
pub(crate) window_failed_requests: Option<i64>,
#[serde(default)]
pub(crate) dns_failures: Option<i64>,
#[serde(default)]
pub(crate) window_dns_failures: Option<i64>,
#[serde(default)]
pub(crate) stream_errors: Option<i64>,
#[serde(default)]
pub(crate) window_stream_errors: Option<i64>,
#[serde(default)]
pub(crate) proxy_metadata: Option<serde_json::Value>,
#[serde(default)]
pub(crate) proxy_version: Option<String>,
apps/aether-gateway/src/tests/control/internal.rs
+17 -4
View File
@@ -52,11 +52,15 @@ async fn gateway_handles_internal_tunnel_heartbeat_locally_with_loopback() {
"heartbeat_id": 77,
"heartbeat_interval": 45,
"active_connections": 5,
"total_requests": 9,
"total_requests": 100,
"avg_latency_ms": 12.5,
"failed_requests": 1,
"dns_failures": 2,
"stream_errors": 3,
"failed_requests": 20,
"dns_failures": 30,
"stream_errors": 40,
"window_total_requests": 9,
"window_failed_requests": 1,
"window_dns_failures": 2,
"window_stream_errors": 3,
"proxy_metadata": {"arch": "arm64"},
"proxy_version": "2.0.0",
}))
@@ -71,6 +75,15 @@ async fn gateway_handles_internal_tunnel_heartbeat_locally_with_loopback() {
assert_eq!(payload["upgrade_to"], "1.2.3");
assert_eq!(payload["remote_config"]["allowed_ports"][0], 443);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let node = repository
.find_proxy_node("node-123")
.await
.expect("node lookup should succeed")
.expect("node should exist");
assert_eq!(node.total_requests, 9);
assert_eq!(node.failed_requests, 1);
assert_eq!(node.dns_failures, 2);
assert_eq!(node.stream_errors, 3);
gateway_handle.abort();
upstream_handle.abort();
apps/aether-gateway/src/tunnel/mod.rs
+18 -4
View File
@@ -72,14 +72,22 @@ struct InternalTunnelHeartbeatRequest {
#[serde(default)]
total_requests: Option<i64>,
#[serde(default)]
window_total_requests: Option<i64>,
#[serde(default)]
avg_latency_ms: Option<f64>,
#[serde(default)]
failed_requests: Option<i64>,
#[serde(default)]
window_failed_requests: Option<i64>,
#[serde(default)]
dns_failures: Option<i64>,
#[serde(default)]
window_dns_failures: Option<i64>,
#[serde(default)]
stream_errors: Option<i64>,
#[serde(default)]
window_stream_errors: Option<i64>,
#[serde(default)]
proxy_metadata: Option<serde_json::Value>,
#[serde(default)]
proxy_version: Option<String>,
@@ -932,11 +940,11 @@ async fn apply_embedded_tunnel_heartbeat(
node_id: node_id.clone(),
heartbeat_interval: payload.heartbeat_interval,
active_connections: payload.active_connections,
total_requests_delta: payload.total_requests,
total_requests_delta: payload.window_total_requests.or(payload.total_requests),
avg_latency_ms: payload.avg_latency_ms,
failed_requests_delta: payload.failed_requests,
dns_failures_delta: payload.dns_failures,
stream_errors_delta: payload.stream_errors,
failed_requests_delta: payload.window_failed_requests.or(payload.failed_requests),
dns_failures_delta: payload.window_dns_failures.or(payload.dns_failures),
stream_errors_delta: payload.window_stream_errors.or(payload.stream_errors),
proxy_metadata: payload.proxy_metadata,
proxy_version: payload.proxy_version,
};
@@ -1009,10 +1017,16 @@ fn parse_embedded_tunnel_heartbeat_request(
.is_some_and(|value| !(5..=600).contains(&value))
|| payload.active_connections.is_some_and(|value| value < 0)
|| payload.total_requests.is_some_and(|value| value < 0)
|| payload.window_total_requests.is_some_and(|value| value < 0)
|| payload.avg_latency_ms.is_some_and(|value| value < 0.0)
|| payload.failed_requests.is_some_and(|value| value < 0)
|| payload
.window_failed_requests
.is_some_and(|value| value < 0)
|| payload.dns_failures.is_some_and(|value| value < 0)
|| payload.window_dns_failures.is_some_and(|value| value < 0)
|| payload.stream_errors.is_some_and(|value| value < 0)
|| payload.window_stream_errors.is_some_and(|value| value < 0)
|| payload
.proxy_version
.as_deref()
apps/aether-tunnel/src/app.rs
+1
View File
@@ -1384,6 +1384,7 @@ mod tests {
upstream_connect_timeout_secs: 30,
upstream_pool_max_idle_per_host: 4,
upstream_pool_idle_timeout_secs: 60,
upstream_client_pool_capacity: crate::config::DEFAULT_UPSTREAM_CLIENT_POOL_CAPACITY,
upstream_tcp_keepalive_secs: 60,
upstream_tcp_nodelay: true,
upstream_proxy_url: None,
apps/aether-tunnel/src/config.rs
+18
View File
@@ -65,6 +65,7 @@ pub const DEFAULT_TUNNEL_SCALE_CHECK_INTERVAL_MS: u64 = 1_000;
pub const DEFAULT_TUNNEL_SCALE_UP_THRESHOLD_PERCENT: u32 = 50;
pub const DEFAULT_TUNNEL_SCALE_DOWN_THRESHOLD_PERCENT: u32 = 35;
pub const DEFAULT_TUNNEL_SCALE_DOWN_GRACE_SECS: u64 = 15;
pub const DEFAULT_UPSTREAM_CLIENT_POOL_CAPACITY: usize = 256;
const AUTO_TUNNEL_CONNECTIONS_REDUNDANT_FLOOR: u64 = 2;
const AUTO_TUNNEL_CONNECTIONS_BASE_CAP: u64 = 4;
// Bias the automatic pool toward a per-device upper band without letting
@@ -524,6 +525,14 @@ pub struct Config {
)]
pub upstream_pool_idle_timeout_secs: u64,
/// Maximum number of keyed upstream HTTP clients retained by the tunnel.
#[arg(
long,
env = "AETHER_TUNNEL_UPSTREAM_CLIENT_POOL_CAPACITY",
default_value_t = DEFAULT_UPSTREAM_CLIENT_POOL_CAPACITY
)]
pub upstream_client_pool_capacity: usize,
/// Upstream TCP keepalive in seconds (0 disables)
#[arg(
long,
@@ -813,6 +822,9 @@ impl Config {
if self.upstream_connect_timeout_secs == 0 {
anyhow::bail!("upstream_connect_timeout_secs must be > 0");
}
if self.upstream_client_pool_capacity == 0 {
anyhow::bail!("upstream_client_pool_capacity must be > 0");
}
if let Some(proxy_url) = normalized_proxy_url(&self.aether_outbound_proxy_url) {
crate::egress_proxy::UpstreamProxyConfig::parse(proxy_url)
.map_err(|err| anyhow::anyhow!("aether_outbound_proxy_url invalid: {err}"))?;
@@ -1044,6 +1056,8 @@ pub struct ConfigFile {
#[serde(skip_serializing_if = "Option::is_none")]
pub upstream_pool_idle_timeout_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub upstream_client_pool_capacity: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub upstream_tcp_keepalive_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub upstream_tcp_nodelay: Option<bool>,
@@ -1227,6 +1241,10 @@ impl ConfigFile {
"AETHER_TUNNEL_UPSTREAM_POOL_IDLE_TIMEOUT",
self.upstream_pool_idle_timeout_secs
);
set!(
"AETHER_TUNNEL_UPSTREAM_CLIENT_POOL_CAPACITY",
self.upstream_client_pool_capacity
);
set!(
"AETHER_TUNNEL_UPSTREAM_TCP_KEEPALIVE",
self.upstream_tcp_keepalive_secs
apps/aether-tunnel/src/tunnel/dispatcher.rs
+32 -1
View File
@@ -239,7 +239,10 @@ where
// Create body channel and spawn handler
let (body_tx, body_rx) = mpsc::channel::<Frame>(64);
streams.insert(frame.stream_id, body_tx);
let request_headers_end_stream = frame.is_end_stream();
if !request_headers_end_stream {
streams.insert(frame.stream_id, body_tx);
}
let state_clone = Arc::clone(&state);
let server_clone = Arc::clone(&server);
@@ -326,8 +329,16 @@ where
// Trigger every 64 frames OR when the count exceeds max_streams.
frames_since_cleanup += 1;
if frames_since_cleanup >= 64 || handler_handles.len() > max_streams {
let closed_streams = prune_closed_stream_senders(&mut streams);
if closed_streams > 0 {
debug!(closed_streams, "removed closed request body stream senders");
}
handler_handles.retain(|h| !h.is_finished());
frames_since_cleanup = 0;
if draining && streams.is_empty() {
info!("tunnel drained after cleanup");
break None;
}
}
};
@@ -397,6 +408,12 @@ fn try_send_stream_error(frame_tx: &FrameSender, stream_id: u32, message: &'stat
}
}
fn prune_closed_stream_senders(streams: &mut HashMap<u32, mpsc::Sender<Frame>>) -> usize {
let before = streams.len();
streams.retain(|_, tx| !tx.is_closed());
before.saturating_sub(streams.len())
}
/// Wait for all active stream handlers to finish (with a timeout).
async fn drain_handlers(handles: Vec<JoinHandle<()>>) {
if handles.is_empty() {
@@ -470,4 +487,18 @@ mod tests {
Bytes::from_static(b"tunnel request body dispatch stalled")
);
}
#[test]
fn prune_closed_stream_senders_drops_streams_with_closed_receivers() {
let (closed_tx, closed_rx) = mpsc::channel::<Frame>(1);
let (open_tx, _open_rx) = mpsc::channel::<Frame>(1);
drop(closed_rx);
let mut streams = HashMap::from([(7, closed_tx), (9, open_tx)]);
let removed = prune_closed_stream_senders(&mut streams);
assert_eq!(removed, 1);
assert!(!streams.contains_key(&7));
assert!(streams.contains_key(&9));
}
}
apps/aether-tunnel/src/tunnel/mod.rs
+1
View File
@@ -532,6 +532,7 @@ mod tests {
upstream_connect_timeout_secs: 30,
upstream_pool_max_idle_per_host: 4,
upstream_pool_idle_timeout_secs: 60,
upstream_client_pool_capacity: crate::config::DEFAULT_UPSTREAM_CLIENT_POOL_CAPACITY,
upstream_tcp_keepalive_secs: 60,
upstream_tcp_nodelay: true,
upstream_proxy_url: None,
apps/aether-tunnel/src/tunnel/stream_handler.rs
+33 -11
View File
@@ -521,6 +521,21 @@ fn prepare_request_body(
}
}
fn prepare_bodyless_request_body(
body_rx: mpsc::Receiver<TunnelFrame>,
follow_redirects: bool,
) -> PreparedRequestBody {
drop(body_rx);
PreparedRequestBody {
first_request_body: Some(empty_request_body()),
replay_body: if follow_redirects {
ReplayableRequestBody::None
} else {
ReplayableRequestBody::NonReplayable
},
}
}
async fn collect_request_body_for_replay(
mut body_rx: mpsc::Receiver<TunnelFrame>,
body_size: Arc<AtomicUsize>,
@@ -1379,17 +1394,7 @@ async fn handle_stream_inner(
0,
)
} else {
PreparedRequestBody {
first_request_body: Some(build_streaming_request_body(
body_rx,
Arc::clone(&request_body_size),
)),
replay_body: if follow_redirects {
ReplayableRequestBody::None
} else {
ReplayableRequestBody::NonReplayable
},
}
prepare_bodyless_request_body(body_rx, follow_redirects)
};
let mut total_dns_ms = 0u64;
@@ -1794,6 +1799,22 @@ mod tests {
assert_eq!(body_size.load(Ordering::Relaxed), 0);
}
#[tokio::test]
async fn bodyless_request_body_completes_without_waiting_for_tunnel_sender() {
let (_tx, rx) = mpsc::channel(4);
let mut prepared = prepare_bodyless_request_body(rx, true);
let mut body = prepared
.first_request_body
.take()
.expect("bodyless request should have an initial body");
let frame = tokio::time::timeout(Duration::from_millis(25), body.frame())
.await
.expect("bodyless request body should not wait for tunnel body frames");
assert!(frame.is_none());
assert!(matches!(prepared.replay_body, ReplayableRequestBody::None));
}
#[tokio::test]
async fn prepare_request_body_streams_immediately_and_replays_after_completion() {
let (tx, rx) = mpsc::channel(4);
@@ -2690,6 +2711,7 @@ mod tests {
upstream_connect_timeout_secs: 30,
upstream_pool_max_idle_per_host: 4,
upstream_pool_idle_timeout_secs: 60,
upstream_client_pool_capacity: crate::config::DEFAULT_UPSTREAM_CLIENT_POOL_CAPACITY,
upstream_tcp_keepalive_secs: 60,
upstream_tcp_nodelay: true,
upstream_proxy_url: None,
apps/aether-tunnel/src/upstream_client.rs
+100 -12
View File
@@ -4,6 +4,7 @@ use std::future::Future;
use std::io;
use std::net::IpAddr;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::sync::Mutex;
use std::task::{Context, Poll};
@@ -65,7 +66,14 @@ pub struct UpstreamClientPoolKey {
pub struct UpstreamClientPool {
config: Arc<Config>,
dns_cache: Arc<DnsCache>,
clients: Arc<Mutex<HashMap<UpstreamClientPoolKey, UpstreamClient>>>,
clients: Arc<Mutex<HashMap<UpstreamClientPoolKey, UpstreamClientPoolEntry>>>,
access_counter: Arc<AtomicU64>,
}
#[derive(Clone)]
struct UpstreamClientPoolEntry {
client: UpstreamClient,
last_used: u64,
}
impl UpstreamClientPool {
@@ -74,18 +82,17 @@ impl UpstreamClientPool {
config,
dns_cache,
clients: Arc::new(Mutex::new(HashMap::new())),
access_counter: Arc::new(AtomicU64::new(0)),
}
}
pub fn get_or_build(&self, key: UpstreamClientPoolKey) -> Result<UpstreamClient, String> {
if let Some(client) = self
.clients
.lock()
.expect("client pool lock")
.get(&key)
.cloned()
{
return Ok(client);
let mut clients = self.clients.lock().expect("client pool lock");
if let Some(entry) = clients.get_mut(&key) {
entry.last_used = self.next_access_id();
return Ok(entry.client.clone());
}
}
validate_proxy_transport_backend(&key.backend)?;
@@ -97,12 +104,45 @@ impl UpstreamClientPool {
Arc::clone(&self.dns_cache),
http1_only,
)?;
self.clients
.lock()
.expect("client pool lock")
.insert(key, client.clone());
let mut clients = self.clients.lock().expect("client pool lock");
if let Some(entry) = clients.get_mut(&key) {
entry.last_used = self.next_access_id();
return Ok(entry.client.clone());
}
evict_lru_client_if_needed(
&mut clients,
self.config.upstream_client_pool_capacity.max(1),
);
clients.insert(
key,
UpstreamClientPoolEntry {
client: client.clone(),
last_used: self.next_access_id(),
},
);
Ok(client)
}
fn next_access_id(&self) -> u64 {
self.access_counter.fetch_add(1, Ordering::Relaxed)
}
}
fn evict_lru_client_if_needed(
clients: &mut HashMap<UpstreamClientPoolKey, UpstreamClientPoolEntry>,
capacity: usize,
) {
if clients.len() < capacity {
return;
}
let Some(oldest_key) = clients
.iter()
.min_by_key(|(_, entry)| entry.last_used)
.map(|(key, _)| key.clone())
else {
return;
};
clients.remove(&oldest_key);
}
pub fn upstream_client_pool_key(
@@ -765,6 +805,44 @@ mod tests {
assert!(error.contains("unsupported transport profile backend"));
}
#[test]
fn upstream_client_pool_evicts_lru_clients_above_capacity() {
let config = Arc::new(
Config::try_parse_from([
"aether-tunnel",
"--aether-url",
"https://aether.example.com",
"--management-token",
"ae_test",
"--node-name",
"tunnel-test",
"--upstream-client-pool-capacity",
"2",
])
.expect("config should parse"),
);
let pool =
UpstreamClientPool::new(config, Arc::new(DnsCache::new(Duration::from_secs(60), 16)));
let key_a = test_pool_key("key-a");
let key_b = test_pool_key("key-b");
let key_c = test_pool_key("key-c");
pool.get_or_build(key_a.clone())
.expect("client A should build");
pool.get_or_build(key_b.clone())
.expect("client B should build");
pool.get_or_build(key_a.clone())
.expect("client A should be reused and become most recent");
pool.get_or_build(key_c.clone())
.expect("client C should build");
let clients = pool.clients.lock().expect("client pool lock");
assert_eq!(clients.len(), 2);
assert!(clients.contains_key(&key_a));
assert!(clients.contains_key(&key_c));
assert!(!clients.contains_key(&key_b));
}
#[test]
fn http_proxy_authorization_header_uses_basic_auth_for_http_proxy() {
assert_eq!(
@@ -777,6 +855,16 @@ mod tests {
);
}
fn test_pool_key(key_id: &str) -> UpstreamClientPoolKey {
upstream_client_pool_key(
Some("provider-1"),
Some("endpoint-1"),
Some(key_id),
None,
false,
)
}
#[tokio::test]
async fn socks5h_target_address_uses_domain_name() {
let request = socks5_target_address("example.com", 443, true)
crates/aether-data/src/repository/proxy_nodes/types.rs
+35 -1
View File
@@ -587,7 +587,8 @@ fn json_string(value: Option<&Value>) -> Option<String> {
fn counter_delta_u64(previous: Option<u64>, current: u64) -> u64 {
match previous {
Some(previous) if current >= previous => current - previous,
Some(_) | None => current,
Some(_) => current,
None => 0,
}
}
@@ -962,6 +963,39 @@ mod tests {
);
}
#[test]
fn builds_tunnel_metrics_sample_uses_first_counter_report_as_baseline() {
let current = json!({
"tunnel_metrics": {
"connect_errors": 12,
"disconnects": 5,
"error_events_total": 7,
"ws_in_bytes": 1_500,
"ws_out_bytes": 2_500,
"ws_in_frames": 15,
"ws_out_frames": 25,
"heartbeat_rtt_last_ms": 44
},
"recent_tunnel_errors": [
{"timestamp_unix_secs": 101, "category": "newer", "message": "new"}
]
});
let sample = build_tunnel_metrics_sample(None, Some(&current), 4, true)
.expect("sample should build");
assert_eq!(sample.samples, 1);
assert_eq!(sample.heartbeat_rtt_ms_sum, 44);
assert_eq!(sample.connect_errors_delta, 0);
assert_eq!(sample.disconnects_delta, 0);
assert_eq!(sample.error_events_delta, 0);
assert_eq!(sample.ws_in_bytes_delta, 0);
assert_eq!(sample.ws_out_bytes_delta, 0);
assert_eq!(sample.ws_in_frames_delta, 0);
assert_eq!(sample.ws_out_frames_delta, 0);
assert!(sample.recent_error_events.is_empty());
}
#[test]
fn preserves_secure_tunnel_metadata_across_heartbeat_metadata_refresh() {
let previous = json!({