Refactored read_chunked_body_with_max, added http_timeout to Lido key retrieval

Author: jclapis

Cargo.lock

@@ -55,11 +55,9 @@ prost = "0.13.4"
 rand = { version = "0.9", features = ["os_rng"] }
 rayon = "1.10.0"
 reqwest = { version = "0.12.4", features = ["json", "stream"] }
-scopeguard = "1.2.0"
 serde = { version = "1.0.202", features = ["derive"] }
 serde_json = "1.0.117"
 serde_yaml = "0.9.33"
-serial_test = "3.2.0"
 sha2 = "0.10.8"
 ssz_types = "0.10"
 tempfile = "3.20.0"

Cargo.toml

@@ -32,7 +32,6 @@ reqwest.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 serde_yaml.workspace = true
-serial_test.workspace = true
 sha2.workspace = true
 ssz_types.workspace = true
 thiserror.workspace = true
@@ -46,6 +45,3 @@ tree_hash.workspace = true
 tree_hash_derive.workspace = true
 unicode-normalization.workspace = true
 url.workspace = true
-
-[dev-dependencies]
-scopeguard.workspace = true

crates/common/Cargo.toml

@@ -78,7 +78,7 @@ pub const PROXY_DIR_SECRETS_DEFAULT: &str = "/proxy_secrets";
 pub const HTTP_TIMEOUT_SECONDS_DEFAULT: u64 = 10;
 
 /// Max content length for Muxer HTTP responses, in bytes
-pub const MUXER_HTTP_MAX_LENGTH: u64 = 1024 * 1024 * 1024 * 10; // 10 MiB
+pub const MUXER_HTTP_MAX_LENGTH: usize = 1024 * 1024 * 10; // 10 MiB
 
 ///////////////////////// MODULES /////////////////////////
 

crates/common/src/config/constants.rs

@@ -8,10 +8,12 @@ use std::{
 use alloy::{
     primitives::{address, Address, U256},
     providers::ProviderBuilder,
-    rpc::types::beacon::BlsPublicKey,
+    rpc::{client::RpcClient, types::beacon::BlsPublicKey},
     sol,
+    transports::http::Http,
 };
 use eyre::{bail, ensure, Context};
+use reqwest::Client;
 use serde::{Deserialize, Serialize};
 use tracing::{debug, info, warn};
 use url::Url;
@@ -193,8 +195,8 @@ impl MuxKeysLoader {
                 }
                 let client = reqwest::ClientBuilder::new().timeout(http_timeout).build()?;
                 let response = client.get(url).send().await?;
-                let pubkeys = safe_read_http_response(response).await?;
-                serde_json::from_str(&pubkeys)
+                let pubkey_bytes = safe_read_http_response(response).await?;
+                serde_json::from_slice(&pubkey_bytes)
                     .wrap_err("failed to fetch mux keys from HTTP endpoint")
             }
 
@@ -204,7 +206,13 @@ impl MuxKeysLoader {
                         bail!("Lido registry requires RPC URL to be set in the PBS config");
                     };
 
-                    fetch_lido_registry_keys(rpc_url, chain, U256::from(*node_operator_id)).await
+                    fetch_lido_registry_keys(
+                        rpc_url,
+                        chain,
+                        U256::from(*node_operator_id),
+                        http_timeout,
+                    )
+                    .await
                 }
                 NORegistry::SSV => {
                     fetch_ssv_pubkeys(chain, U256::from(*node_operator_id), http_timeout).await
@@ -254,10 +262,17 @@ async fn fetch_lido_registry_keys(
     rpc_url: Url,
     chain: Chain,
     node_operator_id: U256,
+    http_timeout: Duration,
 ) -> eyre::Result<Vec<BlsPublicKey>> {
     debug!(?chain, %node_operator_id, "loading operator keys from Lido registry");
 
-    let provider = ProviderBuilder::new().on_http(rpc_url);
+    // Create an RPC provider with HTTP timeout support
+    let client = Client::builder().timeout(http_timeout).build()?;
+    let http = Http::with_client(client, rpc_url);
+    let is_local = http.guess_local();
+    let rpc_client = RpcClient::new(http, is_local);
+    let provider = ProviderBuilder::new().on_client(rpc_client);
+
     let registry_address = lido_registry_address(chain)?;
     let registry = LidoRegistry::new(registry_address, provider);
 
@@ -362,9 +377,8 @@ async fn fetch_ssv_pubkeys_from_url(
     })?;
 
     // Parse the response as JSON
-    let body_string = safe_read_http_response(response).await?;
-    serde_json::from_slice::<SSVResponse>(body_string.as_bytes())
-        .wrap_err("failed to parse SSV response")
+    let body_bytes = safe_read_http_response(response).await?;
+    serde_json::from_slice::<SSVResponse>(&body_bytes).wrap_err("failed to parse SSV response")
 }
 
 #[derive(Deserialize)]
@@ -386,19 +400,17 @@ struct SSVPagination {
 
 #[cfg(test)]
 mod tests {
-    use std::{env, net::SocketAddr};
+    use std::net::SocketAddr;
 
     use alloy::{hex::FromHex, primitives::U256, providers::ProviderBuilder};
     use axum::{response::Response, routing::get};
-    use scopeguard::defer;
-    use serial_test::serial;
     use tokio::{net::TcpListener, task::JoinHandle};
     use url::Url;
 
     use super::*;
-    use crate::config::{
-        CB_TEST_HTTP_DISABLE_CONTENT_LENGTH_ENV, HTTP_TIMEOUT_SECONDS_DEFAULT,
-        MUXER_HTTP_MAX_LENGTH,
+    use crate::{
+        config::{HTTP_TIMEOUT_SECONDS_DEFAULT, MUXER_HTTP_MAX_LENGTH},
+        utils::{set_ignore_content_length, ResponseReadError},
     };
 
     const TEST_HTTP_TIMEOUT: u64 = 2;
@@ -471,25 +483,28 @@ mod tests {
     }
 
     #[tokio::test]
-    #[serial]
     /// Tests that the SSV network fetch is handled properly when the response's
     /// body is too large
     async fn test_ssv_network_fetch_big_data() -> eyre::Result<()> {
         // Start the mock server
         let port = 30101;
-        env::remove_var(CB_TEST_HTTP_DISABLE_CONTENT_LENGTH_ENV);
         let _server_handle = create_mock_server(port).await?;
         let url = format!("http://localhost:{port}/big_data");
         let response = fetch_ssv_pubkeys_from_url(&url, Duration::from_secs(120)).await;
 
         // The response should fail due to content length being too big
-        assert!(response.is_err(), "Expected error due to big content length, but got success");
-        if let Err(e) = response {
-            assert!(
-                e.to_string().contains("content length") &&
-                    e.to_string().contains("exceeds the maximum allowed length"),
-                "Expected content length error, got: {e}",
-            );
+        match response {
+            Ok(_) => {
+                panic!("Expected an error due to big content length, but got a successful response")
+            }
+            Err(e) => match e.downcast_ref::<ResponseReadError>() {
+                Some(ResponseReadError::PayloadTooLarge { max, content_length, raw }) => {
+                    assert_eq!(*max, MUXER_HTTP_MAX_LENGTH);
+                    assert!(*content_length > MUXER_HTTP_MAX_LENGTH);
+                    assert!(raw.is_empty());
+                }
+                _ => panic!("Expected PayloadTooLarge error, got: {}", e),
+            },
         }
 
         // Clean up the server handle
@@ -522,25 +537,29 @@ mod tests {
     }
 
     #[tokio::test]
-    #[serial]
     /// Tests that the SSV network fetch is handled properly when the response's
     /// content-length header is missing
     async fn test_ssv_network_fetch_big_data_without_content_length() -> eyre::Result<()> {
         // Start the mock server
         let port = 30103;
-        env::set_var(CB_TEST_HTTP_DISABLE_CONTENT_LENGTH_ENV, "1");
-        defer! { env::remove_var(CB_TEST_HTTP_DISABLE_CONTENT_LENGTH_ENV); }
+        set_ignore_content_length(true);
         let _server_handle = create_mock_server(port).await?;
         let url = format!("http://localhost:{port}/big_data");
         let response = fetch_ssv_pubkeys_from_url(&url, Duration::from_secs(120)).await;
 
-        // The response should fail due to timeout
-        assert!(response.is_err(), "Expected error due to body size, but got success");
-        if let Err(e) = response {
-            assert!(
-                e.to_string().contains("Response body exceeds the maximum allowed length "),
-                "Expected content length error, got: {e}",
-            );
+        // The response should fail due to the body being too big
+        match response {
+            Ok(_) => {
+                panic!("Expected an error due to excessive data, but got a successful response")
+            }
+            Err(e) => match e.downcast_ref::<ResponseReadError>() {
+                Some(ResponseReadError::PayloadTooLarge { max, content_length, raw }) => {
+                    assert_eq!(*max, MUXER_HTTP_MAX_LENGTH);
+                    assert_eq!(*content_length, 0);
+                    assert!(!raw.is_empty());
+                }
+                _ => panic!("Expected PayloadTooLarge error, got: {}", e),
+            },
         }
 
         // Clean up the server handle
@@ -585,10 +604,12 @@ mod tests {
             .unwrap()
     }
 
-    /// Sends a response with a large body but no content length
+    /// Sends a response with a large body - larger than the maximum allowed.
+    /// Note that hyper overwrites the content-length header automatically, so
+    /// setting it here wouldn't actually change the value that ultimately
+    /// gets sent to the server.
     async fn handle_big_data() -> Response {
-        // Create a response with a large body but no content length
-        let body = "f".repeat(2 * MUXER_HTTP_MAX_LENGTH as usize);
+        let body = "f".repeat(2 * MUXER_HTTP_MAX_LENGTH);
         Response::builder()
             .status(200)
             .header("Content-Type", "application/text")

crates/common/src/config/mux.rs

@@ -1,17 +1,11 @@
-use std::{collections::HashMap, env, path::Path};
+use std::{collections::HashMap, path::Path};
 
 use alloy::rpc::types::beacon::BlsPublicKey;
 use eyre::{bail, Context, Result};
 use serde::de::DeserializeOwned;
 
 use super::JWTS_ENV;
-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";
+use crate::{config::MUXER_HTTP_MAX_LENGTH, types::ModuleId, utils::read_chunked_body_with_max};
 
 pub fn load_env_var(env: &str) -> Result<String> {
     std::env::var(env).wrap_err(format!("{env} is not set"))
@@ -39,40 +33,33 @@ 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(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() {
-        content_length = None;
-    }
-
-    // Break if content length is provided but it's too big
-    if let Some(length) = content_length {
-        if length > MUXER_HTTP_MAX_LENGTH {
-            bail!("Response content length ({length}) exceeds the maximum allowed length ({MUXER_HTTP_MAX_LENGTH} bytes)");
-        }
-    }
+pub async fn safe_read_http_response(response: reqwest::Response) -> Result<Vec<u8>> {
+    // Read the response to a buffer in chunks
+    let status_code = response.status();
+    let response_bytes = read_chunked_body_with_max(response, MUXER_HTTP_MAX_LENGTH)
+        .await
+        .map_err(|e| eyre::Report::new(e));
 
     // Make sure the response is a 200
-    if response.status() != reqwest::StatusCode::OK {
-        bail!("Request failed with status: {}", response.status());
-    }
-
-    // Read the response to a buffer in chunks
-    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
-    match std::str::from_utf8(&bytes) {
-        Ok(s) => Ok(s.to_string()),
-        Err(e) => bail!("Failed to decode response body as UTF-8: {e}"),
+    if status_code != reqwest::StatusCode::OK {
+        match response_bytes {
+            Ok(bytes) => {
+                bail!(
+                    "Request failed with status: {}, body: {}",
+                    status_code,
+                    String::from_utf8_lossy(&bytes)
+                );
+            }
+            Err(e) => {
+                bail!(
+                    "Request failed with status: {} but decoding the response body failed: {}",
+                    status_code,
+                    e
+                );
+            }
+        }
     }
+    response_bytes
 }
 
 /// Removes duplicate entries from a vector of BlsPublicKey