diff --git a/crates/tinytools-std/src/network/gate.rs b/crates/tinytools-std/src/network/gate.rs index 201366b..696148a 100644 --- a/crates/tinytools-std/src/network/gate.rs +++ b/crates/tinytools-std/src/network/gate.rs @@ -207,6 +207,27 @@ pub(crate) fn host_of(url: &str) -> String { .unwrap_or_else(|| "unknown".to_string()) } +/// Reads at most `cap + 1` bytes of `response`'s body, so a caller can tell +/// a body over the cap from one exactly at it without reading the rest. +/// +/// # Errors +/// +/// Returns the transport error if a chunk fails to arrive. +pub(super) async fn read_body_capped( + mut response: reqwest::Response, + cap: usize, +) -> reqwest::Result> { + let limit = cap.saturating_add(1); + let mut body = Vec::new(); + while body.len() < limit + && let Some(chunk) = response.chunk().await? + { + let take = chunk.len().min(limit - body.len()); + body.extend_from_slice(&chunk[..take]); + } + Ok(body) +} + /// Response-size and timeout limits a tool falls back to when it is /// constructed with a `0` (or absent) value. /// diff --git a/crates/tinytools-std/src/network/http_request.rs b/crates/tinytools-std/src/network/http_request.rs index a29074a..c6edb27 100644 --- a/crates/tinytools-std/src/network/http_request.rs +++ b/crates/tinytools-std/src/network/http_request.rs @@ -1,4 +1,4 @@ -use super::gate::{HttpLimits, NetGate, host_of}; +use super::gate::{HttpLimits, NetGate, host_of, read_body_capped}; use async_trait::async_trait; use serde_json::json; use std::sync::Arc; @@ -300,8 +300,8 @@ impl HttpRequestTool { .collect::>() .join(", "); - let response_text = match response.text().await { - Ok(text) => self.truncate_response(&text), + let response_text = match read_body_capped(response, self.max_response_size).await { + Ok(body) => self.truncate_response(&String::from_utf8_lossy(&body)), Err(e) => format!("[Failed to read response body: {e}]"), }; diff --git a/crates/tinytools-std/src/network/http_request_tests.rs b/crates/tinytools-std/src/network/http_request_tests.rs index e395663..a9458dc 100644 --- a/crates/tinytools-std/src/network/http_request_tests.rs +++ b/crates/tinytools-std/src/network/http_request_tests.rs @@ -704,3 +704,35 @@ fn truncate_response_caps_bytes_not_chars() { let marker = "\n\n... [Response truncated due to size limit] ..."; assert!(out.len() <= 10 + marker.len(), "{} bytes", out.len()); } + +// A body larger than the cap is read only up to the cap, not to its end: +// this server never stops sending, so an unbounded read never returns. +#[tokio::test] +async fn an_endless_body_is_read_only_up_to_the_cap() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::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 _ = socket.read(&mut [0u8; 8192]).await; + let _ = socket + .write_all(b"HTTP/1.1 200 OK\r\nConnection: close\r\n\r\n") + .await; + let chunk = vec![b'x'; 65536]; + while socket.write_all(&chunk).await.is_ok() {} + }); + let tool = HttpRequestTool::new(TestNetGate::supervised(), vec![], 1000, 30, DEFAULT_LIMITS); + let target = pinned_to_local(&format!("http://{addr}/big")); + let output = tokio::time::timeout(std::time::Duration::from_secs(5), async { + let response = tool + .execute_request(&target, reqwest::Method::GET, vec![], None) + .await + .unwrap(); + tool.format_response(response).await.unwrap().output() + }) + .await; + assert!(output.is_ok(), "the body read did not stop at the cap"); + let output = output.unwrap(); + assert!(output.contains("[Response truncated due to size limit]")); + assert!(!output.contains(&"x".repeat(1001))); +}