refactor(proxy): preserve original header wire order and add non-streaming body timeout

- Rewrite build_raw_request to emit headers in original
  client-sent sequence instead of hash-map order
- Remove unused OriginalHeaderCases::get_all method
- Add body_timeout to read_decoded_body to prevent
  requests hanging when upstream stalls after headers
This commit is contained in:
YoVinchen
2026-03-29 18:45:23 +08:00
parent 1555dbc55e
commit 5a433b0ff7
3 changed files with 79 additions and 25 deletions
+8 -1
View File
@@ -226,7 +226,14 @@ async fn handle_claude_transform(
}
// 非流式响应转换 (OpenAI/Responses → Anthropic)
let (mut response_headers, _status, body_bytes) = read_decoded_body(response, ctx.tag).await?;
let body_timeout =
if ctx.app_config.auto_failover_enabled && ctx.app_config.non_streaming_timeout > 0 {
std::time::Duration::from_secs(ctx.app_config.non_streaming_timeout as u64)
} else {
std::time::Duration::ZERO
};
let (mut response_headers, _status, body_bytes) =
read_decoded_body(response, ctx.tag, body_timeout).await?;
let body_str = String::from_utf8_lossy(&body_bytes);
+45 -22
View File
@@ -48,16 +48,6 @@ impl OriginalHeaderCases {
Self { cases }
}
/// Look up the original casing for a header name (case-insensitive).
/// Returns an iterator over all original casings for that name.
pub fn get_all<'a>(&'a self, name: &'a str) -> impl Iterator<Item = &'a [u8]> + 'a {
let lower = name.to_ascii_lowercase();
self.cases
.iter()
.filter(move |(k, _)| *k == lower)
.map(|(_, v)| v.as_slice())
}
}
type HyperClient = Client<
@@ -535,20 +525,53 @@ fn build_raw_request(
raw.extend_from_slice(path_and_query.as_bytes());
raw.extend_from_slice(b" HTTP/1.1\r\n");
// Headers with original casing
for name in headers.keys() {
let name_str = name.as_str();
let mut case_iter = original_cases.get_all(name_str);
for value in headers.get_all(name) {
if let Some(orig_name_bytes) = case_iter.next() {
// Headers with original casing, emitted in original wire order.
//
// Strategy:
// 1. Walk `original_cases.cases` in order — this preserves the exact
// header sequence the client sent. For each entry, emit the stored
// original-casing name plus the current value from `headers` (the
// proxy may have rewritten the value, e.g. Authorization).
// Repeated headers with the same name are handled by tracking a
// per-name value cursor so we step through `get_all()` in order.
// 2. After the original headers, append any headers that exist in
// `headers` but were not present in the original request (i.e. added
// by the proxy). These are emitted in lowercase.
//
// This replaces the old `for name in headers.keys()` loop which iterated
// in hash-map order, destroying the original header sequence.
let mut emitted: std::collections::HashSet<String> =
std::collections::HashSet::with_capacity(original_cases.cases.len());
// Per-name cursor: how many values we have already emitted for each name.
let mut value_cursor: std::collections::HashMap<String, usize> =
std::collections::HashMap::with_capacity(original_cases.cases.len());
for (lower_name, orig_name_bytes) in &original_cases.cases {
if let Ok(header_name) = http::header::HeaderName::from_bytes(lower_name.as_bytes()) {
let all_values: Vec<_> = headers.get_all(&header_name).iter().collect();
let cursor = value_cursor.entry(lower_name.clone()).or_insert(0);
if let Some(value) = all_values.get(*cursor) {
raw.extend_from_slice(orig_name_bytes);
} else {
// Header not in original request (added by proxy) — use lowercase
raw.extend_from_slice(name_str.as_bytes());
raw.extend_from_slice(b": ");
raw.extend_from_slice(value.as_bytes());
raw.extend_from_slice(b"\r\n");
*cursor += 1;
emitted.insert(lower_name.clone());
}
raw.extend_from_slice(b": ");
raw.extend_from_slice(value.as_bytes());
raw.extend_from_slice(b"\r\n");
}
}
// Append proxy-added headers (not present in the original request).
for name in headers.keys() {
let lower = name.as_str().to_ascii_lowercase();
if !emitted.contains(&lower) {
for value in headers.get_all(name) {
raw.extend_from_slice(name.as_str().as_bytes());
raw.extend_from_slice(b": ");
raw.extend_from_slice(value.as_bytes());
raw.extend_from_slice(b"\r\n");
}
emitted.insert(lower);
}
}
+26 -2
View File
@@ -76,13 +76,29 @@ pub(crate) fn strip_entity_headers_for_rebuilt_body(headers: &mut HeaderMap) {
}
/// 读取响应体并在需要时解压,确保 headers 与返回 body 一致。
///
/// `body_timeout`: 整包超时。当非零时用 `tokio::time::timeout` 包住 `.bytes()` 调用,
/// 防止上游发完响应头后卡住 body 导致请求永远挂住。
/// 传入 `Duration::ZERO` 表示不启用超时(故障转移关闭时)。
pub(crate) async fn read_decoded_body(
response: ProxyResponse,
tag: &str,
body_timeout: Duration,
) -> Result<(HeaderMap, http::StatusCode, Bytes), ProxyError> {
let mut headers = response.headers().clone();
let status = response.status();
let raw_bytes = response.bytes().await?;
let raw_bytes = if body_timeout.is_zero() {
response.bytes().await?
} else {
tokio::time::timeout(body_timeout, response.bytes())
.await
.map_err(|_| {
ProxyError::Timeout(format!(
"响应体读取超时: {}s(上游发完响应头后 body 未到达)",
body_timeout.as_secs()
))
})??
};
log::debug!(
"[{tag}] 已接收上游响应体: status={}, bytes={}, headers={}",
@@ -184,7 +200,15 @@ pub async fn handle_non_streaming(
state: &ProxyState,
parser_config: &UsageParserConfig,
) -> Result<Response, ProxyError> {
let (response_headers, status, body_bytes) = read_decoded_body(response, ctx.tag).await?;
// 整包超时:仅在故障转移开启且配置值非零时生效
let body_timeout =
if ctx.app_config.auto_failover_enabled && ctx.app_config.non_streaming_timeout > 0 {
Duration::from_secs(ctx.app_config.non_streaming_timeout as u64)
} else {
Duration::ZERO
};
let (response_headers, status, body_bytes) =
read_decoded_body(response, ctx.tag, body_timeout).await?;
log::debug!(
"[{}] 上游响应体内容: {}",