1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
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
}
}
|
