From 7eef9e6c3e0e70fa5495c1bc75c7a008f6dda52f Mon Sep 17 00:00:00 2001 From: argszero Date: Sun, 13 Sep 2026 21:54:22 +0800 Subject: [PATCH] fix(gateway): report a failed upstream body read instead of a fabricated 200 `forward` read the upstream response body with `resp.bytes().await.unwrap_or_default()`, which collapses a body-read failure into an empty vector, and then let the status line alone decide the outcome. A 2xx whose body never arrived was therefore handed to the client as `200 OK` + `content-type: application/json` with **0 bytes** - not a response any client can parse. Measured against a scripted upstream that claims `content-length: 999` and sends 47 bytes before closing: `status=200 body_len=0`. The damage is not just the empty body: `parse_usage` sees no JSON and returns `(0, 0, 0)`, `settle_usage` returns early (its contract is "only settle when tokens > 0"), so the call leaves no `usage_records` row and no `transactions` row, the balance stays unchanged and the key is not cooled. The client gets a fabricated success while the call disappears from both the ledger and the operational counter (`month_calls = COUNT(*) FROM usage_records`) - even though the upstream really spent tokens. This is the same rule the tree already states elsewhere: missing data is not the value zero. The failure is reachable: the non-stream client deliberately keeps a **total** 120 s deadline (`UPSTREAM_TIMEOUT`), so any non-stream generation whose body outlives it lands here, as does a mid-body connection reset or a proxy cutting a large response. The streaming path already handles the identical situation the opposite way - a truncated body reaches the client as an error and is not settled (see `sse_slow_stream_is_cut_by_a_total_deadline`) - and `passthrough` exists to forward an upstream body, not to invent one. Keep the status-based dispatch exactly where it is, but stop discarding the error: log it, and when the status was 2xx answer **502** instead of a fabricated empty 200. A non-2xx keeps today's behaviour byte-for-byte (empty body, then the existing 401/403/429/5xx failover or the 4xx passthrough), so a 5xx whose body is unreadable still switches keys. A 2xx is deliberately **not** retried: the upstream may already have done the work, and retrying could charge the key owner twice. Who pays is unchanged (nobody, before and after) - the point is to stop reporting a failure as a success, and deliberately not to estimate tokens. Tests (2 new, 207 -> 209): a raw scripted TCP upstream (one pre-built HTTP response per connection, `connection: close`) is required here, because the existing axum fakes can only emit well-formed bodies and therefore cannot express "200 whose body never arrives". - `forward_body_read_failure_is_not_a_success` carries its own positive control: a complete 200 + usage must still answer 200, bill exactly once and spend one upstream request; the truncated 200 must answer 502 with no `usage_records` row, an unchanged balance and still exactly one upstream request (no failover on a 2xx). - `unreadable_body_on_a_server_error_still_fails_over` pins the unchanged non-2xx path: a 500 with an unreadable body still fails over, and with every key failed the answer is still 503 with both keys cooled. A/B: reverting only the new error arm (back to `unwrap_or_default()`) leaves exactly `forward_body_read_failure_is_not_a_success` red with `left: 200, right: 502`; relaxing the new arm's condition to every status leaves exactly `unreadable_body_on_a_server_error_still_fails_over` red with `left: 502, right: 503`. Each mutation was applied alone and restored byte-identically. --- src/gateway.rs | 224 ++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 223 insertions(+), 1 deletion(-) diff --git a/src/gateway.rs b/src/gateway.rs index 686f5f3..f0b0918 100644 --- a/src/gateway.rs +++ b/src/gateway.rs @@ -339,7 +339,29 @@ async fn forward( } }; let status = resp.status(); - let bytes = resp.bytes().await.unwrap_or_default().to_vec(); + // 响应体读取失败**不得**折叠成空 body：状态行只说明响应头，不说明 body 到达。 + // 2xx + 空 body 会以 `200 OK` + `application/json` 交给客户端（不是任何客户端能 + // 解析的成功），且 parse_usage 得到 (0,0,0) ⇒ settle_usage 提前返回，该次调用 + // 从账本与运营计数（month_calls = COUNT(*) FROM usage_records）中同时消失 —— 而 + // 上游已真实消耗 token。缺数据 != 数据为 0（rant 2026-08-23T14:05:02 同款口径）。 + // 非 2xx 沿用原分流（401/403/429/5xx 换 key、其它 4xx 透传），仍按空体处理。 + let bytes = match resp.bytes().await { + Ok(b) => b.to_vec(), + Err(e) => { + log::error!( + "上游响应体读取失败 key_id={} status={}: {e}", + key.id, + status + ); + if status.is_success() { + return Err(err_json( + StatusCode::BAD_GATEWAY, + "上游响应读取失败（响应体未完整到达）", + )); + } + Vec::new() + } + }; if status.is_success() { // 成功：解析 usage（按上游出站协议）→ 计量入账 → 粘性 @@ -2608,4 +2630,204 @@ data: {"type":"message_delta","usage":{"output_tokens":30}}"#, ); up.abort(); } + + /// C2086：脚本化 HTTP 上游 —— 每接受一个连接写一条**预先构造**的原始响应后关闭。 + /// 必须用裸 socket：`fake_upstream` 这类 axum 假上游只能发出格式正确的 body， + /// 「状态行 200 但 body 未完整到达」这个形状在它们身上无法表达。 + async fn scripted_http_upstream( + listener: tokio::net::TcpListener, + responses: Vec, + hits: std::sync::Arc, + ) { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let mut idx = 0usize; + loop { + let (mut sock, _) = match listener.accept().await { + Ok(v) => v, + Err(_) => return, + }; + hits.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let mut buf = [0u8; 8192]; + let _ = + tokio::time::timeout(std::time::Duration::from_millis(300), sock.read(&mut buf)) + .await; + let pick = idx.min(responses.len() - 1); + idx += 1; + let _ = sock.write_all(responses[pick].as_bytes()).await; + let _ = sock.flush().await; + drop(sock); + } + } + + /// 构造一条原始响应；`claimed_len` 故意大于实发字节数以模拟 body 未完整到达 + fn scripted_http_response(status: &str, body: &str, claimed_len: Option) -> String { + let n = claimed_len.unwrap_or(body.len()); + format!( + "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {n}\r\nconnection: close\r\n\r\n{body}" + ) + } + + /// C2086 测试状态：单模型定价 + 属主用户，便于同时断言「入账」与「未入账」 + fn scripted_upstream_state(tag: &str, port: u16, keys: usize) -> AppState { + let base = format!("http://127.0.0.1:{port}"); + let st = test_state_eps(tag, "test-c2086", &["openai_chat"], &base); + { + let conn = st.db.lock().unwrap(); + models_row(&conn, "test", "test-model", 10.0, 20.0); + conn.execute( + "INSERT OR IGNORE INTO users (id, email, password_hash, name, role) VALUES (2, 'c2086-owner@t.local', 'x', '分享者', 'user')", + [], + ) + .unwrap(); + conn.execute( + "INSERT OR IGNORE INTO quotas (user_id, balance) VALUES (2, 0)", + [], + ) + .unwrap(); + } + for i in 0..keys { + insert_key(&st, 700 + i as i64, 2, "test-model", "test-c2086"); + } + st + } + + fn consumer_balance(st: &AppState) -> f64 { + let conn = st.db.lock().unwrap(); + conn.query_row("SELECT balance FROM quotas WHERE user_id = 1", [], |r| { + r.get(0) + }) + .unwrap() + } + + /// C2085 缺陷：状态行 200 而响应体未完整到达时，**不得**把它变成「成功的空响应」。 + /// 客户端拿到的必须是显式失败（502），且该次调用不得被记成「零用量成功」 + /// （否则账本与 month_calls 同时漏掉一次真实消耗的调用）。 + /// 本测试自带阳性对照：同样一条上游，body 完整时仍须 200 且恰好入账一次。 + #[tokio::test] + async fn forward_body_read_failure_is_not_a_success() { + const OK_BODY: &str = r#"{"id":"cmpl-c2086","object":"chat.completion","model":"test-model","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":100,"completion_tokens":50}}"#; + + // ---- 阳性对照：完整 200 + usage → 200 且恰好入账一次 ---- + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await + .unwrap(); + let port = listener.local_addr().unwrap().port(); + let hits = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let up = tokio::spawn(scripted_http_upstream( + listener, + vec![scripted_http_response("200 OK", OK_BODY, None)], + hits.clone(), + )); + let st = scripted_upstream_state("c2086ok", port, 1); + let key = login_key(st.clone()).await; + let (s, body) = post_raw( + st.clone(), + "/v1/chat/completions", + r#"{"model":"test-model","messages":[{"role":"user","content":"hi"}]}"#, + Some(&key), + ) + .await; + assert_eq!( + s, + StatusCode::OK, + "阳性对照：body 完整时必须 200，body: {}", + String::from_utf8_lossy(&body) + ); + assert_eq!(usage_records(&st), 1, "阳性对照：完整响应恰好入账一次"); + assert_eq!( + hits.load(std::sync::atomic::Ordering::SeqCst), + 1, + "阳性对照：只应发出一次上游请求" + ); + assert!( + (consumer_balance(&st) - (12471.0 - 0.0144)).abs() < 1e-9, + "阳性对照：消费者应按 usage 扣费，实测 {}", + consumer_balance(&st) + ); + up.abort(); + + // ---- 缺陷臂：content-length 999、实发不足 → 502，不入账、不重试、余额不变 ---- + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await + .unwrap(); + let port = listener.local_addr().unwrap().port(); + let hits = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let up = tokio::spawn(scripted_http_upstream( + listener, + vec![scripted_http_response("200 OK", OK_BODY, Some(999))], + hits.clone(), + )); + let st = scripted_upstream_state("c2086trunc", port, 2); + let key = login_key(st.clone()).await; + let before = consumer_balance(&st); + let (s, body) = post_raw( + st.clone(), + "/v1/chat/completions", + r#"{"model":"test-model","messages":[{"role":"user","content":"hi"}]}"#, + Some(&key), + ) + .await; + assert_eq!( + s, + StatusCode::BAD_GATEWAY, + "body 未完整到达时必须显式失败，而不是伪造 200 + 空 body；实测 body: {}", + String::from_utf8_lossy(&body) + ); + assert!( + String::from_utf8_lossy(&body).contains("上游响应读取失败"), + "502 应说明原因，实测 {}", + String::from_utf8_lossy(&body) + ); + assert_eq!(usage_records(&st), 0, "未完成的响应不得入账"); + assert_eq!(consumer_balance(&st), before, "未完成的响应不得扣费"); + assert_eq!( + hits.load(std::sync::atomic::Ordering::SeqCst), + 1, + "2xx 之后不得重发（上游可能已完成工作，重试可能让 key 属主被双扣）" + ); + up.abort(); + } + + /// C2086 非目标回归：非 2xx 且 body 读不到时，**原分流必须不变** —— + /// 5xx 仍然换 key（可故障转移），全部 key 失败后仍是 503、且不入账。 + #[tokio::test] + async fn unreadable_body_on_a_server_error_still_fails_over() { + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await + .unwrap(); + let port = listener.local_addr().unwrap().port(); + let hits = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let up = tokio::spawn(scripted_http_upstream( + listener, + vec![scripted_http_response( + "500 Internal Server Error", + r#"{"error":{"message":"upstream-boom"}}"#, + Some(999), + )], + hits.clone(), + )); + let st = scripted_upstream_state("c2086failover", port, 2); + let key = login_key(st.clone()).await; + let (s, body) = post_raw( + st.clone(), + "/v1/chat/completions", + r#"{"model":"test-model","messages":[{"role":"user","content":"hi"}]}"#, + Some(&key), + ) + .await; + assert_eq!( + s, + StatusCode::SERVICE_UNAVAILABLE, + "两个 key 都 5xx 后应 503，body: {}", + String::from_utf8_lossy(&body) + ); + assert_eq!( + hits.load(std::sync::atomic::Ordering::SeqCst), + 2, + "5xx 仍须换 key（body 读不到不改变故障转移判定）" + ); + assert_eq!(st.router.cooldown_len(), 2, "失败的 key 仍应进入冷却"); + assert_eq!(usage_records(&st), 0, "失败的调用不入账"); + up.abort(); + } }