1 files changed, 18 insertions, 74 deletions

diff --git a/src/transport/firewall.rs b/src/transport/firewall.rs
index b8e05cd..7b0ec0d 100644
--- a/src/transport/firewall.rs
+++ b/src/transport/firewall.rs
@@ -1,105 +1,49 @@
+//! The debug firewall (`/debug/firewall`, with `--debug-api`): drops messages from and
+//! to peers, by node ID, to test partitions on a real swarm. It works at the bottom of
+//! the node, in the paths layer (`transport::paths::Gate`), so it drops everything a
+//! partition would: WHO and ME, PINGs and their PONGs, as well as the node's messages.
+
 use std::collections::HashSet;
-use std::sync::{Arc, RwLock, RwLockReadGuard};
+use std::sync::{RwLock, RwLockReadGuard};
 
-use anyhow::Result;
-use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
-use utoipa::ToSchema;
 
-use crate::transport::EncodedMessage;
-use crate::transport::{MessageHandler, ParsedMessage, Transport};
+use crate::locks::RwLockExt;
 
-#[derive(Debug, Clone, Default, Serialize, Deserialize, ToSchema)]
-pub struct FirewallRules {
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub(crate) struct FirewallRules {
     /// Messages from these peers are dropped before the application sees them.
     #[serde(default)]
-    pub inbound: HashSet<String>,
+    pub(crate) inbound: HashSet<String>,
     /// Messages to these peers are silently dropped, like a lossy network.
     #[serde(default)]
-    pub outbound: HashSet<String>,
+    pub(crate) outbound: HashSet<String>,
 }
 
 /// Uses a std lock: checks are tiny and the guard is never held across an await.
 #[derive(Default)]
-pub struct Firewall {
+pub(crate) struct Firewall {
     rules: RwLock<FirewallRules>,
 }
 
 impl Firewall {
     fn read(&self) -> RwLockReadGuard<'_, FirewallRules> {
-        self.rules
-            .read()
-            .unwrap_or_else(|poisoned| poisoned.into_inner())
+        self.rules.read_or_recover()
     }
 
-    pub fn rules(&self) -> FirewallRules {
+    pub(crate) fn rules(&self) -> FirewallRules {
         self.read().clone()
     }
 
-    pub fn set_rules(&self, rules: FirewallRules) {
-        *self
-            .rules
-            .write()
-            .unwrap_or_else(|poisoned| poisoned.into_inner()) = rules;
+    pub(crate) fn set_rules(&self, rules: FirewallRules) {
+        *self.rules.write_or_recover() = rules;
     }
 
-    pub fn blocks_inbound(&self, peer: &str) -> bool {
+    pub(crate) fn blocks_inbound(&self, peer: &str) -> bool {
         self.read().inbound.contains(peer)
     }
 
-    pub fn blocks_outbound(&self, peer: &str) -> bool {
+    pub(crate) fn blocks_outbound(&self, peer: &str) -> bool {
         self.read().outbound.contains(peer)
     }
 }
-
-pub struct FirewallTransport {
-    inner: Arc<dyn Transport>,
-    firewall: Arc<Firewall>,
-}
-
-impl FirewallTransport {
-    pub fn new(inner: Arc<dyn Transport>, firewall: Arc<Firewall>) -> Self {
-        Self { inner, firewall }
-    }
-}
-
-#[async_trait]
-impl Transport for FirewallTransport {
-    fn local_id(&self) -> &str {
-        self.inner.local_id()
-    }
-
-    async fn resolve_seed(&self, seed: &str) -> Result<String> {
-        self.inner.resolve_seed(seed).await
-    }
-
-    fn encode(&self, peer: &str, message: &ParsedMessage) -> Result<EncodedMessage> {
-        self.inner.encode(peer, message)
-    }
-
-    async fn send_encoded(&self, message: &EncodedMessage) -> Result<()> {
-        if self.firewall.blocks_outbound(&message.peer) {
-            tracing::trace!(peer = %message.peer, "Firewall dropped outbound message");
-            return Ok(());
-        }
-        self.inner.send_encoded(message).await
-    }
-
-    async fn start_consumer(&self, handler: MessageHandler) -> Result<()> {
-        let firewall = Arc::clone(&self.firewall);
-
-        self.inner
-            .start_consumer(Box::new(move |sender, message, payload_len| {
-                if firewall.blocks_inbound(&sender) {
-                    tracing::trace!(peer = %sender, "Firewall dropped inbound message");
-                    return Box::pin(async { Ok(()) });
-                }
-                handler(sender, message, payload_len)
-            }))
-            .await
-    }
-
-    async fn wait_for_exit(&self) -> std::io::Result<std::process::ExitStatus> {
-        self.inner.wait_for_exit().await
-    }
-}