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 #[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 pub all_addrs: Vec<Multiaddr>,
149 pub tested_addr: Multiaddr,
151 pub client: PeerId,
153 pub data_amount: usize,
155 pub result: Result<(), io::Error>,
157}