protocols/autonat: Add basic behaviour logic

Author: elenaf9

protocols/autonat/src/autonat.rs renamed to protocols/autonat/src/behaviour.rs

@@ -18,49 +18,61 @@
 // FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 // DEALINGS IN THE SOFTWARE.
 
-pub use crate::protocol::DialResponse;
-use crate::protocol::{AutoNatCodec, AutoNatProtocol};
+use crate::protocol::{AutoNatCodec, AutoNatProtocol, DialRequest, DialResponse, ResponseError};
 use libp2p_core::{
     connection::{ConnectionId, ListenerId},
     ConnectedPoint, Multiaddr, PeerId,
 };
 use libp2p_request_response::{
     handler::RequestResponseHandlerEvent, ProtocolSupport, RequestResponse, RequestResponseConfig,
-    RequestResponseEvent, RequestResponseMessage,
+    RequestResponseEvent, RequestResponseMessage, ResponseChannel,
 };
 use libp2p_swarm::{
-    IntoProtocolsHandler, NetworkBehaviour, NetworkBehaviourAction, PollParameters,
+    AddressScore, DialPeerCondition, IntoProtocolsHandler, NetworkBehaviour,
+    NetworkBehaviourAction, PollParameters,
 };
 use std::{
+    collections::{HashMap, HashSet, VecDeque},
     iter,
     task::{Context, Poll},
 };
 
-#[derive(Clone, Debug, Eq, PartialEq)]
-pub enum AutoNatEvent {
-    DialResponse(DialResponse),
-}
+type FiniteAddrScore = u32;
 
 pub struct AutoNat {
     inner: RequestResponse<AutoNatCodec>,
+    local_addresses: HashMap<Multiaddr, FiniteAddrScore>,
+    pending_inbound: HashMap<PeerId, ResponseChannel<DialResponse>>,
+    pending_outbound: HashSet<PeerId>,
+    send_request: VecDeque<PeerId>,
 }
 
-impl AutoNat {
-    pub fn new() -> Self {
+impl Default for AutoNat {
+    fn default() -> Self {
         let protocols = iter::once((AutoNatProtocol, ProtocolSupport::Full));
         let cfg = RequestResponseConfig::default();
         let inner = RequestResponse::new(AutoNatCodec, protocols, cfg);
-        Self { inner }
+        Self {
+            inner,
+            local_addresses: HashMap::default(),
+            pending_inbound: HashMap::default(),
+            pending_outbound: HashSet::default(),
+            send_request: VecDeque::default(),
+        }
     }
+}
 
-    pub fn add_address(&mut self, peer_id: &PeerId, addr: Multiaddr) {
-        self.inner.add_address(peer_id, addr)
+impl AutoNat {
+    pub fn add_local_address(&mut self, address: Multiaddr) {
+        if self.local_addresses.get(&address).is_none() {
+            self.local_addresses.insert(address, 1);
+        }
     }
 }
 
 impl NetworkBehaviour for AutoNat {
     type ProtocolsHandler = <RequestResponse<AutoNatCodec> as NetworkBehaviour>::ProtocolsHandler;
-    type OutEvent = AutoNatEvent;
+    type OutEvent = ();
 
     fn new_handler(&mut self) -> Self::ProtocolsHandler {
         self.inner.new_handler()
@@ -85,7 +97,17 @@ impl NetworkBehaviour for AutoNat {
         endpoint: &ConnectedPoint,
     ) {
         self.inner
-            .inject_connection_established(peer, conn, endpoint)
+            .inject_connection_established(peer, conn, endpoint);
+        if !self.pending_outbound.contains(peer) {
+            self.send_request.push_back(*peer);
+        }
+        if let ConnectedPoint::Dialer { address } = endpoint {
+            if let Some(channel) = self.pending_inbound.remove(peer) {
+                let _ = self
+                    .inner
+                    .send_response(channel, DialResponse::Ok(address.clone()));
+            }
+        }
     }
 
     fn inject_connection_closed(
@@ -96,7 +118,10 @@ impl NetworkBehaviour for AutoNat {
         handler: <Self::ProtocolsHandler as IntoProtocolsHandler>::Handler,
     ) {
         self.inner
-            .inject_connection_closed(peer, conn, endpoint, handler)
+            .inject_connection_closed(peer, conn, endpoint, handler);
+        // Channel can be dropped, as the underlying substream already closed.
+        self.pending_inbound.remove(peer);
+        self.send_request.retain(|p| p != peer);
     }
 
     fn inject_address_change(
@@ -106,16 +131,34 @@ impl NetworkBehaviour for AutoNat {
         old: &ConnectedPoint,
         new: &ConnectedPoint,
     ) {
-        self.inner.inject_address_change(peer, conn, old, new)
+        self.inner.inject_address_change(peer, conn, old, new);
+        if let ConnectedPoint::Listener {
+            local_addr: old_addr,
+            ..
+        } = old
+        {
+            match new {
+                ConnectedPoint::Listener {
+                    local_addr: new_addr,
+                    ..
+                } if old_addr != new_addr => {
+                    self.local_addresses.remove(old_addr);
+                    if !self.local_addresses.contains_key(new_addr) {
+                        self.local_addresses.insert(new_addr.clone(), 1);
+                    }
+                }
+                _ => {}
+            }
+        }
     }
 
     fn inject_event(
         &mut self,
-        peer: PeerId,
+        peer_id: PeerId,
         conn: ConnectionId,
         event: RequestResponseHandlerEvent<AutoNatCodec>,
     ) {
-        self.inner.inject_event(peer, conn, event)
+        self.inner.inject_event(peer_id, conn, event)
     }
 
     fn inject_addr_reach_failure(
@@ -133,7 +176,12 @@ impl NetworkBehaviour for AutoNat {
         handler: Self::ProtocolsHandler,
         error: libp2p_swarm::DialError,
     ) {
-        self.inner.inject_dial_failure(peer_id, handler, error)
+        self.inner.inject_dial_failure(peer_id, handler, error);
+        if let Some(channel) = self.pending_inbound.remove(peer_id) {
+            let _ = self
+                .inner
+                .send_response(channel, DialResponse::Err(ResponseError::DialError));
+        }
     }
 
     fn inject_listen_failure(
@@ -151,11 +199,15 @@ impl NetworkBehaviour for AutoNat {
     }
 
     fn inject_new_listen_addr(&mut self, id: ListenerId, addr: &Multiaddr) {
-        self.inner.inject_new_listen_addr(id, addr)
+        self.inner.inject_new_listen_addr(id, addr);
+        if !self.local_addresses.contains_key(addr) {
+            self.local_addresses.insert(addr.clone(), 0);
+        }
     }
 
     fn inject_expired_listen_addr(&mut self, id: ListenerId, addr: &Multiaddr) {
-        self.inner.inject_expired_listen_addr(id, addr)
+        self.inner.inject_expired_listen_addr(id, addr);
+        self.local_addresses.remove(addr);
     }
 
     fn inject_listener_error(&mut self, id: ListenerId, err: &(dyn std::error::Error + 'static)) {
@@ -167,11 +219,19 @@ impl NetworkBehaviour for AutoNat {
     }
 
     fn inject_new_external_addr(&mut self, addr: &Multiaddr) {
-        self.inner.inject_new_external_addr(addr)
+        self.inner.inject_new_external_addr(addr);
+        match self.local_addresses.get_mut(addr) {
+            Some(score) if *score == 0 => *score = 1,
+            Some(_) => {},
+            None => {
+                self.local_addresses.insert(addr.clone(), 1);
+            }
+        }
     }
 
     fn inject_expired_external_addr(&mut self, addr: &Multiaddr) {
-        self.inner.inject_expired_external_addr(addr)
+        self.inner.inject_expired_external_addr(addr);
+        self.local_addresses.remove(addr);
     }
 
     fn poll(
@@ -180,46 +240,62 @@ impl NetworkBehaviour for AutoNat {
         params: &mut impl PollParameters,
     ) -> Poll<NetworkBehaviourAction<Self::OutEvent, Self::ProtocolsHandler>> {
         loop {
+            if let Some(peer_id) = self.send_request.pop_front() {
+                let mut scores: Vec<(Multiaddr, FiniteAddrScore)> =
+                    self.local_addresses.clone().into_iter().collect();
+                // Sort so that the address with the highest score will be dialed first by the remote.
+                scores.sort_by(|(_, score_a), (_, score_b)| score_b.cmp(score_a));
+                let addrs = scores.into_iter().map(|(a, _)| a).collect();
+                self.inner
+                    .send_request(&peer_id, DialRequest { peer_id, addrs });
+            }
             match self.inner.poll(cx, params) {
                 Poll::Ready(NetworkBehaviourAction::GenerateEvent(
                     RequestResponseEvent::Message { peer, message },
                 )) => match message {
                     RequestResponseMessage::Request {
-                        request_id,
-                        request,
+                        request_id: _,
+                        request: DialRequest { peer_id, addrs },
                         channel,
                     } => {
-                        println!("{} {:?} {:?} {:?}", peer, request_id, request, channel);
+                        for addr in addrs {
+                            self.inner.add_address(&peer, addr)
+                        }
+                        // TODO: Handle if there is already a pending request.
+                        self.pending_inbound.insert(peer_id, channel);
+                        return Poll::Ready(NetworkBehaviourAction::DialPeer {
+                            peer_id: peer,
+                            handler: self.inner.new_handler(),
+                            condition: DialPeerCondition::Always,
+                        });
                     }
                     RequestResponseMessage::Response {
-                        request_id,
+                        request_id: _,
                         response,
                     } => {
-                        println!("{} {:?} {:?}", peer, request_id, response);
+                        self.pending_outbound.remove(&peer);
+                        if let DialResponse::Ok(address) = response {
+                            let score = self.local_addresses.entry(address.clone()).or_insert(1);
+                            *score += 1;
+                            return Poll::Ready(NetworkBehaviourAction::ReportObservedAddr {
+                                address,
+                                score: AddressScore::Finite(*score),
+                            });
+                        }
                     }
                 },
                 Poll::Ready(NetworkBehaviourAction::GenerateEvent(
-                    RequestResponseEvent::ResponseSent { peer, request_id },
-                )) => {
-                    println!("response sent {} {:?}", peer, request_id);
-                }
+                    RequestResponseEvent::ResponseSent { .. },
+                )) => {}
                 Poll::Ready(NetworkBehaviourAction::GenerateEvent(
-                    RequestResponseEvent::OutboundFailure {
-                        peer,
-                        request_id,
-                        error,
-                    },
+                    RequestResponseEvent::OutboundFailure { peer, .. },
                 )) => {
-                    println!("outbound failure {} {:?} {:?}", peer, request_id, error);
+                    self.pending_outbound.remove(&peer);
                 }
                 Poll::Ready(NetworkBehaviourAction::GenerateEvent(
-                    RequestResponseEvent::InboundFailure {
-                        peer,
-                        error,
-                        request_id,
-                    },
+                    RequestResponseEvent::InboundFailure { peer, .. },
                 )) => {
-                    println!("inbound failure {} {:?} {:?}", peer, request_id, error);
+                    self.pending_inbound.remove(&peer);
                 }
                 Poll::Ready(NetworkBehaviourAction::DialAddress { address, handler }) => {
                     return Poll::Ready(NetworkBehaviourAction::DialAddress { address, handler })

protocols/autonat/src/lib.rs

@@ -20,9 +20,9 @@
 
 //! Implementation of the [AutoNAT] protocol.
 
-pub use self::autonat::{AutoNat, AutoNatEvent, DialResponse};
+pub use self::behaviour::AutoNat;
 
-mod autonat;
+mod behaviour;
 mod protocol;
 
 mod structs_proto {