mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
fix: harden tunnel security and timeout handling
This commit is contained in:
@@ -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),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 或显式关闭的旧节点仍按原明文协议工作。
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(¤t_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(),
|
||||
¤t_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,
|
||||
|
||||
Reference in New Issue
Block a user