libp2p_autonat/v2/server/
behaviour.rs

1use std::{
2    collections::{HashMap, VecDeque},
3    io,
4    task::{Context, Poll},
5};
6
7use either::Either;
8use libp2p_core::{transport::PortUse, Endpoint, Multiaddr};
9use libp2p_identity::PeerId;
10use libp2p_swarm::{
11    dial_opts::{DialOpts, PeerCondition},
12    dummy, ConnectionDenied, ConnectionHandler, ConnectionId, DialFailure, FromSwarm,
13    NetworkBehaviour, ToSwarm,
14};
15use rand_core::{OsRng, RngCore};
16
17use crate::v2::server::handler::{
18    dial_back,
19    dial_request::{self, DialBackCommand, DialBackStatus},
20    Handler,
21};
22
23pub struct Behaviour<R = OsRng>
24where
25    R: Clone + Send + RngCore + 'static,
26{
27    dialing_dial_back: HashMap<ConnectionId, DialBackCommand>,
28    pending_events: VecDeque<
29        ToSwarm<
30            <Self as NetworkBehaviour>::ToSwarm,
31            <<Self as NetworkBehaviour>::ConnectionHandler as ConnectionHandler>::FromBehaviour,
32        >,
33    >,
34    rng: R,
35}
36
37impl Default for Behaviour<OsRng> {
38    fn default() -> Self {
39        Self::new(OsRng)
40    }
41}
42
43impl<R> Behaviour<R>
44where
45    R: RngCore + Send + Clone + 'static,
46{
47    pub fn new(rng: R) -> Self {
48        Self {
49            dialing_dial_back: HashMap::new(),
50            pending_events: VecDeque::new(),
51            rng,
52        }
53    }
54}
55
56impl<R> NetworkBehaviour for Behaviour<R>
57where
58    R: RngCore + Send + Clone + 'static,
59{
60    type ConnectionHandler = Handler<R>;
61
62    type ToSwarm = Event;
63
64    fn handle_established_inbound_connection(
65        &mut self,
66        _connection_id: ConnectionId,
67        peer: PeerId,
68        _local_addr: &Multiaddr,
69        remote_addr: &Multiaddr,
70    ) -> Result<<Self as NetworkBehaviour>::ConnectionHandler, ConnectionDenied> {
71        Ok(Either::Right(dial_request::Handler::new(
72            peer,
73            remote_addr.clone(),
74            self.rng.clone(),
75        )))
76    }
77
78    fn handle_established_outbound_connection(
79        &mut self,
80        connection_id: ConnectionId,
81        _peer: PeerId,
82        _addr: &Multiaddr,
83        _role_override: Endpoint,
84        _port_use: PortUse,
85    ) -> Result<<Self as NetworkBehaviour>::ConnectionHandler, ConnectionDenied> {
86        Ok(match self.dialing_dial_back.remove(&connection_id) {
87            Some(cmd) => Either::Left(Either::Left(dial_back::Handler::new(cmd))),
88            None => Either::Left(Either::Right(dummy::ConnectionHandler)),
89        })
90    }
91
92    fn on_swarm_event(&mut self, event: FromSwarm) {
93        if let FromSwarm::DialFailure(DialFailure { connection_id, .. }) = event {
94            if let Some(DialBackCommand { back_channel, .. }) =
95                self.dialing_dial_back.remove(&connection_id)
96            {
97                let dial_back_status = DialBackStatus::DialErr;
98                let _ = back_channel.send(Err(dial_back_status));
99            }
100        }
101    }
102
103    fn on_connection_handler_event(
104        &mut self,
105        peer_id: PeerId,
106        _connection_id: ConnectionId,
107        event: <Handler<R> as ConnectionHandler>::ToBehaviour,
108    ) {
109        match event {
110            Either::Left(Either::Left(Ok(_))) => {}
111            Either::Left(Either::Left(Err(e))) => {
112                tracing::debug!("dial back error: {e:?}");
113            }
114            // TODO: remove when Rust 1.82 is MSRV
115            #[allow(unreachable_patterns)]
116            Either::Left(Either::Right(v)) => libp2p_core::util::unreachable(v),
117            Either::Right(Either::Left(cmd)) => {
118                let addr = cmd.addr.clone();
119                let opts = DialOpts::peer_id(peer_id)
120                    .addresses(Vec::from([addr]))
121                    .condition(PeerCondition::Always)
122                    .allocate_new_port()
123                    .build();
124                let conn_id = opts.connection_id();
125                self.dialing_dial_back.insert(conn_id, cmd);
126                self.pending_events.push_back(ToSwarm::Dial { opts });
127            }
128            Either::Right(Either::Right(status_update)) => self
129                .pending_events
130                .push_back(ToSwarm::GenerateEvent(status_update)),
131        }
132    }
133
134    fn poll(
135        &mut self,
136        _cx: &mut Context<'_>,
137    ) -> Poll<ToSwarm<Self::ToSwarm, <Handler<R> as ConnectionHandler>::FromBehaviour>> {
138        if let Some(event) = self.pending_events.pop_front() {
139            return Poll::Ready(event);
140        }
141        Poll::Pending
142    }
143}
144
145#[derive(Debug)]
146pub struct Event {
147    /// All address that were submitted for testing.
148    pub all_addrs: Vec<Multiaddr>,
149    /// The address that was eventually tested.
150    pub tested_addr: Multiaddr,
151    /// The peer id of the client that submitted addresses for testing.
152    pub client: PeerId,
153    /// The amount of data that was requested by the server and was transmitted.
154    pub data_amount: usize,
155    /// The result of the test.
156    pub result: Result<(), io::Error>,
157}