Fix bugs

95th · Jul 22, 2020 · 27375a5 · 27375a5
1 parent f8e82a4
commit 27375a5
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 16 deletions.
diff --git a/dht/src/contact.rs b/dht/src/contact.rs
@@ -1,7 +1,6 @@
 use crate::id::NodeId;
 use ben::{Encode, Encoder};
 use std::net::SocketAddr;
-use std::time::Instant;
 
 bitflags! {
     pub struct ContactStatus: u8 {
@@ -32,9 +31,6 @@ impl ContactRef<'_> {
 pub struct Contact {
     pub id: NodeId,
     pub addr: SocketAddr,
-    pub last_updated: Instant,
-    pub last_queried: Instant,
-    pub token: Vec<u8>,
     pub status: ContactStatus,
     timeout_count: Option<u8>,
 }
@@ -44,9 +40,6 @@ impl Contact {
         Self {
             id,
             addr,
-            last_updated: Instant::now(),
-            last_queried: Instant::now(),
-            token: vec![],
             timeout_count: None,
             status: ContactStatus::INITIAL,
         }
@@ -59,10 +52,6 @@ impl Contact {
         }
     }
 
-    pub fn touch(&mut self) {
-        self.last_updated = Instant::now();
-    }
-
     pub fn is_pinged(&self) -> bool {
         self.timeout_count.is_some()
     }

diff --git a/dht/src/server/traversal.rs b/dht/src/server/traversal.rs
@@ -54,6 +54,10 @@ impl Traversal {
         }
     }
 
+    /// Handle an incoming response and return `true` if it
+    /// was handled in this traversal.
+    /// Returning `false` means that the response didn't belong
+    /// to this traversal.
     pub fn handle_reply(
         &mut self,
         resp: &Response,

diff --git a/dht/src/server/traversal/bootstrap.rs b/dht/src/server/traversal/bootstrap.rs
@@ -37,7 +37,7 @@ impl BootstrapTraversal {
         if self.nodes.len() < 3 {
             for node in &table.router_nodes {
                 self.nodes.push(TraversalNode {
-                    id: NodeId::gen(),
+                    id: NodeId::new(),
                     addr: *node,
                     status: Status::INITIAL | Status::NO_ID,
                 });
@@ -65,7 +65,13 @@ impl BootstrapTraversal {
     ) -> bool {
         if let Some(req) = self.txns.remove(&resp.txn_id) {
             if req.has_id {
-                table.heard_from(&req.id);
+                if &req.id == resp.id {
+                    table.heard_from(&req.id);
+                } else {
+                    warn!("ID mismatch from {}", addr);
+                    table.failed(&req.id);
+                    return true;
+                }
             }
         } else {
             return false;

diff --git a/dht/src/server/traversal/get_peers.rs b/dht/src/server/traversal/get_peers.rs
@@ -45,7 +45,7 @@ impl GetPeersTraversal {
         if self.nodes.len() < 3 {
             for node in &table.router_nodes {
                 self.nodes.push(TraversalNode {
-                    id: NodeId::gen(),
+                    id: NodeId::new(),
                     addr: *node,
                     status: Status::INITIAL | Status::NO_ID,
                 });
@@ -73,7 +73,13 @@ impl GetPeersTraversal {
     ) -> bool {
         if let Some(req) = self.txns.remove(&resp.txn_id) {
             if req.has_id {
-                table.heard_from(&req.id);
+                if &req.id == resp.id {
+                    table.heard_from(&req.id);
+                } else {
+                    warn!("ID mismatch from {}", addr);
+                    table.failed(&req.id);
+                    return true;
+                }
             }
         } else {
             return false;

diff --git a/dht/src/table.rs b/dht/src/table.rs
@@ -114,14 +114,19 @@ impl RoutingTable {
         self.buckets[idx].live.iter_mut().find(|c| c.id == *id)
     }
 
+    pub fn failed(&mut self, id: &NodeId) {
+        if let Some(c) = self.find_contact(id) {
+            c.timed_out();
+        }
+    }
+
     pub fn heard_from(&mut self, id: &NodeId) {
         let idx = self.find_bucket(id);
         let bucket = &mut self.buckets[idx];
 
         if let Some(c) = bucket.live.iter_mut().find(|c| c.id == *id) {
             c.status = ContactStatus::ALIVE | ContactStatus::QUERIED;
             c.clear_timeout();
-            c.last_updated = Instant::now();
             bucket.last_updated = Instant::now();
         }
     }