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

Improve tunnel proxy diagnostics and request body spooling

Browse Source
This commit is contained in:
fawney19
2026-04-18 21:13:37 +08:00
parent 3363592751
commit f2a3836877
5 changed files with 712 additions and 260 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -913,28 +913,51 @@ fn proxy_reference_matches_node_id(value: Option<&Value>, node_id: &str) -> bool
.is_some_and(|value| value == node_id)
}
fn build_proxy_connectivity_result(
probe_url: &str,
timeout_secs: u64,
success: bool,
latency_ms: Option<u64>,
exit_ip: Option<String>,
error: Option<String>,
) -> Value {
json!({
"success": success,
"latency_ms": latency_ms,
"exit_ip": exit_ip,
"error": error,
"probe_url": probe_url.trim(),
"timeout_secs": timeout_secs,
})
}
async fn test_proxy_node_connectivity(
state: &AdminAppState<'_>,
node: &aether_data::repository::proxy_nodes::StoredProxyNode,
) -> Value {
let probe_url = proxy_connectivity_probe_url();
if node.is_manual {
let Some(proxy_url) = node.proxy_url.as_deref() else {
return json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": "手动节点缺少 proxy_url",
});
return build_proxy_connectivity_result(
&probe_url,
PROXY_CONNECTIVITY_TIMEOUT_SECS,
false,
None,
None,
Some("手动节点缺少 proxy_url".to_string()),
);
};
let endpoint = match parse_manual_proxy_endpoint(proxy_url, "proxy_url") {
Ok(endpoint) => endpoint,
Err(detail) => {
return json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": detail,
});
return build_proxy_connectivity_result(
&probe_url,
PROXY_CONNECTIVITY_TIMEOUT_SECS,
false,
None,
None,
Some(detail),
);
}
};
let proxy_url = proxy_url_with_auth(
@@ -947,83 +970,118 @@ async fn test_proxy_node_connectivity(
}
if !node.tunnel_mode {
return json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": "non-tunnel mode is no longer supported, please upgrade aether-proxy to use tunnel mode",
});
return build_proxy_connectivity_result(
&probe_url,
PROXY_CONNECTIVITY_TIMEOUT_SECS,
false,
None,
None,
Some(
"non-tunnel mode is no longer supported, please upgrade aether-proxy to use tunnel mode"
.to_string(),
),
);
}
if !node.status.eq_ignore_ascii_case("online") || !node.tunnel_connected {
return json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": "tunnel 未连接",
});
return build_proxy_connectivity_result(
&probe_url,
PROXY_CONNECTIVITY_TIMEOUT_SECS,
false,
None,
None,
Some("tunnel 未连接".to_string()),
);
}
let probe_url = proxy_connectivity_probe_url();
match probe_tunnel_proxy_connectivity(state.app(), &node.id, &probe_url).await {
match probe_tunnel_proxy_connectivity(
state.app(),
&node.id,
&probe_url,
PROXY_CONNECTIVITY_TIMEOUT_SECS,
)
.await
{
Ok(result) => {
if let Ok(status) = reqwest::StatusCode::from_u16(result.status) {
if status.is_success() {
return json!({
"success": true,
"latency_ms": result.latency_ms,
"exit_ip": parse_proxy_probe_exit_ip(&result.body),
"error": null,
});
return build_proxy_connectivity_result(
&probe_url,
PROXY_CONNECTIVITY_TIMEOUT_SECS,
true,
Some(result.latency_ms),
parse_proxy_probe_exit_ip(&result.body),
None,
);
}
return json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": sanitize_proxy_error(&format_proxy_probe_status_error(status, &result.body)),
});
return build_proxy_connectivity_result(
&probe_url,
PROXY_CONNECTIVITY_TIMEOUT_SECS,
false,
None,
None,
Some(sanitize_proxy_error(&format_proxy_probe_status_error(
status,
&result.body,
))),
);
}
json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": format!("代理探测返回非法状态码: {}", result.status),
})
build_proxy_connectivity_result(
&probe_url,
PROXY_CONNECTIVITY_TIMEOUT_SECS,
false,
None,
None,
Some(format!("代理探测返回非法状态码: {}", result.status)),
)
}
Err(error) => json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": sanitize_proxy_error(&error),
}),
Err(error) => build_proxy_connectivity_result(
&probe_url,
PROXY_CONNECTIVITY_TIMEOUT_SECS,
false,
None,
None,
Some(sanitize_proxy_error(&error)),
),
}
}
async fn test_manual_proxy_connectivity(proxy_url: &str) -> Value {
let probe_url = proxy_connectivity_probe_url();
test_manual_proxy_connectivity_with_probe_url(proxy_url, &probe_url).await
test_manual_proxy_connectivity_with_probe_url(
proxy_url,
&probe_url,
PROXY_CONNECTIVITY_TIMEOUT_SECS,
)
.await
}
async fn test_manual_proxy_connectivity_with_probe_url(proxy_url: &str, probe_url: &str) -> Value {
async fn test_manual_proxy_connectivity_with_probe_url(
proxy_url: &str,
probe_url: &str,
timeout_secs: u64,
) -> Value {
let started_at = Instant::now();
let proxy = match reqwest::Proxy::all(proxy_url) {
Ok(proxy) => proxy,
Err(error) => {
return json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": sanitize_proxy_error(&format_upstream_request_error(&error)),
});
return build_proxy_connectivity_result(
probe_url,
timeout_secs,
false,
None,
None,
Some(sanitize_proxy_error(&format_upstream_request_error(&error))),
);
}
};
let mut builder = reqwest::Client::builder()
.no_proxy()
.redirect(reqwest::redirect::Policy::none())
.connect_timeout(Duration::from_secs(5))
.timeout(Duration::from_secs(PROXY_CONNECTIVITY_TIMEOUT_SECS))
.timeout(Duration::from_secs(timeout_secs))
.proxy(proxy)
.user_agent("aether-gateway/proxy-connectivity");
if proxy_url
@@ -1036,54 +1094,66 @@ async fn test_manual_proxy_connectivity_with_probe_url(proxy_url: &str, probe_ur
let client = match builder.build() {
Ok(client) => client,
Err(error) => {
return json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": sanitize_proxy_error(&format_upstream_request_error(&error)),
});
return build_proxy_connectivity_result(
probe_url,
timeout_secs,
false,
None,
None,
Some(sanitize_proxy_error(&format_upstream_request_error(&error))),
);
}
};
let response = match client.get(probe_url).send().await {
Ok(response) => response,
Err(error) => {
return json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": sanitize_proxy_error(&format_upstream_request_error(&error)),
});
return build_proxy_connectivity_result(
probe_url,
timeout_secs,
false,
None,
None,
Some(sanitize_proxy_error(&format_upstream_request_error(&error))),
);
}
};
let status = response.status();
let body = match response.text().await {
Ok(body) => body,
Err(error) => {
return json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": sanitize_proxy_error(&format_upstream_request_error(&error)),
});
return build_proxy_connectivity_result(
probe_url,
timeout_secs,
false,
None,
None,
Some(sanitize_proxy_error(&format_upstream_request_error(&error))),
);
}
};
if !status.is_success() {
return json!({
"success": false,
"latency_ms": null,
"exit_ip": null,
"error": sanitize_proxy_error(&format_proxy_probe_status_error(status, &body)),
});
return build_proxy_connectivity_result(
probe_url,
timeout_secs,
false,
None,
None,
Some(sanitize_proxy_error(&format_proxy_probe_status_error(
status, &body,
))),
);
}
json!({
"success": true,
"latency_ms": started_at.elapsed().as_millis() as u64,
"exit_ip": parse_proxy_probe_exit_ip(&body),
"error": null,
})
build_proxy_connectivity_result(
probe_url,
timeout_secs,
true,
Some(started_at.elapsed().as_millis() as u64),
parse_proxy_probe_exit_ip(&body),
None,
)
}
struct TunnelConnectivityProbeResult {
@@ -1096,6 +1166,7 @@ async fn probe_tunnel_proxy_connectivity(
state: &crate::AppState,
node_id: &str,
probe_url: &str,
timeout_secs: u64,
) -> Result<TunnelConnectivityProbeResult, String> {
let trimmed_node_id = node_id.trim();
if trimmed_node_id.is_empty() {
@@ -1103,7 +1174,13 @@ async fn probe_tunnel_proxy_connectivity(
}
if state.tunnel.has_local_proxy(trimmed_node_id) {
return probe_tunnel_proxy_connectivity_locally(state, trimmed_node_id, probe_url).await;
return probe_tunnel_proxy_connectivity_locally(
state,
trimmed_node_id,
probe_url,
timeout_secs,
)
.await;
}
if let Some(owner) = state
@@ -1117,6 +1194,7 @@ async fn probe_tunnel_proxy_connectivity(
state,
trimmed_node_id,
probe_url,
timeout_secs,
&owner.relay_base_url,
&owner.gateway_instance_id,
)
@@ -1130,18 +1208,19 @@ async fn probe_tunnel_proxy_connectivity(
.map_err(|err| format!("clear stale local tunnel attachment failed: {err}"))?;
}
probe_tunnel_proxy_connectivity_locally(state, trimmed_node_id, probe_url).await
probe_tunnel_proxy_connectivity_locally(state, trimmed_node_id, probe_url, timeout_secs).await
}
async fn probe_tunnel_proxy_connectivity_locally(
state: &crate::AppState,
node_id: &str,
probe_url: &str,
timeout_secs: u64,
) -> Result<TunnelConnectivityProbeResult, String> {
let started_at = Instant::now();
let result = state
.tunnel
.probe_node_url_with_response(node_id, probe_url, PROXY_CONNECTIVITY_TIMEOUT_SECS)
.probe_node_url_with_response(node_id, probe_url, timeout_secs)
.await?;
Ok(TunnelConnectivityProbeResult {
status: result.status,
@@ -1154,6 +1233,7 @@ async fn probe_tunnel_proxy_connectivity_via_owner(
state: &crate::AppState,
node_id: &str,
probe_url: &str,
timeout_secs: u64,
relay_base_url: &str,
owner_instance_id: &str,
) -> Result<TunnelConnectivityProbeResult, String> {
@@ -1171,8 +1251,8 @@ async fn probe_tunnel_proxy_connectivity_via_owner(
state.tunnel.local_instance_id(),
)
.header(TUNNEL_RELAY_OWNER_INSTANCE_HEADER, owner_instance_id)
.timeout(Duration::from_secs(PROXY_CONNECTIVITY_TIMEOUT_SECS))
.body(build_tunnel_probe_relay_envelope(probe_url)?)
.timeout(Duration::from_secs(timeout_secs))
.body(build_tunnel_probe_relay_envelope(probe_url, timeout_secs)?)
.send()
.await
.map_err(|error| format!("owner tunnel relay probe failed: {error}"))?;
@@ -1195,12 +1275,15 @@ async fn probe_tunnel_proxy_connectivity_via_owner(
})
}
fn build_tunnel_probe_relay_envelope(probe_url: &str) -> Result<Vec<u8>, String> {
fn build_tunnel_probe_relay_envelope(
probe_url: &str,
timeout_secs: u64,
) -> Result<Vec<u8>, String> {
let meta = crate::tunnel::tunnel_protocol::RequestMeta {
method: "GET".to_string(),
url: probe_url.trim().to_string(),
headers: std::collections::HashMap::new(),
timeout: PROXY_CONNECTIVITY_TIMEOUT_SECS,
timeout: timeout_secs,
follow_redirects: Some(false),
http1_only: false,
};

7
apps/aether-gateway/src/tests/control/admin/proxy_nodes.rs
View File

@@ -943,6 +943,11 @@ async fn gateway_tests_disconnected_tunnel_proxy_nodes_locally() {
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], false);
assert_eq!(payload["error"], "tunnel 未连接");
assert_eq!(
payload["probe_url"],
"https://www.cloudflare.com/cdn-cgi/trace"
);
assert_eq!(payload["timeout_secs"], 10);
gateway_handle.abort();
}
@@ -1072,6 +1077,8 @@ async fn gateway_tests_connected_tunnel_proxy_nodes_with_active_probe() {
assert_eq!(payload["success"], true);
assert!(payload["latency_ms"].is_u64());
assert_eq!(payload["exit_ip"], "203.0.113.10");
assert_eq!(payload["probe_url"], "https://probe.example/cdn-cgi/trace");
assert_eq!(payload["timeout_secs"], 10);
gateway_handle.abort();
}

634
apps/aether-proxy/src/tunnel/stream_handler.rs
View File

@@ -7,6 +7,7 @@ use std::io;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::sync::Mutex;
use std::time::{Duration, Instant};
use aether_runtime::hold_admission_permit_until;
@@ -15,7 +16,7 @@ use futures_util::stream;
use futures_util::StreamExt;
use http_body_util::BodyExt;
use hyper::body::Frame as BodyFrame;
use tokio::sync::mpsc;
use tokio::sync::{mpsc, Notify};
use tracing::{debug, info, warn};
use crate::state::{AppState, ServerContext};
@@ -80,10 +81,10 @@ const REDIRECT_SENSITIVE_HEADERS: &[&str] = &[
"www-authenticate",
];
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone)]
enum ReplayableRequestBody {
None,
Replayable(Bytes),
Pending(Arc<RequestBodyReplayState>),
NonReplayable,
}
@@ -92,6 +93,39 @@ struct PreparedRequestBody {
replay_body: ReplayableRequestBody,
}
#[derive(Debug)]
struct RequestBodyReplayState {
budget_bytes: usize,
state: Mutex<RequestBodyReplayStatus>,
ready: Notify,
}
#[derive(Debug)]
enum RequestBodyReplayStatus {
Collecting {
chunks: Vec<Bytes>,
buffered_len: usize,
},
Ready(Bytes),
Empty,
NonReplayable,
Error(String),
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum ReplayBodyResolution {
Empty,
Replayable(Bytes),
NonReplayable,
}
#[derive(Debug)]
enum SpoolBodyEvent {
Data(Bytes),
Error(String),
End,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RedirectBodyMode {
Empty,
@@ -221,35 +255,134 @@ fn log_stream_failure(ctx: StreamLogContext<'_>, error: &str, duration: Duration
}
impl PreparedRequestBody {
fn streaming(
body_rx: mpsc::Receiver<TunnelFrame>,
body_size: Arc<AtomicUsize>,
) -> PreparedRequestBody {
PreparedRequestBody {
first_request_body: Some(build_streaming_request_body(body_rx, body_size)),
replay_body: ReplayableRequestBody::NonReplayable,
}
}
fn take_first_request_body(&mut self) -> upstream_client::UpstreamRequestBody {
self.first_request_body
.take()
.unwrap_or_else(empty_request_body)
}
}
fn build_redirect_request_body(
&self,
body_mode: RedirectBodyMode,
) -> Option<upstream_client::UpstreamRequestBody> {
match body_mode {
RedirectBodyMode::Empty => Some(empty_request_body()),
RedirectBodyMode::Replay => match &self.replay_body {
ReplayableRequestBody::None => Some(empty_request_body()),
ReplayableRequestBody::Replayable(body) => {
Some(buffered_request_body(body.clone()))
async fn prepare_redirect_request_body(
replay_body: ReplayableRequestBody,
body_mode: RedirectBodyMode,
deadline: Instant,
) -> Result<Option<upstream_client::UpstreamRequestBody>, String> {
match body_mode {
RedirectBodyMode::Empty => Ok(Some(empty_request_body())),
RedirectBodyMode::Replay => match replay_body {
ReplayableRequestBody::None => Ok(Some(empty_request_body())),
ReplayableRequestBody::Pending(state) => {
match state.wait_for_resolution(deadline).await? {
ReplayBodyResolution::Empty => Ok(Some(empty_request_body())),
ReplayBodyResolution::Replayable(body) => Ok(Some(buffered_request_body(body))),
ReplayBodyResolution::NonReplayable => Ok(None),
}
ReplayableRequestBody::NonReplayable => None,
},
}
ReplayableRequestBody::NonReplayable => Ok(None),
},
}
}
impl RequestBodyReplayState {
fn new(budget_bytes: usize) -> Self {
Self {
budget_bytes,
state: Mutex::new(RequestBodyReplayStatus::Collecting {
chunks: Vec::new(),
buffered_len: 0,
}),
ready: Notify::new(),
}
}
fn push_chunk(&self, payload: Bytes) {
let mut notify = false;
{
let mut state = self.state.lock().expect("request body replay state lock");
if let RequestBodyReplayStatus::Collecting {
chunks,
buffered_len,
} = &mut *state
{
let next_len = buffered_len.saturating_add(payload.len());
if next_len > self.budget_bytes {
chunks.clear();
*state = RequestBodyReplayStatus::NonReplayable;
notify = true;
} else {
*buffered_len = next_len;
chunks.push(payload);
}
}
}
if notify {
self.ready.notify_waiters();
}
}
fn finish(&self) {
let notify;
{
let mut state = self.state.lock().expect("request body replay state lock");
let next_state = match std::mem::replace(&mut *state, RequestBodyReplayStatus::Empty) {
RequestBodyReplayStatus::Collecting {
chunks,
buffered_len,
} => {
if buffered_len == 0 {
RequestBodyReplayStatus::Empty
} else {
let mut buffered = BytesMut::with_capacity(buffered_len);
for chunk in chunks {
buffered.extend_from_slice(&chunk);
}
RequestBodyReplayStatus::Ready(buffered.freeze())
}
}
terminal => terminal,
};
notify = !matches!(next_state, RequestBodyReplayStatus::Collecting { .. });
*state = next_state;
}
if notify {
self.ready.notify_waiters();
}
}
fn fail(&self, message: String) {
{
let mut state = self.state.lock().expect("request body replay state lock");
*state = RequestBodyReplayStatus::Error(message);
}
self.ready.notify_waiters();
}
async fn wait_for_resolution(&self, deadline: Instant) -> Result<ReplayBodyResolution, String> {
loop {
let resolution = {
let state = self.state.lock().expect("request body replay state lock");
match &*state {
RequestBodyReplayStatus::Collecting { .. } => None,
RequestBodyReplayStatus::Ready(body) => {
Some(Ok(ReplayBodyResolution::Replayable(body.clone())))
}
RequestBodyReplayStatus::Empty => Some(Ok(ReplayBodyResolution::Empty)),
RequestBodyReplayStatus::NonReplayable => {
Some(Ok(ReplayBodyResolution::NonReplayable))
}
RequestBodyReplayStatus::Error(message) => Some(Err(message.clone())),
}
};
if let Some(resolution) = resolution {
return resolution;
}
let Some(remaining) = remaining_timeout(deadline) else {
return Err("upstream timeout".to_string());
};
tokio::time::timeout(remaining, self.ready.notified())
.await
.map_err(|_| "upstream timeout".to_string())?;
}
}
}
@@ -258,6 +391,29 @@ fn follow_redirects_enabled(meta: &RequestMeta) -> bool {
meta.follow_redirects == Some(true)
}
fn request_likely_has_body(
method: &hyper::Method,
headers: &std::collections::HashMap<String, String>,
) -> bool {
if matches!(
*method,
hyper::Method::GET | hyper::Method::HEAD | hyper::Method::OPTIONS | hyper::Method::TRACE
) {
return headers.iter().any(|(name, value)| {
name.eq_ignore_ascii_case("content-length")
&& value
.trim()
.parse::<u64>()
.ok()
.is_some_and(|value| value > 0)
}) || headers
.keys()
.any(|name| name.eq_ignore_ascii_case("transfer-encoding"));
}
true
}
fn sanitize_upstream_headers(
headers: &std::collections::HashMap<String, String>,
) -> Vec<(String, String)> {
@@ -296,84 +452,38 @@ fn buffered_request_body(body: Bytes) -> upstream_client::UpstreamRequestBody {
upstream_client::stream_request_body(stream::once(async move { Ok(BodyFrame::data(body)) }))
}
async fn prepare_redirect_request_body(
mut body_rx: mpsc::Receiver<TunnelFrame>,
// Drain tunnel body frames on a detached task so the shared dispatcher is no
// longer coupled to upstream body polling. Redirect replay still reuses a full
// in-memory copy when the request body completes within budget.
fn prepare_request_body(
body_rx: mpsc::Receiver<TunnelFrame>,
body_size: Arc<AtomicUsize>,
deadline: Instant,
replay_budget_bytes: usize,
) -> Result<PreparedRequestBody, String> {
let mut prefix_chunks = Vec::new();
let mut buffered_len = 0usize;
let mut finished = false;
) -> PreparedRequestBody {
let (spool_tx, spool_rx) = mpsc::unbounded_channel();
let replay_state = if replay_budget_bytes == 0 {
None
} else {
Some(Arc::new(RequestBodyReplayState::new(replay_budget_bytes)))
};
let replay_body = match replay_state.as_ref() {
Some(state) => ReplayableRequestBody::Pending(Arc::clone(state)),
None => ReplayableRequestBody::NonReplayable,
};
loop {
let Some(frame) = recv_body_frame_with_deadline(&mut body_rx, deadline).await? else {
finished = true;
break;
};
tokio::spawn(spool_request_body(
body_rx,
spool_tx,
replay_state,
body_size,
deadline,
));
match frame.msg_type {
MsgType::RequestBody => {
let end_stream = frame.is_end_stream();
let payload = decompress_if_gzip(&frame)
.map_err(|error| format!("gzip decompress failed: {error}"))?;
if !payload.is_empty() {
body_size.fetch_add(payload.len(), Ordering::Relaxed);
buffered_len = buffered_len.saturating_add(payload.len());
prefix_chunks.push(payload);
}
if end_stream {
finished = true;
break;
}
if buffered_len > replay_budget_bytes {
break;
}
}
MsgType::StreamError => {
let message = String::from_utf8(frame.payload.to_vec())
.unwrap_or_else(|_| "client cancelled request body".to_string());
return Err(message);
}
MsgType::StreamEnd => {
finished = true;
break;
}
_ => continue,
}
PreparedRequestBody {
first_request_body: Some(build_spooled_request_body(spool_rx)),
replay_body,
}
if buffered_len <= replay_budget_bytes && finished {
let total_len: usize = prefix_chunks.iter().map(Bytes::len).sum();
if total_len == 0 {
return Ok(PreparedRequestBody {
first_request_body: Some(empty_request_body()),
replay_body: ReplayableRequestBody::None,
});
}
let mut buffered = BytesMut::with_capacity(total_len);
for chunk in prefix_chunks {
buffered.extend_from_slice(&chunk);
}
let body = buffered.freeze();
return Ok(PreparedRequestBody {
first_request_body: Some(buffered_request_body(body.clone())),
replay_body: ReplayableRequestBody::Replayable(body),
});
}
Ok(PreparedRequestBody {
first_request_body: Some(build_prefixed_request_body(
prefix_chunks,
body_rx,
body_size,
)),
replay_body: ReplayableRequestBody::NonReplayable,
})
}
async fn recv_body_frame_with_deadline(
@@ -392,6 +502,97 @@ fn remaining_timeout(deadline: Instant) -> Option<Duration> {
deadline.checked_duration_since(Instant::now())
}
async fn spool_request_body(
mut body_rx: mpsc::Receiver<TunnelFrame>,
spool_tx: mpsc::UnboundedSender<SpoolBodyEvent>,
replay_state: Option<Arc<RequestBodyReplayState>>,
body_size: Arc<AtomicUsize>,
deadline: Instant,
) {
let mut spool_tx = Some(spool_tx);
loop {
let frame = match recv_body_frame_with_deadline(&mut body_rx, deadline).await {
Ok(frame) => frame,
Err(message) => {
if let Some(state) = &replay_state {
state.fail(message.clone());
}
send_spool_event(&mut spool_tx, SpoolBodyEvent::Error(message));
return;
}
};
let Some(frame) = frame else {
if let Some(state) = &replay_state {
state.finish();
}
send_spool_event(&mut spool_tx, SpoolBodyEvent::End);
return;
};
match frame.msg_type {
MsgType::RequestBody => {
let end_stream = frame.is_end_stream();
let payload = match decompress_if_gzip(&frame) {
Ok(payload) => payload,
Err(error) => {
let message = format!("gzip decompress failed: {error}");
if let Some(state) = &replay_state {
state.fail(message.clone());
}
send_spool_event(&mut spool_tx, SpoolBodyEvent::Error(message));
return;
}
};
if !payload.is_empty() {
body_size.fetch_add(payload.len(), Ordering::Relaxed);
if let Some(state) = &replay_state {
state.push_chunk(payload.clone());
}
send_spool_event(&mut spool_tx, SpoolBodyEvent::Data(payload));
}
if end_stream {
if let Some(state) = &replay_state {
state.finish();
}
send_spool_event(&mut spool_tx, SpoolBodyEvent::End);
return;
}
}
MsgType::StreamError => {
let message = String::from_utf8(frame.payload.to_vec())
.unwrap_or_else(|_| "client cancelled request body".to_string());
if let Some(state) = &replay_state {
state.fail(message.clone());
}
send_spool_event(&mut spool_tx, SpoolBodyEvent::Error(message));
return;
}
MsgType::StreamEnd => {
if let Some(state) = &replay_state {
state.finish();
}
send_spool_event(&mut spool_tx, SpoolBodyEvent::End);
return;
}
_ => continue,
}
}
}
fn send_spool_event(
spool_tx: &mut Option<mpsc::UnboundedSender<SpoolBodyEvent>>,
event: SpoolBodyEvent,
) {
if let Some(sender) = spool_tx.as_ref() {
if sender.send(event).is_err() {
*spool_tx = None;
}
}
}
fn remove_headers_case_insensitive(headers: &mut Vec<(String, String)>, blocked: &[&str]) {
headers.retain(|(name, _)| {
let normalized = name.to_ascii_lowercase();
@@ -433,8 +634,9 @@ fn resolve_redirect<B>(
}
StatusCode::TEMPORARY_REDIRECT | StatusCode::PERMANENT_REDIRECT => match replay_body {
ReplayableRequestBody::NonReplayable => return RedirectDecision::Stop,
ReplayableRequestBody::None => RedirectBodyMode::Empty,
ReplayableRequestBody::Replayable(_) => RedirectBodyMode::Replay,
ReplayableRequestBody::None | ReplayableRequestBody::Pending(_) => {
RedirectBodyMode::Replay
}
},
_ => return RedirectDecision::Stop,
};
@@ -891,41 +1093,37 @@ async fn handle_stream_inner(
let mut current_headers = sanitize_upstream_headers(&meta.headers);
let timeout = Duration::from_secs(meta.timeout.clamp(MIN_TIMEOUT_SECS, MAX_TIMEOUT_SECS));
let request_body_size = Arc::new(AtomicUsize::new(0));
let mut prepared_body = if follow_redirects {
match prepare_redirect_request_body(
let request_has_body = request_likely_has_body(&current_method, &meta.headers);
let mut prepared_body = if request_has_body {
let replay_budget_bytes = if follow_redirects {
state.config.redirect_replay_budget_bytes
} else {
0
};
prepare_request_body(
body_rx,
Arc::clone(&request_body_size),
deadline,
state.config.redirect_replay_budget_bytes,
replay_budget_bytes,
)
.await
{
Ok(body) => body,
Err(message) => {
log_stream_failure(
stream_log_context(
server,
stream_id,
&current_method,
Some(&current_url),
0,
request_body_size.load(Ordering::Relaxed),
),
&message,
Duration::ZERO,
);
send_error(frame_tx, stream_id, &message).await;
return None;
}
}
} else {
PreparedRequestBody::streaming(body_rx, Arc::clone(&request_body_size))
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
},
}
};
let overall_start = Instant::now();
let mut total_dns_ms = 0u64;
let mut redirects_followed = 0usize;
let mut next_body_mode = None::<RedirectBodyMode>;
let mut next_request_body = None::<upstream_client::UpstreamRequestBody>;
loop {
let Some(remaining) = remaining_timeout(deadline) else {
@@ -944,29 +1142,9 @@ async fn handle_stream_inner(
send_error(frame_tx, stream_id, "upstream timeout").await;
return None;
};
let request_body = match next_body_mode.take() {
Some(mode) => match prepared_body.build_redirect_request_body(mode) {
Some(body) => body,
None => {
let error_message = "upstream redirect error: body not replayable";
log_stream_failure(
stream_log_context(
server,
stream_id,
&current_method,
Some(&current_url),
redirects_followed,
request_body_size.load(Ordering::Relaxed),
),
error_message,
overall_start.elapsed(),
);
send_error(frame_tx, stream_id, error_message).await;
return None;
}
},
None => prepared_body.take_first_request_body(),
};
let request_body = next_request_body
.take()
.unwrap_or_else(|| prepared_body.take_first_request_body());
let response_ctx = match execute_upstream_request(
state,
@@ -1030,14 +1208,54 @@ async fn handle_stream_inner(
url,
headers,
body_mode,
} => {
redirects_followed += 1;
current_method = method;
current_url = url;
current_headers = headers;
next_body_mode = Some(body_mode);
continue;
}
} => match prepare_redirect_request_body(
prepared_body.replay_body.clone(),
body_mode,
deadline,
)
.await
{
Ok(Some(body)) => {
redirects_followed += 1;
current_method = method;
current_url = url;
current_headers = headers;
next_request_body = Some(body);
continue;
}
Ok(None) => {
return relay_upstream_response(
server,
stream_id,
&current_method,
&current_url,
frame_tx,
response_ctx.response,
total_dns_ms,
overall_start.elapsed(),
response_ctx.request_timing,
request_body_size.as_ref(),
redirects_followed,
)
.await;
}
Err(message) => {
log_stream_failure(
stream_log_context(
server,
stream_id,
&current_method,
Some(&current_url),
redirects_followed,
request_body_size.load(Ordering::Relaxed),
),
&message,
overall_start.elapsed(),
);
send_error(frame_tx, stream_id, &message).await;
return None;
}
},
RedirectDecision::Error(message) => {
let error_message = format!("upstream redirect error: {message}");
log_stream_failure(
@@ -1096,6 +1314,28 @@ fn build_streaming_request_body(
build_prefixed_request_body(Vec::new(), body_rx, body_size)
}
fn build_spooled_request_body(
spool_rx: mpsc::UnboundedReceiver<SpoolBodyEvent>,
) -> upstream_client::UpstreamRequestBody {
let body_stream = stream::unfold((spool_rx, false), |(mut spool_rx, finished)| async move {
if finished {
return None;
}
match spool_rx.recv().await {
Some(SpoolBodyEvent::Data(payload)) => {
Some((Ok(BodyFrame::data(payload)), (spool_rx, false)))
}
Some(SpoolBodyEvent::Error(message)) => {
Some((Err(io::Error::other(message)), (spool_rx, true)))
}
Some(SpoolBodyEvent::End) | None => None,
}
});
upstream_client::stream_request_body(body_stream)
}
fn build_prefixed_request_body(
prefix_chunks: Vec<Bytes>,
body_rx: mpsc::Receiver<TunnelFrame>,
@@ -1187,6 +1427,15 @@ mod tests {
use crate::target_filter::DnsCache;
use crate::tunnel::client::build_tls_config;
fn completed_replay_body(body: Bytes) -> ReplayableRequestBody {
let state = Arc::new(RequestBodyReplayState::new(body.len().max(1)));
if !body.is_empty() {
state.push_chunk(body);
}
state.finish();
ReplayableRequestBody::Pending(state)
}
#[tokio::test]
async fn streaming_request_body_yields_chunks_and_tracks_size() {
let (tx, rx) = mpsc::channel(4);
@@ -1258,6 +1507,79 @@ mod tests {
assert_eq!(body_size.load(Ordering::Relaxed), 0);
}
#[tokio::test]
async fn prepare_request_body_streams_immediately_and_replays_after_completion() {
let (tx, rx) = mpsc::channel(4);
let body_size = Arc::new(AtomicUsize::new(0));
let prepared = prepare_request_body(
rx,
Arc::clone(&body_size),
Instant::now() + Duration::from_secs(1),
1024,
);
let mut body = prepared
.first_request_body
.expect("first request body should be present");
tx.send(TunnelFrame::new(
1,
MsgType::RequestBody,
0,
Bytes::from_static(b"hello "),
))
.await
.expect("send first chunk");
let first = body
.frame()
.await
.expect("first frame should exist")
.expect("first frame should be ok")
.into_data()
.expect("first data frame");
assert_eq!(first, Bytes::from_static(b"hello "));
tx.send(TunnelFrame::new(
1,
MsgType::RequestBody,
flags::END_STREAM,
Bytes::from_static(b"world"),
))
.await
.expect("send final chunk");
drop(tx);
let second = body
.frame()
.await
.expect("second frame should exist")
.expect("second frame should be ok")
.into_data()
.expect("second data frame");
assert_eq!(second, Bytes::from_static(b"world"));
assert!(body.frame().await.is_none());
let mut replay = prepare_redirect_request_body(
prepared.replay_body.clone(),
RedirectBodyMode::Replay,
Instant::now() + Duration::from_secs(1),
)
.await
.expect("redirect replay should resolve")
.expect("body should be replayable");
let frame = replay
.frame()
.await
.expect("replayed frame should exist")
.expect("replayed frame should be ok");
assert_eq!(
frame.into_data().expect("data frame"),
Bytes::from_static(b"hello world")
);
assert!(replay.frame().await.is_none());
assert_eq!(body_size.load(Ordering::Relaxed), 11);
}
#[test]
fn selects_http1_only_client_when_request_metadata_requires_it() {
let state = sample_state(None, None);
@@ -1289,7 +1611,7 @@ mod tests {
&current_url,
&hyper::Method::POST,
&[("content-type".into(), "application/json".into())],
&ReplayableRequestBody::Replayable(Bytes::from_static(br#"{"ok":true}"#)),
&completed_replay_body(Bytes::from_static(br#"{"ok":true}"#)),
0,
);