Forráskód Böngészése

Consider IP address and port when finding DHT node.

Tatsuhiro Tsujikawa 14 éve
szülő
commit
920583502e
2 módosított fájl, 23 hozzáadás és 1 törlés
  1. 2 1
      src/DHTBucket.cc
  2. 21 0
      test/DHTBucketTest.cc

+ 2 - 1
src/DHTBucket.cc

@@ -232,7 +232,8 @@ SharedHandle<DHTNode> DHTBucket::getNode(const unsigned char* nodeID, const std:
   node->setPort(port);
   std::deque<SharedHandle<DHTNode> >::const_iterator itr =
     std::find_if(nodes_.begin(), nodes_.end(), derefEqual(node));
-  if(itr == nodes_.end()) {
+  if(itr == nodes_.end() ||
+     (*itr)->getIPAddress() != ipaddr || (*itr)->getPort() != port) {
     return SharedHandle<DHTNode>();
   } else {
     return *itr;

+ 21 - 0
test/DHTBucketTest.cc

@@ -24,6 +24,7 @@ class DHTBucketTest:public CppUnit::TestFixture {
   CPPUNIT_TEST(testGetGoodNodes);
   CPPUNIT_TEST(testCacheNode);
   CPPUNIT_TEST(testDropNode);
+  CPPUNIT_TEST(testGetNode);
   CPPUNIT_TEST_SUITE_END();
 public:
   void setUp() {}
@@ -40,6 +41,7 @@ public:
   void testGetGoodNodes();
   void testCacheNode();
   void testDropNode();
+  void testGetNode();
 };
 
 
@@ -451,4 +453,23 @@ void DHTBucketTest::testDropNode()
   CPPUNIT_ASSERT(*cachedNode1 == *bucket.getCachedNodes()[0]);
 }
 
+void DHTBucketTest::testGetNode()
+{
+  unsigned char localNodeID[DHT_ID_LENGTH];
+  memset(localNodeID, 0, DHT_ID_LENGTH);
+  SharedHandle<DHTNode> localNode(new DHTNode(localNodeID));
+  DHTBucket bucket(localNode);
+
+  unsigned char id[DHT_ID_LENGTH];
+  createID(id, 0xf0, 0);
+  SharedHandle<DHTNode> node(new DHTNode(id));
+  node->setIPAddress("192.168.0.1");
+  node->setPort(6881);
+  CPPUNIT_ASSERT(bucket.addNode(node));
+  CPPUNIT_ASSERT(bucket.getNode(id, "192.168.0.1", 6881));
+  CPPUNIT_ASSERT(!bucket.getNode(id, "192.168.0.2", 6881));
+  CPPUNIT_ASSERT(!bucket.getNode(id, "192.168.0.1", 6882));
+  CPPUNIT_ASSERT(!bucket.getNode(localNodeID, "192.168.0.1", 6881));
+}
+
 } // namespace aria2