Просмотр исходного кода

net: AddrsMessage to include both addressis with and without preferred transports

aggstam 2 лет назад
Родитель
Сommit
4ff6176fdd
3 измененных файлов с 81 добавлено и 4 удалено
  1. 63 1
      src/net/hosts.rs
  2. 4 1
      src/net/message.rs
  3. 14 2
      src/net/protocol/protocol_address.rs

+ 63 - 1
src/net/hosts.rs

@@ -296,8 +296,26 @@ impl Hosts {
         urls.iter().map(|&url| url.clone()).collect()
     }
 
+    /// Get up to n random peers that don't match the given transport schemes from the hosts set.
+    pub async fn fetch_n_random_excluding_schemes(&self, schemes: &[String], n: u32) -> Vec<Url> {
+        let n = n as usize;
+        if n == 0 {
+            return vec![]
+        }
+
+        // Retrieve all peers not corresponding to that transport schemes
+        let hosts = self.fetch_exluding_schemes(schemes, None).await;
+        if hosts.is_empty() {
+            return hosts
+        }
+
+        // Grab random ones
+        let urls = hosts.iter().choose_multiple(&mut OsRng, n.min(hosts.len()));
+        urls.iter().map(|&url| url.clone()).collect()
+    }
+
     /// Get up to limit peers that match the given transport schemes from the hosts set.
-    /// If limit was not provided, return all peers.
+    /// If limit was not provided, return all matching peers.
     pub async fn fetch_with_schemes(&self, schemes: &[String], limit: Option<usize>) -> Vec<Url> {
         let addrs = self.addrs.read().await;
         let mut limit = match limit {
@@ -335,6 +353,50 @@ impl Hosts {
 
         ret
     }
+
+    /// Get up to limit peers that don't match the given transport schemes from the hosts set.
+    /// If limit was not provided, return all matching peers.
+    pub async fn fetch_exluding_schemes(
+        &self,
+        schemes: &[String],
+        limit: Option<usize>,
+    ) -> Vec<Url> {
+        let addrs = self.addrs.read().await;
+        let mut limit = match limit {
+            Some(l) => l.min(addrs.len()),
+            None => addrs.len(),
+        };
+        let mut ret = vec![];
+
+        if limit == 0 {
+            return ret
+        }
+
+        for addr in addrs.iter() {
+            if !schemes.contains(&addr.scheme().to_string()) {
+                ret.push(addr.clone());
+                limit -= 1;
+                if limit == 0 {
+                    return ret
+                }
+            }
+        }
+
+        // If we didn't find any, pick some from the quarantine zone
+        if ret.is_empty() {
+            for addr in self.quarantine.read().await.keys() {
+                if !schemes.contains(&addr.scheme().to_string()) {
+                    ret.push(addr.clone());
+                    limit -= 1;
+                    if limit == 0 {
+                        break
+                    }
+                }
+            }
+        }
+
+        ret
+    }
 }
 
 #[cfg(test)]

+ 4 - 1
src/net/message.rs

@@ -59,7 +59,10 @@ impl_p2p_message!(PongMessage, "pong");
 /// Requests address of outbound connecction.
 #[derive(Debug, Clone, SerialEncodable, SerialDecodable)]
 pub struct GetAddrsMessage {
-    /// Maximum number of addresses to receive
+    /// Maximum number of addresses with preferred
+    /// transports to receive. Response vector will
+    /// also containg addresses without the preferred
+    /// transports, so its size will be 2 * max.
     pub max: u32,
     /// Preferred addresses transports
     pub transports: Vec<String>,

+ 14 - 2
src/net/protocol/protocol_address.rs

@@ -94,7 +94,7 @@ impl ProtocolAddress {
             // TODO: We might want to close the channel here if we're getting
             // corrupted addresses.
             // Validate addreses length
-            if addrs_msg.addrs.len() > self.settings.outbound_connections {
+            if addrs_msg.addrs.len() > 2 * self.settings.outbound_connections {
                 continue
             }
 
@@ -128,10 +128,22 @@ impl ProtocolAddress {
                 continue
             }
 
-            let addrs = self
+            // First we grab address with the requested transports
+            let mut addrs = self
                 .hosts
                 .fetch_n_random_with_schemes(&get_addrs_msg.transports, get_addrs_msg.max)
                 .await;
+
+            // Then we grab addresses without the requested transports
+            // to fill a 2 * max length vector.
+            let remain = 2 * get_addrs_msg.max - get_addrs_msg.max;
+            addrs.append(
+                &mut self
+                    .hosts
+                    .fetch_n_random_excluding_schemes(&get_addrs_msg.transports, remain)
+                    .await,
+            );
+
             debug!(
                 target: "net::protocol_address::handle_receive_get_addrs()",
                 "Sending {} addresses to {}", addrs.len(), self.channel.address(),