1 files changed, 105 insertions, 0 deletions
diff --git a/src/firewall.rs b/src/firewall.rs
new file mode 100644
index 0000000..917fe5f
--- /dev/null
+++ b/src/firewall.rs
@@ -0,0 +1,105 @@
+use std::collections::HashSet;
+use std::sync::{Arc, 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};
+
+#[derive(Debug, Clone, Default, Serialize, Deserialize, ToSchema)]
+pub struct FirewallRules {
+    /// Messages from these peers are dropped before the application sees them.
+    #[serde(default)]
+    pub inbound: HashSet<String>,
+    /// Messages to these peers are silently dropped, like a lossy network.
+    #[serde(default)]
+    pub outbound: HashSet<String>,
+}
+
+/// Uses a std lock: checks are tiny and the guard is never held across an await.
+#[derive(Default)]
+pub struct Firewall {
+    rules: RwLock<FirewallRules>,
+}
+
+impl Firewall {
+    fn read(&self) -> RwLockReadGuard<'_, FirewallRules> {
+        self.rules
+            .read()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+    }
+
+    pub 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 fn blocks_inbound(&self, peer: &str) -> bool {
+        self.read().inbound.contains(peer)
+    }
+
+    pub 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_peer(&self, hint: &str) -> Result<String> {
+        self.inner.resolve_peer(hint).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| {
+                if firewall.blocks_inbound(&sender) {
+                    tracing::trace!(peer = %sender, "Firewall dropped inbound message");
+                    return Box::pin(async { Ok(()) });
+                }
+                handler(sender, message)
+            }))
+            .await
+    }
+
+    async fn wait_for_exit(&self) -> std::io::Result<std::process::ExitStatus> {
+        self.inner.wait_for_exit().await
+    }
+}