Skip to content

fix: read body in chunks #164

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ltitanb merged 3 commits into from
Oct 24, 2024
Merged

fix: read body in chunks #164

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c439ebd
read body in chunks
ltitanb Oct 23, 2024
eff3e26
fix test
ltitanb Oct 23, 2024
12dd9ef
add test
ltitanb Oct 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ ethereum_serde_utils = "0.7.0"
# networking
axum = { version = "0.7.5", features = ["macros"] }
axum-extra = { version = "0.9.3", features = ["typed-header"] }
reqwest = { version = "0.12.4", features = ["json"] }
reqwest = { version = "0.12.4", features = ["json", "stream"] }
headers = "0.4.0"

# async / threads
Expand Down
4 changes: 2 additions & 2 deletions crates/common/src/pbs/error.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,8 @@ pub enum PbsError {
#[error("relay response error. Code: {code}, err: {error_msg}")]
RelayResponse { error_msg: String, code: u16 },

#[error("response size exceeds max size: max: {max} got: {got}")]
PayloadTooLarge { max: usize, got: usize },
#[error("response size exceeds max size: max: {max} raw: {raw}")]
PayloadTooLarge { max: usize, raw: String },

#[error("failed validating relay response: {0}")]
Validation(#[from] ValidationError),
Expand Down
2 changes: 0 additions & 2 deletions crates/common/src/pbs/relay.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,6 @@ use super::{
};
use crate::{config::RelayConfig, DEFAULT_REQUEST_TIMEOUT};

pub const MAX_SIZE: usize = 10 * 1024 * 1024;

/// A parsed entry of the relay url in the format: scheme://pubkey@host
#[derive(Debug, Clone)]
pub struct RelayEntry {
Expand Down
9 changes: 9 additions & 0 deletions crates/pbs/src/constants.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,3 +6,12 @@ pub(crate) const GET_HEADER_ENDPOINT_TAG: &str = "get_header";
/// For metrics recorded when a request times out
pub(crate) const TIMEOUT_ERROR_CODE: u16 = 555;
pub(crate) const TIMEOUT_ERROR_CODE_STR: &str = "555";

/// 20 MiB to cover edge cases for heavy blocks and also add a bit of slack for
/// any Ethereum upgrades in the near future
pub(crate) const MAX_SIZE_SUBMIT_BLOCK: usize = 20 * 1024 * 1024;

/// 10 KiB, headers are around 700 bytes + buffer for encoding
pub(crate) const MAX_SIZE_GET_HEADER: usize = 10 * 1024;

pub(crate) const MAX_SIZE_DEFAULT: usize = 1024;
1 change: 1 addition & 0 deletions crates/pbs/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ mod mev_boost;
mod routes;
mod service;
mod state;
mod utils;

pub use api::*;
pub use mev_boost::*;
Expand Down
13 changes: 6 additions & 7 deletions crates/pbs/src/mev_boost/get_header.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ use cb_common::{
pbs::{
error::{PbsError, ValidationError},
GetHeaderParams, GetHeaderResponse, RelayClient, SignedExecutionPayloadHeader,
EMPTY_TX_ROOT_HASH, HEADER_SLOT_UUID_KEY, HEADER_START_TIME_UNIX_MS, MAX_SIZE,
EMPTY_TX_ROOT_HASH, HEADER_SLOT_UUID_KEY, HEADER_START_TIME_UNIX_MS,
},
signature::verify_signed_message,
types::Chain,
Expand All @@ -24,9 +24,12 @@ use tracing::{debug, error, warn, Instrument};
use url::Url;

use crate::{
constants::{GET_HEADER_ENDPOINT_TAG, TIMEOUT_ERROR_CODE, TIMEOUT_ERROR_CODE_STR},
constants::{
GET_HEADER_ENDPOINT_TAG, MAX_SIZE_GET_HEADER, TIMEOUT_ERROR_CODE, TIMEOUT_ERROR_CODE_STR,
},
metrics::{RELAY_HEADER_VALUE, RELAY_LAST_SLOT, RELAY_LATENCY, RELAY_STATUS_CODE},
state::{BuilderApiState, PbsState},
utils::read_chunked_body_with_max,
};

/// Implements https://ethereum.github.io/builder-specs/#/Builder/getHeader
Expand Down Expand Up @@ -247,11 +250,7 @@ async fn send_one_get_header(
let code = res.status();
RELAY_STATUS_CODE.with_label_values(&[code.as_str(), GET_HEADER_ENDPOINT_TAG, &relay.id]).inc();

let response_bytes = res.bytes().await?;
if response_bytes.len() > MAX_SIZE {
return Err(PbsError::PayloadTooLarge { max: MAX_SIZE, got: response_bytes.len() });
}

let response_bytes = read_chunked_body_with_max(res, MAX_SIZE_GET_HEADER).await?;
if !code.is_success() {
return Err(PbsError::RelayResponse {
error_msg: String::from_utf8_lossy(&response_bytes).into_owned(),
Expand Down
10 changes: 4 additions & 6 deletions crates/pbs/src/mev_boost/register_validator.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use std::time::{Duration, Instant};
use alloy::rpc::types::beacon::relay::ValidatorRegistration;
use axum::http::{HeaderMap, HeaderValue};
use cb_common::{
pbs::{error::PbsError, RelayClient, HEADER_START_TIME_UNIX_MS, MAX_SIZE},
pbs::{error::PbsError, RelayClient, HEADER_START_TIME_UNIX_MS},
utils::{get_user_agent_with_version, utcnow_ms},
};
use eyre::bail;
Expand All @@ -12,9 +12,10 @@ use reqwest::header::USER_AGENT;
use tracing::{debug, error, Instrument};

use crate::{
constants::{REGISTER_VALIDATOR_ENDPOINT_TAG, TIMEOUT_ERROR_CODE_STR},
constants::{MAX_SIZE_DEFAULT, REGISTER_VALIDATOR_ENDPOINT_TAG, TIMEOUT_ERROR_CODE_STR},
metrics::{RELAY_LATENCY, RELAY_STATUS_CODE},
state::{BuilderApiState, PbsState},
utils::read_chunked_body_with_max,
};

/// Implements https://ethereum.github.io/builder-specs/#/Builder/registerValidator
Expand Down Expand Up @@ -103,11 +104,8 @@ async fn send_register_validator(
.with_label_values(&[code.as_str(), REGISTER_VALIDATOR_ENDPOINT_TAG, &relay.id])
.inc();

let response_bytes = res.bytes().await?;
if response_bytes.len() > MAX_SIZE {
return Err(PbsError::PayloadTooLarge { max: MAX_SIZE, got: response_bytes.len() });
}
if !code.is_success() {
let response_bytes = read_chunked_body_with_max(res, MAX_SIZE_DEFAULT).await?;
let err = PbsError::RelayResponse {
error_msg: String::from_utf8_lossy(&response_bytes).into_owned(),
code: code.as_u16(),
Expand Down
10 changes: 4 additions & 6 deletions crates/pbs/src/mev_boost/status.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,17 +2,18 @@ use std::time::{Duration, Instant};

use axum::http::HeaderMap;
use cb_common::{
pbs::{error::PbsError, RelayClient, MAX_SIZE},
pbs::{error::PbsError, RelayClient},
utils::get_user_agent_with_version,
};
use futures::future::select_ok;
use reqwest::header::USER_AGENT;
use tracing::{debug, error};

use crate::{
constants::{STATUS_ENDPOINT_TAG, TIMEOUT_ERROR_CODE_STR},
constants::{MAX_SIZE_DEFAULT, STATUS_ENDPOINT_TAG, TIMEOUT_ERROR_CODE_STR},
metrics::{RELAY_LATENCY, RELAY_STATUS_CODE},
state::{BuilderApiState, PbsState},
utils::read_chunked_body_with_max,
};

/// Implements https://ethereum.github.io/builder-specs/#/Builder/status
Expand Down Expand Up @@ -74,11 +75,8 @@ async fn send_relay_check(relay: &RelayClient, headers: HeaderMap) -> Result<(),
let code = res.status();
RELAY_STATUS_CODE.with_label_values(&[code.as_str(), STATUS_ENDPOINT_TAG, &relay.id]).inc();

let response_bytes = res.bytes().await?;
if response_bytes.len() > MAX_SIZE {
return Err(PbsError::PayloadTooLarge { max: MAX_SIZE, got: response_bytes.len() });
}
if !code.is_success() {
let response_bytes = read_chunked_body_with_max(res, MAX_SIZE_DEFAULT).await?;
let err = PbsError::RelayResponse {
error_msg: String::from_utf8_lossy(&response_bytes).into_owned(),
code: code.as_u16(),
Expand Down
13 changes: 5 additions & 8 deletions crates/pbs/src/mev_boost/submit_block.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ use cb_common::{
pbs::{
error::{PbsError, ValidationError},
RelayClient, SignedBlindedBeaconBlock, SubmitBlindedBlockResponse, HEADER_SLOT_UUID_KEY,
HEADER_START_TIME_UNIX_MS, MAX_SIZE,
HEADER_START_TIME_UNIX_MS,
},
utils::{get_user_agent_with_version, utcnow_ms},
};
Expand All @@ -14,9 +14,10 @@ use reqwest::header::USER_AGENT;
use tracing::{debug, warn};

use crate::{
constants::{SUBMIT_BLINDED_BLOCK_ENDPOINT_TAG, TIMEOUT_ERROR_CODE_STR},
constants::{MAX_SIZE_SUBMIT_BLOCK, SUBMIT_BLINDED_BLOCK_ENDPOINT_TAG, TIMEOUT_ERROR_CODE_STR},
metrics::{RELAY_LATENCY, RELAY_STATUS_CODE},
state::{BuilderApiState, PbsState},
utils::read_chunked_body_with_max,
};

/// Implements https://ethereum.github.io/builder-specs/#/Builder/submitBlindedBlock
Expand Down Expand Up @@ -94,18 +95,14 @@ async fn send_submit_block(
.with_label_values(&[code.as_str(), SUBMIT_BLINDED_BLOCK_ENDPOINT_TAG, &relay.id])
.inc();

let response_bytes = res.bytes().await?;

if response_bytes.len() > MAX_SIZE {
return Err(PbsError::PayloadTooLarge { max: MAX_SIZE, got: response_bytes.len() });
}
let response_bytes = read_chunked_body_with_max(res, MAX_SIZE_SUBMIT_BLOCK).await?;
if !code.is_success() {
let err = PbsError::RelayResponse {
error_msg: String::from_utf8_lossy(&response_bytes).into_owned(),
code: code.as_u16(),
};

// we request payload to all relays, but some may have not received it
// we requested the payload from all relays, but some may have not received it
warn!(%err, "failed to get payload (this might be ok if other relays have it)");
return Err(err);
};
Expand Down
25 changes: 25 additions & 0 deletions crates/pbs/src/utils.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
use cb_common::pbs::error::PbsError;
use futures::StreamExt;
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 {
Copy link
Contributor

@fbrv fbrv Oct 23, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

wonder if we can do something like so we avoid variables

pub async fn read_chunked_body_with_max(
    res: Response,
    max_size: usize,
) -> Result<Vec<u8>, PbsError> {
    let response_bytes = res
        .bytes_stream()
        .map_err(|e| PbsError::Reqwest(e))
        .try_fold(Vec::new(), |mut acc, chunk| async move {
            if acc.len() + chunk.len() > max_size {
                Err(PbsError::PayloadTooLarge {
                    max: max_size,
                    raw: String::from_utf8_lossy(&acc).into_owned(),
                })
            } else {
                acc.extend_from_slice(&chunk);
                Ok(acc)
            }
        })
        .await?;

    Ok(response_bytes)
}

Copy link
Collaborator Author

@ltitanb ltitanb Oct 23, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

think this could work but not sure it's more readable?

let chunk = chunk?;
if response_bytes.len() + chunk.len() > max_size {
return Err(PbsError::PayloadTooLarge {
max: max_size,
raw: String::from_utf8_lossy(&response_bytes).into_owned(),
});
}

response_bytes.extend_from_slice(&chunk);
}

Ok(response_bytes)
}