arti-testing: support for conditional TCP failure.

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<Job> {
                 .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<Job> {
             .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,

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.

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<Self, Self::Err> {
+        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<R> {
     #[pin]
     inner: R,
     /// The action to take when we try to make an outbound connection.
-    action: Arc<Mutex<Action>>,
+    action: Arc<Mutex<ConditionalAction>>,
 }
 
 impl<R> BrokenTcpProvider<R> {
     /// 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<R> BrokenTcpProvider<R> {
 
     /// 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
+        }
     }
 }