diff --git a/cli/src/access.rs b/cli/src/access.rs new file mode 100644 index 0000000..922c413 --- /dev/null +++ b/cli/src/access.rs @@ -0,0 +1,57 @@ +// Copyright (c) 2023-2024 Cloudflare, Inc. +// Licensed under the BSD-3-Clause license found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +//! Cloudflare Access support for testing gated endpoints. +//! +//! When the `_MACH_CF_ACCESS_TOKEN` environment variable is set, mach attaches a +//! `cf-access-token` header to its test requests so that Cloudflare +//! Access-gated hosts can be reached. The token is only ever attached to hosts +//! on the allowlist below so that it is never leaked to an unrelated origin +//! returned by a remote configuration. + +use anyhow::Context; +use http::{HeaderMap, HeaderValue}; +use nq_core::ScopedHeaders; + +/// Environment variable holding the Cloudflare Access token. +const ACCESS_TOKEN_ENV: &str = "_MACH_CF_ACCESS_TOKEN"; + +/// Header used to carry the Cloudflare Access token. +const ACCESS_TOKEN_HEADER: &str = "cf-access-token"; + +/// Domain suffixes the access token may be sent to. Matching is done with +/// strict label boundaries by [`ScopedHeaders`]. +const ALLOWED_HOST_SUFFIXES: &[&str] = &["speed.cloudflare.com"]; + +/// Build the scoped headers carrying the Cloudflare Access token, if the +/// `_MACH_CF_ACCESS_TOKEN` environment variable is set to a non-empty value. +/// +/// Returns `Ok(None)` when the variable is unset or blank. Returns an error +/// when the token is not a valid header value, so a misconfigured token fails +/// fast at startup rather than silently sending no header. +pub fn cf_access_scoped_headers() -> anyhow::Result> { + let Ok(token) = std::env::var(ACCESS_TOKEN_ENV) else { + return Ok(None); + }; + + let token = token.trim(); + if token.is_empty() { + return Ok(None); + } + + let mut value = HeaderValue::from_str(token) + .with_context(|| format!("{ACCESS_TOKEN_ENV} is not a valid header value"))?; + // Mark the value sensitive so it is redacted (printed as `Sensitive`) in + // request/response debug logs. + value.set_sensitive(true); + + let mut headers = HeaderMap::new(); + headers.insert(ACCESS_TOKEN_HEADER, value); + + let allowed_host_suffixes = ALLOWED_HOST_SUFFIXES + .iter() + .map(|s| s.to_string()) + .collect(); + + Ok(Some(ScopedHeaders::new(headers, allowed_host_suffixes))) +} diff --git a/cli/src/latency.rs b/cli/src/latency.rs index 144e9ee..0c93961 100644 --- a/cli/src/latency.rs +++ b/cli/src/latency.rs @@ -22,6 +22,7 @@ pub async fn run(url: String, runs: usize) -> anyhow::Result<()> { let result = run_test(&LatencyConfig { url: url.parse()?, runs, + scoped_headers: crate::access::cf_access_scoped_headers()?, }) .await?; diff --git a/cli/src/main.rs b/cli/src/main.rs index 039c258..ada1cbb 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -1,6 +1,7 @@ // Copyright (c) 2023-2024 Cloudflare, Inc. // Licensed under the BSD-3-Clause license found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause +mod access; mod aim_report; pub(crate) mod args; mod latency; diff --git a/cli/src/packet_loss.rs b/cli/src/packet_loss.rs index ce3eb81..4db1a41 100644 --- a/cli/src/packet_loss.rs +++ b/cli/src/packet_loss.rs @@ -26,6 +26,7 @@ pub async fn run(args: PacketLossArgs) -> anyhow::Result<()> { response_wait_time: Duration::from_millis(args.response_wait_time_ms), download_url: args.download_url.parse()?, upload_url: args.upload_url.parse()?, + scoped_headers: crate::access::cf_access_scoped_headers()?, }; info!("fetching TURN server credentials"); @@ -58,10 +59,13 @@ async fn fetch_turn_server_creds( let mut headers = HeaderMap::new(); headers.append(hyper::header::HOST, HeaderValue::from_str(host)?); - let response = Client::default() + let client = Client::default() .new_connection(ConnectionType::h1()) .method("GET") .headers(headers) + .scoped_headers(config.scoped_headers.clone()); + + let response = client .send( request_url.to_string().parse()?, http_body_util::Empty::new(), diff --git a/cli/src/rpm.rs b/cli/src/rpm.rs index 5acb740..5fb77b8 100644 --- a/cli/src/rpm.rs +++ b/cli/src/rpm.rs @@ -25,10 +25,12 @@ use crate::util::pretty_secs_to_ms; pub async fn run(cli_config: RpmArgs) -> anyhow::Result<()> { info!("running responsiveness test"); + let scoped_headers = crate::access::cf_access_scoped_headers()?; + let rpm_urls = match cli_config.config.clone() { Some(endpoint) => { info!("fetching configuration from {endpoint}"); - let urls = get_rpm_config(endpoint).await?.urls; + let urls = get_rpm_config(endpoint, scoped_headers.clone()).await?.urls; info!("retrieved configuration urls: {urls:?}"); urls @@ -53,6 +55,7 @@ pub async fn run(cli_config: RpmArgs) -> anyhow::Result<()> { let rtt_result = crate::latency::run_test(&LatencyConfig { url: rpm_urls.small_https_download_url.parse()?, runs: 20, + scoped_headers: scoped_headers.clone(), }) .await?; info!( @@ -79,6 +82,7 @@ pub async fn run(cli_config: RpmArgs) -> anyhow::Result<()> { max_loaded_connections: cli_config.max_loaded_connections, conn_type: ConnectionType::H2, determine_load_only: false, + scoped_headers, }; info!("running download test"); @@ -156,7 +160,10 @@ pub struct RpmUrls { https_upload_url: String, } -pub async fn get_rpm_config(config_url: String) -> anyhow::Result { +pub async fn get_rpm_config( + config_url: String, + scoped_headers: Option, +) -> anyhow::Result { let shutdown = CancellationToken::new(); let time = Arc::new(TokioTime::new()); let network = Arc::new(TokioNetwork::new( @@ -164,9 +171,12 @@ pub async fn get_rpm_config(config_url: String) -> anyhow::Result anyhow::Result<()> { determine_load_only: true, conn_type: cli_config.conn_type.into(), test_duration: Duration::from_secs(cli_config.duration), + scoped_headers: crate::access::cf_access_scoped_headers()?, ..Default::default() }; diff --git a/cli/src/up_down.rs b/cli/src/up_down.rs index 178b363..ec2bdba 100644 --- a/cli/src/up_down.rs +++ b/cli/src/up_down.rs @@ -26,8 +26,11 @@ pub async fn download(args: DownloadArgs) -> anyhow::Result<()> { let conn_type = args.conn_type.into(); info!("downloading: {}", args.url); - let inflight_body = ThroughputClient::download() + let client = ThroughputClient::download() .new_connection(conn_type) + .scoped_headers(crate::access::cf_access_scoped_headers()?); + + let inflight_body = client .send( args.url.parse().context("parsing download url")?, Arc::clone(&network), @@ -107,9 +110,12 @@ pub async fn upload(args: UploadArgs) -> anyhow::Result<()> { info!("added header: {key}"); } - let inflight_body = ThroughputClient::upload(bytes) + let client = ThroughputClient::upload(bytes) .new_connection(conn_type) .headers(headers) + .scoped_headers(crate::access::cf_access_scoped_headers()?); + + let inflight_body = client .send( args.url.parse()?, Arc::clone(&network), diff --git a/crates/nq-core/src/client.rs b/crates/nq-core/src/client.rs index 27cc1ff..519be61 100644 --- a/crates/nq-core/src/client.rs +++ b/crates/nq-core/src/client.rs @@ -19,7 +19,7 @@ use tokio_util::sync::CancellationToken; use tracing::{Instrument, debug, error, info, trace}; use crate::{ - ConnectionType, EstablishedConnection, Network, OneshotResult, Time, Timestamp, + ConnectionType, EstablishedConnection, Network, OneshotResult, ScopedHeaders, Time, Timestamp, body::{BodyEvent, CountingBody, InflightBody, NqBody, UploadBody, empty}, oneshot_result, }; @@ -49,6 +49,7 @@ pub struct ThroughputClient { connection: Option>>, new_connection_type: Option, headers: Option, + scoped_headers: Option, direction: Direction, } @@ -59,6 +60,7 @@ impl ThroughputClient { connection: None, new_connection_type: None, headers: None, + scoped_headers: None, direction: Direction::Down, } } @@ -69,6 +71,7 @@ impl ThroughputClient { connection: None, new_connection_type: None, headers: None, + scoped_headers: None, direction: Direction::Up(size), } } @@ -91,6 +94,13 @@ impl ThroughputClient { self } + /// Set headers that are only attached when the request's host matches the + /// scope's allowlist. Headers set via [`Self::headers`] take precedence. + pub fn scoped_headers(mut self, scoped_headers: Option) -> Self { + self.scoped_headers = scoped_headers; + self + } + /// Execute a download or upload request against the given [`Uri`]. // #[tracing::instrument(skip(self, network, time, shutdown))] pub fn send( @@ -148,6 +158,10 @@ impl ThroughputClient { } }; + if let Some(scoped_headers) = self.scoped_headers.take() { + scoped_headers.apply(&uri, &mut headers); + } + let mut request = http::Request::builder() .method(method) .uri(uri) @@ -355,6 +369,7 @@ pub struct Client { connection: Option>>, new_connection_type: Option, headers: Option, + scoped_headers: Option, method: Option, } @@ -377,6 +392,13 @@ impl Client { self } + /// Set headers that are only attached when the request's host matches the + /// scope's allowlist. Headers set via [`Self::headers`] take precedence. + pub fn scoped_headers(mut self, scoped_headers: Option) -> Self { + self.scoped_headers = scoped_headers; + self + } + /// Set the method used by the client. pub fn method(mut self, method: &str) -> Self { self.method = Some(method.to_string()); @@ -411,6 +433,10 @@ impl Client { let method: http::Method = self.method.as_deref().unwrap_or("GET").parse()?; + if let Some(scoped_headers) = self.scoped_headers { + scoped_headers.apply(&uri, &mut headers); + } + let mut request = http::Request::builder() .method(method) .uri(uri) diff --git a/crates/nq-core/src/lib.rs b/crates/nq-core/src/lib.rs index 40935e3..55a615b 100644 --- a/crates/nq-core/src/lib.rs +++ b/crates/nq-core/src/lib.rs @@ -13,6 +13,7 @@ mod body; pub mod client; mod connection; mod network; +mod scoped_headers; mod time; mod upgraded; mod util; @@ -21,6 +22,7 @@ pub use crate::{ body::{BodyEvent, CountingBody, NqBody}, connection::{ConnectionManager, ConnectionTiming, ConnectionType, EstablishedConnection}, network::Network, + scoped_headers::ScopedHeaders, time::{Time, Timestamp, TokioTime}, upgraded::ConnectUpgraded, util::{OneshotResult, ResponseFuture, oneshot_result}, diff --git a/crates/nq-core/src/scoped_headers.rs b/crates/nq-core/src/scoped_headers.rs new file mode 100644 index 0000000..0b52f16 --- /dev/null +++ b/crates/nq-core/src/scoped_headers.rs @@ -0,0 +1,230 @@ +// Copyright (c) 2023-2024 Cloudflare, Inc. +// Licensed under the BSD-3-Clause license found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +//! Headers that are only attached to requests whose host matches an allowlist. +//! +//! This is a generic mechanism: callers supply the headers and the list of +//! allowed domain suffixes. It is used to attach sensitive headers (such as an +//! access token) to trusted hosts only, so the header is never sent to an +//! unrelated origin returned by a remote configuration. + +use http::{HeaderMap, Uri}; + +/// A set of headers attached to a request only when the request's host matches +/// one of the allowed domain suffixes. +/// +/// Matching uses strict label boundaries: a host matches a suffix when it is +/// equal to the suffix or ends with `.` followed by the suffix. This prevents +/// look-alike hosts such as `example.com.evil.com` from matching `example.com`. +/// +/// Existing headers are never overwritten, so a header set explicitly by the +/// caller takes precedence over a scoped header with the same name. +#[derive(Debug, Clone, Default)] +pub struct ScopedHeaders { + headers: HeaderMap, + allowed_host_suffixes: Vec, +} + +impl ScopedHeaders { + /// Create a new [`ScopedHeaders`] from the headers to attach and the list + /// of allowed domain suffixes. + /// + /// Suffixes are normalized (lowercased and stripped of any trailing dot) so + /// matching in [`allows_host`](Self::allows_host) is case-insensitive and + /// tolerant of fully-qualified (trailing-dot) hostnames. + pub fn new(headers: HeaderMap, allowed_host_suffixes: Vec) -> Self { + let allowed_host_suffixes = allowed_host_suffixes + .into_iter() + .map(|s| s.trim_end_matches('.').to_ascii_lowercase()) + .collect(); + + Self { + headers, + allowed_host_suffixes, + } + } + + /// Returns whether the given host is covered by an allowed suffix, using + /// strict label-boundary matching. + /// + /// Matching is case-insensitive (DNS hostnames are case-insensitive) and + /// ignores a trailing dot on the host (fully-qualified domain names). + pub fn allows_host(&self, host: &str) -> bool { + let host = host.trim_end_matches('.').to_ascii_lowercase(); + self.allowed_host_suffixes.iter().any(|suffix| { + host == *suffix + || (host.len() > suffix.len() + && host.ends_with(suffix) + && host.as_bytes()[host.len() - suffix.len() - 1] == b'.') + }) + } + + /// Attach the scoped headers to `headers` when the URI's host matches an + /// allowed suffix. Headers already present in `headers` are left untouched. + pub fn apply(&self, uri: &Uri, headers: &mut HeaderMap) { + let Some(host) = uri.host() else { + return; + }; + + if !self.allows_host(host) { + return; + } + + for (name, value) in self.headers.iter() { + if !headers.contains_key(name) { + headers.append(name, value.clone()); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use http::{HeaderName, HeaderValue}; + + const TEST_HEADER: &str = "x-test-token"; + const SUFFIX: &str = "speed.cloudflare.com"; + + /// Build a [`ScopedHeaders`] carrying a single `x-test-token` header scoped + /// to the given suffixes. + fn scoped(suffixes: &[&str]) -> ScopedHeaders { + let mut headers = HeaderMap::new(); + headers.insert( + HeaderName::from_static(TEST_HEADER), + HeaderValue::from_static("secret"), + ); + ScopedHeaders::new( + headers, + suffixes.iter().map(|s| s.to_string()).collect(), + ) + } + + // --- allows_host: matching --- + + #[test] + fn allows_exact_match() { + assert!(scoped(&[SUFFIX]).allows_host("speed.cloudflare.com")); + } + + #[test] + fn allows_subdomain_match() { + assert!(scoped(&[SUFFIX]).allows_host("staging.speed.cloudflare.com")); + } + + #[test] + fn allows_deep_subdomain_match() { + assert!(scoped(&[SUFFIX]).allows_host("a.b.speed.cloudflare.com")); + } + + #[test] + fn allows_any_of_multiple_suffixes() { + let scoped = scoped(&["example.com", SUFFIX]); + assert!(scoped.allows_host("speed.cloudflare.com")); + assert!(scoped.allows_host("www.example.com")); + } + + // --- allows_host: case insensitivity & trailing dot --- + + #[test] + fn allows_case_insensitive_host() { + assert!(scoped(&[SUFFIX]).allows_host("SPEED.Cloudflare.COM")); + } + + #[test] + fn allows_trailing_dot_fqdn() { + assert!(scoped(&[SUFFIX]).allows_host("speed.cloudflare.com.")); + } + + #[test] + fn allows_trailing_dot_subdomain() { + assert!(scoped(&[SUFFIX]).allows_host("staging.speed.cloudflare.com.")); + } + + #[test] + fn allows_suffix_normalized_case_and_trailing_dot() { + // Suffix supplied with mixed case and a trailing dot must still match a + // lowercase host, verifying normalization in `new`. + assert!(scoped(&["SPEED.Cloudflare.com."]).allows_host("speed.cloudflare.com")); + } + + // --- allows_host: rejection (security boundary) --- + + #[test] + fn rejects_look_alike_prefix() { + assert!(!scoped(&[SUFFIX]).allows_host("notspeed.cloudflare.com")); + } + + #[test] + fn rejects_look_alike_suffix() { + assert!(!scoped(&[SUFFIX]).allows_host("speed.cloudflare.com.evil.com")); + } + + #[test] + fn rejects_empty_host() { + assert!(!scoped(&[SUFFIX]).allows_host("")); + } + + #[test] + fn rejects_unrelated_host() { + assert!(!scoped(&[SUFFIX]).allows_host("example.com")); + } + + #[test] + fn rejects_host_shorter_than_suffix() { + assert!(!scoped(&[SUFFIX]).allows_host("cloudflare.com")); + } + + #[test] + fn rejects_partial_trailing_label() { + // Host ends with the suffix bytes but not on a label boundary. + assert!(!scoped(&[SUFFIX]).allows_host("xspeed.cloudflare.com")); + } + + // --- apply --- + + #[test] + fn apply_attaches_header_to_matching_host() { + let uri: Uri = "https://staging.speed.cloudflare.com/config" + .parse() + .unwrap(); + let mut headers = HeaderMap::new(); + scoped(&[SUFFIX]).apply(&uri, &mut headers); + assert_eq!( + headers.get(TEST_HEADER).map(|v| v.as_bytes()), + Some(b"secret".as_slice()) + ); + } + + #[test] + fn apply_does_not_attach_header_to_non_matching_host() { + let uri: Uri = "https://evil.com/config".parse().unwrap(); + let mut headers = HeaderMap::new(); + scoped(&[SUFFIX]).apply(&uri, &mut headers); + assert!(!headers.contains_key(TEST_HEADER)); + } + + #[test] + fn apply_is_noop_when_uri_has_no_host() { + let uri: Uri = "/relative/path".parse().unwrap(); + assert!(uri.host().is_none()); + let mut headers = HeaderMap::new(); + scoped(&[SUFFIX]).apply(&uri, &mut headers); + assert!(headers.is_empty()); + } + + #[test] + fn apply_does_not_overwrite_existing_header() { + let uri: Uri = "https://speed.cloudflare.com/config".parse().unwrap(); + let mut headers = HeaderMap::new(); + headers.insert( + HeaderName::from_static(TEST_HEADER), + HeaderValue::from_static("explicit"), + ); + scoped(&[SUFFIX]).apply(&uri, &mut headers); + assert_eq!( + headers.get(TEST_HEADER).map(|v| v.as_bytes()), + Some(b"explicit".as_slice()) + ); + } +} diff --git a/crates/nq-latency/src/lib.rs b/crates/nq-latency/src/lib.rs index 7a51cab..00bd25e 100644 --- a/crates/nq-latency/src/lib.rs +++ b/crates/nq-latency/src/lib.rs @@ -7,7 +7,7 @@ use anyhow::Context; use http::Request; use http_body_util::BodyExt; use nq_core::Network; -use nq_core::{ConnectionType, Time, Timestamp}; +use nq_core::{ConnectionType, ScopedHeaders, Time, Timestamp}; use nq_stats::TimeSeries; use tokio_util::sync::CancellationToken; use tracing::info; @@ -17,6 +17,9 @@ use url::Url; pub struct LatencyConfig { pub url: Url, pub runs: usize, + /// Headers attached only to requests whose host matches the scope's + /// allowlist. + pub scoped_headers: Option, } impl Default for LatencyConfig { @@ -26,6 +29,7 @@ impl Default for LatencyConfig { .parse() .unwrap(), runs: 20, + scoped_headers: None, } } } @@ -94,11 +98,14 @@ impl Latency { ); // perform a simple GET to do some amount of work + let mut request = Request::get(url.as_str()).body(Default::default())?; + if let Some(scoped_headers) = &self.config.scoped_headers { + let request_uri = request.uri().clone(); + scoped_headers.apply(&request_uri, request.headers_mut()); + } + let response = network - .send_request( - connection, - Request::get(url.as_str()).body(Default::default())?, - ) + .send_request(connection, request) .await .context("GET request failed")?; diff --git a/crates/nq-load-generator/src/lib.rs b/crates/nq-load-generator/src/lib.rs index a487142..07fcbc1 100644 --- a/crates/nq-load-generator/src/lib.rs +++ b/crates/nq-load-generator/src/lib.rs @@ -7,8 +7,8 @@ use anyhow::Context; use http::{HeaderMap, HeaderName, HeaderValue}; use nq_core::client::{Direction, ThroughputClient}; use nq_core::{ - BodyEvent, ConnectionType, EstablishedConnection, Network, OneshotResult, Time, Timestamp, - oneshot_result, + BodyEvent, ConnectionType, EstablishedConnection, Network, OneshotResult, ScopedHeaders, Time, + Timestamp, oneshot_result, }; use nq_stats::CounterSeries; use rand::seq::SliceRandom; @@ -21,6 +21,10 @@ use tracing::Instrument; #[derive(Debug, Deserialize)] pub struct LoadConfig { pub headers: HashMap, + /// Headers attached only to requests whose host matches the scope's + /// allowlist. + #[serde(skip)] + pub scoped_headers: Option, pub download_url: url::Url, pub upload_url: url::Url, pub upload_size: usize, @@ -28,6 +32,7 @@ pub struct LoadConfig { pub struct LoadGenerator { headers: HeaderMap, + scoped_headers: Option, config: LoadConfig, loads: Vec, } @@ -45,6 +50,7 @@ impl LoadGenerator { Ok(Self { headers, + scoped_headers: config.scoped_headers.clone(), config, loads: Vec::new(), }) @@ -66,18 +72,20 @@ impl LoadGenerator { Direction::Up(size) => ThroughputClient::upload(size), }; - let response_fut = client + let client = client .new_connection(conn_type) .headers(self.headers.clone()) - .send( - match direction { - Direction::Up(_) => self.config.upload_url.as_str().parse()?, - Direction::Down => self.config.download_url.as_str().parse()?, - }, - network, - time, - shutdown, - )?; + .scoped_headers(self.scoped_headers.clone()); + + let response_fut = client.send( + match direction { + Direction::Up(_) => self.config.upload_url.as_str().parse()?, + Direction::Down => self.config.download_url.as_str().parse()?, + }, + network, + time, + shutdown, + )?; tracing::debug!("got loaded connection response future"); diff --git a/crates/nq-packetloss/src/lib.rs b/crates/nq-packetloss/src/lib.rs index 1ee1334..7559ad0 100644 --- a/crates/nq-packetloss/src/lib.rs +++ b/crates/nq-packetloss/src/lib.rs @@ -9,7 +9,7 @@ mod webrtc_data_channel; -use nq_core::{ConnectionType, Network, Time, TokioTime, client::Direction}; +use nq_core::{ConnectionType, Network, ScopedHeaders, Time, TokioTime, client::Direction}; use nq_load_generator::{LoadConfig, LoadGenerator, LoadedConnection}; use nq_tokio_network::TokioNetwork; use serde::{Deserialize, Serialize}; @@ -38,6 +38,9 @@ pub struct PacketLossConfig { pub download_url: Url, /// Upload URL to use for for the [`LoadGenerator`] pub upload_url: Url, + /// Headers attached only to requests whose host matches the scope's + /// allowlist. + pub scoped_headers: Option, } impl Default for PacketLossConfig { @@ -53,6 +56,7 @@ impl Default for PacketLossConfig { .parse() .unwrap(), upload_url: "https://h3.speed.cloudflare.com/__up".parse().unwrap(), + scoped_headers: None, } } } @@ -61,6 +65,7 @@ impl PacketLossConfig { pub fn load_config(&self) -> LoadConfig { LoadConfig { headers: HashMap::default(), + scoped_headers: self.scoped_headers.clone(), download_url: self.download_url.clone(), upload_url: self.upload_url.clone(), upload_size: 4_000_000_000, // 4 GB @@ -359,6 +364,7 @@ mod tests { .parse() .unwrap(), upload_url: "https://h3.speed.cloudflare.com/__up".parse().unwrap(), + scoped_headers: None, }; let packet_loss = PacketLoss::new_with_config(config)?; let packet_loss_result = packet_loss @@ -465,6 +471,7 @@ mod tests { .parse() .unwrap(), upload_url: "https://h3.speed.cloudflare.com/__up".parse().unwrap(), + scoped_headers: None, }; let packet_loss = PacketLoss::new_with_config(config)?; diff --git a/crates/nq-rpm/src/lib.rs b/crates/nq-rpm/src/lib.rs index b26e360..ba8ca3d 100644 --- a/crates/nq-rpm/src/lib.rs +++ b/crates/nq-rpm/src/lib.rs @@ -12,7 +12,7 @@ use std::{ use humansize::{DECIMAL, format_size}; use nq_core::{ - ConnectionType, Network, Time, Timestamp, + ConnectionType, Network, ScopedHeaders, Time, Timestamp, client::{Direction, ThroughputClient, wait_for_finish}, }; use nq_load_generator::{LoadConfig, LoadGenerator, LoadedConnection}; @@ -35,12 +35,16 @@ pub struct ResponsivenessConfig { pub max_loaded_connections: usize, pub conn_type: ConnectionType, pub determine_load_only: bool, + /// Headers attached only to requests whose host matches the scope's + /// allowlist. + pub scoped_headers: Option, } impl ResponsivenessConfig { pub fn load_config(&self) -> LoadConfig { LoadConfig { headers: HashMap::default(), + scoped_headers: self.scoped_headers.clone(), download_url: self.large_download_url.clone(), upload_url: self.upload_url.clone(), upload_size: 4_000_000_000, // 4 GB @@ -66,6 +70,7 @@ impl Default for ResponsivenessConfig { max_loaded_connections: 16, conn_type: ConnectionType::H2, determine_load_only: false, + scoped_headers: None, } } } @@ -443,14 +448,16 @@ impl Responsiveness { env: &Env, shutdown: CancellationToken, ) -> anyhow::Result<()> { - let inflight_body_fut = ThroughputClient::download() + let client = ThroughputClient::download() .new_connection(ConnectionType::H2) - .send( - self.config.small_download_url.as_str().parse()?, - Arc::clone(&env.network), - Arc::clone(&env.time), - shutdown, - )?; + .scoped_headers(self.config.scoped_headers.clone()); + + let inflight_body_fut = client.send( + self.config.small_download_url.as_str().parse()?, + Arc::clone(&env.network), + Arc::clone(&env.time), + shutdown, + )?; tokio::spawn(report_err( event_tx.clone(), @@ -511,14 +518,16 @@ impl Responsiveness { return Ok(false); }; - let inflight_body_fut = ThroughputClient::download() + let client = ThroughputClient::download() .with_connection(connection) - .send( - self.config.small_download_url.as_str().parse()?, - Arc::clone(&env.network), - Arc::clone(&env.time), - shutdown, - )?; + .scoped_headers(self.config.scoped_headers.clone()); + + let inflight_body_fut = client.send( + self.config.small_download_url.as_str().parse()?, + Arc::clone(&env.network), + Arc::clone(&env.time), + shutdown, + )?; tokio::spawn(report_err( event_tx.clone(),