From 64c94e15a65b4948fa7d2f3c3475c10f09f4fa8b Mon Sep 17 00:00:00 2001 From: Nina Date: Tue, 28 Jul 2026 17:43:05 +0100 Subject: [PATCH 1/5] draft --- Cargo.lock | 3 + Cargo.toml | 1 + benches/pbs/src/main.rs | 3 +- config.example.toml | 9 + crates/common/src/config/pbs.rs | 11 + crates/common/src/pbs/error.rs | 3 + crates/common/src/pbs/relay.rs | 146 ++++++++++- crates/pbs/Cargo.toml | 3 + crates/pbs/src/mev_boost/get_header.rs | 51 +++- crates/pbs/src/mev_boost/get_header_ws.rs | 293 ++++++++++++++++++++++ crates/pbs/src/mev_boost/mod.rs | 1 + tests/src/utils.rs | 12 +- tests/tests/pbs_cfg_file_update.rs | 7 +- 13 files changed, 520 insertions(+), 23 deletions(-) create mode 100644 crates/pbs/src/mev_boost/get_header_ws.rs diff --git a/Cargo.lock b/Cargo.lock index 1f052de5..1d60ab4e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1828,6 +1828,7 @@ dependencies = [ "async-trait", "axum 0.8.9", "axum-extra", + "bytes", "cb-common", "cb-metrics", "ethereum_serde_utils 0.7.0", @@ -1840,10 +1841,12 @@ dependencies = [ "parking_lot", "prometheus", "reqwest 0.13.2", + "rustls", "serde", "serde_json", "thiserror 2.0.18", "tokio", + "tokio-tungstenite", "tower-http", "tracing", "tree_hash 0.12.1", diff --git a/Cargo.toml b/Cargo.toml index 170cb38a..d2334ae5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -79,6 +79,7 @@ subtle = "2.5" tempfile = "3.20.0" thiserror = "2.0.12" tokio = { version = "1.37.0", features = ["full"] } +tokio-tungstenite = { version = "0.28.0", default-features = false, features = ["connect", "rustls-tls-webpki-roots"] } toml = "0.8.13" tonic = { version = "0.12.3", features = ["channel", "prost", "tls"] } tonic-build = "0.12.3" diff --git a/benches/pbs/src/main.rs b/benches/pbs/src/main.rs index 82f72b5c..4a000765 100644 --- a/benches/pbs/src/main.rs +++ b/benches/pbs/src/main.rs @@ -2,7 +2,7 @@ use std::time::{Duration, Instant}; use alloy::primitives::B256; use cb_common::{ - config::RelayConfig, + config::{GetHeaderTransport, RelayConfig}, pbs::{GetHeaderResponse, RelayClient, RelayEntry}, types::{BlsPublicKey, BlsSecretKey, Chain}, utils::TestRandomSeed, @@ -157,6 +157,7 @@ fn get_mock_validator(bench: BenchConfig) -> RelayClient { id: None, headers: None, get_params: None, + get_header: GetHeaderTransport::Http, enable_timing_games: false, target_first_request_ms: None, frequency_get_header_ms: None, diff --git a/config.example.toml b/config.example.toml index 6804faad..d3a4b446 100644 --- a/config.example.toml +++ b/config.example.toml @@ -92,6 +92,15 @@ headers = { X-MyCustomHeader = "MyCustomValue" } # GET parameters to add to each request URL for this relay # OPTIONAL get_params = { param1 = "value1", param2 = "value2" } +# How to fetch headers from this relay. +# "http" -> one request per get_header, at the relay url above +# { stream = "wss://abc.xyz/stream" } -> websocket stream of bid updates, only for relays that +# support it. The stream stays open for the request +# window and the last bid pushed is the one used, so the +# timing game settings below are ignored. +# `{slot}/{parent_hash}/{pubkey}` is appended to the url. +# OPTIONAL, DEFAULT: "http" +get_header = "http" # Whether to enable timing games, as tuned by `target_first_request_ms` and `frequency_get_header_ms`. # NOTE: if neither `target_first_request_ms` nor `frequency_get_header_ms` is set, this flag has no effect. # diff --git a/crates/common/src/config/pbs.rs b/crates/common/src/config/pbs.rs index 4bb1cff9..08bcaafa 100644 --- a/crates/common/src/config/pbs.rs +++ b/crates/common/src/config/pbs.rs @@ -38,6 +38,14 @@ use crate::{ }, }; +#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum GetHeaderTransport { + #[default] + Http, + Stream(Url), +} + #[derive(Debug, Clone, Deserialize, Serialize)] #[serde(deny_unknown_fields)] pub struct RelayConfig { @@ -50,6 +58,9 @@ pub struct RelayConfig { pub headers: Option>, /// Optional GET parameters to add to each request pub get_params: Option>, + /// How to fetch headers from this relay + #[serde(default)] + pub get_header: GetHeaderTransport, /// Whether to enable timing games #[serde(default = "default_bool::")] pub enable_timing_games: bool, diff --git a/crates/common/src/pbs/error.rs b/crates/common/src/pbs/error.rs index 0be1a40b..9aa6919c 100644 --- a/crates/common/src/pbs/error.rs +++ b/crates/common/src/pbs/error.rs @@ -38,6 +38,9 @@ pub enum PbsError { #[error("SSZ error: {0}")] SszError(#[from] SszValueError), + + #[error("websocket error: {0}")] + WebSocket(String), } impl PbsError { diff --git a/crates/common/src/pbs/relay.rs b/crates/common/src/pbs/relay.rs index 59d963e1..394c6068 100644 --- a/crates/common/src/pbs/relay.rs +++ b/crates/common/src/pbs/relay.rs @@ -12,7 +12,10 @@ use super::{ error::PbsError, }; use crate::{ - DEFAULT_REQUEST_TIMEOUT, config::RelayConfig, pbs::BuilderApiVersion, types::BlsPublicKey, + DEFAULT_REQUEST_TIMEOUT, + config::{GetHeaderTransport, RelayConfig}, + pbs::BuilderApiVersion, + types::BlsPublicKey, }; /// A parsed entry of the relay url in the format: scheme://pubkey@host @@ -49,6 +52,12 @@ impl<'de> Deserialize<'de> for RelayEntry { } } +#[derive(Debug, Clone)] +pub enum GetHeaderRequest { + Http(Url), + Stream(Url), +} + /// A client to interact with a relay, safe to share across threads and cheaply /// cloneable #[derive(Debug, Clone)] @@ -62,7 +71,21 @@ pub struct RelayClient { } impl RelayClient { - pub fn new(config: RelayConfig) -> eyre::Result { + pub fn new(mut config: RelayConfig) -> eyre::Result { + if let GetHeaderTransport::Stream(url) = &mut config.get_header { + eyre::ensure!( + matches!(url.scheme(), "ws" | "wss"), + "get_header stream url must be ws:// or wss://, got {}", + url.scheme() + ); + + // Url::join drops the last segment of a base without a trailing slash + if !url.path().ends_with('/') { + let path = format!("{}/", url.path()); + url.set_path(&path); + } + } + let mut headers = HeaderMap::new(); headers.insert(HEADER_VERSION_KEY, HeaderValue::from_static(HEADER_VERSION_VALUE)); @@ -89,7 +112,11 @@ impl RelayClient { // URL builders pub fn get_url(&self, path: &str) -> Result { - let mut url = self.config.entry.url.join(path).map_err(PbsError::UrlParsing)?; + self.join_url(&self.config.entry.url, path) + } + + fn join_url(&self, base: &Url, path: &str) -> Result { + let mut url = base.join(path).map_err(PbsError::UrlParsing)?; if let Some(get_params) = &self.config.get_params { let mut query_pairs = url.query_pairs_mut(); @@ -120,6 +147,22 @@ impl RelayClient { ) } + pub fn get_header_request( + &self, + slot: u64, + parent_hash: &B256, + validator_pubkey: &BlsPublicKey, + ) -> Result { + Ok(match &self.config.get_header { + GetHeaderTransport::Http => { + GetHeaderRequest::Http(self.get_header_url(slot, parent_hash, validator_pubkey)?) + } + GetHeaderTransport::Stream(base) => GetHeaderRequest::Stream( + self.join_url(base, &format!("{slot}/{parent_hash}/{validator_pubkey}"))?, + ), + }) + } + pub fn get_status_url(&self) -> Result { self.builder_api_url(GET_STATUS_PATH, BuilderApiVersion::V1) } @@ -139,8 +182,11 @@ mod tests { use alloy::primitives::B256; - use super::{RelayClient, RelayEntry}; - use crate::{config::RelayConfig, utils::bls_pubkey_from_hex_unchecked}; + use super::{GetHeaderRequest, RelayClient, RelayEntry}; + use crate::{ + config::{GetHeaderTransport, RelayConfig}, + utils::bls_pubkey_from_hex_unchecked, + }; #[test] fn test_relay_entry() { @@ -226,4 +272,94 @@ mod tests { assert!(url.contains("param1=value1")); assert!(url.contains("param2=value2")); } + + #[test] + fn test_get_header_request() { + let slot = 0; + let parent_hash = B256::ZERO; + let validator_pubkey = bls_pubkey_from_hex_unchecked( + "0xac6e77dfe25ecd6110b8e780608cce0dab71fdd5ebea22a16c0205200f2f8e2e3ad3b71d3499c54ad14d6c21b41a37ae", + ); + let relay_config = r#" + { + "url": "http://0xa1cec75a3f0661e99299274182938151e8433c61a19222347ea1313d839229cb4ce4e3e5aa2bdeb71c8fcf1b084963c2@abc.xyz" + }"#; + let base_config = serde_json::from_str::(relay_config).unwrap(); + + // No ws url configured: plain HTTP endpoint + let relay = RelayClient::new(base_config.clone()).unwrap(); + let GetHeaderRequest::Http(url) = + relay.get_header_request(slot, &parent_hash, &validator_pubkey).unwrap() + else { + panic!("expected http request"); + }; + assert_eq!( + url, + relay.get_header_url(slot, &parent_hash, &validator_pubkey).unwrap(), + "http dispatch must match the plain url builder" + ); + + // A stream url takes over, and a missing trailing slash is normalized so + // the configured path isn't dropped by Url::join + for base in ["wss://abc.xyz/stream", "wss://abc.xyz/stream/"] { + let mut config = base_config.clone(); + config.get_header = GetHeaderTransport::Stream(base.parse().unwrap()); + let relay = RelayClient::new(config).unwrap(); + + let GetHeaderRequest::Stream(url) = + relay.get_header_request(slot, &parent_hash, &validator_pubkey).unwrap() + else { + panic!("expected stream request"); + }; + assert_eq!( + url.to_string(), + format!("wss://abc.xyz/stream/{slot}/{parent_hash}/{validator_pubkey}") + ); + } + + // Non-ws schemes are rejected at construction + let mut config = base_config; + config.get_header = GetHeaderTransport::Stream("https://abc.xyz/stream/".parse().unwrap()); + assert!(RelayClient::new(config).is_err()); + } + + #[test] + fn test_get_header_transport_config() { + let with_transport = |value: &str| { + let relay_config = format!( + r#" + {{ + "url": "http://0xa1cec75a3f0661e99299274182938151e8433c61a19222347ea1313d839229cb4ce4e3e5aa2bdeb71c8fcf1b084963c2@abc.xyz", + "get_header": {value} + }}"# + ); + serde_json::from_str::(&relay_config).map(|config| config.get_header) + }; + + assert_eq!(with_transport(r#""http""#).unwrap(), GetHeaderTransport::Http); + assert_eq!( + with_transport(r#"{ "stream": "wss://abc.xyz/stream" }"#).unwrap(), + GetHeaderTransport::Stream("wss://abc.xyz/stream".parse().unwrap()) + ); + assert!(with_transport(r#""grpc""#).is_err()); + + // Same shapes in the toml the operator actually writes + let toml_config = r#" + url = "http://0xa1cec75a3f0661e99299274182938151e8433c61a19222347ea1313d839229cb4ce4e3e5aa2bdeb71c8fcf1b084963c2@abc.xyz" + get_header = { stream = "wss://abc.xyz/stream" } + "#; + let config = toml::from_str::(toml_config).unwrap(); + assert_eq!( + config.get_header, + GetHeaderTransport::Stream("wss://abc.xyz/stream".parse().unwrap()) + ); + + // Defaults to http when omitted + let relay_config = r#" + { + "url": "http://0xa1cec75a3f0661e99299274182938151e8433c61a19222347ea1313d839229cb4ce4e3e5aa2bdeb71c8fcf1b084963c2@abc.xyz" + }"#; + let config = serde_json::from_str::(relay_config).unwrap(); + assert_eq!(config.get_header, GetHeaderTransport::Http); + } } diff --git a/crates/pbs/Cargo.toml b/crates/pbs/Cargo.toml index 30064c28..401dee32 100644 --- a/crates/pbs/Cargo.toml +++ b/crates/pbs/Cargo.toml @@ -9,6 +9,7 @@ version.workspace = true alloy.workspace = true async-trait.workspace = true axum.workspace = true +bytes.workspace = true axum-extra.workspace = true cb-common.workspace = true cb-metrics.workspace = true @@ -23,9 +24,11 @@ notify.workspace = true parking_lot.workspace = true prometheus.workspace = true reqwest.workspace = true +rustls.workspace = true serde.workspace = true serde_json.workspace = true tokio.workspace = true +tokio-tungstenite.workspace = true tower-http.workspace = true tracing.workspace = true tree_hash.workspace = true diff --git a/crates/pbs/src/mev_boost/get_header.rs b/crates/pbs/src/mev_boost/get_header.rs index dd2b1036..473307b9 100644 --- a/crates/pbs/src/mev_boost/get_header.rs +++ b/crates/pbs/src/mev_boost/get_header.rs @@ -13,8 +13,8 @@ use cb_common::{ constants::APPLICATION_BUILDER_DOMAIN, pbs::{ EMPTY_TX_ROOT_HASH, ExecutionPayloadHeaderRef, ForkName, ForkVersionDecode, GetHeaderInfo, - GetHeaderParams, GetHeaderResponse, HEADER_START_TIME_UNIX_MS, HEADER_TIMEOUT_MS, - RelayClient, SignedBuilderBid, + GetHeaderParams, GetHeaderRequest, GetHeaderResponse, HEADER_START_TIME_UNIX_MS, + HEADER_TIMEOUT_MS, RelayClient, SignedBuilderBid, error::{PbsError, ValidationError}, }, signature::verify_signed_message, @@ -36,6 +36,7 @@ use tracing::{Instrument, debug, error, info, warn}; use tree_hash::TreeHash; use url::Url; +use super::get_header_ws::get_header_ws; use crate::{ constants::{ GET_HEADER_ENDPOINT_TAG, MAX_SIZE_GET_HEADER_RESPONSE, TIMEOUT_ERROR_CODE, @@ -49,11 +50,11 @@ use crate::{ /// Info about an incoming get_header request. /// Sent from get_header to each send_timed_get_header call. #[derive(Clone)] -struct RequestInfo { - params: GetHeaderParams, +pub(super) struct RequestInfo { + pub(super) params: GetHeaderParams, /// Common baseline of headers to send with each request - headers: HeaderMap, + pub(super) headers: HeaderMap, /// The chain the request is for chain: Chain, @@ -179,7 +180,7 @@ pub async fn get_header( let mut handles = Vec::with_capacity(relays.len()); for relay in relays.iter() { handles.push( - send_timed_get_header( + get_header_from_relay( request_info.clone(), relay.clone(), ms_into_slot, @@ -249,15 +250,31 @@ async fn fetch_parent_block( } } -async fn send_timed_get_header( +async fn get_header_from_relay( request_info: Arc, relay: RelayClient, ms_into_slot: u64, - mut timeout_left_ms: u64, + timeout_left_ms: u64, ) -> Result, PbsError> { let params = &request_info.params; - let url = relay.get_header_url(params.slot, ¶ms.parent_hash, ¶ms.pubkey)?; + match relay.get_header_request(params.slot, ¶ms.parent_hash, ¶ms.pubkey)? { + GetHeaderRequest::Stream(url) => { + get_header_ws(&request_info, &relay, url, timeout_left_ms).await + } + GetHeaderRequest::Http(url) => { + send_timed_get_header(request_info, relay, ms_into_slot, url, timeout_left_ms).await + } + } +} + +async fn send_timed_get_header( + request_info: Arc, + relay: RelayClient, + ms_into_slot: u64, + url: Url, + mut timeout_left_ms: u64, +) -> Result, PbsError> { if relay.config.enable_timing_games { if let Some(target_ms) = relay.config.target_first_request_ms { // sleep until target time in slot @@ -380,6 +397,18 @@ async fn send_one_get_header( Some(res) => res, }; + validate_get_header_response(&request_info, &relay, &get_header_response)?; + + Ok((start_request_time, Some(get_header_response))) +} + +/// Validate a decoded header from a relay: header contents, relay signature, +/// and, if enabled, consistency with the parent block. +pub(super) fn validate_get_header_response( + request_info: &RequestInfo, + relay: &RelayClient, + get_header_response: &GetHeaderResponse, +) -> Result<(), PbsError> { // Extract the basic header data needed for validation let header_data = match &get_header_response.data.message.header() { ExecutionPayloadHeaderRef::Bellatrix(_) | @@ -421,7 +450,7 @@ async fn send_one_get_header( if validation.extra_validation_enabled { let parent_block = validation.parent_block.read(); if let Some(parent_block) = parent_block.as_ref() { - extra_validation(parent_block, &get_header_response)?; + extra_validation(parent_block, get_header_response)?; } else { warn!( relay_id = relay.id.as_ref(), @@ -430,7 +459,7 @@ async fn send_one_get_header( } } - Ok((start_request_time, Some(get_header_response))) + Ok(()) } /// Send and decode a full get_header response, with all of the fields. diff --git a/crates/pbs/src/mev_boost/get_header_ws.rs b/crates/pbs/src/mev_boost/get_header_ws.rs new file mode 100644 index 00000000..7354cc49 --- /dev/null +++ b/crates/pbs/src/mev_boost/get_header_ws.rs @@ -0,0 +1,293 @@ +//! Streaming get_header over a websocket, for relays configured with +//! `get_header_ws_url`. One connection per get_header call, dropped when the +//! call returns. +//! +//! The request is the handshake itself: slot / parent_hash / pubkey in the +//! path, deadline and timestamp in headers, same data the HTTP request carries. +//! The relay replies with one binary frame per bid update: +//! +//! ```text +//! u8 message type +//! u8 fork +//! .. SSZ SignedBuilderBid +//! ``` + +use std::{str::FromStr, sync::Once, time::Duration}; + +use alloy::primitives::utils::format_ether; +use axum::http::{HeaderName, HeaderValue, Request, header::USER_AGENT}; +use bytes::Bytes; +use cb_common::{ + pbs::{ + ForkName, ForkVersionDecode, GetHeaderInfo, GetHeaderResponse, HEADER_START_TIME_UNIX_MS, + HEADER_TIMEOUT_MS, HEADER_VERSION_KEY, HEADER_VERSION_VALUE, RelayClient, SignedBuilderBid, + error::PbsError, + }, + utils::utcnow_ms, +}; +use futures::StreamExt; +use rustls::crypto::{CryptoProvider, aws_lc_rs}; +use tokio::time::{Instant, sleep_until, timeout_at}; +use tokio_tungstenite::{ + connect_async_with_config, + tungstenite::{Message, client::IntoClientRequest, protocol::WebSocketConfig}, +}; +use tracing::{debug, info, warn}; +use url::Url; + +use super::get_header::{RequestInfo, validate_get_header_response}; +use crate::{ + constants::{GET_HEADER_ENDPOINT_TAG, MAX_SIZE_GET_HEADER_RESPONSE, TIMEOUT_ERROR_CODE_STR}, + metrics::{RELAY_LATENCY, RELAY_STATUS_CODE}, +}; + +/// Frame prefix: message type + fork. +const FRAME_PREFIX_LEN: usize = 2; + +const MSG_BID: u8 = 0x01; + +fn fork_from_wire(byte: u8) -> Option { + // TODO @nina: I don't see a point of extending a u8 for supporting older forks + // we could rotate these instead, i.e. 0 becomes Hegota, etc + Some(match byte { + 0 => ForkName::Base, + 1 => ForkName::Altair, + 2 => ForkName::Bellatrix, + 3 => ForkName::Capella, + 4 => ForkName::Deneb, + 5 => ForkName::Electra, + 6 => ForkName::Fulu, + 7 => ForkName::Gloas, + _ => return None, + }) +} + +/// Open a stream to the relay, keep the latest bid until the deadline, then +/// validate and return it. +pub(super) async fn get_header_ws( + request_info: &RequestInfo, + relay: &RelayClient, + url: Url, + timeout_ms: u64, +) -> Result, PbsError> { + let deadline = Instant::now() + Duration::from_millis(timeout_ms); + let request = build_handshake_request(request_info, relay, &url, timeout_ms)?; + + let config = WebSocketConfig::default() + .max_message_size(Some(MAX_SIZE_GET_HEADER_RESPONSE)) + .max_frame_size(Some(MAX_SIZE_GET_HEADER_RESPONSE)); + + install_crypto_provider(); + + let start_request = Instant::now(); + let connect = connect_async_with_config(request, Some(config), true); + let (mut stream, _) = match timeout_at(deadline, connect).await { + Ok(Ok(connected)) => connected, + Ok(Err(err)) => { + record_status(TIMEOUT_ERROR_CODE_STR, relay); + return Err(PbsError::WebSocket(format!("connect failed: {err}"))); + } + Err(_) => { + record_status(TIMEOUT_ERROR_CODE_STR, relay); + return Err(PbsError::WebSocket("timed out connecting".to_string())); + } + }; + debug!(relay_id = relay.id.as_ref(), latency = ?start_request.elapsed(), "ws connected"); + + let timer = sleep_until(deadline); + tokio::pin!(timer); + + let mut latest: Option<(ForkName, Bytes)> = None; + let mut first_bid_latency = None; + let mut updates = 0usize; + + loop { + let message = tokio::select! { + biased; + _ = &mut timer => break, + message = stream.next() => message, + }; + + let message = match message { + Some(Ok(message)) => message, + Some(Err(err)) => { + warn!(relay_id = relay.id.as_ref(), %err, "ws stream error"); + break; + } + None => break, + }; + + let payload = match message { + Message::Binary(payload) => payload, + Message::Close(_) => break, + _ => continue, + }; + + match parse_frame(payload) { + Ok((fork, bid)) => { + updates += 1; + first_bid_latency.get_or_insert_with(|| start_request.elapsed()); + latest = Some((fork, bid)); + } + Err(err) => { + warn!(relay_id = relay.id.as_ref(), %err, "invalid ws frame"); + break; + } + } + } + + drop(stream); + + let Some((fork, bid_bytes)) = latest else { + debug!(relay_id = relay.id.as_ref(), "no header from relay"); + record_status("204", relay); + return Ok(None); + }; + + RELAY_LATENCY + .with_label_values(&[GET_HEADER_ENDPOINT_TAG, &relay.id]) + .observe(first_bid_latency.unwrap_or_default().as_secs_f64()); + record_status("200", relay); + + let data = SignedBuilderBid::from_ssz_bytes_by_fork(&bid_bytes, fork).map_err(|err| { + PbsError::SSZDecode { + err: format!("error decoding relay payload from ws stream: {err:?}"), + fork, + } + })?; + let response = GetHeaderResponse { version: fork, data, metadata: Default::default() }; + + info!( + relay_id = relay.id.as_ref(), + header_size_bytes = bid_bytes.len(), + latency = ?start_request.elapsed(), + version = ?fork, + value_eth = format_ether(*response.value()), + block_hash = %response.block_hash(), + updates, + "received new header from ws stream" + ); + + validate_get_header_response(request_info, relay, &response)?; + + Ok(Some(response)) +} + +/// The handshake carries the request: same headers as the HTTP path, minus +/// `Accept` (the stream is SSZ only). +fn build_handshake_request( + request_info: &RequestInfo, + relay: &RelayClient, + url: &Url, + timeout_ms: u64, +) -> Result, PbsError> { + let mut request = url + .as_str() + .into_client_request() + .map_err(|err| PbsError::WebSocket(format!("invalid ws url: {err}")))?; + + let headers = request.headers_mut(); + if let Some(user_agent) = request_info.headers.get(USER_AGENT) { + headers.insert(USER_AGENT, user_agent.clone()); + } + headers.insert(HEADER_VERSION_KEY, HeaderValue::from_static(HEADER_VERSION_VALUE)); + + // The HTTP client bakes these into its default headers, the ws handshake + // needs them explicitly. Validated in RelayClient::new. + for (key, value) in relay.config.headers.iter().flatten() { + let key = HeaderName::from_str(key) + .map_err(|_| PbsError::WebSocket(format!("invalid header name: {key}")))?; + let value = HeaderValue::from_str(value) + .map_err(|_| PbsError::WebSocket(format!("invalid header value for: {key}")))?; + headers.insert(key, value); + } + + headers.insert(HEADER_START_TIME_UNIX_MS, HeaderValue::from(utcnow_ms())); + headers.insert(HEADER_TIMEOUT_MS, HeaderValue::from(timeout_ms)); + + Ok(request) +} + +fn parse_frame(payload: Bytes) -> Result<(ForkName, Bytes), PbsError> { + let &[msg_type, fork_byte] = payload + .first_chunk::() + .ok_or_else(|| PbsError::WebSocket(format!("frame too short: {} bytes", payload.len())))?; + + if msg_type != MSG_BID { + return Err(PbsError::WebSocket(format!("unknown message type: {msg_type}"))); + } + + let fork = fork_from_wire(fork_byte) + .ok_or_else(|| PbsError::WebSocket(format!("unknown fork: {fork_byte}")))?; + + Ok((fork, payload.slice(FRAME_PREFIX_LEN..))) +} + +fn record_status(code: &str, relay: &RelayClient) { + RELAY_STATUS_CODE.with_label_values(&[code, GET_HEADER_ENDPOINT_TAG, &relay.id]).inc(); +} + +/// rustls is built with both `ring` and `aws-lc-rs` here, so the default +/// `ClientConfig::builder()` inside tokio-tungstenite panics unless a +/// process-level provider is installed. Match the signer and use aws-lc-rs. +fn install_crypto_provider() { + static INIT: Once = Once::new(); + INIT.call_once(|| { + if CryptoProvider::get_default().is_none() { + let _ = aws_lc_rs::default_provider().install_default(); + } + }); +} + +#[cfg(test)] +mod tests { + use std::{fs, path::Path}; + + use ssz::Encode; + + use super::*; + + fn bid_frame(fork_byte: u8, bid: &[u8]) -> Bytes { + let mut frame = vec![MSG_BID, fork_byte]; + frame.extend_from_slice(bid); + Bytes::from(frame) + } + + #[test] + fn test_parse_frame() { + assert!(matches!( + parse_frame(bid_frame(6, &[1, 2, 3])), + Ok((ForkName::Fulu, bid)) if bid.as_ref() == [1, 2, 3] + )); + + // Empty bid payload is well-formed at this layer, SSZ decoding rejects it + assert!(matches!(parse_frame(bid_frame(6, &[])), Ok((ForkName::Fulu, _)))); + + for bad in [ + // Truncated prefix + Bytes::from_static(&[]), + Bytes::from_static(&[MSG_BID]), + // Unknown fork + bid_frame(0xff, &[1]), + // Unknown message type + Bytes::from_static(&[0xff, 6]), + ] { + assert!(matches!(parse_frame(bad), Err(PbsError::WebSocket(_)))); + } + } + + #[test] + fn test_decode_streamed_bid() { + let json_bytes = + fs::read(Path::new("../../tests/data/get_header/fulu.json")).expect("file not found"); + let expected: GetHeaderResponse = + serde_json::from_slice(&json_bytes).expect("failed to decode JSON"); + + let frame = bid_frame(6, &expected.data.as_ssz_bytes()); + let (fork, bid_bytes) = parse_frame(frame).unwrap(); + + let data = SignedBuilderBid::from_ssz_bytes_by_fork(&bid_bytes, fork).unwrap(); + assert_eq!(fork, ForkName::Fulu); + assert_eq!(data, expected.data); + } +} diff --git a/crates/pbs/src/mev_boost/mod.rs b/crates/pbs/src/mev_boost/mod.rs index a41b79db..208c863e 100644 --- a/crates/pbs/src/mev_boost/mod.rs +++ b/crates/pbs/src/mev_boost/mod.rs @@ -1,4 +1,5 @@ mod get_header; +mod get_header_ws; mod register_validator; mod reload; mod status; diff --git a/tests/src/utils.rs b/tests/src/utils.rs index 9bc0e362..311a133b 100644 --- a/tests/src/utils.rs +++ b/tests/src/utils.rs @@ -8,11 +8,11 @@ use std::{ use alloy::primitives::{B256, U256}; use cb_common::{ config::{ - COMMIT_BOOST_IMAGE_DEFAULT, CommitBoostConfig, LogsSettings, ModuleKind, - ModuleSigningConfig, PbsConfig, PbsModuleConfig, RelayConfig, ReverseProxyHeaderSetup, - SIGNER_JWT_AUTH_FAIL_LIMIT_DEFAULT, SIGNER_JWT_AUTH_FAIL_TIMEOUT_SECONDS_DEFAULT, - SIGNER_PORT_DEFAULT, SignerConfig, SignerType, StartSignerConfig, StaticModuleConfig, - StaticPbsConfig, TlsMode, + COMMIT_BOOST_IMAGE_DEFAULT, CommitBoostConfig, GetHeaderTransport, LogsSettings, + ModuleKind, ModuleSigningConfig, PbsConfig, PbsModuleConfig, RelayConfig, + ReverseProxyHeaderSetup, SIGNER_JWT_AUTH_FAIL_LIMIT_DEFAULT, + SIGNER_JWT_AUTH_FAIL_TIMEOUT_SECONDS_DEFAULT, SIGNER_PORT_DEFAULT, SignerConfig, + SignerType, StartSignerConfig, StaticModuleConfig, StaticPbsConfig, TlsMode, }, pbs::{RelayClient, RelayEntry}, signer::SignerLoader, @@ -47,6 +47,7 @@ pub fn generate_mock_relay(port: u16, pubkey: BlsPublicKey) -> Result Result<()> { enable_timing_games: false, frequency_get_header_ms: None, get_params: None, + get_header: GetHeaderTransport::Http, headers: None, target_first_request_ms: None, validator_registration_batch_size: None, @@ -142,6 +146,7 @@ async fn test_cfg_file_update() -> Result<()> { enable_timing_games: false, frequency_get_header_ms: None, get_params: None, + get_header: GetHeaderTransport::Http, headers: None, target_first_request_ms: None, validator_registration_batch_size: None, From 2371eb44a999e064330a7419cbf47d010c5bea28 Mon Sep 17 00:00:00 2001 From: Nina Date: Tue, 28 Jul 2026 17:43:26 +0100 Subject: [PATCH 2/5] tests --- Cargo.lock | 2 + tests/Cargo.toml | 2 + tests/src/lib.rs | 1 + tests/src/mock_relay.rs | 50 +++--- tests/src/mock_ws_relay.rs | 171 +++++++++++++++++++ tests/src/utils.rs | 44 ++--- tests/tests/pbs_get_header_ws.rs | 280 +++++++++++++++++++++++++++++++ 7 files changed, 511 insertions(+), 39 deletions(-) create mode 100644 tests/src/mock_ws_relay.rs create mode 100644 tests/tests/pbs_get_header_ws.rs diff --git a/Cargo.lock b/Cargo.lock index 1d60ab4e..66083990 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1898,6 +1898,7 @@ dependencies = [ "eth2", "ethereum_ssz 0.10.3", "eyre", + "futures", "jsonwebtoken", "rcgen", "reqwest 0.13.2", @@ -1905,6 +1906,7 @@ dependencies = [ "serde_json", "tempfile", "tokio", + "tokio-tungstenite", "toml", "tracing", "tracing-subscriber", diff --git a/tests/Cargo.toml b/tests/Cargo.toml index 88b2e377..646f4994 100644 --- a/tests/Cargo.toml +++ b/tests/Cargo.toml @@ -12,6 +12,7 @@ cb-pbs.workspace = true cb-signer.workspace = true eyre.workspace = true ethereum_ssz.workspace = true +futures.workspace = true jsonwebtoken.workspace = true lh_types.workspace = true lh_eth2.workspace = true @@ -21,6 +22,7 @@ serde.workspace = true serde_json.workspace = true tempfile.workspace = true tokio.workspace = true +tokio-tungstenite.workspace = true toml.workspace = true tracing.workspace = true tracing-subscriber.workspace = true diff --git a/tests/src/lib.rs b/tests/src/lib.rs index 42eec95a..c9fc488b 100644 --- a/tests/src/lib.rs +++ b/tests/src/lib.rs @@ -2,5 +2,6 @@ pub mod mock_relay; pub mod mock_ssv_node; pub mod mock_ssv_public; pub mod mock_validator; +pub mod mock_ws_relay; pub mod signer_service; pub mod utils; diff --git a/tests/src/mock_relay.rs b/tests/src/mock_relay.rs index 8f3382bf..d6656ada 100644 --- a/tests/src/mock_relay.rs +++ b/tests/src/mock_relay.rs @@ -9,7 +9,8 @@ use std::{ }; use alloy::{ - eips::eip7594::CELLS_PER_EXT_BLOB, primitives::U256, + eips::eip7594::CELLS_PER_EXT_BLOB, + primitives::{B256, U256}, rpc::types::beacon::relay::ValidatorRegistration, }; use axum::{ @@ -212,6 +213,33 @@ pub fn mock_relay_app_router(state: Arc) -> Router { Router::new().merge(builder_router_v1).merge(builder_router_v2).with_state(state) } +pub fn mock_signed_builder_bid( + chain: Chain, + signer: &BlsSecretKey, + slot: u64, + parent_hash: B256, + value: U256, +) -> SignedBuilderBid { + let mut header = ExecutionPayloadHeaderFulu { + parent_hash: parent_hash.into(), + block_hash: Default::default(), + timestamp: timestamp_of_slot_start_sec(slot, chain), + ..ExecutionPayloadHeaderFulu::test_random() + }; + header.block_hash.0[0] = 1; + + let message = BuilderBid::Fulu(BuilderBidFulu { + header, + blob_kzg_commitments: Default::default(), + execution_requests: ExecutionRequests::default(), + value, + pubkey: signer.public_key().into(), + }); + let signature = sign_builder_root(chain, signer, &message.tree_hash_root()); + + SignedBuilderBid { message, signature } +} + async fn handle_get_header( State(state): State>, Path(GetHeaderParams { parent_hash, slot, .. }): Path, @@ -243,24 +271,8 @@ async fn handle_get_header( let data = match consensus_version_header { ForkName::Fulu => { - let mut header = ExecutionPayloadHeaderFulu { - parent_hash: parent_hash.into(), - block_hash: Default::default(), - timestamp: timestamp_of_slot_start_sec(slot, state.chain), - ..ExecutionPayloadHeaderFulu::test_random() - }; - header.block_hash.0[0] = 1; - - let message = BuilderBid::Fulu(BuilderBidFulu { - header, - blob_kzg_commitments: Default::default(), - execution_requests: ExecutionRequests::default(), - value: bid_value, - pubkey: state.signer.public_key().into(), - }); - let object_root = message.tree_hash_root(); - let signature = sign_builder_root(state.chain, &state.signer, &object_root); - let response = SignedBuilderBid { message, signature }; + let response = + mock_signed_builder_bid(state.chain, &state.signer, slot, parent_hash, bid_value); if content_type == EncodingType::Ssz { response.as_ssz_bytes() } else { diff --git a/tests/src/mock_ws_relay.rs b/tests/src/mock_ws_relay.rs new file mode 100644 index 00000000..fe489a77 --- /dev/null +++ b/tests/src/mock_ws_relay.rs @@ -0,0 +1,171 @@ +use std::{ + str::FromStr, + sync::{ + Arc, Mutex, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, +}; + +use alloy::primitives::{B256, U256}; +use cb_common::{ + pbs::{HEADER_START_TIME_UNIX_MS, HEADER_TIMEOUT_MS, HEADER_VERSION_KEY}, + types::{BlsSecretKey, Chain}, +}; +use futures::SinkExt; +use ssz::Encode; +use tokio::net::{TcpListener, TcpStream}; +use tokio_tungstenite::{ + accept_hdr_async, + tungstenite::{ + Message, + handshake::server::{Request, Response}, + }, +}; +use tracing::debug; + +use crate::mock_relay::mock_signed_builder_bid; + +const MSG_BID: u8 = 0x01; +const FORK_FULU: u8 = 6; + +/// What PBS sent in the handshake, captured for assertions. +#[derive(Debug, Clone)] +pub struct StreamRequest { + pub slot: u64, + pub parent_hash: B256, + pub validator_pubkey: String, + pub timeout_ms: Option, + pub start_time_ms: Option, + pub user_agent: Option, + pub cb_version: Option, +} + +pub struct MockWsRelayState { + pub chain: Chain, + pub signer: BlsSecretKey, + /// One frame pushed per value, in order. The last one is what PBS must + /// return. + bid_values: Vec, + /// Pause between updates + update_interval: Duration, + /// Keep the connection open after the last update, so PBS returns on its + /// own deadline instead of on close + hold_open: bool, + received_connections: AtomicU64, + last_request: Mutex>, +} + +impl MockWsRelayState { + pub fn new(chain: Chain, signer: BlsSecretKey) -> Self { + Self { + chain, + signer, + bid_values: vec![U256::from(10)], + update_interval: Duration::ZERO, + hold_open: false, + received_connections: AtomicU64::new(0), + last_request: Mutex::new(None), + } + } + + pub fn with_bid_values(self, bid_values: Vec) -> Self { + Self { bid_values, ..self } + } + + pub fn with_update_interval(self, update_interval: Duration) -> Self { + Self { update_interval, ..self } + } + + pub fn hold_open(self) -> Self { + Self { hold_open: true, ..self } + } + + pub fn received_connections(&self) -> u64 { + self.received_connections.load(Ordering::Relaxed) + } + + pub fn last_request(&self) -> Option { + self.last_request.lock().unwrap().clone() + } +} + +pub async fn start_mock_ws_relay_service( + state: Arc, + listener: TcpListener, +) -> eyre::Result<()> { + loop { + let (stream, addr) = listener.accept().await?; + let state = state.clone(); + tokio::spawn(async move { + if let Err(err) = serve_stream(state, stream).await { + debug!(%addr, %err, "mock ws relay connection ended"); + } + }); + } +} + +// The handshake callback's Err type is fixed by tungstenite +#[allow(clippy::result_large_err)] +async fn serve_stream(state: Arc, stream: TcpStream) -> eyre::Result<()> { + let mut request = None; + let mut ws = accept_hdr_async(stream, |req: &Request, res: Response| { + request = parse_request(req); + Ok(res) + }) + .await?; + + let request = request.ok_or_else(|| eyre::eyre!("malformed get_header stream request"))?; + state.received_connections.fetch_add(1, Ordering::Relaxed); + *state.last_request.lock().unwrap() = Some(request.clone()); + + for value in &state.bid_values { + let bid = mock_signed_builder_bid( + state.chain, + &state.signer, + request.slot, + request.parent_hash, + *value, + ); + + let mut frame = vec![MSG_BID, FORK_FULU]; + frame.extend_from_slice(&bid.as_ssz_bytes()); + ws.send(Message::Binary(frame.into())).await?; + + if !state.update_interval.is_zero() { + tokio::time::sleep(state.update_interval).await; + } + } + + if state.hold_open { + // Never resolves: PBS has to cut the stream off at its own deadline + std::future::pending::<()>().await; + } + + ws.close(None).await?; + + Ok(()) +} + +/// The request is the handshake: `/{slot}/{parent_hash}/{pubkey}` plus the +/// same headers the HTTP path sends. +fn parse_request(req: &Request) -> Option { + let mut segments = req.uri().path().trim_start_matches('/').split('/'); + let slot = segments.next()?.parse().ok()?; + let parent_hash = B256::from_str(segments.next()?).ok()?; + let validator_pubkey = segments.next()?.to_string(); + + Some(StreamRequest { + slot, + parent_hash, + validator_pubkey, + timeout_ms: header(req, HEADER_TIMEOUT_MS).and_then(|v| v.parse().ok()), + start_time_ms: header(req, HEADER_START_TIME_UNIX_MS).and_then(|v| v.parse().ok()), + user_agent: header(req, "user-agent"), + cb_version: header(req, HEADER_VERSION_KEY), + }) +} + +fn header(req: &Request, name: &str) -> Option { + req.headers().get(name)?.to_str().ok().map(ToOwned::to_owned) +} diff --git a/tests/src/utils.rs b/tests/src/utils.rs index 311a133b..6d9d0dc9 100644 --- a/tests/src/utils.rs +++ b/tests/src/utils.rs @@ -39,11 +39,13 @@ pub fn setup_test_env() { }); } -pub fn generate_mock_relay(port: u16, pubkey: BlsPublicKey) -> Result { - let entry = - RelayEntry { id: format!("mock_{port}"), pubkey, url: get_local_address(port).parse()? }; - let config = RelayConfig { - entry, +fn mock_relay_config(port: u16, pubkey: BlsPublicKey) -> Result { + Ok(RelayConfig { + entry: RelayEntry { + id: format!("mock_{port}"), + pubkey, + url: get_local_address(port).parse()?, + }, id: None, headers: None, get_params: None, @@ -52,8 +54,11 @@ pub fn generate_mock_relay(port: u16, pubkey: BlsPublicKey) -> Result Result { + RelayClient::new(mock_relay_config(port, pubkey)?) } pub fn generate_mock_relay_with_batch_size( @@ -61,19 +66,18 @@ pub fn generate_mock_relay_with_batch_size( pubkey: BlsPublicKey, batch_size: usize, ) -> Result { - let entry = - RelayEntry { id: format!("mock_{port}"), pubkey, url: get_local_address(port).parse()? }; - let config = RelayConfig { - entry, - id: None, - headers: None, - get_params: None, - get_header: GetHeaderTransport::Http, - enable_timing_games: false, - target_first_request_ms: None, - frequency_get_header_ms: None, - validator_registration_batch_size: Some(batch_size), - }; + let mut config = mock_relay_config(port, pubkey)?; + config.validator_registration_batch_size = Some(batch_size); + RelayClient::new(config) +} + +pub fn generate_mock_stream_relay( + port: u16, + stream_url: Url, + pubkey: BlsPublicKey, +) -> Result { + let mut config = mock_relay_config(port, pubkey)?; + config.get_header = GetHeaderTransport::Stream(stream_url); RelayClient::new(config) } diff --git a/tests/tests/pbs_get_header_ws.rs b/tests/tests/pbs_get_header_ws.rs new file mode 100644 index 00000000..ca027505 --- /dev/null +++ b/tests/tests/pbs_get_header_ws.rs @@ -0,0 +1,280 @@ +use std::{ + path::PathBuf, + sync::Arc, + time::{Duration, Instant}, +}; + +use alloy::primitives::{B256, U256}; +use cb_common::{ + pbs::{GetHeaderResponse, HEADER_VERSION_VALUE}, + signature::sign_builder_root, + signer::random_secret, + types::{BlsPublicKeyBytes, BlsSecretKey, Chain, KnownChain}, + utils::timestamp_of_slot_start_sec, + wire::EncodingType, +}; +use cb_pbs::{DefaultBuilderApi, PbsService, PbsState}; +use cb_tests::{ + mock_relay::{MockRelayState, start_mock_relay_service_with_listener}, + mock_validator::MockValidator, + mock_ws_relay::{MockWsRelayState, start_mock_ws_relay_service}, + utils::{ + generate_mock_relay, generate_mock_stream_relay, get_free_listener, get_pbs_config, + setup_test_env, to_pbs_config, + }, +}; +use eyre::Result; +use lh_types::ForkName; +use reqwest::StatusCode; +use tree_hash::TreeHash; +use url::Url; + +fn request_slot() -> u64 { + KnownChain::Hoodi.fulu_fork_slot() + 1 +} + +/// Start a streaming relay on a free port and return the client PBS should use +/// plus the mock state. +async fn start_stream_relay( + state: MockWsRelayState, + pubkey: cb_common::types::BlsPublicKey, +) -> Result<(Arc, cb_common::pbs::RelayClient)> { + let listener = get_free_listener().await; + let port = listener.local_addr()?.port(); + let state = Arc::new(state); + tokio::spawn(start_mock_ws_relay_service(state.clone(), listener)); + + let stream_url: Url = format!("ws://127.0.0.1:{port}/").parse()?; + let relay = generate_mock_stream_relay(port, stream_url, pubkey)?; + + Ok((state, relay)) +} + +/// Boot PBS on a free port with the given relays and header timeout. +async fn start_pbs( + chain: Chain, + relays: Vec, + timeout_get_header_ms: u64, +) -> Result { + let listener = get_free_listener().await; + let port = listener.local_addr()?.port(); + + let mut pbs_config = get_pbs_config(port); + pbs_config.timeout_get_header_ms = timeout_get_header_ms; + + let config = to_pbs_config(chain, pbs_config, relays); + let state = PbsState::new(config, PathBuf::new()); + drop(listener); + tokio::spawn(PbsService::run::<(), DefaultBuilderApi>(state)); + + // leave some time to start servers + tokio::time::sleep(Duration::from_millis(100)).await; + + MockValidator::new(port) +} + +async fn get_header_json( + validator: &MockValidator, +) -> Result<(StatusCode, Option)> { + let res = validator.do_get_header(None, vec![EncodingType::Json], ForkName::Fulu).await?; + let code = res.status(); + if code != StatusCode::OK { + return Ok((code, None)); + } + + Ok((code, Some(serde_json::from_slice(&res.bytes().await?)?))) +} + +fn assert_bid(res: &GetHeaderResponse, chain: Chain, signer: &BlsSecretKey, value: U256) { + assert_eq!(*res.data.message.value(), value); + assert_eq!(res.data.message.header().parent_hash().0, B256::ZERO); + assert_eq!(res.data.message.header().block_hash().0[0], 1); + assert_eq!(*res.data.message.pubkey(), BlsPublicKeyBytes::from(signer.public_key())); + assert_eq!( + res.data.message.header().timestamp(), + timestamp_of_slot_start_sec(request_slot(), chain) + ); + assert_eq!( + res.data.signature, + sign_builder_root(chain, signer, &res.data.message.tree_hash_root()) + ); +} + +/// The relay's last word wins, even when an earlier update paid more, and a +/// close from the relay ends the wait before the deadline. +#[tokio::test] +async fn test_get_header_ws_returns_latest_bid() -> Result<()> { + setup_test_env(); + let signer = random_secret(); + let pubkey = signer.public_key(); + let chain = Chain::Hoodi; + let timeout_ms = 1_000; + + let (relay_state, relay) = start_stream_relay( + MockWsRelayState::new(chain, signer.clone()).with_bid_values(vec![ + U256::from(30), + U256::from(20), + U256::from(10), + ]), + pubkey, + ) + .await?; + + let validator = start_pbs(chain, vec![relay], timeout_ms).await?; + + let started = Instant::now(); + let (code, res) = get_header_json(&validator).await?; + assert_eq!(code, StatusCode::OK); + // The relay closed, so PBS must not have sat on the deadline + assert!(started.elapsed() < Duration::from_millis(timeout_ms)); + + // Last update, not the highest one + assert_bid(&res.unwrap(), chain, &signer, U256::from(10)); + + assert_eq!(relay_state.received_connections(), 1); + Ok(()) +} + +/// The handshake carries the same request data as the HTTP call +#[tokio::test] +async fn test_get_header_ws_handshake_carries_request() -> Result<()> { + setup_test_env(); + let signer = random_secret(); + let pubkey = signer.public_key(); + let chain = Chain::Hoodi; + let timeout_ms = 1_000; + + let (relay_state, relay) = + start_stream_relay(MockWsRelayState::new(chain, signer.clone()), pubkey).await?; + let validator = start_pbs(chain, vec![relay], timeout_ms).await?; + + let sent_at = cb_common::utils::utcnow_ms(); + let (code, _) = get_header_json(&validator).await?; + assert_eq!(code, StatusCode::OK); + + let request = relay_state.last_request().expect("relay saw no request"); + assert_eq!(request.slot, request_slot()); + assert_eq!(request.parent_hash, B256::ZERO); + assert!(request.validator_pubkey.starts_with("0x")); + + // No timeout header from the caller, so PBS passes its own budget through + assert_eq!(request.timeout_ms, Some(timeout_ms)); + let start_time_ms = request.start_time_ms.expect("missing start time header"); + assert!((sent_at..sent_at + timeout_ms).contains(&start_time_ms)); + + assert!(request.user_agent.is_some_and(|ua| ua.contains("commit-boost"))); + assert_eq!(request.cb_version.as_deref(), Some(HEADER_VERSION_VALUE)); + Ok(()) +} + +/// A relay that keeps the stream open is cut off at the PBS deadline, and the +/// last update received up to that point is the one returned. +#[tokio::test] +async fn test_get_header_ws_returns_at_deadline() -> Result<()> { + setup_test_env(); + let signer = random_secret(); + let pubkey = signer.public_key(); + let chain = Chain::Hoodi; + let timeout_ms = 400; + + let (_relay_state, relay) = start_stream_relay( + MockWsRelayState::new(chain, signer.clone()) + .with_bid_values(vec![U256::from(10), U256::from(20)]) + .with_update_interval(Duration::from_millis(50)) + .hold_open(), + pubkey, + ) + .await?; + + let validator = start_pbs(chain, vec![relay], timeout_ms).await?; + + let started = Instant::now(); + let (code, res) = get_header_json(&validator).await?; + let elapsed = started.elapsed(); + + assert_eq!(code, StatusCode::OK); + assert_bid(&res.unwrap(), chain, &signer, U256::from(20)); + + // Held open, so PBS waited out its full budget and no longer + assert!(elapsed >= Duration::from_millis(timeout_ms), "returned early: {elapsed:?}"); + assert!(elapsed < Duration::from_millis(2 * timeout_ms), "returned late: {elapsed:?}"); + Ok(()) +} + +/// A stream that never delivers a bid is a 204, same as an HTTP relay with no +/// header for the slot. +#[tokio::test] +async fn test_get_header_ws_no_bid_returns_204() -> Result<()> { + setup_test_env(); + let signer = random_secret(); + let pubkey = signer.public_key(); + let chain = Chain::Hoodi; + + let (relay_state, relay) = start_stream_relay( + MockWsRelayState::new(chain, signer).with_bid_values(vec![]).hold_open(), + pubkey, + ) + .await?; + + let validator = start_pbs(chain, vec![relay], 300).await?; + + let (code, _) = get_header_json(&validator).await?; + assert_eq!(code, StatusCode::NO_CONTENT); + assert_eq!(relay_state.received_connections(), 1); + Ok(()) +} + +/// An unreachable stream relay fails that relay only, it doesn't fail the call +#[tokio::test] +async fn test_get_header_ws_unreachable_relay_returns_204() -> Result<()> { + setup_test_env(); + let signer = random_secret(); + let pubkey = signer.public_key(); + let chain = Chain::Hoodi; + + // Take a port and immediately give it back, so nothing is listening + let listener = get_free_listener().await; + let port = listener.local_addr()?.port(); + drop(listener); + + let stream_url: Url = format!("ws://127.0.0.1:{port}/").parse()?; + let relay = generate_mock_stream_relay(port, stream_url, pubkey)?; + let validator = start_pbs(chain, vec![relay], 300).await?; + + let (code, _) = get_header_json(&validator).await?; + assert_eq!(code, StatusCode::NO_CONTENT); + Ok(()) +} + +/// A streamed bid competes in the same auction as an HTTP one +#[tokio::test] +async fn test_get_header_ws_wins_auction_against_http() -> Result<()> { + setup_test_env(); + let signer = random_secret(); + let pubkey = signer.public_key(); + let chain = Chain::Hoodi; + + let http_listener = get_free_listener().await; + let http_port = http_listener.local_addr()?.port(); + let http_state = + Arc::new(MockRelayState::new(chain, signer.clone()).with_bid_value(U256::from(10))); + let http_relay = generate_mock_relay(http_port, pubkey.clone())?; + tokio::spawn(start_mock_relay_service_with_listener(http_state.clone(), http_listener)); + + let (stream_state, stream_relay) = start_stream_relay( + MockWsRelayState::new(chain, signer.clone()).with_bid_values(vec![U256::from(50)]), + pubkey, + ) + .await?; + + let validator = start_pbs(chain, vec![http_relay, stream_relay], 1_000).await?; + + let (code, res) = get_header_json(&validator).await?; + assert_eq!(code, StatusCode::OK); + assert_bid(&res.unwrap(), chain, &signer, U256::from(50)); + + // Both transports were actually queried + assert_eq!(http_state.received_get_header(), 1); + assert_eq!(stream_state.received_connections(), 1); + Ok(()) +} From 703730f8a36a7a8925f00b1c7c7d16b93c97860d Mon Sep 17 00:00:00 2001 From: Nina Date: Tue, 28 Jul 2026 18:41:15 +0100 Subject: [PATCH 3/5] fixes --- Cargo.lock | 2 +- Cargo.toml | 1 + crates/common/src/pbs/error.rs | 9 +- crates/common/src/pbs/relay.rs | 77 ++++++++++---- crates/pbs/Cargo.toml | 2 +- crates/pbs/src/mev_boost/get_header.rs | 2 - crates/pbs/src/mev_boost/get_header_ws.rs | 121 +++++++++++++--------- tests/src/mock_ws_relay.rs | 15 +++ tests/tests/pbs_get_header_ws.rs | 25 +++++ 9 files changed, 182 insertions(+), 72 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 66083990..27645840 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1828,7 +1828,6 @@ dependencies = [ "async-trait", "axum 0.8.9", "axum-extra", - "bytes", "cb-common", "cb-metrics", "ethereum_serde_utils 0.7.0", @@ -1853,6 +1852,7 @@ dependencies = [ "types", "url", "uuid 1.23.1", + "webpki-roots 1.0.7", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index d2334ae5..6b68e99c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -94,6 +94,7 @@ typenum = "1.17.0" unicode-normalization = "0.1.24" url = { version = "2.5.0", features = ["serde"] } uuid = { version = "1.8.0", features = ["fast-rng", "serde", "v4"] } +webpki-roots = "1.0" [patch.crates-io] blstrs_plus = { git = "https://github.com/Commit-Boost/blstrs" } diff --git a/crates/common/src/pbs/error.rs b/crates/common/src/pbs/error.rs index 9aa6919c..a5cc19a4 100644 --- a/crates/common/src/pbs/error.rs +++ b/crates/common/src/pbs/error.rs @@ -41,11 +41,18 @@ pub enum PbsError { #[error("websocket error: {0}")] WebSocket(String), + + #[error("websocket timed out")] + WebSocketTimeout, } impl PbsError { pub fn is_timeout(&self) -> bool { - matches!(self, PbsError::Reqwest(err) if err.is_timeout()) + match self { + PbsError::Reqwest(err) => err.is_timeout(), + PbsError::WebSocketTimeout => true, + _ => false, + } } /// Extract the HTTP status code from relay-originated errors. diff --git a/crates/common/src/pbs/relay.rs b/crates/common/src/pbs/relay.rs index 394c6068..81954fb2 100644 --- a/crates/common/src/pbs/relay.rs +++ b/crates/common/src/pbs/relay.rs @@ -66,24 +66,20 @@ pub struct RelayClient { pub id: Arc, /// HTTP client to send requests pub client: reqwest::Client, + /// Baseline headers for the get_header stream handshake. + stream_headers: Arc, /// Configuration of the relay pub config: Arc, } impl RelayClient { - pub fn new(mut config: RelayConfig) -> eyre::Result { - if let GetHeaderTransport::Stream(url) = &mut config.get_header { + pub fn new(config: RelayConfig) -> eyre::Result { + if let GetHeaderTransport::Stream(url) = &config.get_header { eyre::ensure!( matches!(url.scheme(), "ws" | "wss"), "get_header stream url must be ws:// or wss://, got {}", url.scheme() ); - - // Url::join drops the last segment of a base without a trailing slash - if !url.path().ends_with('/') { - let path = format!("{}/", url.path()); - url.set_path(&path); - } } let mut headers = HeaderMap::new(); @@ -98,35 +94,48 @@ impl RelayClient { } } + let stream_headers = match config.get_header { + GetHeaderTransport::Http => HeaderMap::new(), + GetHeaderTransport::Stream(_) => headers.clone(), + }; + let client = reqwest::Client::builder() .default_headers(headers) .timeout(DEFAULT_REQUEST_TIMEOUT) .build()?; - Ok(Self { id: Arc::new(config.id().to_owned()), client, config: Arc::new(config) }) + Ok(Self { + id: Arc::new(config.id().to_owned()), + client, + stream_headers: Arc::new(stream_headers), + config: Arc::new(config), + }) } pub fn pubkey(&self) -> &BlsPublicKey { &self.config.entry.pubkey } + pub fn stream_headers(&self) -> &HeaderMap { + &self.stream_headers + } + // URL builders pub fn get_url(&self, path: &str) -> Result { - self.join_url(&self.config.entry.url, path) + let mut url = self.config.entry.url.join(path).map_err(PbsError::UrlParsing)?; + self.append_get_params(&mut url); + Ok(url) } - fn join_url(&self, base: &Url, path: &str) -> Result { - let mut url = base.join(path).map_err(PbsError::UrlParsing)?; - + fn append_get_params(&self, url: &mut Url) { if let Some(get_params) = &self.config.get_params { let mut query_pairs = url.query_pairs_mut(); for (key, value) in get_params { query_pairs.append_pair(key, value); } } - - Ok(url) } + pub fn builder_api_url( &self, path: &str, @@ -157,9 +166,22 @@ impl RelayClient { GetHeaderTransport::Http => { GetHeaderRequest::Http(self.get_header_url(slot, parent_hash, validator_pubkey)?) } - GetHeaderTransport::Stream(base) => GetHeaderRequest::Stream( - self.join_url(base, &format!("{slot}/{parent_hash}/{validator_pubkey}"))?, - ), + GetHeaderTransport::Stream(base) => { + let mut url = base.clone(); + // Push segments rather than Url::join: join drops both the last + // segment of a base without a trailing slash and its query + url.path_segments_mut() + .map_err(|_| PbsError::WebSocket("stream url cannot be a base".to_string()))? + .pop_if_empty() + .extend([ + slot.to_string(), + parent_hash.to_string(), + validator_pubkey.to_string(), + ]); + + self.append_get_params(&mut url); + GetHeaderRequest::Stream(url) + } }) } @@ -299,8 +321,8 @@ mod tests { "http dispatch must match the plain url builder" ); - // A stream url takes over, and a missing trailing slash is normalized so - // the configured path isn't dropped by Url::join + // A stream url takes over, and a missing trailing slash doesn't drop the + // configured path for base in ["wss://abc.xyz/stream", "wss://abc.xyz/stream/"] { let mut config = base_config.clone(); config.get_header = GetHeaderTransport::Stream(base.parse().unwrap()); @@ -317,6 +339,21 @@ mod tests { ); } + // An auth token on the stream url survives the path append + let mut config = base_config.clone(); + config.get_header = + GetHeaderTransport::Stream("wss://abc.xyz/stream?token=abc".parse().unwrap()); + let relay = RelayClient::new(config).unwrap(); + let GetHeaderRequest::Stream(url) = + relay.get_header_request(slot, &parent_hash, &validator_pubkey).unwrap() + else { + panic!("expected stream request"); + }; + assert_eq!( + url.to_string(), + format!("wss://abc.xyz/stream/{slot}/{parent_hash}/{validator_pubkey}?token=abc") + ); + // Non-ws schemes are rejected at construction let mut config = base_config; config.get_header = GetHeaderTransport::Stream("https://abc.xyz/stream/".parse().unwrap()); diff --git a/crates/pbs/Cargo.toml b/crates/pbs/Cargo.toml index 401dee32..7497aa93 100644 --- a/crates/pbs/Cargo.toml +++ b/crates/pbs/Cargo.toml @@ -9,7 +9,6 @@ version.workspace = true alloy.workspace = true async-trait.workspace = true axum.workspace = true -bytes.workspace = true axum-extra.workspace = true cb-common.workspace = true cb-metrics.workspace = true @@ -34,4 +33,5 @@ tracing.workspace = true tree_hash.workspace = true url.workspace = true uuid.workspace = true +webpki-roots.workspace = true thiserror.workspace = true diff --git a/crates/pbs/src/mev_boost/get_header.rs b/crates/pbs/src/mev_boost/get_header.rs index 473307b9..2903b9fc 100644 --- a/crates/pbs/src/mev_boost/get_header.rs +++ b/crates/pbs/src/mev_boost/get_header.rs @@ -402,8 +402,6 @@ async fn send_one_get_header( Ok((start_request_time, Some(get_header_response))) } -/// Validate a decoded header from a relay: header contents, relay signature, -/// and, if enabled, consistency with the parent block. pub(super) fn validate_get_header_response( request_info: &RequestInfo, relay: &RelayClient, diff --git a/crates/pbs/src/mev_boost/get_header_ws.rs b/crates/pbs/src/mev_boost/get_header_ws.rs index 7354cc49..978af766 100644 --- a/crates/pbs/src/mev_boost/get_header_ws.rs +++ b/crates/pbs/src/mev_boost/get_header_ws.rs @@ -1,6 +1,6 @@ //! Streaming get_header over a websocket, for relays configured with -//! `get_header_ws_url`. One connection per get_header call, dropped when the -//! call returns. +//! `get_header = { stream = "wss://..." }`. One connection per get_header +//! call, dropped when the call returns. //! //! The request is the handshake itself: slot / parent_hash / pubkey in the //! path, deadline and timestamp in headers, same data the HTTP request carries. @@ -12,25 +12,26 @@ //! .. SSZ SignedBuilderBid //! ``` -use std::{str::FromStr, sync::Once, time::Duration}; +use std::{ + sync::{Arc, OnceLock}, + time::Duration, +}; use alloy::primitives::utils::format_ether; -use axum::http::{HeaderName, HeaderValue, Request, header::USER_AGENT}; -use bytes::Bytes; +use axum::http::{HeaderValue, Request, header::USER_AGENT}; use cb_common::{ pbs::{ ForkName, ForkVersionDecode, GetHeaderInfo, GetHeaderResponse, HEADER_START_TIME_UNIX_MS, - HEADER_TIMEOUT_MS, HEADER_VERSION_KEY, HEADER_VERSION_VALUE, RelayClient, SignedBuilderBid, - error::PbsError, + HEADER_TIMEOUT_MS, RelayClient, SignedBuilderBid, error::PbsError, }, utils::utcnow_ms, }; use futures::StreamExt; -use rustls::crypto::{CryptoProvider, aws_lc_rs}; +use rustls::{ClientConfig, RootCertStore, crypto::aws_lc_rs}; use tokio::time::{Instant, sleep_until, timeout_at}; use tokio_tungstenite::{ - connect_async_with_config, - tungstenite::{Message, client::IntoClientRequest, protocol::WebSocketConfig}, + Connector, connect_async_tls_with_config, + tungstenite::{Bytes, Message, client::IntoClientRequest, protocol::WebSocketConfig}, }; use tracing::{debug, info, warn}; use url::Url; @@ -77,10 +78,13 @@ pub(super) async fn get_header_ws( .max_message_size(Some(MAX_SIZE_GET_HEADER_RESPONSE)) .max_frame_size(Some(MAX_SIZE_GET_HEADER_RESPONSE)); - install_crypto_provider(); - let start_request = Instant::now(); - let connect = connect_async_with_config(request, Some(config), true); + let connect = connect_async_tls_with_config( + request, + Some(config), + true, + Some(Connector::Rustls(tls_config().clone())), + ); let (mut stream, _) = match timeout_at(deadline, connect).await { Ok(Ok(connected)) => connected, Ok(Err(err)) => { @@ -89,17 +93,20 @@ pub(super) async fn get_header_ws( } Err(_) => { record_status(TIMEOUT_ERROR_CODE_STR, relay); - return Err(PbsError::WebSocket("timed out connecting".to_string())); + return Err(PbsError::WebSocketTimeout); } }; - debug!(relay_id = relay.id.as_ref(), latency = ?start_request.elapsed(), "ws connected"); + let connect_latency = start_request.elapsed(); + debug!(relay_id = relay.id.as_ref(), ?connect_latency, "ws connected"); let timer = sleep_until(deadline); tokio::pin!(timer); let mut latest: Option<(ForkName, Bytes)> = None; - let mut first_bid_latency = None; + let mut first_bid_latency: Option = None; let mut updates = 0usize; + let mut invalid_frames = 0usize; + let mut stream_error = None; loop { let message = tokio::select! { @@ -112,6 +119,7 @@ pub(super) async fn get_header_ws( Some(Ok(message)) => message, Some(Err(err)) => { warn!(relay_id = relay.id.as_ref(), %err, "ws stream error"); + stream_error = Some(PbsError::WebSocket(format!("stream error: {err}"))); break; } None => break, @@ -126,12 +134,16 @@ pub(super) async fn get_header_ws( match parse_frame(payload) { Ok((fork, bid)) => { updates += 1; - first_bid_latency.get_or_insert_with(|| start_request.elapsed()); + if first_bid_latency.is_none() { + first_bid_latency = Some(start_request.elapsed()); + } latest = Some((fork, bid)); } Err(err) => { - warn!(relay_id = relay.id.as_ref(), %err, "invalid ws frame"); - break; + invalid_frames += 1; + if invalid_frames == 1 { + warn!(relay_id = relay.id.as_ref(), %err, "invalid ws frame, skipping"); + } } } } @@ -139,14 +151,21 @@ pub(super) async fn get_header_ws( drop(stream); let Some((fork, bid_bytes)) = latest else { - debug!(relay_id = relay.id.as_ref(), "no header from relay"); + if let Some(err) = stream_error { + record_status(TIMEOUT_ERROR_CODE_STR, relay); + return Err(err); + } + + debug!(relay_id = relay.id.as_ref(), ?connect_latency, invalid_frames, "no header"); record_status("204", relay); return Ok(None); }; - RELAY_LATENCY - .with_label_values(&[GET_HEADER_ENDPOINT_TAG, &relay.id]) - .observe(first_bid_latency.unwrap_or_default().as_secs_f64()); + if let Some(first_bid_latency) = first_bid_latency { + RELAY_LATENCY + .with_label_values(&[GET_HEADER_ENDPOINT_TAG, &relay.id]) + .observe(first_bid_latency.as_secs_f64()); + } record_status("200", relay); let data = SignedBuilderBid::from_ssz_bytes_by_fork(&bid_bytes, fork).map_err(|err| { @@ -157,24 +176,29 @@ pub(super) async fn get_header_ws( })?; let response = GetHeaderResponse { version: fork, data, metadata: Default::default() }; + let start_validate = Instant::now(); + let validated = validate_get_header_response(request_info, relay, &response); + let validate_latency = start_validate.elapsed(); + info!( relay_id = relay.id.as_ref(), header_size_bytes = bid_bytes.len(), - latency = ?start_request.elapsed(), + ?connect_latency, + ?first_bid_latency, + ?validate_latency, version = ?fork, value_eth = format_ether(*response.value()), block_hash = %response.block_hash(), updates, + invalid_frames, "received new header from ws stream" ); - validate_get_header_response(request_info, relay, &response)?; + validated?; Ok(Some(response)) } -/// The handshake carries the request: same headers as the HTTP path, minus -/// `Accept` (the stream is SSZ only). fn build_handshake_request( request_info: &RequestInfo, relay: &RelayClient, @@ -190,16 +214,9 @@ fn build_handshake_request( if let Some(user_agent) = request_info.headers.get(USER_AGENT) { headers.insert(USER_AGENT, user_agent.clone()); } - headers.insert(HEADER_VERSION_KEY, HeaderValue::from_static(HEADER_VERSION_VALUE)); - - // The HTTP client bakes these into its default headers, the ws handshake - // needs them explicitly. Validated in RelayClient::new. - for (key, value) in relay.config.headers.iter().flatten() { - let key = HeaderName::from_str(key) - .map_err(|_| PbsError::WebSocket(format!("invalid header name: {key}")))?; - let value = HeaderValue::from_str(value) - .map_err(|_| PbsError::WebSocket(format!("invalid header value for: {key}")))?; - headers.insert(key, value); + + for (key, value) in relay.stream_headers() { + headers.insert(key, value.clone()); } headers.insert(HEADER_START_TIME_UNIX_MS, HeaderValue::from(utcnow_ms())); @@ -227,16 +244,26 @@ fn record_status(code: &str, relay: &RelayClient) { RELAY_STATUS_CODE.with_label_values(&[code, GET_HEADER_ENDPOINT_TAG, &relay.id]).inc(); } -/// rustls is built with both `ring` and `aws-lc-rs` here, so the default -/// `ClientConfig::builder()` inside tokio-tungstenite panics unless a -/// process-level provider is installed. Match the signer and use aws-lc-rs. -fn install_crypto_provider() { - static INIT: Once = Once::new(); - INIT.call_once(|| { - if CryptoProvider::get_default().is_none() { - let _ = aws_lc_rs::default_provider().install_default(); - } - }); +/// One TLS config for every stream connection. Left to tokio-tungstenite it is +/// rebuilt per connect, which reparses the root store and, worse, gives each +/// connection its own session cache: every slot then pays a full handshake +/// instead of a resumed one. The provider is named explicitly because rustls is +/// built with both `ring` and `aws-lc-rs` here, so the default builder needs a +/// process-wide install to pick one. +fn tls_config() -> &'static Arc { + static CONFIG: OnceLock> = OnceLock::new(); + CONFIG.get_or_init(|| { + let mut roots = RootCertStore::empty(); + roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned()); + + Arc::new( + ClientConfig::builder_with_provider(Arc::new(aws_lc_rs::default_provider())) + .with_safe_default_protocol_versions() + .expect("aws-lc-rs supports tls 1.2 and 1.3") + .with_root_certificates(roots) + .with_no_client_auth(), + ) + }) } #[cfg(test)] diff --git a/tests/src/mock_ws_relay.rs b/tests/src/mock_ws_relay.rs index fe489a77..39044113 100644 --- a/tests/src/mock_ws_relay.rs +++ b/tests/src/mock_ws_relay.rs @@ -52,6 +52,9 @@ pub struct MockWsRelayState { /// Keep the connection open after the last update, so PBS returns on its /// own deadline instead of on close hold_open: bool, + /// Precede each bid with frames PBS can't parse, which it must skip rather + /// than treat as the end of the stream + unknown_frames: bool, received_connections: AtomicU64, last_request: Mutex>, } @@ -64,6 +67,7 @@ impl MockWsRelayState { bid_values: vec![U256::from(10)], update_interval: Duration::ZERO, hold_open: false, + unknown_frames: false, received_connections: AtomicU64::new(0), last_request: Mutex::new(None), } @@ -81,6 +85,10 @@ impl MockWsRelayState { Self { hold_open: true, ..self } } + pub fn with_unknown_frames(self) -> Self { + Self { unknown_frames: true, ..self } + } + pub fn received_connections(&self) -> u64 { self.received_connections.load(Ordering::Relaxed) } @@ -120,6 +128,13 @@ async fn serve_stream(state: Arc, stream: TcpStream) -> eyre:: *state.last_request.lock().unwrap() = Some(request.clone()); for value in &state.bid_values { + if state.unknown_frames { + // Unknown message type, unknown fork, truncated prefix + for frame in [vec![0x7f, FORK_FULU, 1], vec![MSG_BID, 0xff, 1], vec![MSG_BID]] { + ws.send(Message::Binary(frame.into())).await?; + } + } + let bid = mock_signed_builder_bid( state.chain, &state.signer, diff --git a/tests/tests/pbs_get_header_ws.rs b/tests/tests/pbs_get_header_ws.rs index ca027505..072572a4 100644 --- a/tests/tests/pbs_get_header_ws.rs +++ b/tests/tests/pbs_get_header_ws.rs @@ -135,6 +135,31 @@ async fn test_get_header_ws_returns_latest_bid() -> Result<()> { Ok(()) } +/// Frames PBS can't parse are skipped, not treated as the end of the stream: +/// the updates after them still count. +#[tokio::test] +async fn test_get_header_ws_skips_unknown_frames() -> Result<()> { + setup_test_env(); + let signer = random_secret(); + let pubkey = signer.public_key(); + let chain = Chain::Hoodi; + + let (_relay_state, relay) = start_stream_relay( + MockWsRelayState::new(chain, signer.clone()) + .with_bid_values(vec![U256::from(10), U256::from(20)]) + .with_unknown_frames(), + pubkey, + ) + .await?; + + let validator = start_pbs(chain, vec![relay], 1_000).await?; + + let (code, res) = get_header_json(&validator).await?; + assert_eq!(code, StatusCode::OK); + assert_bid(&res.unwrap(), chain, &signer, U256::from(20)); + Ok(()) +} + /// The handshake carries the same request data as the HTTP call #[tokio::test] async fn test_get_header_ws_handshake_carries_request() -> Result<()> { From 7fbcf5630fe2cc1722ef7d077509470a79ab6a83 Mon Sep 17 00:00:00 2001 From: Nina Date: Wed, 29 Jul 2026 13:20:11 +0100 Subject: [PATCH 4/5] fixes --- config.example.toml | 11 +- crates/common/src/config/pbs.rs | 4 +- crates/common/src/pbs/constants.rs | 2 + crates/common/src/pbs/relay.rs | 117 ++++++++++++---------- crates/pbs/src/constants.rs | 4 + crates/pbs/src/mev_boost/get_header.rs | 2 +- crates/pbs/src/mev_boost/get_header_ws.rs | 39 +++++--- tests/src/mock_ws_relay.rs | 12 ++- tests/src/utils.rs | 8 +- tests/tests/pbs_get_header_ws.rs | 9 +- 10 files changed, 111 insertions(+), 97 deletions(-) diff --git a/config.example.toml b/config.example.toml index d3a4b446..209089a7 100644 --- a/config.example.toml +++ b/config.example.toml @@ -93,12 +93,11 @@ headers = { X-MyCustomHeader = "MyCustomValue" } # OPTIONAL get_params = { param1 = "value1", param2 = "value2" } # How to fetch headers from this relay. -# "http" -> one request per get_header, at the relay url above -# { stream = "wss://abc.xyz/stream" } -> websocket stream of bid updates, only for relays that -# support it. The stream stays open for the request -# window and the last bid pushed is the one used, so the -# timing game settings below are ignored. -# `{slot}/{parent_hash}/{pubkey}` is appended to the url. +# "http" -> one request per get_header, at the relay url above +# "stream" -> websocket stream of bid updates, only for relays that support it. Connects to +# ws(s):///eth/v1/builder/header_stream/{slot}/{parent_hash}/{pubkey}. +# The stream stays open for the whole request window and the last bid pushed is the +# one used. # OPTIONAL, DEFAULT: "http" get_header = "http" # Whether to enable timing games, as tuned by `target_first_request_ms` and `frequency_get_header_ms`. diff --git a/crates/common/src/config/pbs.rs b/crates/common/src/config/pbs.rs index 08bcaafa..c261db59 100644 --- a/crates/common/src/config/pbs.rs +++ b/crates/common/src/config/pbs.rs @@ -38,12 +38,12 @@ use crate::{ }, }; -#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize, Serialize)] +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize, Serialize)] #[serde(rename_all = "snake_case")] pub enum GetHeaderTransport { #[default] Http, - Stream(Url), + Stream, } #[derive(Debug, Clone, Deserialize, Serialize)] diff --git a/crates/common/src/pbs/constants.rs b/crates/common/src/pbs/constants.rs index bbe20b0d..d0bd3e41 100644 --- a/crates/common/src/pbs/constants.rs +++ b/crates/common/src/pbs/constants.rs @@ -4,6 +4,8 @@ pub const BUILDER_V1_API_PATH: &str = "/eth/v1/builder"; pub const BUILDER_V2_API_PATH: &str = "/eth/v2/builder"; pub const GET_HEADER_PATH: &str = "/header/{slot}/{parent_hash}/{pubkey}"; + +pub const GET_HEADER_STREAM_PATH: &str = "/header_stream"; pub const GET_STATUS_PATH: &str = "/status"; pub const REGISTER_VALIDATOR_PATH: &str = "/validators"; pub const SUBMIT_BLOCK_PATH: &str = "/blinded_blocks"; diff --git a/crates/common/src/pbs/relay.rs b/crates/common/src/pbs/relay.rs index 81954fb2..12d438cd 100644 --- a/crates/common/src/pbs/relay.rs +++ b/crates/common/src/pbs/relay.rs @@ -8,7 +8,9 @@ use url::Url; use super::{ HEADER_VERSION_KEY, HEADER_VERSION_VALUE, - constants::{GET_STATUS_PATH, REGISTER_VALIDATOR_PATH, SUBMIT_BLOCK_PATH}, + constants::{ + GET_HEADER_STREAM_PATH, GET_STATUS_PATH, REGISTER_VALIDATOR_PATH, SUBMIT_BLOCK_PATH, + }, error::PbsError, }; use crate::{ @@ -58,6 +60,22 @@ pub enum GetHeaderRequest { Stream(Url), } +fn stream_url(entry: &Url) -> eyre::Result { + let scheme = match entry.scheme() { + "http" | "ws" => "ws", + "https" | "wss" => "wss", + other => eyre::bail!("get_header stream needs an http(s) relay url, got {other}"), + }; + + let mut url = entry.clone(); + url.set_scheme(scheme).map_err(|_| eyre::eyre!("cannot use {scheme} for {entry}"))?; + url.set_username("").map_err(|_| eyre::eyre!("cannot strip credentials from {entry}"))?; + url.set_password(None).map_err(|_| eyre::eyre!("cannot strip credentials from {entry}"))?; + url.set_path(&format!("{}{GET_HEADER_STREAM_PATH}", BuilderApiVersion::V1.path())); + + Ok(url) +} + /// A client to interact with a relay, safe to share across threads and cheaply /// cloneable #[derive(Debug, Clone)] @@ -66,6 +84,8 @@ pub struct RelayClient { pub id: Arc, /// HTTP client to send requests pub client: reqwest::Client, + /// Base url of the get_header stream, `Some` only when the relay streams. + stream_url: Option, /// Baseline headers for the get_header stream handshake. stream_headers: Arc, /// Configuration of the relay @@ -74,13 +94,10 @@ pub struct RelayClient { impl RelayClient { pub fn new(config: RelayConfig) -> eyre::Result { - if let GetHeaderTransport::Stream(url) = &config.get_header { - eyre::ensure!( - matches!(url.scheme(), "ws" | "wss"), - "get_header stream url must be ws:// or wss://, got {}", - url.scheme() - ); - } + let stream_url = match config.get_header { + GetHeaderTransport::Http => None, + GetHeaderTransport::Stream => Some(stream_url(&config.entry.url)?), + }; let mut headers = HeaderMap::new(); headers.insert(HEADER_VERSION_KEY, HeaderValue::from_static(HEADER_VERSION_VALUE)); @@ -96,7 +113,7 @@ impl RelayClient { let stream_headers = match config.get_header { GetHeaderTransport::Http => HeaderMap::new(), - GetHeaderTransport::Stream(_) => headers.clone(), + GetHeaderTransport::Stream => headers.clone(), }; let client = reqwest::Client::builder() @@ -107,6 +124,7 @@ impl RelayClient { Ok(Self { id: Arc::new(config.id().to_owned()), client, + stream_url, stream_headers: Arc::new(stream_headers), config: Arc::new(config), }) @@ -162,22 +180,13 @@ impl RelayClient { parent_hash: &B256, validator_pubkey: &BlsPublicKey, ) -> Result { - Ok(match &self.config.get_header { - GetHeaderTransport::Http => { + Ok(match &self.stream_url { + None => { GetHeaderRequest::Http(self.get_header_url(slot, parent_hash, validator_pubkey)?) } - GetHeaderTransport::Stream(base) => { + Some(base) => { let mut url = base.clone(); - // Push segments rather than Url::join: join drops both the last - // segment of a base without a trailing slash and its query - url.path_segments_mut() - .map_err(|_| PbsError::WebSocket("stream url cannot be a base".to_string()))? - .pop_if_empty() - .extend([ - slot.to_string(), - parent_hash.to_string(), - validator_pubkey.to_string(), - ]); + url.set_path(&format!("{}/{slot}/{parent_hash}/{validator_pubkey}", base.path())); self.append_get_params(&mut url); GetHeaderRequest::Stream(url) @@ -308,7 +317,7 @@ mod tests { }"#; let base_config = serde_json::from_str::(relay_config).unwrap(); - // No ws url configured: plain HTTP endpoint + // Default transport: plain HTTP endpoint let relay = RelayClient::new(base_config.clone()).unwrap(); let GetHeaderRequest::Http(url) = relay.get_header_request(slot, &parent_hash, &validator_pubkey).unwrap() @@ -321,28 +330,29 @@ mod tests { "http dispatch must match the plain url builder" ); - // A stream url takes over, and a missing trailing slash doesn't drop the - // configured path - for base in ["wss://abc.xyz/stream", "wss://abc.xyz/stream/"] { - let mut config = base_config.clone(); - config.get_header = GetHeaderTransport::Stream(base.parse().unwrap()); - let relay = RelayClient::new(config).unwrap(); - - let GetHeaderRequest::Stream(url) = - relay.get_header_request(slot, &parent_hash, &validator_pubkey).unwrap() - else { - panic!("expected stream request"); - }; - assert_eq!( - url.to_string(), - format!("wss://abc.xyz/stream/{slot}/{parent_hash}/{validator_pubkey}") - ); - } + // Streaming: the relay url over ws, at the fixed stream path, with the + // pubkey credentials dropped + let mut config = base_config.clone(); + config.get_header = GetHeaderTransport::Stream; + let relay = RelayClient::new(config).unwrap(); + let GetHeaderRequest::Stream(url) = + relay.get_header_request(slot, &parent_hash, &validator_pubkey).unwrap() + else { + panic!("expected stream request"); + }; + assert_eq!( + url.to_string(), + format!( + "ws://abc.xyz/eth/v1/builder/header_stream/{slot}/{parent_hash}/{validator_pubkey}" + ) + ); - // An auth token on the stream url survives the path append + // https relays stream over wss, and a port and get_params carry over let mut config = base_config.clone(); - config.get_header = - GetHeaderTransport::Stream("wss://abc.xyz/stream?token=abc".parse().unwrap()); + config.entry.url = + format!("https://{}@abc.xyz:4444/relay-api", config.entry.pubkey).parse().unwrap(); + config.get_header = GetHeaderTransport::Stream; + config.get_params = Some(HashMap::from([("token".to_string(), "abc".to_string())])); let relay = RelayClient::new(config).unwrap(); let GetHeaderRequest::Stream(url) = relay.get_header_request(slot, &parent_hash, &validator_pubkey).unwrap() @@ -351,12 +361,15 @@ mod tests { }; assert_eq!( url.to_string(), - format!("wss://abc.xyz/stream/{slot}/{parent_hash}/{validator_pubkey}?token=abc") + format!( + "wss://abc.xyz:4444/eth/v1/builder/header_stream/{slot}/{parent_hash}/{validator_pubkey}?token=abc" + ) ); - // Non-ws schemes are rejected at construction + // A relay url we can't stream over is rejected at construction let mut config = base_config; - config.get_header = GetHeaderTransport::Stream("https://abc.xyz/stream/".parse().unwrap()); + config.entry.url = "unix:/tmp/relay.sock".parse().unwrap(); + config.get_header = GetHeaderTransport::Stream; assert!(RelayClient::new(config).is_err()); } @@ -374,22 +387,16 @@ mod tests { }; assert_eq!(with_transport(r#""http""#).unwrap(), GetHeaderTransport::Http); - assert_eq!( - with_transport(r#"{ "stream": "wss://abc.xyz/stream" }"#).unwrap(), - GetHeaderTransport::Stream("wss://abc.xyz/stream".parse().unwrap()) - ); + assert_eq!(with_transport(r#""stream""#).unwrap(), GetHeaderTransport::Stream); assert!(with_transport(r#""grpc""#).is_err()); // Same shapes in the toml the operator actually writes let toml_config = r#" url = "http://0xa1cec75a3f0661e99299274182938151e8433c61a19222347ea1313d839229cb4ce4e3e5aa2bdeb71c8fcf1b084963c2@abc.xyz" - get_header = { stream = "wss://abc.xyz/stream" } + get_header = "stream" "#; let config = toml::from_str::(toml_config).unwrap(); - assert_eq!( - config.get_header, - GetHeaderTransport::Stream("wss://abc.xyz/stream".parse().unwrap()) - ); + assert_eq!(config.get_header, GetHeaderTransport::Stream); // Defaults to http when omitted let relay_config = r#" diff --git a/crates/pbs/src/constants.rs b/crates/pbs/src/constants.rs index 0b0faeb3..301c9b34 100644 --- a/crates/pbs/src/constants.rs +++ b/crates/pbs/src/constants.rs @@ -8,6 +8,10 @@ pub const RELOAD_ENDPOINT_TAG: &str = "reload"; pub const TIMEOUT_ERROR_CODE: u16 = 555; pub const TIMEOUT_ERROR_CODE_STR: &str = "555"; +/// For metrics recorded when the relay connection itself failed rather than the +/// deadline expiring: refused, dns, tls, or a stream that broke mid-window +pub const TRANSPORT_ERROR_CODE_STR: &str = "556"; + /// 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 const MAX_SIZE_SUBMIT_BLOCK_RESPONSE: usize = 20 * 1024 * 1024; diff --git a/crates/pbs/src/mev_boost/get_header.rs b/crates/pbs/src/mev_boost/get_header.rs index 2903b9fc..b9d9767f 100644 --- a/crates/pbs/src/mev_boost/get_header.rs +++ b/crates/pbs/src/mev_boost/get_header.rs @@ -607,7 +607,7 @@ fn decode_json_payload(response_bytes: &[u8]) -> Result Result { diff --git a/crates/pbs/src/mev_boost/get_header_ws.rs b/crates/pbs/src/mev_boost/get_header_ws.rs index 978af766..368d9192 100644 --- a/crates/pbs/src/mev_boost/get_header_ws.rs +++ b/crates/pbs/src/mev_boost/get_header_ws.rs @@ -21,8 +21,8 @@ use alloy::primitives::utils::format_ether; use axum::http::{HeaderValue, Request, header::USER_AGENT}; use cb_common::{ pbs::{ - ForkName, ForkVersionDecode, GetHeaderInfo, GetHeaderResponse, HEADER_START_TIME_UNIX_MS, - HEADER_TIMEOUT_MS, RelayClient, SignedBuilderBid, error::PbsError, + ForkName, GetHeaderInfo, GetHeaderResponse, HEADER_START_TIME_UNIX_MS, HEADER_TIMEOUT_MS, + RelayClient, error::PbsError, }, utils::utcnow_ms, }; @@ -31,15 +31,21 @@ use rustls::{ClientConfig, RootCertStore, crypto::aws_lc_rs}; use tokio::time::{Instant, sleep_until, timeout_at}; use tokio_tungstenite::{ Connector, connect_async_tls_with_config, - tungstenite::{Bytes, Message, client::IntoClientRequest, protocol::WebSocketConfig}, + tungstenite::{ + Bytes, Error as WsError, Message, client::IntoClientRequest, protocol::WebSocketConfig, + }, }; use tracing::{debug, info, warn}; use url::Url; use super::get_header::{RequestInfo, validate_get_header_response}; use crate::{ - constants::{GET_HEADER_ENDPOINT_TAG, MAX_SIZE_GET_HEADER_RESPONSE, TIMEOUT_ERROR_CODE_STR}, + constants::{ + GET_HEADER_ENDPOINT_TAG, MAX_SIZE_GET_HEADER_RESPONSE, TIMEOUT_ERROR_CODE_STR, + TRANSPORT_ERROR_CODE_STR, + }, metrics::{RELAY_LATENCY, RELAY_STATUS_CODE}, + mev_boost::get_header::decode_ssz_payload, }; /// Frame prefix: message type + fork. @@ -88,7 +94,13 @@ pub(super) async fn get_header_ws( let (mut stream, _) = match timeout_at(deadline, connect).await { Ok(Ok(connected)) => connected, Ok(Err(err)) => { - record_status(TIMEOUT_ERROR_CODE_STR, relay); + let rejected = match &err { + WsError::Http(res) => Some(res.status()), + _ => None, + }; + let status = rejected.as_ref().map_or(TRANSPORT_ERROR_CODE_STR, |code| code.as_str()); + + record_status(status, relay); return Err(PbsError::WebSocket(format!("connect failed: {err}"))); } Err(_) => { @@ -152,7 +164,7 @@ pub(super) async fn get_header_ws( let Some((fork, bid_bytes)) = latest else { if let Some(err) = stream_error { - record_status(TIMEOUT_ERROR_CODE_STR, relay); + record_status(TRANSPORT_ERROR_CODE_STR, relay); return Err(err); } @@ -166,15 +178,8 @@ pub(super) async fn get_header_ws( .with_label_values(&[GET_HEADER_ENDPOINT_TAG, &relay.id]) .observe(first_bid_latency.as_secs_f64()); } - record_status("200", relay); - let data = SignedBuilderBid::from_ssz_bytes_by_fork(&bid_bytes, fork).map_err(|err| { - PbsError::SSZDecode { - err: format!("error decoding relay payload from ws stream: {err:?}"), - fork, - } - })?; - let response = GetHeaderResponse { version: fork, data, metadata: Default::default() }; + let response = decode_ssz_payload(&bid_bytes, fork)?; let start_validate = Instant::now(); let validated = validate_get_header_response(request_info, relay, &response); @@ -196,6 +201,8 @@ pub(super) async fn get_header_ws( validated?; + record_status("200", relay); + Ok(Some(response)) } @@ -313,8 +320,8 @@ mod tests { let frame = bid_frame(6, &expected.data.as_ssz_bytes()); let (fork, bid_bytes) = parse_frame(frame).unwrap(); - let data = SignedBuilderBid::from_ssz_bytes_by_fork(&bid_bytes, fork).unwrap(); + let decoded = decode_ssz_payload(&bid_bytes, fork).unwrap(); assert_eq!(fork, ForkName::Fulu); - assert_eq!(data, expected.data); + assert_eq!(decoded.data, expected.data); } } diff --git a/tests/src/mock_ws_relay.rs b/tests/src/mock_ws_relay.rs index 39044113..e1b46103 100644 --- a/tests/src/mock_ws_relay.rs +++ b/tests/src/mock_ws_relay.rs @@ -9,7 +9,10 @@ use std::{ use alloy::primitives::{B256, U256}; use cb_common::{ - pbs::{HEADER_START_TIME_UNIX_MS, HEADER_TIMEOUT_MS, HEADER_VERSION_KEY}, + pbs::{ + BuilderApiVersion, GET_HEADER_STREAM_PATH, HEADER_START_TIME_UNIX_MS, HEADER_TIMEOUT_MS, + HEADER_VERSION_KEY, + }, types::{BlsSecretKey, Chain}, }; use futures::SinkExt; @@ -162,10 +165,11 @@ async fn serve_stream(state: Arc, stream: TcpStream) -> eyre:: Ok(()) } -/// The request is the handshake: `/{slot}/{parent_hash}/{pubkey}` plus the -/// same headers the HTTP path sends. +/// The request is the handshake: the fixed stream path followed by +/// `/{slot}/{parent_hash}/{pubkey}`, plus the same headers the HTTP path sends. fn parse_request(req: &Request) -> Option { - let mut segments = req.uri().path().trim_start_matches('/').split('/'); + let prefix = format!("{}{GET_HEADER_STREAM_PATH}/", BuilderApiVersion::V1.path()); + let mut segments = req.uri().path().strip_prefix(&prefix)?.split('/'); let slot = segments.next()?.parse().ok()?; let parent_hash = B256::from_str(segments.next()?).ok()?; let validator_pubkey = segments.next()?.to_string(); diff --git a/tests/src/utils.rs b/tests/src/utils.rs index 6d9d0dc9..0559877a 100644 --- a/tests/src/utils.rs +++ b/tests/src/utils.rs @@ -71,13 +71,9 @@ pub fn generate_mock_relay_with_batch_size( RelayClient::new(config) } -pub fn generate_mock_stream_relay( - port: u16, - stream_url: Url, - pubkey: BlsPublicKey, -) -> Result { +pub fn generate_mock_stream_relay(port: u16, pubkey: BlsPublicKey) -> Result { let mut config = mock_relay_config(port, pubkey)?; - config.get_header = GetHeaderTransport::Stream(stream_url); + config.get_header = GetHeaderTransport::Stream; RelayClient::new(config) } diff --git a/tests/tests/pbs_get_header_ws.rs b/tests/tests/pbs_get_header_ws.rs index 072572a4..7f5392c6 100644 --- a/tests/tests/pbs_get_header_ws.rs +++ b/tests/tests/pbs_get_header_ws.rs @@ -27,7 +27,6 @@ use eyre::Result; use lh_types::ForkName; use reqwest::StatusCode; use tree_hash::TreeHash; -use url::Url; fn request_slot() -> u64 { KnownChain::Hoodi.fulu_fork_slot() + 1 @@ -44,10 +43,7 @@ async fn start_stream_relay( let state = Arc::new(state); tokio::spawn(start_mock_ws_relay_service(state.clone(), listener)); - let stream_url: Url = format!("ws://127.0.0.1:{port}/").parse()?; - let relay = generate_mock_stream_relay(port, stream_url, pubkey)?; - - Ok((state, relay)) + Ok((state, generate_mock_stream_relay(port, pubkey)?)) } /// Boot PBS on a free port with the given relays and header timeout. @@ -262,8 +258,7 @@ async fn test_get_header_ws_unreachable_relay_returns_204() -> Result<()> { let port = listener.local_addr()?.port(); drop(listener); - let stream_url: Url = format!("ws://127.0.0.1:{port}/").parse()?; - let relay = generate_mock_stream_relay(port, stream_url, pubkey)?; + let relay = generate_mock_stream_relay(port, pubkey)?; let validator = start_pbs(chain, vec![relay], 300).await?; let (code, _) = get_header_json(&validator).await?; From fcc6e6e77dff4703f2bb484e63de5ff2811d0d87 Mon Sep 17 00:00:00 2001 From: Nina Date: Wed, 29 Jul 2026 17:22:32 +0100 Subject: [PATCH 5/5] api key --- benches/pbs/src/main.rs | 1 + config.example.toml | 4 +++ crates/common/src/config/pbs.rs | 3 +++ crates/common/src/pbs/constants.rs | 1 + crates/common/src/pbs/relay.rs | 21 ++++++++++++--- tests/src/mock_ws_relay.rs | 6 +++-- tests/src/utils.rs | 12 +++++++++ tests/tests/pbs_cfg_file_update.rs | 2 ++ tests/tests/pbs_get_header_ws.rs | 41 ++++++++++++++++++++++++++++-- 9 files changed, 84 insertions(+), 7 deletions(-) diff --git a/benches/pbs/src/main.rs b/benches/pbs/src/main.rs index 4a000765..aca59054 100644 --- a/benches/pbs/src/main.rs +++ b/benches/pbs/src/main.rs @@ -158,6 +158,7 @@ fn get_mock_validator(bench: BenchConfig) -> RelayClient { headers: None, get_params: None, get_header: GetHeaderTransport::Http, + api_key_env: None, enable_timing_games: false, target_first_request_ms: None, frequency_get_header_ms: None, diff --git a/config.example.toml b/config.example.toml index 209089a7..1cd29586 100644 --- a/config.example.toml +++ b/config.example.toml @@ -100,6 +100,10 @@ get_params = { param1 = "value1", param2 = "value2" } # one used. # OPTIONAL, DEFAULT: "http" get_header = "http" +# Name of the environment variable holding this relay's API key. It is sent as `X-Api-Key` when opening a get_header stream, for relays that require one. +# Only used when `get_header = "stream"`. +# OPTIONAL +# api_key_env = "EXAMPLE_RELAY_API_KEY" # Whether to enable timing games, as tuned by `target_first_request_ms` and `frequency_get_header_ms`. # NOTE: if neither `target_first_request_ms` nor `frequency_get_header_ms` is set, this flag has no effect. # diff --git a/crates/common/src/config/pbs.rs b/crates/common/src/config/pbs.rs index c261db59..79e8779d 100644 --- a/crates/common/src/config/pbs.rs +++ b/crates/common/src/config/pbs.rs @@ -61,6 +61,9 @@ pub struct RelayConfig { /// How to fetch headers from this relay #[serde(default)] pub get_header: GetHeaderTransport, + /// Name of the env var holding this relay's API key, sent as + /// `HEADER_API_KEY` when opening a get_header stream. + pub api_key_env: Option, /// Whether to enable timing games #[serde(default = "default_bool::")] pub enable_timing_games: bool, diff --git a/crates/common/src/pbs/constants.rs b/crates/common/src/pbs/constants.rs index d0bd3e41..66ab42f0 100644 --- a/crates/common/src/pbs/constants.rs +++ b/crates/common/src/pbs/constants.rs @@ -19,6 +19,7 @@ pub const HEADER_VERSION_KEY: &str = "X-CommitBoost-Version"; pub const HEADER_VERSION_VALUE: &str = COMMIT_BOOST_VERSION; pub const HEADER_START_TIME_UNIX_MS: &str = "Date-Milliseconds"; pub const HEADER_TIMEOUT_MS: &str = "X-Timeout-Ms"; +pub const HEADER_API_KEY: &str = "X-Api-Key"; pub const HEADER_CONSENSUS_VERSION: &str = "Eth-Consensus-Version"; pub const DEFAULT_PBS_JWT_KEY: &str = "DEFAULT_PBS"; diff --git a/crates/common/src/pbs/relay.rs b/crates/common/src/pbs/relay.rs index 12d438cd..8e43cbe8 100644 --- a/crates/common/src/pbs/relay.rs +++ b/crates/common/src/pbs/relay.rs @@ -7,7 +7,7 @@ use serde::{Deserialize, Serialize}; use url::Url; use super::{ - HEADER_VERSION_KEY, HEADER_VERSION_VALUE, + HEADER_API_KEY, HEADER_VERSION_KEY, HEADER_VERSION_VALUE, constants::{ GET_HEADER_STREAM_PATH, GET_STATUS_PATH, REGISTER_VALIDATOR_PATH, SUBMIT_BLOCK_PATH, }, @@ -15,7 +15,7 @@ use super::{ }; use crate::{ DEFAULT_REQUEST_TIMEOUT, - config::{GetHeaderTransport, RelayConfig}, + config::{GetHeaderTransport, RelayConfig, load_env_var}, pbs::BuilderApiVersion, types::BlsPublicKey, }; @@ -111,11 +111,26 @@ impl RelayClient { } } - let stream_headers = match config.get_header { + let mut stream_headers = match config.get_header { GetHeaderTransport::Http => HeaderMap::new(), GetHeaderTransport::Stream => headers.clone(), }; + if let Some(env) = &config.api_key_env { + match config.get_header { + GetHeaderTransport::Stream => { + let key = load_env_var(env)?; + eyre::ensure!(!key.is_empty(), "{env} is empty"); + + let mut value = HeaderValue::from_str(&key) + .map_err(|_| eyre::eyre!("{env} is not a valid header value"))?; + value.set_sensitive(true); + stream_headers.insert(HEADER_API_KEY, value); + } + GetHeaderTransport::Http => {} + } + } + let client = reqwest::Client::builder() .default_headers(headers) .timeout(DEFAULT_REQUEST_TIMEOUT) diff --git a/tests/src/mock_ws_relay.rs b/tests/src/mock_ws_relay.rs index e1b46103..34470fe6 100644 --- a/tests/src/mock_ws_relay.rs +++ b/tests/src/mock_ws_relay.rs @@ -10,8 +10,8 @@ use std::{ use alloy::primitives::{B256, U256}; use cb_common::{ pbs::{ - BuilderApiVersion, GET_HEADER_STREAM_PATH, HEADER_START_TIME_UNIX_MS, HEADER_TIMEOUT_MS, - HEADER_VERSION_KEY, + BuilderApiVersion, GET_HEADER_STREAM_PATH, HEADER_API_KEY, HEADER_START_TIME_UNIX_MS, + HEADER_TIMEOUT_MS, HEADER_VERSION_KEY, }, types::{BlsSecretKey, Chain}, }; @@ -42,6 +42,7 @@ pub struct StreamRequest { pub start_time_ms: Option, pub user_agent: Option, pub cb_version: Option, + pub api_key: Option, } pub struct MockWsRelayState { @@ -182,6 +183,7 @@ fn parse_request(req: &Request) -> Option { start_time_ms: header(req, HEADER_START_TIME_UNIX_MS).and_then(|v| v.parse().ok()), user_agent: header(req, "user-agent"), cb_version: header(req, HEADER_VERSION_KEY), + api_key: header(req, HEADER_API_KEY), }) } diff --git a/tests/src/utils.rs b/tests/src/utils.rs index 0559877a..6b5e4267 100644 --- a/tests/src/utils.rs +++ b/tests/src/utils.rs @@ -50,6 +50,7 @@ fn mock_relay_config(port: u16, pubkey: BlsPublicKey) -> Result { headers: None, get_params: None, get_header: GetHeaderTransport::Http, + api_key_env: None, enable_timing_games: false, target_first_request_ms: None, frequency_get_header_ms: None, @@ -77,6 +78,17 @@ pub fn generate_mock_stream_relay(port: u16, pubkey: BlsPublicKey) -> Result Result { + let mut config = mock_relay_config(port, pubkey)?; + config.get_header = GetHeaderTransport::Stream; + config.api_key_env = Some(api_key_env.to_string()); + RelayClient::new(config) +} + pub fn get_pbs_config(port: u16) -> PbsConfig { PbsConfig { host: Ipv4Addr::UNSPECIFIED, diff --git a/tests/tests/pbs_cfg_file_update.rs b/tests/tests/pbs_cfg_file_update.rs index c576a6a6..21876362 100644 --- a/tests/tests/pbs_cfg_file_update.rs +++ b/tests/tests/pbs_cfg_file_update.rs @@ -93,6 +93,7 @@ async fn test_cfg_file_update() -> Result<()> { frequency_get_header_ms: None, get_params: None, get_header: GetHeaderTransport::Http, + api_key_env: None, headers: None, target_first_request_ms: None, validator_registration_batch_size: None, @@ -147,6 +148,7 @@ async fn test_cfg_file_update() -> Result<()> { frequency_get_header_ms: None, get_params: None, get_header: GetHeaderTransport::Http, + api_key_env: None, headers: None, target_first_request_ms: None, validator_registration_batch_size: None, diff --git a/tests/tests/pbs_get_header_ws.rs b/tests/tests/pbs_get_header_ws.rs index 7f5392c6..12b1cec2 100644 --- a/tests/tests/pbs_get_header_ws.rs +++ b/tests/tests/pbs_get_header_ws.rs @@ -19,8 +19,8 @@ use cb_tests::{ mock_validator::MockValidator, mock_ws_relay::{MockWsRelayState, start_mock_ws_relay_service}, utils::{ - generate_mock_relay, generate_mock_stream_relay, get_free_listener, get_pbs_config, - setup_test_env, to_pbs_config, + generate_mock_relay, generate_mock_stream_relay, generate_mock_stream_relay_with_api_key, + get_free_listener, get_pbs_config, setup_test_env, to_pbs_config, }, }; use eyre::Result; @@ -156,6 +156,43 @@ async fn test_get_header_ws_skips_unknown_frames() -> Result<()> { Ok(()) } +#[tokio::test] +async fn test_get_header_ws_sends_api_key() -> Result<()> { + setup_test_env(); + let signer = random_secret(); + let pubkey = signer.public_key(); + let chain = Chain::Hoodi; + + let listener = get_free_listener().await; + let port = listener.local_addr()?.port(); + let relay_state = Arc::new(MockWsRelayState::new(chain, signer)); + tokio::spawn(start_mock_ws_relay_service(relay_state.clone(), listener)); + + // SAFETY: single-threaded setup, before PBS reads it + unsafe { std::env::set_var("TEST_RELAY_API_KEY", "secret-key") }; + let relay = generate_mock_stream_relay_with_api_key(port, pubkey, "TEST_RELAY_API_KEY")?; + + let validator = start_pbs(chain, vec![relay], 1_000).await?; + + let (code, _) = get_header_json(&validator).await?; + assert_eq!(code, StatusCode::OK); + + let request = relay_state.last_request().expect("relay saw no request"); + assert_eq!(request.api_key.as_deref(), Some("secret-key")); + Ok(()) +} + +#[tokio::test] +async fn test_stream_relay_requires_api_key_env() -> Result<()> { + setup_test_env(); + let pubkey = random_secret().public_key(); + + assert!( + generate_mock_stream_relay_with_api_key(1234, pubkey, "TEST_RELAY_API_KEY_UNSET").is_err() + ); + Ok(()) +} + /// The handshake carries the same request data as the HTTP call #[tokio::test] async fn test_get_header_ws_handshake_carries_request() -> Result<()> {