subspace_networking/protocols/
autonat_wrapper.rs

1use crate::utils::is_global_address_or_dns;
2use libp2p::autonat::{Behaviour as Autonat, Config as AutonatConfig, Event as AutonatEvent};
3use libp2p::core::transport::PortUse;
4use libp2p::core::Endpoint;
5use libp2p::multiaddr::Protocol;
6use libp2p::swarm::{
7    ConnectionDenied, ConnectionId, FromSwarm, NetworkBehaviour, THandler, THandlerInEvent,
8    THandlerOutEvent, ToSwarm,
9};
10use libp2p::{Multiaddr, PeerId};
11use std::collections::HashSet;
12use std::task::{Context, Poll};
13use tracing::debug;
14
15pub(crate) struct Config {
16    pub(crate) inner_config: AutonatConfig,
17    pub(crate) local_peer_id: PeerId,
18    pub(crate) servers: Vec<Multiaddr>,
19}
20
21pub(crate) struct Behaviour {
22    inner: Autonat,
23    private_ips_enabled: bool,
24    listen_addresses: HashSet<Multiaddr>,
25}
26
27impl Behaviour {
28    pub(crate) fn new(config: Config) -> Self {
29        let mut inner = Autonat::new(config.local_peer_id, config.inner_config.clone());
30
31        for server in config.servers {
32            let maybe_peer_id = server.iter().find_map(|protocol| {
33                if let Protocol::P2p(peer_id) = protocol {
34                    Some(peer_id)
35                } else {
36                    None
37                }
38            });
39            if let Some(peer_id) = maybe_peer_id {
40                inner.add_server(peer_id, Some(server));
41            }
42        }
43
44        Self {
45            inner,
46            private_ips_enabled: !config.inner_config.only_global_ips,
47            listen_addresses: Default::default(),
48        }
49    }
50
51    fn address_corresponds_to_listening_addresses(&self, addr: &Multiaddr) -> bool {
52        let Some(candidate_protocol) = addr.iter().find_map(|protocol| match protocol {
53            tcp @ Protocol::Tcp(_) => Some(tcp),
54            _ => None,
55        }) else {
56            return false;
57        };
58
59        let address_result = self
60            .listen_addresses
61            .iter()
62            .any(|addr| addr.iter().any(|protocol| protocol == candidate_protocol));
63
64        debug!(
65            %address_result,
66            ?addr,
67            listen_addresses=?self.listen_addresses,
68            "Address candidate corresponds to listening addresses."
69        );
70
71        address_result
72    }
73
74    pub(crate) fn public_address(&self) -> Option<&Multiaddr> {
75        self.inner.public_address()
76    }
77
78    pub(crate) fn confidence(&self) -> usize {
79        self.inner.confidence()
80    }
81}
82
83impl NetworkBehaviour for Behaviour {
84    type ConnectionHandler = <Autonat as NetworkBehaviour>::ConnectionHandler;
85    type ToSwarm = AutonatEvent;
86
87    fn handle_established_inbound_connection(
88        &mut self,
89        connection_id: ConnectionId,
90        peer: PeerId,
91        local_addr: &Multiaddr,
92        remote_addr: &Multiaddr,
93    ) -> Result<THandler<Self>, ConnectionDenied> {
94        self.inner.handle_established_inbound_connection(
95            connection_id,
96            peer,
97            local_addr,
98            remote_addr,
99        )
100    }
101
102    fn handle_established_outbound_connection(
103        &mut self,
104        connection_id: ConnectionId,
105        peer: PeerId,
106        addr: &Multiaddr,
107        role_override: Endpoint,
108        port_use: PortUse,
109    ) -> Result<THandler<Self>, ConnectionDenied> {
110        self.inner.handle_established_outbound_connection(
111            connection_id,
112            peer,
113            addr,
114            role_override,
115            port_use,
116        )
117    }
118
119    fn on_swarm_event(&mut self, event: FromSwarm) {
120        match event {
121            new_listen_addr_event @ FromSwarm::NewListenAddr(_) => {
122                if let FromSwarm::NewListenAddr(addr) = new_listen_addr_event {
123                    //TODO: handle listener address change
124                    self.listen_addresses.insert(addr.addr.clone());
125
126                    if self.private_ips_enabled || is_global_address_or_dns(addr.addr) {
127                        self.inner.on_swarm_event(new_listen_addr_event);
128                    } else {
129                        debug!(addr=?addr.addr, "Skipped listening address in AutonatWrapper.");
130                    }
131                }
132            }
133            new_external_addr_event @ FromSwarm::NewExternalAddrCandidate(_) => {
134                if let FromSwarm::NewExternalAddrCandidate(addr) = new_external_addr_event {
135                    if self.address_corresponds_to_listening_addresses(addr.addr) {
136                        self.inner.on_swarm_event(new_external_addr_event);
137                    } else {
138                        debug!(
139                            addr=?addr.addr,
140                            "Skipped external address candidate in AutonatWrapper."
141                        );
142                    }
143                }
144            }
145            event => {
146                self.inner.on_swarm_event(event);
147            }
148        }
149    }
150
151    fn on_connection_handler_event(
152        &mut self,
153        peer_id: PeerId,
154        connection_id: ConnectionId,
155        event: THandlerOutEvent<Self>,
156    ) {
157        self.inner
158            .on_connection_handler_event(peer_id, connection_id, event)
159    }
160
161    fn poll(
162        &mut self,
163        cx: &mut Context<'_>,
164    ) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
165        self.inner.poll(cx)
166    }
167}