mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-07-24 12:44:18 +08:00
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:
@@ -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);
|
||||
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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!(
|
||||
"[{}] 上游响应体内容: {}",
|
||||
|
||||
Reference in New Issue
Block a user