From 05d77e4e4512ea68ecb4d8bd05d15ceea45b7030 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Sat, 10 Oct 2026 18:24:53 +0800 Subject: [PATCH] feat(gateway): decode compressed request bodies and upstream responses Clients may send a request body with `Content-Encoding: gzip`, `deflate` or `zstd`; the gateway only ever read the raw bytes and handed them to the JSON parser, so such a request failed on every route. Upstreams that answer compressed (some do even without `Accept-Encoding`) reached the client with the compressed bytes and a `content-encoding` header the client may not understand, and the gateway's own readers (usage, guards, plugins, the session log) saw garbage. Inbound: `server::intake` decodes the body right after reading it, on every gateway route, and removes `Content-Encoding` from the request, so everything downstream (parsing, conversion, passthrough, the Codex backend's own zstd re-compression) sees an ordinary uncompressed request and nothing about the client's encoding reaches the upstream. The size limit applies to the decoded size with a bounded reader: a few hundred bytes that inflate to more than the cap are refused with 413 like an oversized plain body, regardless of what the compressed data or `Content-Length` claim. An unknown encoding is refused before the body is read; a body that does not decode is a 400 in the client's format. Outbound: the new `inflate` module wraps the upstream response body in a streaming decoder (gzip/deflate/zstd) and drops `content-encoding` and the stale `content-length`. It is layered *outside* the traffic meter, not delegated to reqwest's decompression features: the traffic counters report wire bytes, and enabling reqwest's gzip/zstd would silently turn them into decoded bytes (the existing guard test keeps standing). SSE is decoded frame by frame, so streaming still arrives event by event. `br` is not decoded (no decoder in the dependency tree) and passes through as before. The traffic end-to-end test now expects the zstd answer to reach the client decoded while the received-byte count stays the wire size. Co-Authored-By: Claude Fable 5.1 --- Cargo.lock | 1 + crates/tw-api/msg-codes.txt | 3 + crates/tw-gateway/Cargo.toml | 2 + crates/tw-gateway/src/inflate.rs | 252 ++++++++++++++ crates/tw-gateway/src/lib.rs | 1 + crates/tw-gateway/src/server.rs | 7 +- crates/tw-gateway/src/server/intake.rs | 181 +++++++++- crates/tw-gateway/src/server/pipeline/hop.rs | 17 +- crates/tw-gateway/src/traffic.rs | 5 +- crates/tw-gateway/tests/compressed.rs | 327 +++++++++++++++++++ crates/tw-gateway/tests/traffic.rs | 13 +- 11 files changed, 783 insertions(+), 26 deletions(-) create mode 100644 crates/tw-gateway/src/inflate.rs create mode 100644 crates/tw-gateway/tests/compressed.rs diff --git a/Cargo.lock b/Cargo.lock index 30c8835a..f91cf37d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3283,6 +3283,7 @@ dependencies = [ "blake3", "bytes", "chrono", + "flate2", "futures", "http", "http-body-util", diff --git a/crates/tw-api/msg-codes.txt b/crates/tw-api/msg-codes.txt index 52d372f6..a86ef981 100644 --- a/crates/tw-api/msg-codes.txt +++ b/crates/tw-api/msg-codes.txt @@ -374,9 +374,12 @@ gw.probe.key_rejected gw.probe.request_failed gw.probe.timeout gw.request.aborted +gw.request.bad_encoding gw.request.body_declared_too_large gw.request.body_over_limit gw.request.body_unreadable +gw.request.decoded_body_over_limit +gw.request.unsupported_encoding gw.route.all_selected_disabled gw.route.denied gw.route.failed diff --git a/crates/tw-gateway/Cargo.toml b/crates/tw-gateway/Cargo.toml index 3ddbe4c1..29240ee1 100644 --- a/crates/tw-gateway/Cargo.toml +++ b/crates/tw-gateway/Cargo.toml @@ -53,6 +53,8 @@ thiserror = { workspace = true } tokio = { workspace = true } tracing = { workspace = true } zstd = { workspace = true } +# 压缩过的请求体和上游回答：gzip、deflate 用它解，zstd 用上面那个 +flate2 = { workspace = true } [dev-dependencies] # 假的 Bedrock 要回 eventstream 的二进制帧 diff --git a/crates/tw-gateway/src/inflate.rs b/crates/tw-gateway/src/inflate.rs new file mode 100644 index 00000000..c77fd557 --- /dev/null +++ b/crates/tw-gateway/src/inflate.rs @@ -0,0 +1,252 @@ +//! 传输编码：压缩过的请求体和上游回答在这里解开。 +//! +//! 两头各一处。**进来的**请求体由 `server::intake` 读完整个再解（[`Coding::decode`]，解出来的 +//! 有上限）；**收回来的**上游回答边收边解（[`inflated`]），流式回答一块一块地解，SSE 照样一条 +//! 一条地到。 +//! +//! 只认 gzip、deflate、zstd：HTTP 里常见的就是这三种，而 gzip、zstd 的解码库本来就在依赖树里 +//! （价目表、Codex 后端的请求体）。`br` 要另引一个库，没见哪家上游在没问过 `Accept-Encoding` +//! 时主动发它，不认：带着它的回答原样交给客户端，和从前一样。 +//! +//! **上游回答的解压不交给 reqwest**（它的 gzip、zstd 特性都没开）：流量按线上的字节数 +//! （[`crate::traffic`]），计数套在解压之前，开了自动解压就数不到压缩着的那些了。网关也不发 +//! `Accept-Encoding`：没问就压缩着回的上游不多，问了反而要为每一次 SSE 多解一层。 + +use std::io::{self, Read, Write}; + +use bytes::Bytes; +use futures::StreamExt; +use http_body_util::BodyExt; + +/// 一个 `Content-Encoding` 值说的是哪种。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Coding { + Gzip, + Deflate, + Zstd, +} + +/// 整个解开时出的事 +pub enum Decode { + /// 解出来的超过了上限 + OverLimit, + /// 数据和编码对不上、解到一半断了 + Failed(String), +} + +impl Coding { + /// `identity` 和没写一样是 `None`；大小写不论；`x-gzip` 是 gzip 的旧名。叠了几层的 + /// （`gzip, zstd`）和不认识的一律是 `None`，由调用的那一头决定拒还是原样 + pub fn parse(value: &str) -> Option { + match value.trim().to_ascii_lowercase().as_str() { + "gzip" | "x-gzip" => Some(Self::Gzip), + "deflate" => Some(Self::Deflate), + "zstd" => Some(Self::Zstd), + _ => None, + } + } + + pub fn name(self) -> &'static str { + match self { + Self::Gzip => "gzip", + Self::Deflate => "deflate", + Self::Zstd => "zstd", + } + } + + /// 解开 `raw`，解出来的最多 `max` 字节。**读到第 `max + 1` 个字节就停**：解出来的有多大， + /// 压缩数据自己说了不算（zstd 的帧头可以写一个假的长度），只信实际解出来的 + pub fn decode(self, raw: &[u8], max: usize) -> Result { + let reader: Box = match self { + Self::Gzip => Box::new(flate2::read::MultiGzDecoder::new(raw)), + // HTTP 的 deflate 是 zlib 封装的（RFC 9110 §8.4.1.2） + Self::Deflate => Box::new(flate2::read::ZlibDecoder::new(raw)), + Self::Zstd => Box::new( + zstd::stream::read::Decoder::new(raw).map_err(|e| Decode::Failed(e.to_string()))?, + ), + }; + let mut out = Vec::new(); + reader + .take(max as u64 + 1) + .read_to_end(&mut out) + .map_err(|e| Decode::Failed(e.to_string()))?; + if out.len() > max { + return Err(Decode::OverLimit); + } + Ok(Bytes::from(out)) + } +} + +/// 边收边解：喂一块压缩数据，吐出此刻解得出的那些 +enum Unpacker { + Gzip(flate2::write::MultiGzDecoder>), + Deflate(flate2::write::ZlibDecoder>), + Zstd(zstd::stream::write::Decoder<'static, Vec>), +} + +impl Unpacker { + fn new(coding: Coding) -> io::Result { + Ok(match coding { + Coding::Gzip => Self::Gzip(flate2::write::MultiGzDecoder::new(Vec::new())), + Coding::Deflate => Self::Deflate(flate2::write::ZlibDecoder::new(Vec::new())), + Coding::Zstd => Self::Zstd(zstd::stream::write::Decoder::new(Vec::new())?), + }) + } + + /// 喂一块。**喂完就 flush**：SSE 的一条事件压缩后多半不满一个压缩块，不 flush 的话它会 + /// 留在解码器里等下一块，客户端看到的就晚了一条 + fn feed(&mut self, chunk: &[u8]) -> io::Result { + match self { + Self::Gzip(d) => { + d.write_all(chunk)?; + d.flush()?; + Ok(Bytes::from(std::mem::take(d.get_mut()))) + } + Self::Deflate(d) => { + d.write_all(chunk)?; + d.flush()?; + Ok(Bytes::from(std::mem::take(d.get_mut()))) + } + Self::Zstd(d) => { + d.write_all(chunk)?; + d.flush()?; + Ok(Bytes::from(std::mem::take(d.get_mut()))) + } + } + } + + /// 收完了：把最后那点吐出来。gzip 在这里核对尾部的校验和 + fn finish(self) -> io::Result { + Ok(Bytes::from(match self { + Self::Gzip(d) => d.finish()?, + Self::Deflate(d) => d.finish()?, + Self::Zstd(mut d) => { + d.flush()?; + d.into_inner() + } + })) + } +} + +/// 上游的回答压缩着（`Content-Encoding` 是认得的那三种之一）就套一层解码：读它的人拿到的是 +/// 解开的字节，`Content-Encoding` 摘掉，`Content-Length`（压缩着的长度）也摘掉。没压缩的、 +/// 认不得的原样返回。 +/// +/// **套在 [`crate::traffic::metered`] 之后**：流量数的是线上的字节。解到一半对不上，流就在那里 +/// 断掉，读的那一头当成上游断了 +pub fn inflated(r: reqwest::Response) -> reqwest::Response { + let coding = r + .headers() + .get(http::header::CONTENT_ENCODING) + .and_then(|v| v.to_str().ok()) + .and_then(Coding::parse); + let Some(coding) = coding else { + return r; + }; + let r: http::Response = r.into(); + let (mut parts, body) = r.into_parts(); + parts.headers.remove(http::header::CONTENT_ENCODING); + parts.headers.remove(http::header::CONTENT_LENGTH); + let body = reqwest::Body::wrap_stream(unpacked(coding, body.into_data_stream())); + reqwest::Response::from(http::Response::from_parts(parts, body)) +} + +type Chunk = Result>; + +/// 压缩着的字节流解成解开的字节流。空块不吐：读 SSE 的那一头把一个空块当成什么都没来 +fn unpacked( + coding: Coding, + mut packed: impl futures::Stream> + Unpin + Send + 'static, +) -> impl futures::Stream + Send + 'static { + async_stream::try_stream! { + let mut unpacker = Unpacker::new(coding)?; + while let Some(chunk) = packed.next().await { + let out = unpacker.feed(&chunk?)?; + if !out.is_empty() { + yield out; + } + } + let rest = unpacker.finish()?; + if !rest.is_empty() { + yield rest; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn gz_frames(parts: &[&[u8]]) -> Vec> { + // 一个 gzip 流分几段发：每段 flush 一次，像上游边生成边压缩那样 + let mut e = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + let mut frames = Vec::new(); + for p in parts { + e.write_all(p).unwrap(); + e.flush().unwrap(); + frames.push(std::mem::take(e.get_mut())); + } + frames.push(e.finish().unwrap()); + frames + } + + fn response(encoding: &'static str, frames: Vec>) -> reqwest::Response { + let stream = futures::stream::iter( + frames + .into_iter() + .map(|f| Ok::<_, std::io::Error>(Bytes::from(f))), + ); + let r = http::Response::builder() + .header("content-encoding", encoding) + .header("content-length", "999") + .header("content-type", "text/event-stream") + .body(reqwest::Body::wrap_stream(stream)) + .unwrap(); + reqwest::Response::from(r) + } + + /// 每一段压缩数据到了，它里面的那条事件就解得出来，不等整个流结束 + #[tokio::test] + async fn each_frame_comes_out_as_it_arrives() { + let frames = gz_frames(&[b"event: a\n\n", b"event: b\n\n"]); + let r = inflated(response("gzip", frames)); + assert!(r.headers().get("content-encoding").is_none()); + assert!(r.headers().get("content-length").is_none()); + assert_eq!( + r.headers().get("content-type").unwrap(), + "text/event-stream" + ); + let mut s = r.bytes_stream(); + assert_eq!(&s.next().await.unwrap().unwrap()[..], b"event: a\n\n"); + assert_eq!(&s.next().await.unwrap().unwrap()[..], b"event: b\n\n"); + assert!(s.next().await.is_none()); + } + + #[tokio::test] + async fn zstd_and_deflate_are_decoded_whole() { + let z = zstd::stream::encode_all(&b"{\"ok\":true}"[..], 3).unwrap(); + let got = inflated(response("zstd", vec![z])).bytes().await.unwrap(); + assert_eq!(&got[..], b"{\"ok\":true}"); + let mut e = flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default()); + e.write_all(b"{\"ok\":true}").unwrap(); + let got = inflated(response("deflate", vec![e.finish().unwrap()])) + .bytes() + .await + .unwrap(); + assert_eq!(&got[..], b"{\"ok\":true}"); + } + + /// 认不得的编码原样交出去，头也留着 + #[tokio::test] + async fn an_unknown_encoding_is_left_alone() { + let r = inflated(response("br", vec![b"xyz".to_vec()])); + assert_eq!(r.headers().get("content-encoding").unwrap(), "br"); + assert_eq!(&r.bytes().await.unwrap()[..], b"xyz"); + } + + #[tokio::test] + async fn a_corrupt_stream_errors_instead_of_handing_over_garbage() { + let r = inflated(response("gzip", vec![b"not gzip at all".to_vec()])); + assert!(r.bytes().await.is_err()); + } +} diff --git a/crates/tw-gateway/src/lib.rs b/crates/tw-gateway/src/lib.rs index 59ff9edd..9bdf4d7d 100644 --- a/crates/tw-gateway/src/lib.rs +++ b/crates/tw-gateway/src/lib.rs @@ -29,6 +29,7 @@ pub mod glm; pub mod guard; pub mod health; pub mod hint; +pub mod inflate; pub mod key_limits; pub mod l1; pub mod l3; diff --git a/crates/tw-gateway/src/server.rs b/crates/tw-gateway/src/server.rs index 58fc80c3..b96db449 100644 --- a/crates/tw-gateway/src/server.rs +++ b/crates/tw-gateway/src/server.rs @@ -101,7 +101,7 @@ async fn passthrough( RawQuery(query): RawQuery, // **必须排在 `body` 前面。**`body` 拿走的是整个请求 crate::ws::MaybeUpgrade(upgrade): crate::ws::MaybeUpgrade, - headers: HeaderMap, + mut headers: HeaderMap, // **这里还没读。**过完来源和密钥检查才读（见 `intake`） body: axum::body::Body, ) -> Result { @@ -188,8 +188,9 @@ async fn passthrough( // 交给管线里每一条 `return Err` 各自去报的话，漏掉一条的后果不是没报， // 而是被 Drop 报成「客户端取消」—— 一次策略拒绝会记到客户端头上。 let mut ending: Option = None; - // **到这里才读体**：来源、密钥、方法都过了。超了上限的还没有开始事件，不记一行 - let body = intake::read(&headers, body, intake::MAX_BODY) + // **到这里才读体**：来源、密钥、方法都过了。超了上限的还没有开始事件，不记一行。 + // 压缩过的在这里解开（`Content-Encoding` 一并摘掉），往下就是一个没压缩的请求 + let body = intake::read(&mut headers, body, intake::MAX_BODY) .await .map_err(|e| e.in_dialect(dialect))?; // 手动中止的开关（见 `crate::abort`）。管线发出开始事件时登记上它 diff --git a/crates/tw-gateway/src/server/intake.rs b/crates/tw-gateway/src/server/intake.rs index 857f2f16..665c1a55 100644 --- a/crates/tw-gateway/src/server/intake.rs +++ b/crates/tw-gateway/src/server/intake.rs @@ -4,6 +4,11 @@ //! `Bytes` 提取器：提取器在处理函数之前就把整个体读进内存，一个不在放行网段里、 //! 没带密钥的连接也能先塞满上限再被拒；它自带的 2 MB 默认上限还会把贴了几张截图的 //! 对话拒成 413，连压缩请求一起，会话就再也压不回来了（Core#295）。 +//! +//! **压缩过的请求体在这里解开**（`Content-Encoding: gzip`、`deflate`、`zstd`）：读完、解完， +//! 后面的每一条路 —— 解析、转换、直通、Codex 后端的再压缩 —— 拿到的都是解开的字节，和客户端 +//! 没压缩时一个样；`Content-Encoding` 头一并摘掉，上游不会看到它。上限按**解开之后**的大小算： +//! 一小段压缩数据能解出几百兆，声明的长度和压缩的长度都不作数。 use axum::body::Body; use axum::http::{HeaderMap, header}; @@ -11,19 +16,22 @@ use bytes::Bytes; use http_body_util::{BodyExt, LengthLimitError, Limited}; use crate::error::GatewayError; +use crate::inflate::{Coding, Decode}; use tw_types::msg; /// 256 MiB。大到能装下几张 4K 图的 base64（膨胀 33%），小到失控的 /// 客户端打不爆内存。 pub(super) const MAX_BODY: usize = 256 * 1024 * 1024; -/// 读完请求体，最多 `max` 字节。 +/// 读完请求体，最多 `max` 字节；压缩过的解开，解开的也最多 `max` 字节。 /// /// **`Content-Length` 只用来提前回绝**：声明的已经超了，不读就拒，老实的客户端不用先 /// 传完几百兆再挨一个 413。真正卡住的是读的时候按实际收到的字节计数 —— 声明的长度 /// 是对方自己写的，所以报错里只说「声明的」。 +/// +/// 解开之后 `Content-Encoding` 从 `headers` 里摘掉：从这里往下的请求就是一个没压缩的请求。 pub(super) async fn read( - headers: &HeaderMap, + headers: &mut HeaderMap, body: Body, max: usize, ) -> Result { @@ -37,22 +45,64 @@ pub(super) async fn read( "The request declares a {size}-byte body, over the {max}-byte limit." ))); } - match Limited::new(body, max).collect().await { - Ok(c) => Ok(c.to_bytes()), + // 认不得的编码在读体之前就拒：不用先收完几百兆再说不认识 + let coding = coding_of(headers)?; + let raw = match Limited::new(body, max).collect().await { + Ok(c) => c.to_bytes(), Err(e) if e.downcast_ref::().is_some() => { - Err(GatewayError::too_large(msg!( + return Err(GatewayError::too_large(msg!( "gw.request.body_over_limit", max = max => "The request body is over the {max}-byte limit." - ))) + ))); } // 客户端传到一半断了。它多半已经听不到这句了，但这一行得有个结局 - Err(e) => Err(GatewayError::request(msg!( - "gw.request.body_unreadable", detail = e => - "The request body could not be read: {detail}" + Err(e) => { + return Err(GatewayError::request(msg!( + "gw.request.body_unreadable", detail = e => + "The request body could not be read: {detail}" + ))); + } + }; + let Some(coding) = coding else { + return Ok(raw); + }; + headers.remove(header::CONTENT_ENCODING); + // 解压是纯 CPU 的活，几百兆要几百毫秒：不占 async 线程 + let decoded = tokio::task::spawn_blocking(move || coding.decode(&raw, max)) + .await + .unwrap_or_else(|e| Err(Decode::Failed(e.to_string()))); + match decoded { + Ok(b) => Ok(b), + Err(Decode::OverLimit) => Err(GatewayError::too_large(msg!( + "gw.request.decoded_body_over_limit", max = max => + "The request body is over the {max}-byte limit once decoded." + ))), + Err(Decode::Failed(detail)) => Err(GatewayError::request(msg!( + "gw.request.bad_encoding", encoding = coding.name(), detail = detail => + "The request body could not be decoded as {encoding}: {detail}" ))), } } +/// 这个请求的体压没压：没压是 `Ok(None)`，压成认不得的样子是 400。 +/// +/// **只认 gzip、deflate、zstd**（见 [`crate::inflate`]）；叠了几层的（`gzip, zstd`）也不认 +fn coding_of(headers: &HeaderMap) -> Result, GatewayError> { + let Some(v) = headers.get(header::CONTENT_ENCODING) else { + return Ok(None); + }; + let text = String::from_utf8_lossy(v.as_bytes()).trim().to_string(); + if text.is_empty() || text.eq_ignore_ascii_case("identity") { + return Ok(None); + } + Coding::parse(&text).map(Some).ok_or_else(|| { + GatewayError::request(msg!( + "gw.request.unsupported_encoding", encoding = text => + "The request body is encoded as {encoding}; the gateway decodes only gzip, deflate and zstd." + )) + }) +} + #[cfg(test)] mod tests { use super::*; @@ -66,7 +116,7 @@ mod tests { #[tokio::test] async fn a_body_at_the_limit_is_read_whole() { - let b = read(&HeaderMap::new(), chunked(&[b"12", b"345"]), 5) + let b = read(&mut HeaderMap::new(), chunked(&[b"12", b"345"]), 5) .await .unwrap(); assert_eq!(&b[..], b"12345"); @@ -74,7 +124,7 @@ mod tests { #[tokio::test] async fn a_body_without_a_length_stops_at_the_limit() { - let e = read(&HeaderMap::new(), chunked(&[b"123", b"456"]), 5) + let e = read(&mut HeaderMap::new(), chunked(&[b"123", b"456"]), 5) .await .unwrap_err(); assert_eq!(e.detail.code, "gw.request.body_over_limit"); @@ -87,7 +137,7 @@ mod tests { h.insert(header::CONTENT_LENGTH, HeaderValue::from_static("6")); // 体是一条永远不结束的流：读了就挂住，拒得掉说明根本没读 let e = read( - &h, + &mut h, Body::from_stream(futures::stream::pending::>()), 5, ) @@ -101,7 +151,112 @@ mod tests { async fn a_declared_length_that_understates_the_body_does_not_lift_the_limit() { let mut h = HeaderMap::new(); h.insert(header::CONTENT_LENGTH, HeaderValue::from_static("1")); - let e = read(&h, chunked(&[b"123", b"456"]), 5).await.unwrap_err(); + let e = read(&mut h, chunked(&[b"123", b"456"]), 5) + .await + .unwrap_err(); assert_eq!(e.detail.code, "gw.request.body_over_limit"); } + + fn gz(data: &[u8]) -> Vec { + use std::io::Write; + let mut e = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + e.write_all(data).unwrap(); + e.finish().unwrap() + } + + fn with_encoding(v: &'static str) -> HeaderMap { + let mut h = HeaderMap::new(); + h.insert(header::CONTENT_ENCODING, HeaderValue::from_static(v)); + h + } + + #[tokio::test] + async fn a_gzip_body_is_decoded_and_the_encoding_header_is_dropped() { + let packed = gz(b"12345"); + let mut h = with_encoding("gzip"); + let b = read(&mut h, Body::from(packed), 1024).await.unwrap(); + assert_eq!(&b[..], b"12345"); + assert!(!h.contains_key(header::CONTENT_ENCODING)); + } + + #[tokio::test] + async fn deflate_and_zstd_bodies_are_decoded() { + use std::io::Write; + let mut e = flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default()); + e.write_all(b"12345").unwrap(); + let b = read( + &mut with_encoding("deflate"), + Body::from(e.finish().unwrap()), + 1024, + ) + .await + .unwrap(); + assert_eq!(&b[..], b"12345"); + let z = zstd::stream::encode_all(&b"12345"[..], 3).unwrap(); + let b = read(&mut with_encoding("zstd"), Body::from(z), 1024) + .await + .unwrap(); + assert_eq!(&b[..], b"12345"); + } + + /// 一小段压缩数据解出一大片：按解开之后的算，和没压缩的超限一样是 413。压缩着的 + /// 那几十字节本身在上限之内，拦下它的只能是解开之后的长度 + #[tokio::test] + async fn a_body_that_decodes_to_more_than_the_limit_is_too_large() { + let packed = gz(&[b'x'; 10_000]); + assert!(packed.len() < 100); + let e = read(&mut with_encoding("gzip"), Body::from(packed), 100) + .await + .unwrap_err(); + assert_eq!(e.detail.code, "gw.request.decoded_body_over_limit"); + assert_eq!(e.source, crate::error::Source::TooLarge); + } + + #[tokio::test] + async fn a_body_that_is_not_what_its_encoding_says_is_a_bad_request() { + let e = read( + &mut with_encoding("gzip"), + Body::from(&b"this is not gzip"[..]), + 1024, + ) + .await + .unwrap_err(); + assert_eq!(e.detail.code, "gw.request.bad_encoding"); + assert_eq!(e.detail.args["encoding"], "gzip"); + assert_eq!(e.source, crate::error::Source::Request); + // 解到一半断了的也是 + let mut packed = gz(&[b'x'; 10_000]); + packed.truncate(packed.len() / 2); + let e = read(&mut with_encoding("gzip"), Body::from(packed), 1 << 20) + .await + .unwrap_err(); + assert_eq!(e.detail.code, "gw.request.bad_encoding"); + } + + /// 认不得的编码在读体之前就拒：体是一条永远不结束的流 + #[tokio::test] + async fn an_unknown_encoding_is_refused_without_reading_the_body() { + let e = read( + &mut with_encoding("br"), + Body::from_stream(futures::stream::pending::>()), + 5, + ) + .await + .unwrap_err(); + assert_eq!(e.detail.code, "gw.request.unsupported_encoding"); + assert_eq!(e.detail.args["encoding"], "br"); + // 叠了几层的也不认 + let e = read(&mut with_encoding("gzip, zstd"), Body::from(gz(b"1")), 5) + .await + .unwrap_err(); + assert_eq!(e.detail.code, "gw.request.unsupported_encoding"); + } + + #[tokio::test] + async fn identity_means_not_encoded() { + let b = read(&mut with_encoding("identity"), Body::from(&b"123"[..]), 5) + .await + .unwrap(); + assert_eq!(&b[..], b"123"); + } } diff --git a/crates/tw-gateway/src/server/pipeline/hop.rs b/crates/tw-gateway/src/server/pipeline/hop.rs index 530dc6c1..ec673082 100644 --- a/crates/tw-gateway/src/server/pipeline/hop.rs +++ b/crates/tw-gateway/src/server/pipeline/hop.rs @@ -1758,8 +1758,9 @@ async fn send( } /// 发出去一次：请求体（`size` 字节）记进这个请求的流量，响应套上一层、读到的每一块都记上 -/// （见 [`crate::traffic`]）。**连不上的退回**：地址不通、代理拒绝、握手失败时上游一个字节都 -/// 没收到。别的错误（超时、发到一半断了）照记：请求体多半已经出去了 +/// （见 [`crate::traffic`]）；压缩着回的在计数之外再套一层解开（见 [`crate::inflate`]）， +/// 后面读它的每一处拿到的都是解开的字节。**连不上的退回**：地址不通、代理拒绝、握手失败时 +/// 上游一个字节都没收到。别的错误（超时、发到一半断了）照记：请求体多半已经出去了 async fn dispatch( traffic: &std::sync::Arc, size: usize, @@ -1767,7 +1768,9 @@ async fn dispatch( ) -> reqwest::Result { traffic.sending(size); match sending.await { - Ok(r) => Ok(crate::traffic::metered(r, traffic)), + Ok(r) => Ok(crate::inflate::inflated(crate::traffic::metered( + r, traffic, + ))), Err(e) => { if e.is_connect() { traffic.unsent(size); @@ -1826,7 +1829,13 @@ async fn bedrock_refusal( so it is not passed on." ), }; - for h in ["content-length", "content-type", "transfer-encoding"] { + // 正文换成了我们自己的一句话：说明原来那份正文长什么样的头都不能留 + for h in [ + "content-length", + "content-type", + "content-encoding", + "transfer-encoding", + ] { headers.remove(h); } // 前缀只在交给客户端的这一份上（和网关自己的错误一样，见 `crate::error`）：记录里的 diff --git a/crates/tw-gateway/src/traffic.rs b/crates/tw-gateway/src/traffic.rs index e646c724..468cbed3 100644 --- a/crates/tw-gateway/src/traffic.rs +++ b/crates/tw-gateway/src/traffic.rs @@ -13,8 +13,9 @@ //! 第一段内容的、把回答交给客户端的，各走各的路，数的是同一个数，不会漏一处、也不会数两遍。 //! //! 解压之前怎么数得到：**网关的 HTTP 客户端不解压**（reqwest 没开 gzip、brotli、zstd、 -//! deflate 这几个特性），读到的就是线上的那些字节，上游压缩过的也还压缩着。哪天开了自动 -//! 解压，这里数到的就成了解压之后的 —— `tests::the_client_does_not_decompress` 守着这一条。 +//! deflate 这几个特性），读到的就是线上的那些字节，上游压缩过的也还压缩着；解压是网关自己 +//! 的一层（[`crate::inflate`]），套在计数的外面。哪天开了 reqwest 的自动解压，这里数到的就成了 +//! 解压之后的 —— `tests::the_client_does_not_decompress` 守着这一条。 //! //! WebSocket 不走 HTTP 客户端：帧的载荷由那条路自己记（见 [`crate::ending::Ending::count`]、 //! [`crate::ending::Ending::sending`]）。 diff --git a/crates/tw-gateway/tests/compressed.rs b/crates/tw-gateway/tests/compressed.rs new file mode 100644 index 00000000..3e664c82 --- /dev/null +++ b/crates/tw-gateway/tests/compressed.rs @@ -0,0 +1,327 @@ +//! 压缩过的请求体和上游回答，端到端（见 `tw_gateway::inflate`、`server::intake`）。 +//! +//! - 客户端压缩着发（`Content-Encoding: gzip` / `zstd`）：直通和转换两条路，上游收到的都是 +//! 解开的请求，不带 `Content-Encoding` +//! - 压坏了的是 400，按客户端的格式回，上游什么都收不到 +//! - 上游压缩着回（整包和 SSE）：客户端收到的是解开的，不带 `Content-Encoding`；流式的 +//! 一条一条到，不等整个流结束 +//! +//! 解开之后超限的在 `server::intake` 的单元测试里（上限是 256 MiB，端到端要真解出那么多）。 + +use std::io::Write; +use std::net::SocketAddr; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use axum::Router; +use axum::extract::{OriginalUri, State}; +use axum::http::HeaderMap; +use futures::StreamExt; +use serde_json::{Value, json}; +use tw_config::{Client, Config, Listen, Protocol, Provider}; + +#[derive(Debug, Default, Clone)] +struct Seen { + uri: String, + headers: HeaderMap, + body: Vec, +} + +/// 假上游怎么回 +#[derive(Clone)] +enum Reply { + /// 没压缩的整包 + Plain(&'static str), + /// gzip 压缩的整包 + Gzip(&'static str), + /// gzip 压缩的 SSE，每条事件一段、段与段之间停一会儿，像边生成边压缩那样 + GzipStream(&'static [&'static str]), +} + +const ANSWER: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5","content":[{"type":"text","text":"hello there"}],"stop_reason":"end_turn","usage":{"input_tokens":12,"output_tokens":3}}"#; + +const STREAM: &[&str] = &[ + "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-sonnet-4-5\",\"content\":[],\"usage\":{\"input_tokens\":12,\"output_tokens\":0}}}\n\n", + "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n", + "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello there\"}}\n\n", + "event: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}\n\n", + "event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":3}}\n\n", + "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", +]; + +fn gzip(data: &[u8]) -> Vec { + let mut e = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + e.write_all(data).unwrap(); + e.finish().unwrap() +} + +/// 一个 gzip 流分几段：每段 flush 一次，拿到那一刻压出来的字节 +fn gzip_frames(parts: &[&str]) -> Vec> { + let mut e = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + let mut frames = Vec::new(); + for p in parts { + e.write_all(p.as_bytes()).unwrap(); + e.flush().unwrap(); + frames.push(std::mem::take(e.get_mut())); + } + frames.push(e.finish().unwrap()); + frames +} + +async fn upstream(reply: Reply) -> (SocketAddr, Arc>) { + let seen = Arc::new(Mutex::new(Seen::default())); + let app = Router::new() + .fallback( + move |State(s): State>>, + OriginalUri(uri): OriginalUri, + headers: HeaderMap, + body: bytes::Bytes| { + let reply = reply.clone(); + async move { + *s.lock().unwrap() = Seen { + uri: uri.to_string(), + headers, + body: body.to_vec(), + }; + let b = axum::response::Response::builder(); + match reply { + Reply::Plain(text) => b + .header("content-type", "application/json") + .body(axum::body::Body::from(text)) + .unwrap(), + Reply::Gzip(text) => { + let packed = gzip(text.as_bytes()); + b.header("content-type", "application/json") + .header("content-encoding", "gzip") + .header("content-length", packed.len()) + .body(axum::body::Body::from(packed)) + .unwrap() + } + Reply::GzipStream(events) => { + let frames = gzip_frames(events); + let stream = async_stream::stream! { + for f in frames { + yield Ok::<_, std::io::Error>(bytes::Bytes::from(f)); + tokio::time::sleep(Duration::from_millis(150)).await; + } + }; + b.header("content-type", "text/event-stream") + .header("content-encoding", "gzip") + .body(axum::body::Body::from_stream(stream)) + .unwrap() + } + } + } + }, + ) + .with_state(seen.clone()); + let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = l.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(l, app).await.unwrap() }); + (addr, seen) +} + +async fn gateway(up: SocketAddr) -> SocketAddr { + let cfg = Config { + version: 1, + listen: Listen::default(), + clients: vec![Client { + name: "c".into(), + key: "tw-k".into(), + ..Default::default() + }], + providers: vec![Provider { + name: "up".into(), + base_url: format!("http://{up}"), + key: Some("sk-upstream".into()), + protocol: Some(Protocol::Anthropic), + ..Default::default() + }], + ..Default::default() + }; + let state = tw_gateway::AppState::new(cfg).unwrap(); + let addr = tw_gateway::serve(state, ([127, 0, 0, 1], 0).into()) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + addr +} + +/// 发一个请求体（已经按 `encoding` 压好的），交回响应。测试用的客户端不解压，看到的就是 +/// 网关发出的字节 +async fn post( + gw: SocketAddr, + path: &str, + encoding: Option<&str>, + body: Vec, +) -> reqwest::Response { + let mut r = reqwest::Client::new() + .post(format!("http://{gw}{path}")) + .header("x-api-key", "tw-k") + .header("content-type", "application/json"); + if let Some(e) = encoding { + r = r.header("content-encoding", e); + } + r.body(body).send().await.unwrap() +} + +fn claude_request() -> String { + json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi there"}] + }) + .to_string() +} + +#[tokio::test] +async fn a_gzip_request_body_reaches_the_upstream_decoded_on_the_passthrough_route() { + let (up, seen) = upstream(Reply::Plain(ANSWER)).await; + let gw = gateway(up).await; + let plain = claude_request(); + let r = post(gw, "/v1/messages", Some("gzip"), gzip(plain.as_bytes())).await; + assert_eq!(r.status(), 200, "{}", r.text().await.unwrap()); + + let s = seen.lock().unwrap().clone(); + assert_eq!(s.uri, "/v1/messages"); + assert_eq!( + String::from_utf8_lossy(&s.body), + plain, + "上游收到的该是解开的、一个字节都没改的请求体" + ); + assert!( + s.headers.get("content-encoding").is_none(), + "客户端的 Content-Encoding 不该到上游：{:?}", + s.headers + ); + assert_eq!(s.headers.get("content-type").unwrap(), "application/json"); +} + +#[tokio::test] +async fn a_zstd_request_body_on_a_converted_route_is_decoded_before_conversion() { + let (up, seen) = upstream(Reply::Plain(ANSWER)).await; + let gw = gateway(up).await; + let chat = json!({ + "model": "claude-sonnet-4-5", + "messages": [{"role": "system", "content": "简短"}, {"role": "user", "content": "hi there"}] + }) + .to_string(); + let packed = zstd::stream::encode_all(chat.as_bytes(), 3).unwrap(); + let r = post(gw, "/v1/chat/completions", Some("zstd"), packed).await; + assert_eq!(r.status(), 200, "{}", r.text().await.unwrap()); + + let s = seen.lock().unwrap().clone(); + assert_eq!(s.uri, "/v1/messages", "转换过的请求发到 Anthropic 的路径"); + assert!( + s.headers.get("content-encoding").is_none(), + "{:?}", + s.headers + ); + let sent: Value = serde_json::from_slice(&s.body).expect("上游收到的该是 JSON"); + assert_eq!(sent["system"][0]["text"], "简短"); + assert_eq!(sent["messages"][0]["content"][0]["text"], "hi there"); +} + +#[tokio::test] +async fn a_compressed_count_tokens_request_is_decoded_too() { + let (up, seen) = upstream(Reply::Plain(r#"{"input_tokens":12}"#)).await; + let gw = gateway(up).await; + let plain = claude_request(); + let r = post( + gw, + "/v1/messages/count_tokens", + Some("gzip"), + gzip(plain.as_bytes()), + ) + .await; + assert_eq!(r.status(), 200, "{}", r.text().await.unwrap()); + let s = seen.lock().unwrap().clone(); + assert_eq!(s.uri, "/v1/messages/count_tokens"); + assert_eq!(String::from_utf8_lossy(&s.body), plain); + assert!(s.headers.get("content-encoding").is_none()); +} + +#[tokio::test] +async fn a_corrupt_compressed_body_is_a_400_in_the_clients_format() { + let (up, seen) = upstream(Reply::Plain(ANSWER)).await; + let gw = gateway(up).await; + let r = post( + gw, + "/v1/messages", + Some("gzip"), + b"this is not gzip".to_vec(), + ) + .await; + assert_eq!(r.status(), 400); + assert_eq!(r.headers().get("x-thinkwatch-error").unwrap(), "request"); + let body: Value = r.json().await.unwrap(); + assert_eq!(body["type"], "error", "{body}"); + assert_eq!(body["error"]["type"], "invalid_request_error", "{body}"); + let message = body["error"]["message"].as_str().unwrap(); + assert!(message.contains("gzip"), "{message}"); + assert!(seen.lock().unwrap().body.is_empty(), "上游不该收到任何东西"); +} + +#[tokio::test] +async fn an_encoding_the_gateway_does_not_decode_is_refused() { + let (up, seen) = upstream(Reply::Plain(ANSWER)).await; + let gw = gateway(up).await; + let r = post(gw, "/v1/messages", Some("br"), b"whatever".to_vec()).await; + assert_eq!(r.status(), 400); + let text = r.text().await.unwrap(); + assert!(text.contains("br"), "{text}"); + assert!(seen.lock().unwrap().body.is_empty()); +} + +#[tokio::test] +async fn an_upstream_that_answers_gzip_json_reaches_the_client_decoded() { + let (up, _) = upstream(Reply::Gzip(ANSWER)).await; + let gw = gateway(up).await; + let r = post(gw, "/v1/messages", None, claude_request().into_bytes()).await; + assert_eq!(r.status(), 200); + assert!( + r.headers().get("content-encoding").is_none(), + "解开了就不能再说它压缩着：{:?}", + r.headers() + ); + assert_eq!(r.headers().get("content-type").unwrap(), "application/json"); + let body = r.bytes().await.unwrap(); + assert_eq!(String::from_utf8_lossy(&body), ANSWER); +} + +#[tokio::test] +async fn an_upstream_that_streams_gzip_sse_reaches_the_client_decoded_event_by_event() { + let (up, _) = upstream(Reply::GzipStream(STREAM)).await; + let gw = gateway(up).await; + let r = post(gw, "/v1/messages", None, claude_request().into_bytes()).await; + assert_eq!(r.status(), 200); + assert!( + r.headers().get("content-encoding").is_none(), + "{:?}", + r.headers() + ); + assert_eq!( + r.headers().get("content-type").unwrap(), + "text/event-stream" + ); + + let mut chunks: Vec<(Duration, bytes::Bytes)> = Vec::new(); + let started = std::time::Instant::now(); + let mut s = r.bytes_stream(); + while let Some(c) = s.next().await { + chunks.push((started.elapsed(), c.unwrap())); + } + let text: String = chunks + .iter() + .map(|(_, c)| String::from_utf8_lossy(c).into_owned()) + .collect(); + assert_eq!(text, STREAM.concat(), "解开的 SSE 该和上游压之前的一模一样"); + // 上游每隔 150 毫秒发一段；第一段到客户端的时候最后一段还没发出来。整个流攒到结束才 + // 交的话，第一块要等到 5 × 150 毫秒之后 + let (first, _) = &chunks[0]; + assert!( + *first < Duration::from_millis(450), + "第一条事件 {first:?} 才到：流被攒起来了" + ); + assert!(chunks.len() >= 2, "{} 块", chunks.len()); +} diff --git a/crates/tw-gateway/tests/traffic.rs b/crates/tw-gateway/tests/traffic.rs index aed0e4ce..110cc4dc 100644 --- a/crates/tw-gateway/tests/traffic.rs +++ b/crates/tw-gateway/tests/traffic.rs @@ -1,7 +1,7 @@ //! 一个请求在网关和上游之间走了多少流量、从哪个出口出去的，端到端（见 `tw_gateway::traffic`）。 //! //! - 故障转移：**每一跳都算**，连不上的那一跳没发出去、不算；失败的那一跳读过的错误正文算 -//! - 收到的按线上的样子数，**解压之前**：上游压缩过的回答数压缩着的字节 +//! - 收到的按线上的样子数，**解压之前**：上游压缩过的回答数压缩着的字节，交给客户端的是解开的 //! - 走代理的上游：结局的出口、尝试链上每一跳的代理都是它的名字；直连的没有 use std::net::SocketAddr; @@ -73,7 +73,7 @@ fn cfg(providers: Vec) -> Config { } } -/// 发一个请求，交回状态码和**原样**的正文（测试用的客户端不解压） +/// 发一个请求，交回状态码和**原样**的正文（测试用的客户端不解压：看到的就是网关发出的字节） async fn post(gw: SocketAddr) -> (u16, bytes::Bytes) { let r = reqwest::Client::new() .post(format!("http://{gw}/v1/messages")) @@ -109,7 +109,8 @@ async fn routed_and_ended(rx: &mut tokio::sync::broadcast::Receiver) -> ( /// 三家：第一家连不上，第二家回 503，第三家回一个 zstd 压缩过的整包回答。 /// /// **发出去的是两跳的请求体**：连不上的那一跳一个字节都没出去。**收回来的是 503 的错误正文 -/// 加上压缩着的回答** —— 解压之后的长度不是流量。全都直连：没有出口 +/// 加上压缩着的回答** —— 解压之后的长度不是流量；回答本身解开了交给客户端（见 +/// `tw_gateway::inflate`），数的还是线上的字节。全都直连：没有出口 #[tokio::test] async fn every_hop_counts_and_what_arrived_is_counted_before_decompression() { let dead = format!("http://127.0.0.1:{}", spare_port()); @@ -144,7 +145,11 @@ async fn every_hop_counts_and_what_arrived_is_counted_before_decompression() { let (status, body) = post(gw).await; assert_eq!(status, 200); - assert_eq!(&body[..], &packed[..], "回答原样交给了客户端"); + assert_eq!( + String::from_utf8_lossy(&body), + ANSWER, + "回答解开了交给客户端" + ); let (routed, end) = routed_and_ended(&mut rx).await; let Event::RequestRouted {