Merge pull request #210 from tox-rs/kbucket

Add parameter for eviction strategy

Author: kpp

src/lib.rs

@@ -27,7 +27,6 @@ extern crate cookie_factory;
 extern crate sodiumoxide;
 
 #[cfg(test)]
-#[macro_use]
 extern crate quickcheck;
 
 extern crate tokio;

src/toxcore/dht/dht_friend.rs

@@ -8,6 +8,7 @@ use std::net::SocketAddr;
 use toxcore::time::*;
 use toxcore::dht::kbucket::*;
 use toxcore::crypto_core::*;
+use toxcore::dht::packed_node::*;
 use toxcore::dht::server::hole_punching::*;
 
 /// Number of bootstrap nodes each friend has.
@@ -79,6 +80,16 @@ impl DhtFriend {
 
         addrs
     }
+
+    /// Try to add a node to the friend's close nodes list.
+    pub fn try_add_to_close(&mut self, node: &PackedNode) -> bool {
+        self.close_nodes.try_add(&self.pk, node, /* evict */ true)
+    }
+
+    /// Check if a node can be added to the friend's close nodes list.
+    pub fn can_add_to_close(&self, node: &PackedNode) -> bool {
+        self.close_nodes.can_add(&self.pk, node, /* evict */ true)
+    }
 }
 
 #[cfg(test)]
@@ -91,16 +102,15 @@ mod tests {
     use tokio_timer::clock::*;
 
     use toxcore::dht::dht_node::*;
-    use toxcore::dht::packed_node::*;
     use toxcore::time::ConstNow;
 
     #[test]
     fn addr_is_unknown() {
         let pk = gen_keypair().0;
         let mut friend = DhtFriend::new(pk);
 
-        assert!(friend.close_nodes.try_add(&pk, &PackedNode::new("192.168.1.1:12345".parse().unwrap(), &gen_keypair().0)));
-        assert!(friend.close_nodes.try_add(&pk, &PackedNode::new("192.168.1.2:12345".parse().unwrap(), &gen_keypair().0)));
+        assert!(friend.try_add_to_close(&PackedNode::new("192.168.1.1:12345".parse().unwrap(), &gen_keypair().0)));
+        assert!(friend.try_add_to_close(&PackedNode::new("192.168.1.2:12345".parse().unwrap(), &gen_keypair().0)));
 
         assert!(!friend.is_addr_known())
     }
@@ -110,10 +120,10 @@ mod tests {
         let pk = gen_keypair().0;
         let mut friend = DhtFriend::new(pk);
 
-        assert!(friend.close_nodes.try_add(&pk, &PackedNode::new("192.168.1.1:12345".parse().unwrap(), &gen_keypair().0)));
-        assert!(friend.close_nodes.try_add(&pk, &PackedNode::new("192.168.1.2:12345".parse().unwrap(), &gen_keypair().0)));
+        assert!(friend.try_add_to_close(&PackedNode::new("192.168.1.1:12345".parse().unwrap(), &gen_keypair().0)));
+        assert!(friend.try_add_to_close(&PackedNode::new("192.168.1.2:12345".parse().unwrap(), &gen_keypair().0)));
 
-        assert!(friend.close_nodes.try_add(&pk, &PackedNode::new("192.168.1.3:12345".parse().unwrap(), &pk)));
+        assert!(friend.try_add_to_close(&PackedNode::new("192.168.1.3:12345".parse().unwrap(), &pk)));
 
         assert!(friend.is_addr_known())
     }
@@ -135,7 +145,7 @@ mod tests {
         ];
 
         for (&node, &addr) in nodes.iter().zip(addrs.iter()) {
-            friend.close_nodes.try_add(&pk, &node);
+            friend.try_add_to_close(&node);
             let dht_node = friend.close_nodes.get_node_mut(&pk, &node.pk).unwrap();
             dht_node.update_returned_addr(addr);
         }
@@ -167,7 +177,7 @@ mod tests {
         ];
 
         for (&node, &addr) in nodes.iter().zip(addrs.iter()) {
-            friend.close_nodes.try_add(&pk, &node);
+            friend.try_add_to_close(&node);
             let dht_node = friend.close_nodes.get_node_mut(&pk, &node.pk).unwrap();
             dht_node.update_returned_addr(addr);
         }
@@ -181,4 +191,25 @@ mod tests {
             assert!(friend.get_returned_addrs().is_empty());
         });
     }
+
+    #[test]
+    fn can_and_try_add_to_close() {
+        let pk = PublicKey([0; PUBLICKEYBYTES]);
+        let mut friend = DhtFriend::new(pk);
+
+        for i in 0 .. 8 {
+            let addr = SocketAddr::new("1.2.3.4".parse().unwrap(), 12345 + i as u16);
+            let node = PackedNode::new(addr, &PublicKey([i + 2; PUBLICKEYBYTES]));
+            assert!(friend.try_add_to_close(&node));
+        }
+
+        let closer_node = PackedNode::new(
+            "1.2.3.5:12345".parse().unwrap(),
+            &PublicKey([1; PUBLICKEYBYTES])
+        );
+
+        // should add a new closer node with eviction
+        assert!(friend.can_add_to_close(&closer_node));
+        assert!(friend.try_add_to_close(&closer_node));
+    }
 }

src/toxcore/dht/dht_node.rs

@@ -213,9 +213,6 @@ impl DhtNode {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use quickcheck::quickcheck;
-
-    use toxcore::dht::kbucket::*;
 
     #[test]
     fn dht_node_clonable() {
@@ -226,27 +223,4 @@ mod tests {
         let dht_node = DhtNode::new(pn);
         let _ = dht_node.clone();
     }
-
-    #[test]
-    fn dht_node_bucket_try_add_test() {
-        fn with_nodes(n1: PackedNode, n2: PackedNode, n3: PackedNode,
-                      n4: PackedNode, n5: PackedNode, n6: PackedNode,
-                      n7: PackedNode, n8: PackedNode) {
-            let pk = PublicKey([0; PUBLICKEYBYTES]);
-            let mut bucket = Bucket::new(BUCKET_DEFAULT_SIZE);
-            assert_eq!(true, bucket.try_add(&pk, &n1));
-            assert_eq!(true, bucket.try_add(&pk, &n2));
-            assert_eq!(true, bucket.try_add(&pk, &n3));
-            assert_eq!(true, bucket.try_add(&pk, &n4));
-            assert_eq!(true, bucket.try_add(&pk, &n5));
-            assert_eq!(true, bucket.try_add(&pk, &n6));
-            assert_eq!(true, bucket.try_add(&pk, &n7));
-            assert_eq!(true, bucket.try_add(&pk, &n8));
-
-            // updating bucket
-            assert_eq!(true, bucket.try_add(&pk, &n1));
-        }
-        quickcheck(with_nodes as fn(PackedNode, PackedNode, PackedNode, PackedNode,
-                                    PackedNode, PackedNode, PackedNode, PackedNode));
-    }
 }