diff --git a/Cargo.lock b/Cargo.lock index 1f052de5..27645840 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1840,16 +1840,19 @@ 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", "types", "url", "uuid 1.23.1", + "webpki-roots 1.0.7", ] [[package]] @@ -1895,6 +1898,7 @@ dependencies = [ "eth2", "ethereum_ssz 0.10.3", "eyre", + "futures", "jsonwebtoken", "rcgen", "reqwest 0.13.2", @@ -1902,6 +1906,7 @@ dependencies = [ "serde_json", "tempfile", "tokio", + "tokio-tungstenite", "toml", "tracing", "tracing-subscriber", diff --git a/Cargo.toml b/Cargo.toml index 170cb38a..6b68e99c 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" @@ -93,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/benches/pbs/src/main.rs b/benches/pbs/src/main.rs index 82f72b5c..aca59054 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,8 @@ fn get_mock_validator(bench: BenchConfig) -> RelayClient { id: None, 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 6804faad..1cd29586 100644 --- a/config.example.toml +++ b/config.example.toml @@ -92,6 +92,18 @@ 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" -> 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" +# 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 4bb1cff9..79e8779d 100644 --- a/crates/common/src/config/pbs.rs +++ b/crates/common/src/config/pbs.rs @@ -38,6 +38,14 @@ use crate::{ }, }; +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum GetHeaderTransport { + #[default] + Http, + Stream, +} + #[derive(Debug, Clone, Deserialize, Serialize)] #[serde(deny_unknown_fields)] pub struct RelayConfig { @@ -50,6 +58,12 @@ 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, + /// 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 bbe20b0d..66ab42f0 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"; @@ -17,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/error.rs b/crates/common/src/pbs/error.rs index 0be1a40b..a5cc19a4 100644 --- a/crates/common/src/pbs/error.rs +++ b/crates/common/src/pbs/error.rs @@ -38,11 +38,21 @@ pub enum PbsError { #[error("SSZ error: {0}")] SszError(#[from] SszValueError), + + #[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 59d963e1..8e43cbe8 100644 --- a/crates/common/src/pbs/relay.rs +++ b/crates/common/src/pbs/relay.rs @@ -7,12 +7,17 @@ use serde::{Deserialize, Serialize}; use url::Url; use super::{ - HEADER_VERSION_KEY, HEADER_VERSION_VALUE, - constants::{GET_STATUS_PATH, REGISTER_VALIDATOR_PATH, SUBMIT_BLOCK_PATH}, + HEADER_API_KEY, HEADER_VERSION_KEY, HEADER_VERSION_VALUE, + constants::{ + GET_HEADER_STREAM_PATH, GET_STATUS_PATH, REGISTER_VALIDATOR_PATH, SUBMIT_BLOCK_PATH, + }, error::PbsError, }; use crate::{ - DEFAULT_REQUEST_TIMEOUT, config::RelayConfig, pbs::BuilderApiVersion, types::BlsPublicKey, + DEFAULT_REQUEST_TIMEOUT, + config::{GetHeaderTransport, RelayConfig, load_env_var}, + pbs::BuilderApiVersion, + types::BlsPublicKey, }; /// A parsed entry of the relay url in the format: scheme://pubkey@host @@ -49,6 +54,28 @@ impl<'de> Deserialize<'de> for RelayEntry { } } +#[derive(Debug, Clone)] +pub enum GetHeaderRequest { + Http(Url), + 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)] @@ -57,12 +84,21 @@ 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 pub config: Arc, } impl RelayClient { pub fn new(config: RelayConfig) -> eyre::Result { + 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)); @@ -75,31 +111,64 @@ impl RelayClient { } } + 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) .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_url, + 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 { let mut url = self.config.entry.url.join(path).map_err(PbsError::UrlParsing)?; + self.append_get_params(&mut url); + Ok(url) + } + 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, @@ -120,6 +189,26 @@ impl RelayClient { ) } + pub fn get_header_request( + &self, + slot: u64, + parent_hash: &B256, + validator_pubkey: &BlsPublicKey, + ) -> Result { + Ok(match &self.stream_url { + None => { + GetHeaderRequest::Http(self.get_header_url(slot, parent_hash, validator_pubkey)?) + } + Some(base) => { + let mut url = base.clone(); + url.set_path(&format!("{}/{slot}/{parent_hash}/{validator_pubkey}", base.path())); + + self.append_get_params(&mut url); + GetHeaderRequest::Stream(url) + } + }) + } + pub fn get_status_url(&self) -> Result { self.builder_api_url(GET_STATUS_PATH, BuilderApiVersion::V1) } @@ -139,8 +228,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 +318,107 @@ 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(); + + // 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() + 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" + ); + + // 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}" + ) + ); + + // https relays stream over wss, and a port and get_params carry over + let mut config = base_config.clone(); + 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() + else { + panic!("expected stream request"); + }; + assert_eq!( + url.to_string(), + format!( + "wss://abc.xyz:4444/eth/v1/builder/header_stream/{slot}/{parent_hash}/{validator_pubkey}?token=abc" + ) + ); + + // A relay url we can't stream over is rejected at construction + let mut config = base_config; + config.entry.url = "unix:/tmp/relay.sock".parse().unwrap(); + config.get_header = GetHeaderTransport::Stream; + 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""#).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" + "#; + let config = toml::from_str::(toml_config).unwrap(); + assert_eq!(config.get_header, GetHeaderTransport::Stream); + + // 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..7497aa93 100644 --- a/crates/pbs/Cargo.toml +++ b/crates/pbs/Cargo.toml @@ -23,12 +23,15 @@ 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 url.workspace = true uuid.workspace = true +webpki-roots.workspace = true thiserror.workspace = true 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 dd2b1036..b9d9767f 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,16 @@ 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))) +} + +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 +448,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 +457,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. @@ -580,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 new file mode 100644 index 00000000..368d9192 --- /dev/null +++ b/crates/pbs/src/mev_boost/get_header_ws.rs @@ -0,0 +1,327 @@ +//! Streaming get_header over a websocket, for relays configured with +//! `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. +//! The relay replies with one binary frame per bid update: +//! +//! ```text +//! u8 message type +//! u8 fork +//! .. SSZ SignedBuilderBid +//! ``` + +use std::{ + sync::{Arc, OnceLock}, + time::Duration, +}; + +use alloy::primitives::utils::format_ether; +use axum::http::{HeaderValue, Request, header::USER_AGENT}; +use cb_common::{ + pbs::{ + ForkName, GetHeaderInfo, GetHeaderResponse, HEADER_START_TIME_UNIX_MS, HEADER_TIMEOUT_MS, + RelayClient, error::PbsError, + }, + utils::utcnow_ms, +}; +use futures::StreamExt; +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, 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, + TRANSPORT_ERROR_CODE_STR, + }, + metrics::{RELAY_LATENCY, RELAY_STATUS_CODE}, + mev_boost::get_header::decode_ssz_payload, +}; + +/// 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)); + + let start_request = Instant::now(); + 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)) => { + 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(_) => { + record_status(TIMEOUT_ERROR_CODE_STR, relay); + return Err(PbsError::WebSocketTimeout); + } + }; + 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: Option = None; + let mut updates = 0usize; + let mut invalid_frames = 0usize; + let mut stream_error = None; + + 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"); + stream_error = Some(PbsError::WebSocket(format!("stream error: {err}"))); + break; + } + None => break, + }; + + let payload = match message { + Message::Binary(payload) => payload, + Message::Close(_) => break, + _ => continue, + }; + + match parse_frame(payload) { + Ok((fork, bid)) => { + updates += 1; + if first_bid_latency.is_none() { + first_bid_latency = Some(start_request.elapsed()); + } + latest = Some((fork, bid)); + } + Err(err) => { + invalid_frames += 1; + if invalid_frames == 1 { + warn!(relay_id = relay.id.as_ref(), %err, "invalid ws frame, skipping"); + } + } + } + } + + drop(stream); + + let Some((fork, bid_bytes)) = latest else { + if let Some(err) = stream_error { + record_status(TRANSPORT_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); + }; + + 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()); + } + + let response = decode_ssz_payload(&bid_bytes, fork)?; + + 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(), + ?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" + ); + + validated?; + + record_status("200", relay); + + Ok(Some(response)) +} + +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()); + } + + for (key, value) in relay.stream_headers() { + headers.insert(key, value.clone()); + } + + 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(); +} + +/// 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)] +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 decoded = decode_ssz_payload(&bid_bytes, fork).unwrap(); + assert_eq!(fork, ForkName::Fulu); + assert_eq!(decoded.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/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..34470fe6 --- /dev/null +++ b/tests/src/mock_ws_relay.rs @@ -0,0 +1,192 @@ +use std::{ + str::FromStr, + sync::{ + Arc, Mutex, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, +}; + +use alloy::primitives::{B256, U256}; +use cb_common::{ + pbs::{ + BuilderApiVersion, GET_HEADER_STREAM_PATH, HEADER_API_KEY, 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 api_key: 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, + /// 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>, +} + +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, + unknown_frames: 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 with_unknown_frames(self) -> Self { + Self { unknown_frames: 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 { + 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, + 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: 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 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(); + + 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), + api_key: header(req, HEADER_API_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 9bc0e362..6b5e4267 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, @@ -39,20 +39,27 @@ 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, + get_header: GetHeaderTransport::Http, + api_key_env: None, enable_timing_games: false, target_first_request_ms: None, frequency_get_header_ms: None, validator_registration_batch_size: None, - }; - RelayClient::new(config) + }) +} + +pub fn generate_mock_relay(port: u16, pubkey: BlsPublicKey) -> Result { + RelayClient::new(mock_relay_config(port, pubkey)?) } pub fn generate_mock_relay_with_batch_size( @@ -60,18 +67,25 @@ 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, - 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, pubkey: BlsPublicKey) -> Result { + let mut config = mock_relay_config(port, pubkey)?; + config.get_header = GetHeaderTransport::Stream; + RelayClient::new(config) +} + +pub fn generate_mock_stream_relay_with_api_key( + port: u16, + pubkey: BlsPublicKey, + api_key_env: &str, +) -> 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) } diff --git a/tests/tests/pbs_cfg_file_update.rs b/tests/tests/pbs_cfg_file_update.rs index 37dd4eb3..21876362 100644 --- a/tests/tests/pbs_cfg_file_update.rs +++ b/tests/tests/pbs_cfg_file_update.rs @@ -2,7 +2,10 @@ use std::{net::Ipv4Addr, sync::Arc, time::Duration}; use alloy::primitives::U256; use cb_common::{ - config::{CommitBoostConfig, LogsSettings, PbsConfig, RelayConfig, StaticPbsConfig}, + config::{ + CommitBoostConfig, GetHeaderTransport, LogsSettings, PbsConfig, RelayConfig, + StaticPbsConfig, + }, pbs::RelayEntry, signer::random_secret, types::Chain, @@ -89,6 +92,8 @@ async fn test_cfg_file_update() -> Result<()> { enable_timing_games: false, 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, @@ -142,6 +147,8 @@ async fn test_cfg_file_update() -> Result<()> { enable_timing_games: false, 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 new file mode 100644 index 00000000..12b1cec2 --- /dev/null +++ b/tests/tests/pbs_get_header_ws.rs @@ -0,0 +1,337 @@ +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, generate_mock_stream_relay_with_api_key, + 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; + +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)); + + Ok((state, generate_mock_stream_relay(port, pubkey)?)) +} + +/// 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(()) +} + +/// 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(()) +} + +#[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<()> { + 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 relay = generate_mock_stream_relay(port, 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(()) +}