mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-06-16 13:34:04 +08:00
fix(proxy): address code review feedback on response handling
Fixes from PR #1714 code review: - Extract `read_decoded_body()` and `strip_entity_headers_for_rebuilt_body()` in response_processor to properly clean content-encoding/content-length headers after decompression - Reuse `read_decoded_body()` in handlers.rs for Claude transform path, ensuring compressed responses are decoded before format conversion - Make `build_proxy_url_from_config()` public so forwarder can pass proxy URL to the hyper raw write path - Add `has_system_proxy_env()` utility with test coverage - Add 50ms backoff after accept() failures in server.rs to prevent tight-loop CPU spin on transient socket errors
This commit is contained in:
@@ -18,7 +18,10 @@ use super::{
|
||||
streaming_responses::create_anthropic_sse_stream_from_responses, transform,
|
||||
transform_responses,
|
||||
},
|
||||
response_processor::{create_logged_passthrough_stream, process_response, SseUsageCollector},
|
||||
response_processor::{
|
||||
create_logged_passthrough_stream, process_response, read_decoded_body,
|
||||
strip_entity_headers_for_rebuilt_body, SseUsageCollector,
|
||||
},
|
||||
server::ProxyState,
|
||||
types::*,
|
||||
usage::parser::TokenUsage,
|
||||
@@ -211,9 +214,7 @@ async fn handle_claude_transform(
|
||||
}
|
||||
|
||||
// 非流式响应转换 (OpenAI/Responses → Anthropic)
|
||||
let response_headers = response.headers().clone();
|
||||
|
||||
let body_bytes = response.bytes().await?;
|
||||
let (mut response_headers, _status, body_bytes) = read_decoded_body(response, ctx.tag).await?;
|
||||
|
||||
let body_str = String::from_utf8_lossy(&body_bytes);
|
||||
|
||||
@@ -266,13 +267,10 @@ async fn handle_claude_transform(
|
||||
|
||||
// 构建响应
|
||||
let mut builder = axum::response::Response::builder().status(status);
|
||||
strip_entity_headers_for_rebuilt_body(&mut response_headers);
|
||||
|
||||
for (key, value) in response_headers.iter() {
|
||||
if key.as_str().to_lowercase() != "content-length"
|
||||
&& key.as_str().to_lowercase() != "transfer-encoding"
|
||||
{
|
||||
builder = builder.header(key, value);
|
||||
}
|
||||
builder = builder.header(key, value);
|
||||
}
|
||||
|
||||
builder = builder.header("content-type", "application/json");
|
||||
|
||||
@@ -280,6 +280,28 @@ fn system_proxy_points_to_loopback() -> bool {
|
||||
.any(|value| proxy_points_to_loopback(&value))
|
||||
}
|
||||
|
||||
/// 检查当前环境是否存在可用的系统代理配置。
|
||||
#[allow(dead_code)]
|
||||
pub fn has_system_proxy_env() -> bool {
|
||||
const KEYS: [&str; 6] = [
|
||||
"HTTP_PROXY",
|
||||
"http_proxy",
|
||||
"HTTPS_PROXY",
|
||||
"https_proxy",
|
||||
"ALL_PROXY",
|
||||
"all_proxy",
|
||||
];
|
||||
|
||||
if system_proxy_points_to_loopback() {
|
||||
return false;
|
||||
}
|
||||
|
||||
KEYS.iter()
|
||||
.filter_map(|key| env::var(key).ok())
|
||||
.map(|value| value.trim().to_string())
|
||||
.any(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn proxy_points_to_loopback(value: &str) -> bool {
|
||||
fn host_is_loopback(host: &str) -> bool {
|
||||
if host.eq_ignore_ascii_case("localhost") {
|
||||
@@ -337,7 +359,7 @@ pub fn mask_url(url: &str) -> String {
|
||||
/// 根据供应商单独代理配置构建代理 URL
|
||||
///
|
||||
/// 将 ProviderProxyConfig 转换为代理 URL 字符串
|
||||
fn build_proxy_url_from_config(config: &ProviderProxyConfig) -> Option<String> {
|
||||
pub fn build_proxy_url_from_config(config: &ProviderProxyConfig) -> Option<String> {
|
||||
let proxy_type = config.proxy_type.as_deref().unwrap_or("http");
|
||||
let host = config.proxy_host.as_deref()?;
|
||||
let port = config.proxy_port?;
|
||||
@@ -544,4 +566,19 @@ mod tests {
|
||||
std::env::remove_var(key);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_has_system_proxy_env_ignores_own_loopback_proxy() {
|
||||
let _guard = env_lock().lock().unwrap();
|
||||
set_proxy_port(15721);
|
||||
std::env::remove_var("HTTP_PROXY");
|
||||
|
||||
std::env::set_var("HTTP_PROXY", "http://127.0.0.1:15721");
|
||||
assert!(!has_system_proxy_env());
|
||||
|
||||
std::env::set_var("HTTP_PROXY", "http://127.0.0.1:7890");
|
||||
assert!(has_system_proxy_env());
|
||||
|
||||
std::env::remove_var("HTTP_PROXY");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -68,6 +68,52 @@ fn get_content_encoding(headers: &HeaderMap) -> Option<String> {
|
||||
.filter(|s| !s.is_empty() && s != "identity")
|
||||
}
|
||||
|
||||
/// 移除在重建响应体后会失真的实体头。
|
||||
pub(crate) fn strip_entity_headers_for_rebuilt_body(headers: &mut HeaderMap) {
|
||||
headers.remove(axum::http::header::CONTENT_ENCODING);
|
||||
headers.remove(axum::http::header::CONTENT_LENGTH);
|
||||
headers.remove(axum::http::header::TRANSFER_ENCODING);
|
||||
}
|
||||
|
||||
/// 读取响应体并在需要时解压,确保 headers 与返回 body 一致。
|
||||
pub(crate) async fn read_decoded_body(
|
||||
response: ProxyResponse,
|
||||
tag: &str,
|
||||
) -> Result<(HeaderMap, http::StatusCode, Bytes), ProxyError> {
|
||||
let mut headers = response.headers().clone();
|
||||
let status = response.status();
|
||||
let raw_bytes = response.bytes().await?;
|
||||
|
||||
log::debug!(
|
||||
"[{tag}] 已接收上游响应体: status={}, bytes={}, headers={}",
|
||||
status.as_u16(),
|
||||
raw_bytes.len(),
|
||||
format_headers(&headers)
|
||||
);
|
||||
|
||||
let mut body_bytes = raw_bytes.clone();
|
||||
let mut decoded = false;
|
||||
|
||||
if let Some(encoding) = get_content_encoding(&headers) {
|
||||
log::debug!("[{tag}] 解压非流式响应: content-encoding={encoding}");
|
||||
match decompress_body(&encoding, &raw_bytes) {
|
||||
Ok(decompressed) => {
|
||||
body_bytes = Bytes::from(decompressed);
|
||||
decoded = true;
|
||||
}
|
||||
Err(e) => {
|
||||
log::warn!("[{tag}] 解压失败 ({encoding}): {e},使用原始数据");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if decoded {
|
||||
strip_entity_headers_for_rebuilt_body(&mut headers);
|
||||
}
|
||||
|
||||
Ok((headers, status, body_bytes))
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// 公共接口
|
||||
// ============================================================================
|
||||
@@ -138,32 +184,7 @@ pub async fn handle_non_streaming(
|
||||
state: &ProxyState,
|
||||
parser_config: &UsageParserConfig,
|
||||
) -> Result<Response, ProxyError> {
|
||||
let response_headers = response.headers().clone();
|
||||
let status = response.status();
|
||||
|
||||
// 读取响应体
|
||||
let raw_bytes = response.bytes().await?;
|
||||
log::debug!(
|
||||
"[{}] 已接收上游响应体: status={}, bytes={}, headers={}",
|
||||
ctx.tag,
|
||||
status.as_u16(),
|
||||
raw_bytes.len(),
|
||||
format_headers(&response_headers)
|
||||
);
|
||||
|
||||
// 手动解压(reqwest 自动解压已禁用以透传 accept-encoding)
|
||||
let body_bytes: Bytes = if let Some(encoding) = get_content_encoding(&response_headers) {
|
||||
log::debug!("[{}] 解压非流式响应: content-encoding={encoding}", ctx.tag);
|
||||
match decompress_body(&encoding, &raw_bytes) {
|
||||
Ok(decompressed) => Bytes::from(decompressed),
|
||||
Err(e) => {
|
||||
log::warn!("[{}] 解压失败 ({encoding}): {e},使用原始数据", ctx.tag);
|
||||
raw_bytes
|
||||
}
|
||||
}
|
||||
} else {
|
||||
raw_bytes
|
||||
};
|
||||
let (response_headers, status, body_bytes) = read_decoded_body(response, ctx.tag).await?;
|
||||
|
||||
log::debug!(
|
||||
"[{}] 上游响应体内容: {}",
|
||||
@@ -670,6 +691,22 @@ mod tests {
|
||||
assert_eq!(super::strip_sse_field("id:1", "data"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_entity_headers_removes_encoding_and_length_headers() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert("content-encoding", "gzip".parse().unwrap());
|
||||
headers.insert("content-length", "123".parse().unwrap());
|
||||
headers.insert("transfer-encoding", "chunked".parse().unwrap());
|
||||
headers.insert("content-type", "application/json".parse().unwrap());
|
||||
|
||||
strip_entity_headers_for_rebuilt_body(&mut headers);
|
||||
|
||||
assert!(!headers.contains_key("content-encoding"));
|
||||
assert!(!headers.contains_key("content-length"));
|
||||
assert!(!headers.contains_key("transfer-encoding"));
|
||||
assert_eq!(headers.get("content-type").unwrap(), "application/json");
|
||||
}
|
||||
|
||||
fn build_state(db: Arc<Database>) -> ProxyState {
|
||||
ProxyState {
|
||||
db: db.clone(),
|
||||
|
||||
@@ -133,18 +133,44 @@ impl ProxyServer {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
log::error!("[{SRV}] accept 失败: {e}", SRV = log_srv::ACCEPT_ERR);
|
||||
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let app = app.clone();
|
||||
tokio::spawn(async move {
|
||||
// Peek raw TCP bytes to capture original header casing
|
||||
// before hyper parses (and lowercases) the header names.
|
||||
let original_cases = {
|
||||
let mut peek_buf = vec![0u8; 8192];
|
||||
match stream.peek(&mut peek_buf).await {
|
||||
Ok(n) => {
|
||||
let cases = super::hyper_client::OriginalHeaderCases::from_raw_bytes(&peek_buf[..n]);
|
||||
log::debug!(
|
||||
"[ProxyServer] Peeked {} bytes, captured {} header casings",
|
||||
n, cases.cases.len()
|
||||
);
|
||||
cases
|
||||
}
|
||||
Err(e) => {
|
||||
log::debug!("[ProxyServer] peek failed (non-fatal): {e}");
|
||||
super::hyper_client::OriginalHeaderCases::default()
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// service_fn 将 axum Router(tower::Service)桥接到 hyper
|
||||
let service = hyper::service::service_fn(move |req: hyper::Request<hyper::body::Incoming>| {
|
||||
let mut router = app.clone();
|
||||
let cases = original_cases.clone();
|
||||
async move {
|
||||
// 将 hyper::body::Incoming 转为 axum::body::Body,保留 extensions
|
||||
let (parts, body) = req.into_parts();
|
||||
let (mut parts, body) = req.into_parts();
|
||||
|
||||
// Insert our own header case map alongside hyper's internal one
|
||||
parts.extensions.insert(cases);
|
||||
|
||||
let body = axum::body::Body::new(body);
|
||||
let axum_req = http::Request::from_parts(parts, body);
|
||||
<Router as tower::Service<http::Request<axum::body::Body>>>::call(&mut router, axum_req).await
|
||||
|
||||
Reference in New Issue
Block a user