mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-08-03 11:01:22 +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,
|
streaming_responses::create_anthropic_sse_stream_from_responses, transform,
|
||||||
transform_responses,
|
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,
|
server::ProxyState,
|
||||||
types::*,
|
types::*,
|
||||||
usage::parser::TokenUsage,
|
usage::parser::TokenUsage,
|
||||||
@@ -211,9 +214,7 @@ async fn handle_claude_transform(
|
|||||||
}
|
}
|
||||||
|
|
||||||
// 非流式响应转换 (OpenAI/Responses → Anthropic)
|
// 非流式响应转换 (OpenAI/Responses → Anthropic)
|
||||||
let response_headers = response.headers().clone();
|
let (mut response_headers, _status, body_bytes) = read_decoded_body(response, ctx.tag).await?;
|
||||||
|
|
||||||
let body_bytes = response.bytes().await?;
|
|
||||||
|
|
||||||
let body_str = String::from_utf8_lossy(&body_bytes);
|
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);
|
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() {
|
for (key, value) in response_headers.iter() {
|
||||||
if key.as_str().to_lowercase() != "content-length"
|
builder = builder.header(key, value);
|
||||||
&& key.as_str().to_lowercase() != "transfer-encoding"
|
|
||||||
{
|
|
||||||
builder = builder.header(key, value);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
builder = builder.header("content-type", "application/json");
|
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))
|
.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 proxy_points_to_loopback(value: &str) -> bool {
|
||||||
fn host_is_loopback(host: &str) -> bool {
|
fn host_is_loopback(host: &str) -> bool {
|
||||||
if host.eq_ignore_ascii_case("localhost") {
|
if host.eq_ignore_ascii_case("localhost") {
|
||||||
@@ -337,7 +359,7 @@ pub fn mask_url(url: &str) -> String {
|
|||||||
/// 根据供应商单独代理配置构建代理 URL
|
/// 根据供应商单独代理配置构建代理 URL
|
||||||
///
|
///
|
||||||
/// 将 ProviderProxyConfig 转换为代理 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 proxy_type = config.proxy_type.as_deref().unwrap_or("http");
|
||||||
let host = config.proxy_host.as_deref()?;
|
let host = config.proxy_host.as_deref()?;
|
||||||
let port = config.proxy_port?;
|
let port = config.proxy_port?;
|
||||||
@@ -544,4 +566,19 @@ mod tests {
|
|||||||
std::env::remove_var(key);
|
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")
|
.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,
|
state: &ProxyState,
|
||||||
parser_config: &UsageParserConfig,
|
parser_config: &UsageParserConfig,
|
||||||
) -> Result<Response, ProxyError> {
|
) -> Result<Response, ProxyError> {
|
||||||
let response_headers = response.headers().clone();
|
let (response_headers, status, body_bytes) = read_decoded_body(response, ctx.tag).await?;
|
||||||
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
|
|
||||||
};
|
|
||||||
|
|
||||||
log::debug!(
|
log::debug!(
|
||||||
"[{}] 上游响应体内容: {}",
|
"[{}] 上游响应体内容: {}",
|
||||||
@@ -670,6 +691,22 @@ mod tests {
|
|||||||
assert_eq!(super::strip_sse_field("id:1", "data"), None);
|
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 {
|
fn build_state(db: Arc<Database>) -> ProxyState {
|
||||||
ProxyState {
|
ProxyState {
|
||||||
db: db.clone(),
|
db: db.clone(),
|
||||||
|
|||||||
@@ -133,18 +133,44 @@ impl ProxyServer {
|
|||||||
Ok(v) => v,
|
Ok(v) => v,
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
log::error!("[{SRV}] accept 失败: {e}", SRV = log_srv::ACCEPT_ERR);
|
log::error!("[{SRV}] accept 失败: {e}", SRV = log_srv::ACCEPT_ERR);
|
||||||
|
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
let app = app.clone();
|
let app = app.clone();
|
||||||
tokio::spawn(async move {
|
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
|
// service_fn 将 axum Router(tower::Service)桥接到 hyper
|
||||||
let service = hyper::service::service_fn(move |req: hyper::Request<hyper::body::Incoming>| {
|
let service = hyper::service::service_fn(move |req: hyper::Request<hyper::body::Incoming>| {
|
||||||
let mut router = app.clone();
|
let mut router = app.clone();
|
||||||
|
let cases = original_cases.clone();
|
||||||
async move {
|
async move {
|
||||||
// 将 hyper::body::Incoming 转为 axum::body::Body,保留 extensions
|
// 将 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 body = axum::body::Body::new(body);
|
||||||
let axum_req = http::Request::from_parts(parts, 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
|
<Router as tower::Service<http::Request<axum::body::Body>>>::call(&mut router, axum_req).await
|
||||||
|
|||||||
Reference in New Issue
Block a user