From 92caef37f8c80e0ee31ab9acad117662bc169f78 Mon Sep 17 00:00:00 2001 From: Jasper Date: Mon, 24 Aug 2026 05:08:24 +0000 Subject: [PATCH] fix(providers): bound non-streaming JSON responses (#11109) Signed-off-by: Jasper Hugo --- Cargo.lock | 1 + crates/goose-providers/Cargo.toml | 2 + crates/goose-providers/src/http_status.rs | 247 +++++++++++++++++- crates/goose-providers/src/openai.rs | 83 +++++- .../goose-providers/src/openai_compatible.rs | 76 +++++- 5 files changed, 388 insertions(+), 21 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7d5478b4d..c633f5358 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5270,6 +5270,7 @@ dependencies = [ "async-trait", "chrono", "env-lock", + "flate2", "futures", "goose-local-inference", "goose-provider-types", diff --git a/crates/goose-providers/Cargo.toml b/crates/goose-providers/Cargo.toml index e233b3476..9608264e5 100644 --- a/crates/goose-providers/Cargo.toml +++ b/crates/goose-providers/Cargo.toml @@ -59,6 +59,8 @@ tokio = { workspace = true, features = ["io-util", "macros", "net", "rt-multi-th tokio-stream = { workspace = true } env-lock = { workspace = true } wiremock.workspace = true +flate2 = "1.1.9" +reqwest = { workspace = true, features = ["gzip"] } [[example]] name = "streaming" diff --git a/crates/goose-providers/src/http_status.rs b/crates/goose-providers/src/http_status.rs index 44ed287d2..a7ec2c969 100644 --- a/crates/goose-providers/src/http_status.rs +++ b/crates/goose-providers/src/http_status.rs @@ -8,10 +8,14 @@ use std::time::{Duration, SystemTime}; use crate::errors::ProviderError; use chrono::{DateTime, NaiveDateTime, TimeZone, Utc}; +use futures::TryStreamExt; use reqwest::header::{HeaderMap, RETRY_AFTER}; use reqwest::{Response, StatusCode}; +use serde::de::DeserializeOwned; use serde_json::Value; +pub const MAX_PROVIDER_JSON_RESPONSE_BYTES: usize = 16 * 1024 * 1024; + /// Strip credentials and sensitive query parameters from a URL for safe /// inclusion in error messages and logs. Drops userinfo (`user:pass@`) and /// all query parameters (which may contain API keys like `?key=...`). @@ -326,22 +330,82 @@ pub async fn send_bounded( Ok(response) } -pub async fn read_error_body(response: Response) -> Option { - match response.extensions().get::().copied() { - Some(ResponseDeadline(deadline)) => tokio::time::timeout_at(deadline, response.text()) - .await - .ok() - .and_then(Result::ok), - None => response.text().await.ok(), +async fn read_response_body_with_limit( + response: Response, + limit: usize, +) -> Result, ProviderError> { + let deadline = response.extensions().get::().copied(); + let read = async move { + let mut stream = response.bytes_stream(); + let mut body = Vec::new(); + + while let Some(chunk) = stream.try_next().await.map_err(|e| { + ProviderError::RequestFailed(format!("Failed to read response body: {e}")) + })? { + if chunk.len() > limit.saturating_sub(body.len()) { + return Err(ProviderError::RequestFailed(format!( + "Provider response body exceeds the {limit} byte limit" + ))); + } + body.try_reserve(chunk.len()).map_err(|_| { + ProviderError::RequestFailed("Failed to allocate response body".to_string()) + })?; + body.extend_from_slice(&chunk); + } + Ok(body) + }; + + match deadline { + Some(ResponseDeadline(deadline)) => { + tokio::time::timeout_at(deadline, read).await.map_err(|_| { + ProviderError::NetworkError( + "Response body timed out — check your network connection and try again." + .to_string(), + ) + })? + } + None => read.await, } } +pub async fn read_error_body(response: Response) -> Option { + read_response_body_with_limit(response, MAX_PROVIDER_JSON_RESPONSE_BYTES) + .await + .ok() + .map(|body| String::from_utf8_lossy(&body).into_owned()) +} + +pub async fn read_json_response( + response: Response, +) -> Result { + read_json_response_with_limit(response, MAX_PROVIDER_JSON_RESPONSE_BYTES).await +} + +async fn read_json_response_with_limit( + response: Response, + limit: usize, +) -> Result { + let body = read_response_body_with_limit(response, limit).await?; + serde_json::from_slice(&body) + .map_err(|e| ProviderError::RequestFailed(format!("Response body is not valid JSON: {e}"))) +} + pub async fn handle_status(response: Response) -> Result { + handle_status_with_limit(response, MAX_PROVIDER_JSON_RESPONSE_BYTES).await +} + +async fn handle_status_with_limit( + response: Response, + limit: usize, +) -> Result { let status = response.status(); if !status.is_success() { let url = sanitize_url(response.url().as_str()); let headers = response.headers().clone(); - let body = read_error_body(response).await.unwrap_or_default(); + let body = read_response_body_with_limit(response, limit) + .await + .unwrap_or_default(); + let body = String::from_utf8_lossy(&body); let payload = serde_json::from_str::(&body).ok(); let mut err = map_http_error_to_provider_error(status, payload.clone(), &url); if let ProviderError::RateLimitExceeded { details, .. } = &err { @@ -356,17 +420,26 @@ pub async fn handle_status(response: Response) -> Result Result { - let response = handle_status(response).await?; + handle_response_with_limit(response, MAX_PROVIDER_JSON_RESPONSE_BYTES).await +} - response.json::().await.map_err(|e| { - ProviderError::RequestFailed(format!("Response body is not valid JSON: {}", e)) - }) +async fn handle_response_with_limit( + response: Response, + limit: usize, +) -> Result { + let response = handle_status_with_limit(response, limit).await?; + read_json_response_with_limit(response, limit).await } #[cfg(test)] mod tests { use super::*; + use flate2::{write::GzEncoder, Compression}; use serde_json::json; + use std::io::Write; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + use wiremock::{Mock, MockServer, ResponseTemplate}; fn empty_headers() -> HeaderMap { HeaderMap::new() @@ -378,6 +451,156 @@ mod tests { h } + async fn response_from_raw(raw_response: Vec) -> Response { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 4096]; + let _ = socket.read(&mut request).await; + socket.write_all(&raw_response).await.unwrap(); + }); + + reqwest::Client::new() + .get(format!("http://{addr}")) + .send() + .await + .unwrap() + } + + #[tokio::test] + async fn bounded_json_accepts_body_without_content_length() { + let response = response_from_raw( + b"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\nconnection: close\r\n\r\n{\"ok\":true}" + .to_vec(), + ) + .await; + + let value: Value = read_json_response_with_limit(response, 64).await.unwrap(); + assert_eq!(value, json!({"ok": true})); + } + + #[tokio::test] + async fn bounded_json_rejects_oversized_chunked_body() { + let body = format!("{{\"value\":\"{}\"}}", "a".repeat(64)); + let raw = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ntransfer-encoding: chunked\r\n\r\n{:x}\r\n{}\r\n0\r\n\r\n", + body.len(), + body + ) + .into_bytes(); + let response = response_from_raw(raw).await; + + let err = read_json_response_with_limit::(response, 64) + .await + .unwrap_err(); + assert!(err.to_string().contains("64 byte limit"), "got: {err}"); + } + + #[tokio::test] + async fn bounded_handle_response_rejects_oversized_success_body() { + let body = format!("{{\"value\":\"{}\"}}", "a".repeat(64)); + let raw = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ntransfer-encoding: chunked\r\n\r\n{:x}\r\n{}\r\n0\r\n\r\n", + body.len(), + body + ) + .into_bytes(); + let response = response_from_raw(raw).await; + + let err = handle_response_with_limit(response, 64).await.unwrap_err(); + assert!(err.to_string().contains("64 byte limit"), "got: {err}"); + } + + #[tokio::test] + async fn bounded_status_preserves_authentication_for_oversized_error_body() { + let body = format!("{{\"error\":{{\"message\":\"{}\"}}}}", "a".repeat(64)); + let raw = format!( + "HTTP/1.1 401 Unauthorized\r\ncontent-type: application/json\r\ntransfer-encoding: chunked\r\n\r\n{:x}\r\n{}\r\n0\r\n\r\n", + body.len(), + body + ) + .into_bytes(); + let response = response_from_raw(raw).await; + + let err = handle_status_with_limit(response, 64).await.unwrap_err(); + assert!( + matches!(err, ProviderError::Authentication(_)), + "got: {err}" + ); + } + + #[tokio::test] + async fn bounded_status_preserves_retry_after_for_oversized_error_body() { + let body = format!("{{\"error\":{{\"message\":\"{}\"}}}}", "a".repeat(64)); + let raw = format!( + "HTTP/1.1 429 Too Many Requests\r\ncontent-type: application/json\r\nretry-after: 17\r\ntransfer-encoding: chunked\r\n\r\n{:x}\r\n{}\r\n0\r\n\r\n", + body.len(), + body + ) + .into_bytes(); + let response = response_from_raw(raw).await; + + let err = handle_status_with_limit(response, 64).await.unwrap_err(); + assert!( + matches!( + err, + ProviderError::RateLimitExceeded { + retry_delay: Some(delay), + .. + } if delay == Duration::from_secs(17) + ), + "got: {err}" + ); + } + + #[tokio::test] + async fn bounded_json_accepts_many_small_chunks() { + let body = br#"{"ok":true}"#; + let mut raw = + b"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ntransfer-encoding: chunked\r\n\r\n" + .to_vec(); + for byte in body { + raw.extend_from_slice(b"1\r\n"); + raw.push(*byte); + raw.extend_from_slice(b"\r\n"); + } + raw.extend_from_slice(b"0\r\n\r\n"); + + let response = response_from_raw(raw).await; + let value: Value = read_json_response_with_limit(response, 64).await.unwrap(); + assert_eq!(value, json!({"ok": true})); + } + + #[tokio::test] + async fn bounded_json_limits_decompressed_body() { + let server = MockServer::start().await; + let body = format!("{{\"value\":\"{}\"}}", "a".repeat(128)); + let mut encoder = GzEncoder::new(Vec::new(), Compression::default()); + encoder.write_all(body.as_bytes()).unwrap(); + let compressed = encoder.finish().unwrap(); + + Mock::given(wiremock::matchers::method("GET")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "application/json") + .insert_header("content-encoding", "gzip") + .set_body_bytes(compressed), + ) + .mount(&server) + .await; + + let response = reqwest::Client::new() + .get(server.uri()) + .send() + .await + .unwrap(); + let err = read_json_response_with_limit::(response, 64) + .await + .unwrap_err(); + assert!(err.to_string().contains("64 byte limit"), "got: {err}"); + } + fn error_payload(fields: [(&str, Value); N]) -> Value { let mut error = json!({ "message": "invalid request" }); let error = error.as_object_mut().unwrap(); diff --git a/crates/goose-providers/src/openai.rs b/crates/goose-providers/src/openai.rs index dff3338a5..57653cb59 100644 --- a/crates/goose-providers/src/openai.rs +++ b/crates/goose-providers/src/openai.rs @@ -15,6 +15,7 @@ use crate::formats::openai_responses::{ create_responses_request_for_model, get_responses_usage, responses_api_to_message, ResponsesApiResponse, }; +use crate::http_status::read_json_response; use crate::images::ImageFormat; use crate::openai_compatible::{ handle_response_openai_compat, handle_status, stream_openai_compat, stream_responses_compat, @@ -324,9 +325,7 @@ impl OpenAiProvider { if self.supports_streaming { stream_responses_compat(response, log) } else { - let json: serde_json::Value = response.json().await.map_err(|e| { - ProviderError::RequestFailed(format!("Failed to parse JSON: {}", e)) - })?; + let json: serde_json::Value = read_json_response(response).await?; let parsed: ResponsesApiResponse = serde_json::from_value(json.clone()).map_err(|e| { ProviderError::ExecutionError(format!( @@ -817,9 +816,7 @@ impl Provider for OpenAiProvider { if self.supports_streaming { stream_openai_compat(response, log) } else { - let json: serde_json::Value = response.json().await.map_err(|e| { - ProviderError::RequestFailed(format!("Failed to parse JSON: {}", e)) - })?; + let json: serde_json::Value = read_json_response(response).await?; let message = response_to_message(&json).map_err(|e| { ProviderError::RequestFailed(format!("Failed to parse message: {}", e)) @@ -1601,6 +1598,80 @@ mod tests { assert!(!err.is_endpoint_not_found(), "got: {:?}", err); } + #[tokio::test] + async fn nonstreaming_chat_accepts_legitimate_response() { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + }))) + .mount(&server) + .await; + + let mut provider = make_provider_with_custom_models( + &server.uri(), + "v1/chat/completions", + vec!["test-model".to_string()], + ); + provider.supports_streaming = false; + + let _stream = provider + .stream(&ModelConfig::new("test-model"), "", &[], &[]) + .await + .expect("legitimate non-streaming response should be accepted"); + } + + #[tokio::test] + async fn nonstreaming_chat_rejects_oversized_response_body() { + use crate::http_status::MAX_PROVIDER_JSON_RESPONSE_BYTES; + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": "a".repeat(MAX_PROVIDER_JSON_RESPONSE_BYTES + 1) + }, + "finish_reason": "stop" + }] + }))) + .mount(&server) + .await; + + let mut provider = make_provider_with_custom_models( + &server.uri(), + "v1/chat/completions", + vec!["test-model".to_string()], + ); + provider.supports_streaming = false; + + let err = match provider + .stream(&ModelConfig::new("test-model"), "", &[], &[]) + .await + { + Ok(_) => panic!("oversized response should be rejected"), + Err(err) => err, + }; + assert!( + err.to_string().contains("response body exceeds"), + "got: {err}" + ); + } + #[tokio::test] async fn fetch_supported_models_accepts_payload_with_extra_fields() { use wiremock::matchers::{method, path}; diff --git a/crates/goose-providers/src/openai_compatible.rs b/crates/goose-providers/src/openai_compatible.rs index 7d5e7cfde..6654104b6 100644 --- a/crates/goose-providers/src/openai_compatible.rs +++ b/crates/goose-providers/src/openai_compatible.rs @@ -1,4 +1,5 @@ use crate::conversation::token_usage::{CostSource, ProviderUsage}; +use crate::http_status::read_json_response; use crate::images::ImageFormat; use anyhow::Error; use async_stream::try_stream; @@ -126,9 +127,7 @@ impl OpenAiCompatibleProvider { if self.supports_streaming { stream_openai_compat(response, log) } else { - let json = response.json().await.map_err(|e| { - ProviderError::RequestFailed(format!("Failed to parse JSON: {}", e)) - })?; + let json = read_json_response(response).await?; let message = response_to_message(&json).map_err(|e| { ProviderError::RequestFailed(format!("Failed to parse message: {}", e)) })?; @@ -381,4 +380,75 @@ mod tests { assert_eq!(payload.get("stream"), None); assert_eq!(payload.get("stream_options"), None); } + + #[tokio::test] + async fn nonstreaming_completion_accepts_legitimate_response() { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "choices": [{ + "message": {"role": "assistant", "content": "hello"} + }] + }))) + .mount(&server) + .await; + + let provider = OpenAiCompatibleProvider::new( + "test".to_string(), + ApiClient::new_with_tls(server.uri(), crate::api_client::AuthMethod::NoAuth, None) + .unwrap(), + String::new(), + ) + .with_supports_streaming(false); + + let _stream = provider + .stream(&ModelConfig::new("test-model"), "", &[], &[]) + .await + .expect("legitimate non-streaming response should be accepted"); + } + + #[tokio::test] + async fn nonstreaming_completion_rejects_oversized_response_body() { + use crate::http_status::MAX_PROVIDER_JSON_RESPONSE_BYTES; + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "choices": [{ + "message": { + "role": "assistant", + "content": "a".repeat(MAX_PROVIDER_JSON_RESPONSE_BYTES + 1) + } + }] + }))) + .mount(&server) + .await; + + let provider = OpenAiCompatibleProvider::new( + "test".to_string(), + ApiClient::new_with_tls(server.uri(), crate::api_client::AuthMethod::NoAuth, None) + .unwrap(), + String::new(), + ) + .with_supports_streaming(false); + + let err = match provider + .stream(&ModelConfig::new("test-model"), "", &[], &[]) + .await + { + Ok(_) => panic!("oversized response should be rejected"), + Err(err) => err, + }; + assert!( + err.to_string().contains("response body exceeds"), + "got: {err}" + ); + } }