Skip to main content

brainwires_network/routing/
broadcast.rs

1use anyhow::Result;
2use async_trait::async_trait;
3
4use super::peer_table::PeerTable;
5use super::traits::{Router, RoutingStrategy};
6use crate::network::MessageEnvelope;
7use crate::transport::TransportAddress;
8
9/// Broadcast router.
10///
11/// Returns the transport addresses of all known peers (excluding the
12/// sender). Used for broadcast and peer-discovery announcements.
13#[derive(Debug, Default)]
14pub struct BroadcastRouter;
15
16impl BroadcastRouter {
17    /// Create a new broadcast router.
18    pub fn new() -> Self {
19        Self
20    }
21}
22
23#[async_trait]
24impl Router for BroadcastRouter {
25    async fn route(
26        &self,
27        envelope: &MessageEnvelope,
28        peers: &PeerTable,
29    ) -> Result<Vec<TransportAddress>> {
30        let mut addrs = Vec::new();
31
32        for peer_id in peers.all_peer_ids() {
33            // Don't send back to the sender
34            if *peer_id == envelope.sender {
35                continue;
36            }
37            if let Some(peer_addrs) = peers.addresses(peer_id) {
38                addrs.extend_from_slice(peer_addrs);
39            }
40        }
41
42        Ok(addrs)
43    }
44
45    fn strategy(&self) -> RoutingStrategy {
46        RoutingStrategy::Broadcast
47    }
48}
49
50#[cfg(test)]
51mod tests {
52    use super::*;
53    use crate::identity::AgentIdentity;
54    use crate::network::Payload;
55    use uuid::Uuid;
56
57    #[tokio::test]
58    async fn broadcast_reaches_all_except_sender() {
59        let router = BroadcastRouter::new();
60        let mut peers = PeerTable::new();
61
62        let sender = AgentIdentity::new("sender");
63        let sender_id = sender.id;
64        let peer_a = AgentIdentity::new("a");
65        let peer_b = AgentIdentity::new("b");
66
67        peers.upsert(
68            sender,
69            vec![TransportAddress::Tcp("127.0.0.1:1000".parse().unwrap())],
70        );
71        peers.upsert(
72            peer_a,
73            vec![TransportAddress::Tcp("127.0.0.1:2000".parse().unwrap())],
74        );
75        peers.upsert(
76            peer_b,
77            vec![TransportAddress::Tcp("127.0.0.1:3000".parse().unwrap())],
78        );
79
80        let env = MessageEnvelope::broadcast(sender_id, Payload::Text("ping".into()));
81        let addrs = router.route(&env, &peers).await.unwrap();
82
83        // Should get addresses for peer_a and peer_b, not sender
84        assert_eq!(addrs.len(), 2);
85        assert!(!addrs.contains(&TransportAddress::Tcp("127.0.0.1:1000".parse().unwrap())));
86    }
87
88    #[tokio::test]
89    async fn broadcast_empty_peers() {
90        let router = BroadcastRouter::new();
91        let peers = PeerTable::new();
92
93        let env = MessageEnvelope::broadcast(Uuid::new_v4(), Payload::Text("ping".into()));
94        let addrs = router.route(&env, &peers).await.unwrap();
95        assert!(addrs.is_empty());
96    }
97}