Moved read_chunked_body_with_max to cb-common

Author: jclapis

Cargo.lock

@@ -23,6 +23,7 @@ ethereum_serde_utils.workspace = true
 ethereum_ssz.workspace = true
 ethereum_ssz_derive.workspace = true
 eyre.workspace = true
+futures.workspace = true
 jsonwebtoken.workspace = true
 pbkdf2.workspace = true
 rand.workspace = true

crates/common/Cargo.toml

@@ -1,12 +1,15 @@
 use std::{collections::HashMap, env, path::Path};
 
 use alloy::rpc::types::beacon::BlsPublicKey;
-use bytes::{BufMut, BytesMut};
 use eyre::{bail, Context, Result};
 use serde::de::DeserializeOwned;
 
 use super::JWTS_ENV;
-use crate::{config::MUXER_HTTP_MAX_LENGTH, types::ModuleId};
+use crate::{
+    config::MUXER_HTTP_MAX_LENGTH,
+    types::ModuleId,
+    utils::{read_chunked_body_with_max, ResponseReadError},
+};
 
 pub const CB_TEST_HTTP_DISABLE_CONTENT_LENGTH_ENV: &str = "CB_TEST_HTTP_DISABLE_CONTENT_LENGTH";
 
@@ -36,7 +39,7 @@ pub fn load_jwt_secrets() -> Result<HashMap<ModuleId, String>> {
 
 /// Reads an HTTP response safely, erroring out if it failed or if the body is
 /// too large.
-pub async fn safe_read_http_response(mut response: reqwest::Response) -> Result<String> {
+pub async fn safe_read_http_response(response: reqwest::Response) -> Result<String> {
     // Get the content length from the response headers
     let mut content_length = response.content_length();
     if env::var(CB_TEST_HTTP_DISABLE_CONTENT_LENGTH_ENV).is_ok() {
@@ -56,18 +59,16 @@ pub async fn safe_read_http_response(mut response: reqwest::Response) -> Result<
     }
 
     // Read the response to a buffer in chunks
-    let mut buffer = BytesMut::with_capacity(1024);
-    while let Some(chunk) = response.chunk().await? {
-        if buffer.len() > MUXER_HTTP_MAX_LENGTH as usize {
-            bail!(
-                "Response body exceeds the maximum allowed length ({MUXER_HTTP_MAX_LENGTH} bytes)"
-            );
-        }
-        buffer.put(chunk);
-    }
+    let result = read_chunked_body_with_max(response, MUXER_HTTP_MAX_LENGTH as usize).await;
+    let bytes = match result {
+        Ok(bytes) => Ok(bytes),
+        Err(ResponseReadError::PayloadTooLarge { max: _, raw: _ }) => bail!(
+            "Response body exceeds the maximum allowed length ({MUXER_HTTP_MAX_LENGTH} bytes)"
+        ),
+        Err(ResponseReadError::ChunkError { inner }) => Err(inner),
+    }?;
 
     // Convert the buffer to a string
-    let bytes = buffer.freeze();
     match std::str::from_utf8(&bytes) {
         Ok(s) => Ok(s.to_string()),
         Err(e) => bail!("Failed to decode response body as UTF-8: {e}"),

crates/common/src/config/utils.rs

@@ -9,11 +9,13 @@ use alloy::{
 };
 use axum::http::HeaderValue;
 use blst::min_pk::{PublicKey, Signature};
+use futures::StreamExt;
 use rand::{distr::Alphanumeric, Rng};
-use reqwest::header::HeaderMap;
+use reqwest::{header::HeaderMap, Error as ReqwestError, Response};
 use serde::{de::DeserializeOwned, Serialize};
 use serde_json::Value;
 use ssz::{Decode, Encode};
+use thiserror::Error;
 use tracing::Level;
 use tracing_appender::{non_blocking::WorkerGuard, rolling::Rotation};
 use tracing_subscriber::{
@@ -31,6 +33,46 @@ use crate::{
 
 const MILLIS_PER_SECOND: u64 = 1_000;
 
+#[derive(Debug, Error)]
+pub enum ResponseReadError {
+    #[error("response size exceeds max size: max: {max} raw: {raw}")]
+    PayloadTooLarge { max: usize, raw: String },
+
+    #[error("error reading chunk from response: {inner:?}")]
+    ChunkError { inner: ReqwestError },
+}
+
+/// Reads the body of a response as a chunked stream, ensuring the size does not
+/// exceed `max_size`.
+pub async fn read_chunked_body_with_max(
+    res: Response,
+    max_size: usize,
+) -> Result<Vec<u8>, ResponseReadError> {
+    let mut stream = res.bytes_stream();
+    let mut response_bytes = Vec::new();
+
+    while let Some(chunk) = stream.next().await {
+        let chunk = match chunk {
+            Ok(c) => c,
+            Err(e) => {
+                return Err(ResponseReadError::ChunkError { inner: e });
+            }
+        };
+        if response_bytes.len() + chunk.len() > max_size {
+            // avoid spamming logs if the message is too large
+            response_bytes.truncate(1024);
+            return Err(ResponseReadError::PayloadTooLarge {
+                max: max_size,
+                raw: String::from_utf8_lossy(&response_bytes).into_owned(),
+            });
+        }
+
+        response_bytes.extend_from_slice(&chunk);
+    }
+
+    Ok(response_bytes)
+}
+
 pub fn timestamp_of_slot_start_sec(slot: u64, chain: Chain) -> u64 {
     chain.genesis_time_sec() + slot * chain.slot_time_sec()
 }

crates/common/src/utils.rs

@@ -1,29 +1,21 @@
-use cb_common::pbs::error::PbsError;
-use futures::StreamExt;
+use cb_common::{
+    pbs::error::PbsError,
+    utils::{read_chunked_body_with_max as read_chunked_body_with_max_impl, ResponseReadError},
+};
 use reqwest::Response;
 
 pub async fn read_chunked_body_with_max(
     res: Response,
     max_size: usize,
 ) -> Result<Vec<u8>, PbsError> {
-    let mut stream = res.bytes_stream();
-    let mut response_bytes = Vec::new();
-
-    while let Some(chunk) = stream.next().await {
-        let chunk = chunk?;
-        if response_bytes.len() + chunk.len() > max_size {
-            // avoid spamming logs if the message is too large
-            response_bytes.truncate(1024);
-            return Err(PbsError::PayloadTooLarge {
-                max: max_size,
-                raw: String::from_utf8_lossy(&response_bytes).into_owned(),
-            });
+    let result = read_chunked_body_with_max_impl(res, max_size).await;
+    match result {
+        Ok(bytes) => Ok(bytes),
+        Err(ResponseReadError::PayloadTooLarge { max, raw }) => {
+            Err(PbsError::PayloadTooLarge { max, raw })
         }
-
-        response_bytes.extend_from_slice(&chunk);
+        Err(ResponseReadError::ChunkError { inner }) => Err(PbsError::Reqwest(inner)),
     }
-
-    Ok(response_bytes)
 }
 
 const GAS_LIMIT_ADJUSTMENT_FACTOR: u64 = 1024;