From d3a9d5d3ebff98e342dff1df0ebad86f346719f5 Mon Sep 17 00:00:00 2001 From: tv Date: Mon, 19 Oct 2026 23:44:20 +0200 Subject: initial commit --- src/firewall.rs | 105 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 105 insertions(+) create mode 100644 src/firewall.rs (limited to 'src/firewall.rs') 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, + /// Messages to these peers are silently dropped, like a lossy network. + #[serde(default)] + pub outbound: HashSet, +} + +/// Uses a std lock: checks are tiny and the guard is never held across an await. +#[derive(Default)] +pub struct Firewall { + rules: RwLock, +} + +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, + firewall: Arc, +} + +impl FirewallTransport { + pub fn new(inner: Arc, firewall: Arc) -> 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 { + self.inner.resolve_peer(hint).await + } + + fn encode(&self, peer: &str, message: &ParsedMessage) -> Result { + 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 { + self.inner.wait_for_exit().await + } +} -- cgit v1.3.1