libp2p_autonat/v2/server/
behaviour.rs1use 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 Either::Left(Either::Right(v)) => libp2p_core::util::unreachable(v),
115 Either::Right(Either::Left(cmd)) => {
116 let addr = cmd.addr.clone();
117 let opts = DialOpts::peer_id(peer_id)
118 .addresses(Vec::from([addr]))
119 .condition(PeerCondition::Always)
120 .allocate_new_port()
121 .build();
122 let conn_id = opts.connection_id();
123 self.dialing_dial_back.insert(conn_id, cmd);
124 self.pending_events.push_back(ToSwarm::Dial { opts });
125 }
126 Either::Right(Either::Right(status_update)) => self
127 .pending_events
128 .push_back(ToSwarm::GenerateEvent(status_update)),
129 }
130 }
131
132 fn poll(
133 &mut self,
134 _cx: &mut Context<'_>,
135 ) -> Poll<ToSwarm<Self::ToSwarm, <Handler<R> as ConnectionHandler>::FromBehaviour>> {
136 if let Some(event) = self.pending_events.pop_front() {
137 return Poll::Ready(event);
138 }
139 Poll::Pending
140 }
141}
142
143#[derive(Debug)]
144pub struct Event {
145 pub all_addrs: Vec<Multiaddr>,
147 pub tested_addr: Multiaddr,
149 pub client: PeerId,
151 pub data_amount: usize,
153 pub result: Result<(), io::Error>,
155}