diff --git a/crates/arti-testing/src/config.rs b/crates/arti-testing/src/config.rs index de7f760d5..2d15903ff 100644 --- a/crates/arti-testing/src/config.rs +++ b/crates/arti-testing/src/config.rs @@ -1,5 +1,6 @@ //! Reading configuration and command line issues in arti-testing. +use crate::rt::badtcp::ConditionalAction; use crate::{Action, Job, TcpBreakage}; use anyhow::{anyhow, Result}; @@ -69,6 +70,13 @@ pub(crate) fn parse_cmdline() -> Result { .value_name("none|timeout|error|blackhole") .global(true), ) + .arg( + Arg::with_name("tcp-failure-on") + .long("tcp-failure-on") + .takes_value(true) + .value_name("all|v4|v6|non443") + .global(true), + ) .arg( Arg::with_name("tcp-failure-stage") .long("tcp-failure-stage") @@ -148,6 +156,11 @@ pub(crate) fn parse_cmdline() -> Result { .value_of("tcp-failure-delay") .map(|d| d.parse().map(Duration::from_secs)) .transpose()?; + let when = matches + .value_of("tcp-failure-on") + .unwrap_or("all") + .parse()?; + let action = ConditionalAction { action, when }; TcpBreakage { action, diff --git a/crates/arti-testing/src/main.rs b/crates/arti-testing/src/main.rs index 36285cd0f..fe21868cb 100644 --- a/crates/arti-testing/src/main.rs +++ b/crates/arti-testing/src/main.rs @@ -36,7 +36,7 @@ //! o by timing out //! - sporadically //! o by succeeding and black-holing data. -//! - depending on address / port / family +//! o depending on address / port / family //! o Install this after a delay //! - make TLS fail //! - With wrong cert @@ -170,7 +170,7 @@ impl FromStr for BreakageStage { #[derive(Debug, Clone)] struct TcpBreakage { /// What kind of breakage to install (if any) - action: rt::badtcp::Action, + action: rt::badtcp::ConditionalAction, /// What stage to apply the breakage at. stage: BreakageStage, /// Delay (if any) after the start of the stage to apply breakage @@ -274,8 +274,10 @@ impl Job { /// XXXX Eventually this should come up with some kind of result that's meaningful. async fn run_job(&self) -> Result<()> { let runtime = PreferredRuntime::current()?; - let broken_tcp = - rt::badtcp::BrokenTcpProvider::new(runtime.clone(), rt::badtcp::Action::Work); + let broken_tcp = rt::badtcp::BrokenTcpProvider::new( + runtime.clone(), + rt::badtcp::ConditionalAction::default(), + ); // We put the counting TCP provider outside the one that breaks: we want // to know how many attempts to connect there are, and BrokenTcpProvider // eats the attempts that it fails without passing them down the stack. diff --git a/crates/arti-testing/src/rt/badtcp.rs b/crates/arti-testing/src/rt/badtcp.rs index a2d59d7da..9bd1635d8 100644 --- a/crates/arti-testing/src/rt/badtcp.rs +++ b/crates/arti-testing/src/rt/badtcp.rs @@ -30,6 +30,31 @@ pub(crate) enum Action { Blackhole, } +/// When should an Action apply? +#[derive(Debug, Clone)] +pub(crate) enum ActionPat { + /// always apply + Always, + /// Apply to all ipv4 + V4, + /// apply to all ipv6 + V6, + /// apply to all ports but 443 + Non443, +} + +/// An Action plus a set of conditions when it applies. +/// +/// (When the action doesn't apply, connections will just `Action::Work`. +#[derive(Debug, Clone)] +pub(crate) struct ConditionalAction { + /// The underlying action + pub(crate) action: Action, + + /// When should the action apply? + pub(crate) when: ActionPat, +} + impl FromStr for Action { type Err = anyhow::Error; @@ -44,6 +69,41 @@ impl FromStr for Action { } } +impl FromStr for ActionPat { + type Err = anyhow::Error; + + fn from_str(s: &str) -> Result { + Ok(match s { + "all" => ActionPat::Always, + "v4" => ActionPat::V4, + "v6" => ActionPat::V6, + "non443" => ActionPat::Non443, + _ => return Err(anyhow!("unrecognized tcp breakage condition {:?}", s)), + }) + } +} + +impl ConditionalAction { + fn applies_to(&self, addr: &SocketAddr) -> bool { + match (addr, &self.when) { + (_, ActionPat::Always) => true, + (SocketAddr::V4(_), ActionPat::V4) => true, + (SocketAddr::V6(_), ActionPat::V6) => true, + (sa, ActionPat::Non443) if sa.port() != 443 => true, + (_, _) => false, + } + } +} + +impl Default for ConditionalAction { + fn default() -> Self { + Self { + action: Action::Work, + when: ActionPat::Always, + } + } +} + /// A TcpProvider that can make its connections fail. #[derive(Debug, Clone)] #[pin_project] @@ -52,13 +112,13 @@ pub(crate) struct BrokenTcpProvider { #[pin] inner: R, /// The action to take when we try to make an outbound connection. - action: Arc>, + action: Arc>, } impl BrokenTcpProvider { /// Construct a new BrokenTcpProvider which responds to all outbound /// connections by taking the specified action. - pub(crate) fn new(inner: R, action: Action) -> Self { + pub(crate) fn new(inner: R, action: ConditionalAction) -> Self { Self { inner, action: Arc::new(Mutex::new(action)), @@ -67,13 +127,18 @@ impl BrokenTcpProvider { /// Cause the provider to respond to all outbound connection attempts /// with the specified action. - pub(crate) fn set_action(&self, action: Action) { + pub(crate) fn set_action(&self, action: ConditionalAction) { *self.action.lock().expect("Lock poisoned") = action; } /// Return the action to take for a connection to `addr`. - fn get_action(&self, _addr: &SocketAddr) -> Action { - self.action.lock().expect("Lock poisoned").clone() + fn get_action(&self, addr: &SocketAddr) -> Action { + let action = self.action.lock().expect("Lock poisoned"); + if action.applies_to(addr) { + action.action.clone() + } else { + Action::Work + } } }