From 726fdd915ffe8e43c360fdec5faefeacbeaa40c2 Mon Sep 17 00:00:00 2001 From: kriss39 Date: Sat, 12 Sep 2026 14:33:13 +0400 Subject: [PATCH] fix(spammer): return errors for invalid inputs --- crates/spammer/src/main.rs | 23 +++++++++++++-------- crates/spammer/src/rate_limiter.rs | 33 +++++++++++++++++++++++------- crates/spammer/src/spammer.rs | 4 ++-- 3 files changed, 43 insertions(+), 17 deletions(-) diff --git a/crates/spammer/src/main.rs b/crates/spammer/src/main.rs index 64bc3ae8..60b1a68b 100644 --- a/crates/spammer/src/main.rs +++ b/crates/spammer/src/main.rs @@ -129,7 +129,7 @@ async fn main() -> Result<()> { let target_ws_urls = match cli.command { TargetCommand::Ws { targets } => { let target_nodes = targets.unwrap_or_else(|| vec![DEFAULT_WS_TARGET.to_string()]); - ws_urls_from_strings(target_nodes) + ws_urls_from_strings(target_nodes)? } TargetCommand::Nodes { nodes_path, @@ -230,20 +230,20 @@ fn write_atomic(contents: &str, path: &Path, label: &str) -> Result<()> { } // Build the WebSocket URLs of the target nodes from the list of IP addresses and ports -fn ws_urls_from_strings(target_nodes: Vec) -> Vec<(String, Url)> { +fn ws_urls_from_strings(target_nodes: Vec) -> Result> { target_nodes .into_iter() - .map(|s| (s.clone(), ws_url_from_str(s))) + .map(|s| ws_url_from_str(s.clone()).map(|url| (s, url))) .collect() } -fn ws_url_from_str(ip_port: String) -> Url { +fn ws_url_from_str(ip_port: String) -> Result { let url_str = if !ip_port.starts_with("ws://") { format!("ws://{ip_port}") } else { ip_port }; - Url::parse(&url_str).unwrap() + Url::parse(&url_str).wrap_err_with(|| format!("Invalid WebSocket target: {url_str}")) } // Read the file with node metadata and obtain the WebSocket URLs of the target nodes @@ -313,22 +313,29 @@ mod tests { #[test] fn ws_url_from_str_adds_ws_scheme_if_missing() { - let url = ws_url_from_str("127.0.0.1:8546".to_string()); + let url = ws_url_from_str("127.0.0.1:8546".to_string()).unwrap(); assert_eq!(url.as_str(), "ws://127.0.0.1:8546/"); } #[test] fn ws_url_from_str_parses_endpoint() { - let url = ws_url_from_str("ws://127.0.0.1:8546".to_string()); + let url = ws_url_from_str("ws://127.0.0.1:8546".to_string()).unwrap(); assert_eq!(url.as_str(), "ws://127.0.0.1:8546/"); } + #[test] + fn ws_url_from_str_rejects_invalid_endpoint() { + let err = ws_url_from_str("http://[::1".to_string()).unwrap_err(); + assert!(err.to_string().contains("Invalid WebSocket target")); + } + #[test] fn ws_urls_from_strings_parses_endpoints() { let urls = ws_urls_from_strings(vec![ "127.0.0.1:8546".to_string(), "ws://127.0.0.1:9546".to_string(), - ]); + ]) + .unwrap(); assert_eq!(urls.len(), 2); assert_eq!(urls[0].0, "127.0.0.1:8546"); diff --git a/crates/spammer/src/rate_limiter.rs b/crates/spammer/src/rate_limiter.rs index 26e428ee..90db7f68 100644 --- a/crates/spammer/src/rate_limiter.rs +++ b/crates/spammer/src/rate_limiter.rs @@ -17,6 +17,7 @@ use std::num::NonZeroU32; use std::sync::atomic::{AtomicU64, Ordering}; +use color_eyre::eyre::{self, Result}; use governor::{Jitter, Quota}; /// Token-bucket rate limiter for transaction sending. @@ -32,22 +33,26 @@ pub(crate) struct RateLimiter { } impl RateLimiter { - pub fn new(tps: u64, max_num_txs: u64, num_senders: usize) -> Self { - let tps_u32 = u32::try_from(tps).expect("TPS must fit in u32"); - let tps_nz = NonZeroU32::new(tps_u32).expect("TPS must be > 0"); + pub fn new(tps: u64, max_num_txs: u64, num_senders: usize) -> Result { + let tps_u32 = u32::try_from(tps) + .map_err(|_| eyre::eyre!("TPS must fit in u32, got {tps}"))?; + let tps_nz = + NonZeroU32::new(tps_u32).ok_or_else(|| eyre::eyre!("TPS must be greater than 0"))?; let burst = (tps / num_senders.max(1) as u64).max(1); - let burst_nz = NonZeroU32::new(u32::try_from(burst).expect("burst must fit in u32")) - .expect("burst must be > 0"); + let burst_u32 = u32::try_from(burst) + .map_err(|_| eyre::eyre!("burst must fit in u32, got {burst}"))?; + let burst_nz = NonZeroU32::new(burst_u32) + .ok_or_else(|| eyre::eyre!("burst must be greater than 0"))?; let quota = Quota::per_second(tps_nz).allow_burst(burst_nz); let limiter = governor::RateLimiter::direct(quota); // Uniformly random jitter up to half the interval let jitter = Jitter::up_to(quota.replenish_interval() / 2); - Self { + Ok(Self { limiter, jitter, max_num_txs, total_counter: AtomicU64::new(0), - } + }) } /// Wait until the rate limiter permits the next send. @@ -62,3 +67,17 @@ impl RateLimiter { prev < self.max_num_txs } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_tps_above_u32_max() { + let result = RateLimiter::new(u64::from(u32::MAX) + 1, 0, 1); + let Err(err) = result else { + panic!("expected oversized TPS to be rejected"); + }; + assert!(err.to_string().contains("TPS must fit in u32")); + } +} diff --git a/crates/spammer/src/spammer.rs b/crates/spammer/src/spammer.rs index d9f636d2..c9d427e1 100644 --- a/crates/spammer/src/spammer.rs +++ b/crates/spammer/src/spammer.rs @@ -267,7 +267,7 @@ impl Spammer { config.max_rate, config.max_num_txs, config.num_generators, - )); + )?); // Create transaction generators and senders let (tx_generators, tx_senders, tx_ack_receivers) = if config.fire_and_forget { @@ -689,7 +689,7 @@ impl Spammer { config.max_rate, config.max_num_txs, num_generators, - )); + )?); let mut tx_generators = Vec::new(); let mut tx_senders = Vec::new();