feat: impl NetworkGroup for peer and multiaddr

nervosnetwork · Dec 18, 2018 · e1e5750 · e1e5750
1 parent f61ddaf
commit e1e5750
Show file tree

Hide file tree

Showing 6 changed files with 65 additions and 6 deletions.
diff --git a/network/src/lib.rs b/network/src/lib.rs
@@ -9,6 +9,7 @@ mod identify_service;
 mod memory_peer_store;
 mod network;
 mod network_config;
+mod network_group;
 mod network_service;
 mod outgoing_service;
 mod peer_store;
@@ -25,7 +26,9 @@ pub use crate::errors::{Error, ErrorKind};
 pub use crate::network::{Network, PeerInfo, SessionInfo};
 pub use crate::network_config::NetworkConfig;
 pub use crate::network_service::NetworkService;
-pub use libp2p::{core::Endpoint, multiaddr::AddrComponent, Multiaddr, PeerId};
+pub use libp2p::{
+    core::Endpoint, multiaddr::AddrComponent, multiaddr::ToMultiaddr, Multiaddr, PeerId,
+};
 
 pub type TimerToken = usize;
 pub type ProtocolId = [u8; 3];

diff --git a/network/src/network.rs b/network/src/network.rs
@@ -147,7 +147,7 @@ impl Network {
                 .map(|addr| addr.to_owned())
                 .collect::<Vec<_>>()
         });
-        addrs.unwrap_or(Vec::new())
+        addrs.unwrap_or_default()
     }
 
     pub(crate) fn peers(&self) -> impl Iterator<Item = PeerId> {

diff --git a/network/src/network_group.rs b/network/src/network_group.rs
@@ -0,0 +1,49 @@
+use libp2p::core::{AddrComponent, Multiaddr};
+use std::net::IpAddr;
+
+pub type Group = Vec<u8>;
+
+pub trait NetworkGroup {
+    fn network_group(&self) -> Group;
+}
+
+fn extract_ip_addr(addr: &Multiaddr) -> Option<IpAddr> {
+    for addr_component in addr {
+        match addr_component {
+            AddrComponent::IP4(ipv4) => return Some(IpAddr::V4(ipv4)),
+            AddrComponent::IP6(ipv6) => return Some(IpAddr::V6(ipv6)),
+            _ => (),
+        }
+    }
+    None
+}
+
+impl NetworkGroup for Multiaddr {
+    fn network_group(&self) -> Group {
+        if let Some(ip_addr) = extract_ip_addr(self) {
+            if ip_addr.is_loopback() {
+                // Local NetworkGroup
+                return vec![1];
+            }
+            // TODO uncomment after ip feature stable
+            // if !ip_addr.is_global() {
+            //     // Global NetworkGroup
+            //     return vec![2]
+            // }
+
+            // IPv4 NetworkGroup
+            if let IpAddr::V4(ipv4) = ip_addr {
+                return ipv4.octets()[0..2].to_vec();
+            }
+            // IPv6 NetworkGroup
+            if let IpAddr::V6(ipv6) = ip_addr {
+                if let Some(ipv4) = ipv6.to_ipv4() {
+                    return ipv4.octets()[0..2].to_vec();
+                }
+                return ipv6.octets()[0..4].to_vec();
+            }
+        }
+        // Can't group addr
+        vec![0]
+    }
+}
diff --git a/network/src/peers_registry.rs b/network/src/peers_registry.rs
@@ -1,3 +1,4 @@
+use crate::network_group::{Group, NetworkGroup};
 use crate::peer_store::PeerStore;
 use crate::{Error, ErrorKind, PeerId, PeerIndex, ProtocolId};
 use bytes::Bytes;
@@ -102,7 +103,7 @@ impl PeerConnection {
     pub fn new(connected_addr: Multiaddr, endpoint_role: Endpoint) -> Self {
         PeerConnection {
             endpoint_role,
-            connected_addr: connected_addr,
+            connected_addr,
             pinger_loader: UniqueConnec::empty(),
             identify_info: None,
             ckb_protocols: Vec::with_capacity(1),
@@ -121,6 +122,12 @@ impl PeerConnection {
     pub fn is_incoming(&self) -> bool {
         !self.is_outgoing()
     }
+
+    #[allow(dead_code)]
+    #[inline]
+    fn network_group(&self) -> Group {
+        self.connected_addr.network_group()
+    }
 }
 
 pub struct ConnectionStatus {

diff --git a/network/src/ping_service.rs b/network/src/ping_service.rs
@@ -125,7 +125,7 @@ impl<T: Send> ProtocolService<T> for PingService {
                 for peer_id in network.peers() {
                     let peer_id = peer_id.clone();
                     // only ping first address?
-                    if let Some(addr) = network.get_peer_remote_addresses(&peer_id).get(0) {
+                    if let Some(addr) = network.get_peer_addresses(&peer_id).get(0) {
                         if let Some(pinger_loader) = network.get_peer_pinger(&peer_id) {
                             let ping_future = pinger_loader
                                 .dial(&swarm_controller, &addr, transport.clone())

diff --git a/sync/src/synchronizer/mod.rs b/sync/src/synchronizer/mod.rs
@@ -716,7 +716,7 @@ mod tests {
     use ckb_db::memorydb::MemoryKeyValueDB;
     use ckb_network::{
         random_peer_id, CKBProtocolContext, Endpoint, Error as NetworkError, PeerIndex, PeerInfo,
-        ProtocolId, SessionInfo, Severity, TimerToken,
+        ProtocolId, SessionInfo, Severity, TimerToken, ToMultiaddr,
     };
     use ckb_notify::{NotifyController, NotifyService};
     use ckb_protocol::{Block as FbsBlock, Headers as FbsHeaders};
@@ -1034,7 +1034,7 @@ mod tests {
                 peer_id: random_peer_id().unwrap(),
                 endpoint_role: Endpoint::Dialer,
                 last_ping_time: None,
-                remote_addresses: vec![],
+                connected_addr: "/ip4/127.0.0.1".to_multiaddr().unwrap(),
                 identify_info: None,
             },
             protocol_version: None,