mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-07-28 08:44:41 +08:00
fix: decompress body before forward and support zstd (#3817)
* fix(proxy): decompress Codex request body before forward, support zstd Codex Desktop sends zstd-compressed request bodies when authenticated against the Codex backend, which broke local proxy routing because the handlers parsed the raw bytes with serde_json directly. Reworked on top of current main so it preserves the response_processor behavior that landed after this PR was first opened: - Extract content-encoding helpers into a shared proxy::content_encoding module. decompress_body keeps returning Option<Vec<u8>> so unknown encodings stay pass-through with their content-encoding header intact, and keeps the deflate zlib-then-raw fallback (RFC 9110). - Add zstd/zst support (zstd 0.13) and disable reqwest's auto zstd decompression via .no_zstd() for parity with gzip/br/deflate. - Decompress the request body before JSON parsing in the three Codex handlers (chat_completions / responses / responses_compact) and strip the stale content-encoding / content-length / transfer-encoding headers so the forwarder regenerates them. - Support stacked codings (e.g. "gzip, zstd") by decoding in reverse order and merge repeated Content-Encoding headers via get_all. Fixes #3764 Fixes #3696 Co-authored-by: chenx-dust <16610294+chenx-dust@users.noreply.github.com> * fix(proxy): decompress upstream error bodies before reading them The forwarder error branch consumes non-2xx responses via String::from_utf8 directly, bypassing read_decoded_body. reqwest has no auto-decompression feature enabled, so a compressed error body (gzip/br/deflate/zstd) arrives as raw bytes, fails from_utf8, and gets dropped, hiding upstream rate-limit and auth details from the client. Decode the error body with the shared proxy::content_encoding helper, mirroring the success path. Falls back to the raw bytes when the encoding is unsupported or decoding fails. Co-authored-by: chenx-dust <16610294+chenx-dust@users.noreply.github.com> --------- Co-authored-by: Jason <farion1231@gmail.com> Co-authored-by: chenx-dust <16610294+chenx-dust@users.noreply.github.com>
This commit is contained in:
Chenx Dust
@@ -3,6 +3,7 @@
|
||||
//! 统一处理流式和非流式 API 响应
|
||||
|
||||
use super::{
|
||||
content_encoding::{decompress_body, get_content_encoding},
|
||||
forwarder::ActiveConnectionGuard,
|
||||
handler_config::{StreamUsageEventFilter, UsageParserConfig},
|
||||
handler_context::{RequestContext, StreamingTimeoutConfig},
|
||||
@@ -19,7 +20,6 @@ use bytes::Bytes;
|
||||
use futures::stream::{Stream, StreamExt};
|
||||
use serde_json::Value;
|
||||
use std::{
|
||||
io::Read,
|
||||
sync::{
|
||||
atomic::{AtomicBool, Ordering},
|
||||
Arc,
|
||||
@@ -29,60 +29,9 @@ use std::{
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
// ============================================================================
|
||||
// 响应解压
|
||||
// 响应头处理
|
||||
// ============================================================================
|
||||
|
||||
/// 根据 content-encoding 解压响应体字节
|
||||
///
|
||||
/// reqwest 自动解压已禁用（为了透传 accept-encoding），需要手动解压。
|
||||
/// 返回 `Ok(None)` 表示编码不受支持、原样透传——此时调用方必须保留
|
||||
/// content-encoding 头，否则下游（诊断/客户端）会把压缩字节误当明文。
|
||||
fn decompress_body(content_encoding: &str, body: &[u8]) -> Result<Option<Vec<u8>>, std::io::Error> {
|
||||
match content_encoding {
|
||||
"gzip" | "x-gzip" => {
|
||||
let mut decoder = flate2::read::GzDecoder::new(body);
|
||||
let mut decompressed = Vec::new();
|
||||
decoder.read_to_end(&mut decompressed)?;
|
||||
Ok(Some(decompressed))
|
||||
}
|
||||
"deflate" => {
|
||||
// RFC 9110: deflate 指 zlib 包裹格式；但部分上游发 raw deflate 流。
|
||||
// 先按规范尝试 zlib，失败再回退 raw —— 否则合规上游必然解压失败，
|
||||
// 原始压缩字节会被 fail-open 透传给 JSON 解析（#2234 形态 C 之一）。
|
||||
let mut decompressed = Vec::new();
|
||||
let mut zlib = flate2::read::ZlibDecoder::new(body);
|
||||
match zlib.read_to_end(&mut decompressed) {
|
||||
Ok(_) => Ok(Some(decompressed)),
|
||||
Err(zlib_err) => {
|
||||
log::debug!("deflate 按 zlib 解压失败（{zlib_err}），回退 raw deflate");
|
||||
let mut decompressed = Vec::new();
|
||||
let mut raw = flate2::read::DeflateDecoder::new(body);
|
||||
raw.read_to_end(&mut decompressed)?;
|
||||
Ok(Some(decompressed))
|
||||
}
|
||||
}
|
||||
}
|
||||
"br" => {
|
||||
let mut decompressed = Vec::new();
|
||||
brotli::BrotliDecompress(&mut std::io::Cursor::new(body), &mut decompressed)?;
|
||||
Ok(Some(decompressed))
|
||||
}
|
||||
_ => {
|
||||
log::warn!("未知的 content-encoding: {content_encoding}，跳过解压");
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 从响应头提取 content-encoding（忽略 identity 和 chunked）
|
||||
fn get_content_encoding(headers: &HeaderMap) -> Option<String> {
|
||||
headers
|
||||
.get("content-encoding")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(|s| s.trim().to_lowercase())
|
||||
.filter(|s| !s.is_empty() && s != "identity")
|
||||
}
|
||||
|
||||
/// RFC 2616 / RFC 7230 中定义的不应被代理继续转发的响应头。
|
||||
const HOP_BY_HOP_RESPONSE_HEADERS: &[&str] = &[
|
||||
"connection",
|
||||
@@ -878,40 +827,6 @@ mod tests {
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
#[test]
|
||||
fn decompress_body_deflate_handles_zlib_wrapped_per_rfc9110() {
|
||||
// RFC 9110 规范的 deflate = zlib 包裹格式（合规上游发的就是这个）
|
||||
let payload = br#"{"ok":true}"#;
|
||||
let mut encoder =
|
||||
flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default());
|
||||
std::io::Write::write_all(&mut encoder, payload).unwrap();
|
||||
let compressed = encoder.finish().unwrap();
|
||||
|
||||
let decompressed = decompress_body("deflate", &compressed).unwrap().unwrap();
|
||||
assert_eq!(decompressed, payload);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decompress_body_deflate_falls_back_to_raw_stream() {
|
||||
// 部分上游违规发 raw deflate 流，保持兼容
|
||||
let payload = br#"{"ok":true}"#;
|
||||
let mut encoder =
|
||||
flate2::write::DeflateEncoder::new(Vec::new(), flate2::Compression::default());
|
||||
std::io::Write::write_all(&mut encoder, payload).unwrap();
|
||||
let compressed = encoder.finish().unwrap();
|
||||
|
||||
let decompressed = decompress_body("deflate", &compressed).unwrap().unwrap();
|
||||
assert_eq!(decompressed, payload);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decompress_body_unknown_encoding_returns_none_to_keep_headers() {
|
||||
// 未知编码必须返回 None（而非伪装成"已解码"），否则 content-encoding
|
||||
// 头被剥掉，下游诊断会把压缩字节误报成明文
|
||||
let result = decompress_body("zstd", b"\x28\xb5\x2f\xfd").unwrap();
|
||||
assert!(result.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_sse_field_accepts_optional_space() {
|
||||
assert_eq!(
|
||||
|
||||
Reference in New Issue
Block a user