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

fix: harden tunnel security and timeout handling

Browse Source
This commit is contained in:
fawney19
2026-05-23 14:17:04 +08:00
parent 74f7348529
commit 6447fda852
19 changed files with 571 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files

177
apps/aether-gateway/src/execution_runtime/transport.rs
View File

@@ -38,6 +38,9 @@ use crate::{AppState, GatewayError};
const HUB_RELAY_CONTENT_TYPE: &str = "application/vnd.aether.tunnel-envelope";
const HUB_RELAY_ERROR_HEADER: &str = "x-aether-tunnel-error";
const TUNNEL_RELAY_PATH_PREFIX: &str = "/api/internal/tunnel/relay";
const DEFAULT_TUNNEL_TIMEOUT_MS: u64 = 60_000;
const MIN_TUNNEL_TIMEOUT_SECS: u64 = 1;
const MAX_TUNNEL_TIMEOUT_SECS: u64 = 300;
pub(crate) fn format_upstream_request_error(err: &reqwest::Error) -> String {
let mut kinds = Vec::new();
if err.is_connect() {
@@ -176,6 +179,12 @@ struct RelayRequestMeta {
method: String,
url: String,
headers: BTreeMap<String, String>,
#[serde(default, skip_serializing_if = "is_false")]
stream: bool,
#[serde(skip_serializing_if = "Option::is_none")]
request_timeout_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
stream_first_byte_timeout_ms: Option<u64>,
timeout: u64,
#[serde(skip_serializing_if = "Option::is_none")]
follow_redirects: Option<bool>,
@@ -195,6 +204,13 @@ pub(crate) struct ExecutionTransportControls {
accept_invalid_certs: bool,
}
#[derive(Debug, Clone, Copy)]
struct TunnelTimeoutMetadata {
request_timeout_ms: Option<u64>,
stream_first_byte_timeout_ms: Option<u64>,
legacy_timeout_secs: u64,
}
pub(crate) enum DirectUpstreamResponse {
Reqwest(reqwest::Response),
BrowserWreq(wreq::Response),
@@ -579,6 +595,7 @@ fn build_direct_tunnel_request_meta(
headers: &HeaderMap,
transport_controls: ExecutionTransportControls,
) -> tunnel_protocol::RequestMeta {
let timeout_metadata = resolve_tunnel_timeout_metadata(plan);
tunnel_protocol::RequestMeta {
provider_id: Some(plan.provider_id.clone()),
endpoint_id: Some(plan.endpoint_id.clone()),
@@ -586,7 +603,10 @@ fn build_direct_tunnel_request_meta(
method: plan.method.clone(),
url: plan.url.clone(),
headers: header_map_to_string_map(headers).into_iter().collect(),
timeout: resolve_relay_timeout_seconds(plan),
stream: plan.stream,
request_timeout_ms: timeout_metadata.request_timeout_ms,
stream_first_byte_timeout_ms: timeout_metadata.stream_first_byte_timeout_ms,
timeout: timeout_metadata.legacy_timeout_secs,
follow_redirects: transport_controls.follow_redirects,
http1_only: transport_controls.http1_only,
transport_profile: plan.transport_profile.clone(),
@@ -753,7 +773,8 @@ async fn send_via_tunnel_relay(
) -> Result<reqwest::Response, ExecutionRuntimeTransportError> {
let client = build_relay_client(plan.timeouts.as_ref())?;
let relay_url = build_relay_url(plan.proxy.as_ref(), node_id);
let timeout_secs = resolve_relay_timeout_seconds(plan);
let timeout_metadata = resolve_tunnel_timeout_metadata(plan);
let timeout_secs = timeout_metadata.legacy_timeout_secs;
let envelope = build_relay_envelope(
RelayRequestMeta {
provider_id: plan.provider_id.clone(),
@@ -762,6 +783,9 @@ async fn send_via_tunnel_relay(
method: method.as_str().to_string(),
url: plan.url.clone(),
headers: header_map_to_string_map(&headers),
stream: plan.stream,
request_timeout_ms: timeout_metadata.request_timeout_ms,
stream_first_byte_timeout_ms: timeout_metadata.stream_first_byte_timeout_ms,
timeout: timeout_secs,
follow_redirects: transport_controls.follow_redirects,
http1_only: transport_controls.http1_only,
@@ -791,15 +815,22 @@ async fn send_via_tunnel_relay(
.request(reqwest::Method::POST, relay_url)
.header(reqwest::header::CONTENT_TYPE, HUB_RELAY_CONTENT_TYPE)
.body(envelope);
if let Some(timeout) = total_timeout {
request = request.timeout(timeout);
if !plan.stream {
if let Some(timeout) = total_timeout {
request = request.timeout(timeout);
}
}
let first_byte_timeout = if plan.stream {
resolve_tunnel_first_byte_timeout(plan)
} else {
None
};
let started_at = Instant::now();
let response = request
.send()
let response = send_relay_request(request, first_byte_timeout)
.await
.map_err(|err| ExecutionRuntimeTransportError::RelayError(err.to_string()))?;
.map_err(ExecutionRuntimeTransportError::RelayError)?;
let elapsed_ms = started_at.elapsed().as_millis() as u64;
let status_code = response.status().as_u16();
let proxy_timing = response
@@ -869,6 +900,21 @@ async fn send_via_tunnel_relay(
Ok(response)
}
async fn send_relay_request(
request: reqwest::RequestBuilder,
first_byte_timeout: Option<Duration>,
) -> Result<reqwest::Response, String> {
if let Some(timeout) = first_byte_timeout {
return match tokio::time::timeout(timeout, request.send()).await {
Ok(Ok(response)) => Ok(response),
Ok(Err(error)) => Err(error.to_string()),
Err(_) => Err("tunnel relay first byte timeout".to_string()),
};
}
request.send().await.map_err(|err| err.to_string())
}
pub(crate) fn build_request_body(
plan: &ExecutionPlan,
) -> Result<Vec<u8>, ExecutionRuntimeTransportError> {
@@ -965,18 +1011,49 @@ fn resolve_tunnel_base_url_from_proxy(proxy: &ProxySnapshot) -> Option<String> {
}
fn resolve_relay_timeout_seconds(plan: &ExecutionPlan) -> u64 {
let ms = plan
.timeouts
resolve_tunnel_timeout_metadata(plan).legacy_timeout_secs
}
fn resolve_tunnel_first_byte_timeout(plan: &ExecutionPlan) -> Option<Duration> {
plan.stream.then(|| {
Duration::from_millis(
resolve_selected_tunnel_timeout_ms(plan).unwrap_or(DEFAULT_TUNNEL_TIMEOUT_MS),
)
})
}
fn resolve_tunnel_timeout_metadata(plan: &ExecutionPlan) -> TunnelTimeoutMetadata {
TunnelTimeoutMetadata {
request_timeout_ms: plan
.timeouts
.as_ref()
.and_then(|timeouts| timeouts.total_ms),
stream_first_byte_timeout_ms: plan
.timeouts
.as_ref()
.and_then(|timeouts| timeouts.first_byte_ms),
legacy_timeout_secs: timeout_ms_to_secs(
resolve_selected_tunnel_timeout_ms(plan).unwrap_or(DEFAULT_TUNNEL_TIMEOUT_MS),
),
}
}
fn resolve_selected_tunnel_timeout_ms(plan: &ExecutionPlan) -> Option<u64> {
plan.timeouts
.as_ref()
.and_then(|timeouts| {
timeouts
.read_ms
.or(timeouts.total_ms)
.or(timeouts.connect_ms)
if plan.stream {
timeouts.first_byte_ms.or(timeouts.total_ms)
} else {
timeouts.total_ms.or(timeouts.first_byte_ms)
}
})
.unwrap_or(60_000);
.map(|value| value.max(1))
}
fn timeout_ms_to_secs(ms: u64) -> u64 {
let secs = ms.div_ceil(1_000);
secs.clamp(1, 300)
secs.clamp(MIN_TUNNEL_TIMEOUT_SECS, MAX_TUNNEL_TIMEOUT_SECS)
}
fn resolve_tunnel_node_id(proxy: Option<&ProxySnapshot>) -> Option<String> {
@@ -1522,12 +1599,12 @@ mod tests {
use tokio::sync::watch;
use super::{
build_browser_wreq_client, build_client, build_execution_response_body,
build_request_headers, execute_sync_plan, record_manual_proxy_request_failure,
record_manual_proxy_request_outcome, record_manual_proxy_request_success,
record_manual_proxy_stream_error, resolve_execution_transport_controls,
response_body_is_json, DirectSyncExecutionRuntime, ExecutionRuntimeTransportError,
ExecutionTransportControls,
build_browser_wreq_client, build_client, build_direct_tunnel_request_meta,
build_execution_response_body, build_request_headers, execute_sync_plan,
record_manual_proxy_request_failure, record_manual_proxy_request_outcome,
record_manual_proxy_request_success, record_manual_proxy_stream_error,
resolve_execution_transport_controls, response_body_is_json, DirectSyncExecutionRuntime,
ExecutionRuntimeTransportError, ExecutionTransportControls,
};
use crate::constants::{
EXECUTION_RUNTIME_LOOP_GUARD_HEADER, EXECUTION_RUNTIME_LOOP_GUARD_VIA_TOKEN,
@@ -1629,6 +1706,64 @@ mod tests {
.is_none());
}
#[test]
fn tunnel_request_meta_uses_total_timeout_for_non_stream_requests() {
let plan = tunnel_timeout_plan(false);
let meta = build_direct_tunnel_request_meta(
&plan,
&reqwest::header::HeaderMap::new(),
ExecutionTransportControls::default(),
);
assert!(!meta.stream);
assert_eq!(meta.request_timeout_ms, Some(90_000));
assert_eq!(meta.stream_first_byte_timeout_ms, Some(12_345));
assert_eq!(meta.timeout, 90);
}
#[test]
fn tunnel_request_meta_uses_first_byte_timeout_for_stream_requests() {
let plan = tunnel_timeout_plan(true);
let meta = build_direct_tunnel_request_meta(
&plan,
&reqwest::header::HeaderMap::new(),
ExecutionTransportControls::default(),
);
assert!(meta.stream);
assert_eq!(meta.request_timeout_ms, Some(90_000));
assert_eq!(meta.stream_first_byte_timeout_ms, Some(12_345));
assert_eq!(meta.timeout, 13);
}
fn tunnel_timeout_plan(stream: bool) -> ExecutionPlan {
ExecutionPlan {
request_id: "req-timeout".into(),
candidate_id: None,
provider_name: Some("provider".into()),
provider_id: "prov-1".into(),
endpoint_id: "ep-1".into(),
key_id: "key-1".into(),
method: "POST".into(),
url: "https://example.com/chat".into(),
headers: BTreeMap::from([("content-type".into(), "application/json".into())]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({"model": "gpt-4.1"})),
stream,
client_api_format: "openai:chat".into(),
provider_api_format: "openai:chat".into(),
model_name: Some("gpt-4.1".into()),
proxy: None,
transport_profile: None,
timeouts: Some(ExecutionTimeouts {
total_ms: Some(90_000),
first_byte_ms: Some(12_345),
..ExecutionTimeouts::default()
}),
}
}
fn tunnel_proxy_snapshot(base_url: String) -> ProxySnapshot {
ProxySnapshot {
enabled: Some(true),

3
apps/aether-gateway/src/handlers/admin/system/proxy_nodes.rs
View File

@@ -1374,6 +1374,9 @@ fn build_tunnel_probe_relay_envelope(
method: "GET".to_string(),
url: probe_url.trim().to_string(),
headers: std::collections::HashMap::new(),
stream: false,
request_timeout_ms: None,
stream_first_byte_timeout_ms: None,
timeout: timeout_secs,
follow_redirects: Some(false),
http1_only: false,

3
apps/aether-gateway/src/tunnel/embedded/hub.rs
View File

@@ -1307,6 +1307,9 @@ mod tests {
method: "GET".to_string(),
url: "https://example.com".to_string(),
headers: HashMap::new(),
stream: false,
request_timeout_ms: None,
stream_first_byte_timeout_ms: None,
timeout: 30,
follow_redirects: None,
http1_only: false,

25
apps/aether-gateway/src/tunnel/embedded/local_relay.rs
View File

@@ -24,6 +24,8 @@ use super::AppState;
pub const TUNNEL_ERROR_HEADER: &str = "x-aether-tunnel-error";
const MAX_RELAY_META_LEN: usize = 256 * 1024;
const MIN_RELAY_TIMEOUT_MS: u64 = 1;
const MAX_RELAY_TIMEOUT_MS: u64 = 300_000;
struct StreamGuard {
hub: std::sync::Arc<super::hub::HubRouter>,
@@ -92,7 +94,7 @@ pub(crate) async fn open_direct_relay_stream(
return Err(format!("connect: {error}"));
}
let wait_timeout = Duration::from_secs(meta.timeout.clamp(5, 300));
let wait_timeout = relay_header_timeout(&meta);
let response_head = match stream.wait_headers(wait_timeout).await {
Ok(response) => response,
Err(error) => {
@@ -148,6 +150,19 @@ fn map_request_admission_error(error: super::RequestAdmissionError) -> String {
}
}
fn relay_header_timeout(meta: &protocol::RequestMeta) -> Duration {
let timeout_ms = if meta.stream {
meta.stream_first_byte_timeout_ms
.or(meta.request_timeout_ms)
.unwrap_or_else(|| meta.timeout.saturating_mul(1_000))
} else {
meta.request_timeout_ms
.or(meta.stream_first_byte_timeout_ms)
.unwrap_or_else(|| meta.timeout.saturating_mul(1_000))
};
Duration::from_millis(timeout_ms.clamp(MIN_RELAY_TIMEOUT_MS, MAX_RELAY_TIMEOUT_MS))
}
pub async fn relay_request(
Path(node_id): Path<String>,
State(state): State<AppState>,
@@ -322,7 +337,7 @@ pub async fn relay_request(
finished: false,
};
let wait_timeout = Duration::from_secs(meta.timeout.clamp(5, 300));
let wait_timeout = relay_header_timeout(&meta);
let response_head = match stream.wait_headers(wait_timeout).await {
Ok(response) => response,
Err(error) => {
@@ -639,6 +654,9 @@ mod tests {
method: "GET".to_string(),
url: "https://example.com/health".to_string(),
headers: HashMap::new(),
stream: false,
request_timeout_ms: None,
stream_first_byte_timeout_ms: None,
timeout: 30,
follow_redirects: None,
http1_only: false,
@@ -754,6 +772,9 @@ mod tests {
method: "GET".to_string(),
url: "https://example.com/headers".to_string(),
headers: HashMap::new(),
stream: false,
request_timeout_ms: None,
stream_first_byte_timeout_ms: None,
timeout: 30,
follow_redirects: None,
http1_only: false,

3
apps/aether-gateway/src/tunnel/mod.rs
View File

@@ -519,6 +519,9 @@ impl EmbeddedTunnelState {
method: "GET".to_string(),
url: url.trim().to_string(),
headers: HashMap::new(),
stream: false,
request_timeout_ms: None,
stream_first_byte_timeout_ms: None,
timeout: timeout_secs,
follow_redirects: Some(false),
http1_only: false,

2
apps/aether-tunnel/README.md
View File

@@ -229,7 +229,7 @@ node_name = "jp-proxy-02"
tunnel_encryption_key = "base64-32-bytes"
```
`tunnel_security = "non_tls_required"` 是非 TLS secure tunnel 的 MVP 配置面：它要求同时提供当前 `[[servers]]` 条目的 `tunnel_encryption_key`，后续握手使用 `node_name` / `X-Node-Id` 查找对应 PSK，不引入 `tunnel_encryption_key_id`。`wss://` 仍是推荐方案；`ws:// + secure tunnel` 只保护 Aether ↔ tunnel 之间的 token 和 payload，不等价于 HTTPS 伪装，也不覆盖 tunnel ↔ origin/provider 这段链路。
`tunnel_security = "non_tls_required"` 是非 TLS secure tunnel 的 MVP 配置面：它要求同时提供当前 `[[servers]]` 条目的 `tunnel_encryption_key`，后续握手使用 `node_name` / `X-Node-Id` 查找对应 PSK，不引入 `tunnel_encryption_key_id`。`wss://` 仍是推荐方案；`ws:// + secure tunnel` 只加密注册完成后的 WebSocket tunnel frame，不保护安装脚本、注册请求、`management_token` 或 PSK 的首次分发；这些 bootstrap 凭据仍必须通过 HTTPS 或其他可信通道交付。它不等价于 HTTPS 伪装，也不覆盖 tunnel ↔ origin/provider 这段链路。
如果 `aether_url` 使用 `http://` 且当前 `[[servers]]` 条目提供了 `tunnel_encryption_key`，省略 `tunnel_security` 时运行时会自动按 `non_tls_required` 生效；显式配置 `tunnel_security = "off"` 会关闭该自动推断。secure tunnel 会在 WebSocket tunnel 上加密所有二进制 tunnel frame；未配置 key 或显式关闭的旧节点仍按原明文协议工作。

7
apps/aether-tunnel/src/app.rs
View File

@@ -221,6 +221,13 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
entry.tunnel_encryption_key.as_deref(),
);
if tunnel_security == crate::config::TunnelSecurity::NonTlsRequired {
if entry.aether_url.trim_start().starts_with("http://") {
warn!(
server = %label,
url = %entry.aether_url,
"secure tunnel frame encryption starts after registration; deliver install and registration credentials over HTTPS or another trusted bootstrap channel"
);
}
let key = entry
.tunnel_encryption_key
.as_deref()

3
apps/aether-tunnel/src/tunnel/mod.rs
View File

@@ -392,6 +392,9 @@ mod tests {
method: "GET".to_string(),
url: "http://127.0.0.1:80/blocked".to_string(),
headers: std::collections::HashMap::new(),
stream: false,
request_timeout_ms: None,
stream_first_byte_timeout_ms: None,
timeout: 5,
follow_redirects: None,
http1_only: false,

213
apps/aether-tunnel/src/tunnel/stream_handler.rs
View File

@@ -39,10 +39,10 @@ const SLOW_STREAM_LOG_THRESHOLD: Duration = Duration::from_secs(2);
const SUCCESS_LOG_SAMPLE_MODULO: u32 = 256;
const REQUEST_BODY_SPOOL_QUEUE_CAPACITY: usize = 64;
/// Minimum allowed upstream request timeout (seconds).
const MIN_TIMEOUT_SECS: u64 = 5;
/// Maximum allowed upstream request timeout (seconds).
const MAX_TIMEOUT_SECS: u64 = 300;
/// Minimum allowed upstream request timeout (milliseconds).
const MIN_TIMEOUT_MS: u64 = 1;
/// Maximum allowed upstream request timeout (milliseconds).
const MAX_TIMEOUT_MS: u64 = 300_000;
/// Match reqwest's default redirect budget so direct execution and tunnel relay
/// fail at the same point instead of diverging after a different number of hops.
const MAX_REDIRECTS: usize = 10;
@@ -95,6 +95,12 @@ struct PreparedRequestBody {
replay_body: ReplayableRequestBody,
}
#[derive(Debug, Clone, Copy)]
struct RequestTimeouts {
first_byte_timeout: Duration,
response_body_timeout: Option<Duration>,
}
#[derive(Debug)]
struct RequestBodyReplayState {
budget_bytes: usize,
@@ -588,6 +594,45 @@ fn remaining_timeout(deadline: Instant) -> Option<Duration> {
deadline.checked_duration_since(Instant::now())
}
fn resolve_request_timeouts(meta: &RequestMeta) -> RequestTimeouts {
let first_byte_timeout = if meta.stream {
meta.stream_first_byte_timeout_ms
.or(meta.request_timeout_ms)
.map(timeout_duration_from_ms)
.unwrap_or_else(|| timeout_duration_from_legacy_secs(meta.timeout))
} else {
meta.request_timeout_ms
.or(meta.stream_first_byte_timeout_ms)
.map(timeout_duration_from_ms)
.unwrap_or_else(|| timeout_duration_from_legacy_secs(meta.timeout))
};
let response_body_timeout = if meta.stream {
None
} else {
Some(
meta.request_timeout_ms
.or(meta.stream_first_byte_timeout_ms)
.map(timeout_duration_from_ms)
.unwrap_or_else(|| timeout_duration_from_legacy_secs(meta.timeout)),
)
};
RequestTimeouts {
first_byte_timeout,
response_body_timeout,
}
}
fn timeout_duration_from_ms(ms: u64) -> Duration {
Duration::from_millis(ms.clamp(MIN_TIMEOUT_MS, MAX_TIMEOUT_MS))
}
fn timeout_duration_from_legacy_secs(secs: u64) -> Duration {
let ms = secs.saturating_mul(1_000);
timeout_duration_from_ms(ms)
}
async fn spool_request_body(
mut body_rx: mpsc::Receiver<TunnelFrame>,
mut spool_tx: mpsc::Sender<SpoolBodyEvent>,
@@ -888,7 +933,7 @@ async fn relay_upstream_response<B>(
redirect_count: usize,
request_body_mode: &'static str,
emit_proxy_timing_header: bool,
deadline: Instant,
response_body_deadline: Option<Instant>,
) -> Option<Duration>
where
B: hyper::body::Body<Data = Bytes> + Send + Unpin + 'static,
@@ -956,28 +1001,8 @@ where
let mut stream = response.into_body().into_data_stream();
loop {
let Some(remaining) = remaining_timeout(deadline) else {
server.metrics.stream_errors.fetch_add(1, Ordering::Release);
let error_message = "upstream response body timeout".to_string();
log_stream_failure(
stream_log_context(
server,
stream_id,
method,
Some(request_url),
redirect_count,
request_body_size.load(Ordering::Relaxed),
),
&error_message,
total_elapsed,
);
send_error(frame_tx, stream_id, &error_message).await;
return Some(total_elapsed);
};
let chunk_result = match tokio::time::timeout(remaining, stream.next()).await {
Ok(chunk_result) => chunk_result,
Err(_) => {
let chunk_result = if let Some(deadline) = response_body_deadline {
let Some(remaining) = remaining_timeout(deadline) else {
server.metrics.stream_errors.fetch_add(1, Ordering::Release);
let error_message = "upstream response body timeout".to_string();
log_stream_failure(
@@ -994,7 +1019,31 @@ where
);
send_error(frame_tx, stream_id, &error_message).await;
return Some(total_elapsed);
};
match tokio::time::timeout(remaining, stream.next()).await {
Ok(chunk_result) => chunk_result,
Err(_) => {
server.metrics.stream_errors.fetch_add(1, Ordering::Release);
let error_message = "upstream response body timeout".to_string();
log_stream_failure(
stream_log_context(
server,
stream_id,
method,
Some(request_url),
redirect_count,
request_body_size.load(Ordering::Relaxed),
),
&error_message,
total_elapsed,
);
send_error(frame_tx, stream_id, &error_message).await;
return Some(total_elapsed);
}
}
} else {
stream.next().await
};
let Some(chunk_result) = chunk_result else {
@@ -1272,16 +1321,19 @@ async fn handle_stream_inner(
}
}
let deadline = Instant::now()
+ Duration::from_secs(meta.timeout.clamp(MIN_TIMEOUT_SECS, MAX_TIMEOUT_SECS));
let overall_start = Instant::now();
let request_timeouts = resolve_request_timeouts(&meta);
let first_byte_deadline = overall_start + request_timeouts.first_byte_timeout;
let response_body_deadline = request_timeouts
.response_body_timeout
.map(|timeout| overall_start + timeout);
let follow_redirects = follow_redirects_enabled(&meta);
let mut current_headers = sanitize_upstream_headers(&meta.headers);
let timeout = Duration::from_secs(meta.timeout.clamp(MIN_TIMEOUT_SECS, MAX_TIMEOUT_SECS));
let first_byte_timeout = request_timeouts.first_byte_timeout;
let request_body_size = Arc::new(AtomicUsize::new(0));
let request_has_body = request_likely_has_body(&current_method, &meta.headers);
let replay_budget_bytes = state.config.redirect_replay_budget_bytes;
let can_buffer_redirect_body = request_has_body && follow_redirects && replay_budget_bytes > 0;
let overall_start = Instant::now();
let request_body_mode = if can_buffer_redirect_body {
"buffered_fixed"
} else if request_has_body {
@@ -1293,7 +1345,7 @@ async fn handle_stream_inner(
let buffered_body = match collect_request_body_for_replay(
body_rx,
Arc::clone(&request_body_size),
deadline,
first_byte_deadline,
replay_budget_bytes,
)
.await
@@ -1321,7 +1373,12 @@ async fn handle_stream_inner(
replay_body: replay_body_from_buffered(buffered_body, replay_budget_bytes),
}
} else if request_has_body {
prepare_request_body(body_rx, Arc::clone(&request_body_size), deadline, 0)
prepare_request_body(
body_rx,
Arc::clone(&request_body_size),
first_byte_deadline,
0,
)
} else {
PreparedRequestBody {
first_request_body: Some(build_streaming_request_body(
@@ -1341,7 +1398,7 @@ async fn handle_stream_inner(
let mut next_request_body = None::<upstream_client::UpstreamRequestBody>;
loop {
let Some(remaining) = remaining_timeout(deadline) else {
let Some(remaining) = remaining_timeout(first_byte_deadline) else {
log_stream_failure(
stream_log_context(
server,
@@ -1369,7 +1426,7 @@ async fn handle_stream_inner(
current_method.clone(),
&current_headers,
request_body,
remaining.min(timeout),
remaining.min(first_byte_timeout),
meta.http1_only,
)
.await
@@ -1419,7 +1476,7 @@ async fn handle_stream_inner(
redirects_followed,
request_body_mode,
state.config.emit_proxy_timing_header,
deadline,
response_body_deadline,
)
.await;
}
@@ -1431,7 +1488,7 @@ async fn handle_stream_inner(
} => match prepare_redirect_request_body(
prepared_body.replay_body.clone(),
body_mode,
deadline,
first_byte_deadline,
)
.await
{
@@ -1459,7 +1516,7 @@ async fn handle_stream_inner(
redirects_followed,
request_body_mode,
state.config.emit_proxy_timing_header,
deadline,
response_body_deadline,
)
.await;
}
@@ -1515,7 +1572,7 @@ async fn handle_stream_inner(
redirects_followed,
request_body_mode,
state.config.emit_proxy_timing_header,
deadline,
response_body_deadline,
)
.await;
}
@@ -1850,6 +1907,34 @@ mod tests {
);
}
#[test]
fn stream_request_timeouts_use_first_byte_without_response_body_deadline() {
let mut meta = sample_request_meta();
meta.stream = true;
meta.request_timeout_ms = Some(90_000);
meta.stream_first_byte_timeout_ms = Some(12_345);
let timeouts = resolve_request_timeouts(&meta);
assert_eq!(timeouts.first_byte_timeout, Duration::from_millis(12_345));
assert!(timeouts.response_body_timeout.is_none());
}
#[test]
fn non_stream_request_timeouts_use_total_for_response_body_deadline() {
let mut meta = sample_request_meta();
meta.request_timeout_ms = Some(90_000);
meta.stream_first_byte_timeout_ms = Some(12_345);
let timeouts = resolve_request_timeouts(&meta);
assert_eq!(timeouts.first_byte_timeout, Duration::from_millis(90_000));
assert_eq!(
timeouts.response_body_timeout,
Some(Duration::from_millis(90_000))
);
}
#[test]
fn resolve_redirect_changes_post_to_get_for_302() {
let current_url = url::Url::parse("https://redirect.test/start").expect("url");
@@ -2073,7 +2158,7 @@ mod tests {
0,
"empty",
true,
Instant::now(),
Some(Instant::now()),
)
.await;
@@ -2086,6 +2171,51 @@ mod tests {
assert_eq!(server.metrics.stream_errors.load(Ordering::Acquire), 1);
}
#[tokio::test]
async fn stream_response_body_without_total_deadline_allows_late_chunk() {
let state = sample_state(None, None);
let server = sample_server(&state);
let (frame_tx, sent, writer_handle) = spawn_test_writer();
let request_url = url::Url::parse("https://example.com/stream").expect("url");
let request_body_size = AtomicUsize::new(0);
let body = Body::from_stream(stream::once(async {
tokio::time::sleep(Duration::from_millis(10)).await;
Ok::<Bytes, std::convert::Infallible>(Bytes::from_static(b"late"))
}));
let response = Response::builder()
.status(StatusCode::OK)
.body(body)
.expect("response");
relay_upstream_response(
&server,
14,
&hyper::Method::GET,
&request_url,
&frame_tx,
response,
0,
Duration::ZERO,
upstream_client::RequestTiming::default(),
&request_body_size,
0,
"empty",
true,
None,
)
.await;
let result = collect_stream_result(frame_tx, sent, writer_handle).await;
assert!(
result.error.is_none(),
"unexpected error: {:?}",
result.error
);
assert_eq!(result.response.expect("response metadata").status, 200);
assert_eq!(result.body, Bytes::from_static(b"late"));
assert_eq!(server.metrics.stream_errors.load(Ordering::Acquire), 0);
}
#[tokio::test]
async fn follows_redirects_when_explicitly_enabled_for_replayable_post_requests() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
@@ -2427,6 +2557,9 @@ mod tests {
method: "GET".to_string(),
url: "https://example.com/ok".to_string(),
headers: HashMap::new(),
stream: false,
request_timeout_ms: None,
stream_first_byte_timeout_ms: None,
timeout: 30,
follow_redirects: None,
http1_only: false,

6
crates/aether-contracts/src/tunnel.rs
View File

@@ -185,6 +185,12 @@ pub struct RequestMeta {
pub method: String,
pub url: String,
pub headers: std::collections::HashMap<String, String>,
#[serde(default, skip_serializing_if = "is_false")]
pub stream: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request_timeout_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub stream_first_byte_timeout_ms: Option<u64>,
#[serde(default = "default_timeout", deserialize_with = "deserialize_timeout")]
pub timeout: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]

72
crates/aether-contracts/src/tunnel_security.rs
View File

@@ -39,6 +39,8 @@ pub enum TunnelSecurityError {
MissingEncryptedFlag,
#[error("secure tunnel frame payload is too short")]
PayloadTooShort,
#[error("secure tunnel frame sequence is not the expected next value")]
UnexpectedSequence,
#[error("secure tunnel frame encryption failed")]
Encrypt,
#[error("secure tunnel frame decryption failed")]
@@ -51,6 +53,7 @@ pub struct SecureFrameCodec {
seal_prefix: [u8; 4],
open_prefix: [u8; 4],
next_sequence: AtomicU64,
next_open_sequence: AtomicU64,
}
impl SecureFrameCodec {
@@ -90,6 +93,7 @@ impl SecureFrameCodec {
seal_prefix,
open_prefix,
next_sequence: AtomicU64::new(0),
next_open_sequence: AtomicU64::new(0),
})
}
@@ -132,6 +136,10 @@ impl SecureFrameCodec {
let mut payload = frame.payload.clone();
let sequence = payload.get_u64();
let expected_sequence = self.next_open_sequence.load(Ordering::Relaxed);
if sequence != expected_sequence {
return Err(TunnelSecurityError::UnexpectedSequence);
}
let nonce_bytes = nonce_bytes(self.open_prefix, sequence);
let nonce = Nonce::from_slice(&nonce_bytes);
let clear_flags = frame.flags & !FLAG_ENCRYPTED;
@@ -146,6 +154,8 @@ impl SecureFrameCodec {
},
)
.map_err(|_| TunnelSecurityError::Decrypt)?;
self.next_open_sequence
.store(expected_sequence.wrapping_add(1), Ordering::Relaxed);
Ok(Frame::new(
frame.stream_id,
@@ -243,6 +253,68 @@ mod tests {
));
}
#[test]
fn secure_frame_rejects_replayed_sequence() {
let client = SecureFrameCodec::new(&test_key(), "session-1", TunnelSecurityRole::Client)
.expect("client codec");
let server = SecureFrameCodec::new(&test_key(), "session-1", TunnelSecurityRole::Server)
.expect("server codec");
let encrypted = client
.encrypt_frame(Frame::new(
1,
MsgType::RequestBody,
0,
Bytes::from_static(b"secret"),
))
.expect("encrypt");
let wire = Frame::decode(encrypted).expect("wire frame");
server.decrypt_frame(wire.clone()).expect("first decrypt");
assert!(matches!(
server.decrypt_frame(wire),
Err(TunnelSecurityError::UnexpectedSequence)
));
}
#[test]
fn secure_frame_rejects_out_of_order_sequence_without_advancing() {
let client = SecureFrameCodec::new(&test_key(), "session-1", TunnelSecurityRole::Client)
.expect("client codec");
let server = SecureFrameCodec::new(&test_key(), "session-1", TunnelSecurityRole::Server)
.expect("server codec");
let first = Frame::decode(
client
.encrypt_frame(Frame::new(
1,
MsgType::RequestBody,
0,
Bytes::from_static(b"first"),
))
.expect("encrypt first"),
)
.expect("first wire frame");
let second = Frame::decode(
client
.encrypt_frame(Frame::new(
1,
MsgType::RequestBody,
0,
Bytes::from_static(b"second"),
))
.expect("encrypt second"),
)
.expect("second wire frame");
assert!(matches!(
server.decrypt_frame(second),
Err(TunnelSecurityError::UnexpectedSequence)
));
assert_eq!(
server.decrypt_frame(first).expect("first decrypt").payload,
Bytes::from_static(b"first")
);
}
#[test]
fn secure_frame_uses_session_in_key_derivation() {
let session_a = "node-1:connection-a";

15
crates/aether-data/src/repository/proxy_nodes/memory.rs
View File

@@ -9,11 +9,12 @@ use uuid::Uuid;
use super::types::{
bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample,
log_reported_tunnel_error_event, normalize_proxy_metadata,
reconcile_remote_config_after_heartbeat, ProxyNodeEventQuery, ProxyNodeHeartbeatMutation,
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary,
ProxyNodeMetricsStep, ProxyNodeReadRepository, ProxyNodeRegistrationMutation,
ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation,
ProxyNodeWriteRepository, StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
preserve_proxy_metadata_tunnel_security, reconcile_remote_config_after_heartbeat,
ProxyNodeEventQuery, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary, ProxyNodeMetricsStep,
ProxyNodeReadRepository, ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation,
ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository,
StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
StoredProxyNodeMetricsBucket, TunnelMetricsSample, PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
};
use crate::DataLayerError;
@@ -590,6 +591,10 @@ impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
mutation.proxy_metadata.as_ref(),
mutation.proxy_version.as_deref(),
);
let normalized_proxy_metadata = preserve_proxy_metadata_tunnel_security(
previous_proxy_metadata.as_ref(),
normalized_proxy_metadata,
);
if let Some(value) = normalized_proxy_metadata {
node.proxy_metadata = Some(value);
}

20
crates/aether-data/src/repository/proxy_nodes/mysql.rs
View File

@@ -4,11 +4,12 @@ use sqlx::{mysql::MySqlRow, Row};
use super::types::{
bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample,
log_reported_tunnel_error_event, normalize_proxy_metadata,
reconcile_remote_config_after_heartbeat, ProxyNodeEventQuery, ProxyNodeHeartbeatMutation,
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary,
ProxyNodeMetricsStep, ProxyNodeReadRepository, ProxyNodeRegistrationMutation,
ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation,
ProxyNodeWriteRepository, StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
preserve_proxy_metadata_tunnel_security, reconcile_remote_config_after_heartbeat,
ProxyNodeEventQuery, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary, ProxyNodeMetricsStep,
ProxyNodeReadRepository, ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation,
ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository,
StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
StoredProxyNodeMetricsBucket, PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
};
use crate::driver::mysql::MysqlPool;
@@ -779,10 +780,15 @@ WHERE is_manual = 0
if let Some(value) = mutation.avg_latency_ms {
node.avg_latency_ms = Some(value);
}
if let Some(value) = normalize_proxy_metadata(
let normalized_proxy_metadata = normalize_proxy_metadata(
mutation.proxy_metadata.as_ref(),
mutation.proxy_version.as_deref(),
) {
);
let normalized_proxy_metadata = preserve_proxy_metadata_tunnel_security(
previous_proxy_metadata.as_ref(),
normalized_proxy_metadata,
);
if let Some(value) = normalized_proxy_metadata {
node.proxy_metadata = Some(value);
}
if let Some(value) = mutation.total_requests_delta.filter(|value| *value > 0) {

15
crates/aether-data/src/repository/proxy_nodes/postgres.rs
View File

@@ -6,11 +6,12 @@ use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
use super::types::{
bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample,
log_reported_tunnel_error_event, normalize_proxy_metadata,
reconcile_remote_config_after_heartbeat, ProxyNodeEventQuery, ProxyNodeHeartbeatMutation,
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary,
ProxyNodeMetricsStep, ProxyNodeReadRepository, ProxyNodeRegistrationMutation,
ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation,
ProxyNodeWriteRepository, StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
preserve_proxy_metadata_tunnel_security, reconcile_remote_config_after_heartbeat,
ProxyNodeEventQuery, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary, ProxyNodeMetricsStep,
ProxyNodeReadRepository, ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation,
ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository,
StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
StoredProxyNodeMetricsBucket, TunnelMetricsSample, PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
};
use crate::{
@@ -1174,6 +1175,10 @@ impl ProxyNodeWriteRepository for SqlxProxyNodeRepository {
mutation.proxy_metadata.as_ref(),
mutation.proxy_version.as_deref(),
);
let normalized_proxy_metadata = preserve_proxy_metadata_tunnel_security(
existing.proxy_metadata.as_ref(),
normalized_proxy_metadata,
);
sqlx::query(APPLY_HEARTBEAT_SQL)
.bind(&mutation.node_id)

20
crates/aether-data/src/repository/proxy_nodes/sqlite.rs
View File

@@ -4,11 +4,12 @@ use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite};
use super::types::{
bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample,
log_reported_tunnel_error_event, normalize_proxy_metadata,
reconcile_remote_config_after_heartbeat, ProxyNodeEventQuery, ProxyNodeHeartbeatMutation,
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary,
ProxyNodeMetricsStep, ProxyNodeReadRepository, ProxyNodeRegistrationMutation,
ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation,
ProxyNodeWriteRepository, StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
preserve_proxy_metadata_tunnel_security, reconcile_remote_config_after_heartbeat,
ProxyNodeEventQuery, ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation,
ProxyNodeManualUpdateMutation, ProxyNodeMetricsCleanupSummary, ProxyNodeMetricsStep,
ProxyNodeReadRepository, ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation,
ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation, ProxyNodeWriteRepository,
StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
StoredProxyNodeMetricsBucket, PROXY_NODE_EVENT_TYPE_TUNNEL_ERROR,
};
use crate::driver::sqlite::SqlitePool;
@@ -770,10 +771,15 @@ WHERE is_manual = 0
if let Some(value) = mutation.avg_latency_ms {
node.avg_latency_ms = Some(value);
}
if let Some(value) = normalize_proxy_metadata(
let normalized_proxy_metadata = normalize_proxy_metadata(
mutation.proxy_metadata.as_ref(),
mutation.proxy_version.as_deref(),
) {
);
let normalized_proxy_metadata = preserve_proxy_metadata_tunnel_security(
previous_proxy_metadata.as_ref(),
normalized_proxy_metadata,
);
if let Some(value) = normalized_proxy_metadata {
node.proxy_metadata = Some(value);
}
if let Some(value) = mutation.total_requests_delta.filter(|value| *value > 0) {

70
crates/aether-data/src/repository/proxy_nodes/types.rs
View File

@@ -485,6 +485,34 @@ pub fn normalize_proxy_metadata(
}
}
pub fn preserve_proxy_metadata_tunnel_security(
previous_proxy_metadata: Option<&Value>,
next_proxy_metadata: Option<Value>,
) -> Option<Value> {
let Some(tunnel_security) = previous_proxy_metadata
.and_then(|value| value.get("tunnel_security"))
.filter(|value| value.is_object())
.cloned()
else {
return next_proxy_metadata;
};
match next_proxy_metadata {
Some(Value::Object(mut metadata)) => {
metadata
.entry("tunnel_security".to_string())
.or_insert(tunnel_security);
Some(Value::Object(metadata))
}
Some(value) => Some(value),
None => {
let mut metadata = serde_json::Map::new();
metadata.insert("tunnel_security".to_string(), tunnel_security);
Some(Value::Object(metadata))
}
}
}
fn extract_tunnel_metrics_counters(
proxy_metadata: Option<&Value>,
) -> Option<TunnelMetricsCounters> {
@@ -772,10 +800,10 @@ mod tests {
use super::{
bucket_start_unix_secs, build_tunnel_error_event_detail, build_tunnel_metrics_sample,
normalize_proxy_node_scheduling_state, proxy_node_accepts_new_tunnels,
proxy_reported_version, reconcile_remote_config_after_heartbeat,
remote_config_scheduling_state, remote_config_upgrade_target, ProxyNodeMetricsStep,
StoredProxyNode,
normalize_proxy_node_scheduling_state, preserve_proxy_metadata_tunnel_security,
proxy_node_accepts_new_tunnels, proxy_reported_version,
reconcile_remote_config_after_heartbeat, remote_config_scheduling_state,
remote_config_upgrade_target, ProxyNodeMetricsStep, StoredProxyNode,
};
#[test]
@@ -934,6 +962,40 @@ mod tests {
);
}
#[test]
fn preserves_secure_tunnel_metadata_across_heartbeat_metadata_refresh() {
let previous = json!({
"version": "1.0.0",
"tunnel_security": {
"mode": "non_tls_required",
"encryption_key": "BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc="
}
});
let next = json!({
"version": "1.0.1",
"tunnel_metrics": {"connect_successes": 1}
});
let merged = preserve_proxy_metadata_tunnel_security(Some(&previous), Some(next))
.expect("metadata should remain present");
assert_eq!(
merged
.pointer("/tunnel_security/mode")
.and_then(|v| v.as_str()),
Some("non_tls_required")
);
assert_eq!(
merged
.pointer("/tunnel_security/encryption_key")
.and_then(|v| v.as_str()),
Some("BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc=")
);
assert_eq!(
merged.pointer("/tunnel_metrics/connect_successes"),
Some(&json!(1))
);
}
#[test]
fn maps_timestamps_to_metric_buckets() {
assert_eq!(

3
crates/aether-testkit/src/bin/capacity_curve_baseline.rs
View File

@@ -609,6 +609,9 @@ fn relay_envelope() -> Vec<u8> {
"content-type".to_string(),
"application/json".to_string(),
)]),
stream: false,
request_timeout_ms: None,
stream_first_byte_timeout_ms: None,
timeout: 30,
follow_redirects: None,
http1_only: false,

3
crates/aether-testkit/src/bin/gateway_tunnel_stream_baseline.rs
View File

@@ -117,6 +117,9 @@ fn relay_envelope() -> Vec<u8> {
"content-type".to_string(),
"application/json".to_string(),
)]),
stream: true,
request_timeout_ms: None,
stream_first_byte_timeout_ms: None,
timeout: 30,
follow_redirects: None,
http1_only: false,

3
crates/aether-testkit/src/bin/multi_instance_owner_relay_baseline.rs
View File

@@ -253,6 +253,9 @@ fn relay_envelope() -> Vec<u8> {
"content-type".to_string(),
"application/json".to_string(),
)]),
stream: false,
request_timeout_ms: None,
stream_first_byte_timeout_ms: None,
timeout: 30,
follow_redirects: None,
http1_only: false,