Browse Source

Make assertion hold

If we have peer which is already used, make addAndCheckoutPeer() fail.
If peer is in unused list, first remove the old entry, and then add
new one.  This would make assertion in ~DefaultPeerStorage hold.
Tatsuhiro Tsujikawa 9 năm trước cách đây
mục cha
commit
7a31f9bf29
2 tập tin đã thay đổi với 51 bổ sung1 xóa
  1. 17 1
      src/DefaultPeerStorage.cc
  2. 34 0
      test/DefaultPeerStorageTest.cc

+ 17 - 1
src/DefaultPeerStorage.cc

@@ -164,8 +164,24 @@ std::shared_ptr<Peer>
 DefaultPeerStorage::addAndCheckoutPeer(const std::shared_ptr<Peer>& peer,
                                        cuid_t cuid)
 {
+  if (isPeerAlreadyAdded(peer)) {
+    auto it = std::find_if(std::begin(unusedPeers_), std::end(unusedPeers_),
+                           [&peer](const std::shared_ptr<Peer>& p) {
+                             return p->getIPAddress() == peer->getIPAddress() &&
+                                    p->getOrigPort() == peer->getOrigPort();
+                           });
+    if (it == std::end(unusedPeers_)) {
+      // peer is in usedPeers_.
+      return nullptr;
+    }
+
+    unusedPeers_.erase(it);
+  }
+  else {
+    addUniqPeer(peer);
+  }
+
   unusedPeers_.push_front(peer);
-  addUniqPeer(peer);
 
   return checkoutPeer(cuid);
 }

+ 34 - 0
test/DefaultPeerStorageTest.cc

@@ -19,6 +19,7 @@ class DefaultPeerStorageTest : public CppUnit::TestFixture {
   CPPUNIT_TEST(testCountAllPeer);
   CPPUNIT_TEST(testDeleteUnusedPeer);
   CPPUNIT_TEST(testAddPeer);
+  CPPUNIT_TEST(testAddAndCheckoutPeer);
   CPPUNIT_TEST(testIsPeerAvailable);
   CPPUNIT_TEST(testCheckoutPeer);
   CPPUNIT_TEST(testReturnPeer);
@@ -42,6 +43,7 @@ public:
   void testCountAllPeer();
   void testDeleteUnusedPeer();
   void testAddPeer();
+  void testAddAndCheckoutPeer();
   void testIsPeerAvailable();
   void testCheckoutPeer();
   void testReturnPeer();
@@ -121,6 +123,38 @@ void DefaultPeerStorageTest::testAddPeer()
   CPPUNIT_ASSERT(!ps.addPeer(peer1));
 }
 
+void DefaultPeerStorageTest::testAddAndCheckoutPeer()
+{
+  DefaultPeerStorage ps;
+  auto btRuntime = std::make_shared<BtRuntime>();
+
+  ps.setBtRuntime(btRuntime);
+
+  auto peer1 = std::make_shared<Peer>("192.168.0.1", 6889);
+  // peer2 and peer3 have the same address and port deliberately.
+  auto peer2 = std::make_shared<Peer>("192.168.0.2", 6889);
+  auto peer3 = std::make_shared<Peer>("192.168.0.2", 6889);
+
+  CPPUNIT_ASSERT(ps.addPeer(peer1));
+  CPPUNIT_ASSERT(ps.addPeer(peer2));
+  CPPUNIT_ASSERT_EQUAL((size_t)2, ps.getUnusedPeers().size());
+
+  auto usedPeer = ps.checkoutPeer(1);
+
+  CPPUNIT_ASSERT_EQUAL((size_t)1, ps.getUnusedPeers().size());
+  CPPUNIT_ASSERT_EQUAL(peer1->getIPAddress(), usedPeer->getIPAddress());
+  // Since peer1 is already used, this fails.
+  CPPUNIT_ASSERT(!ps.addAndCheckoutPeer(peer1, 2));
+  CPPUNIT_ASSERT_EQUAL((size_t)1, ps.getUnusedPeers().size());
+
+  // Since peer2 is not used yet, this succeeds.
+  usedPeer = ps.addAndCheckoutPeer(peer3, 2);
+
+  CPPUNIT_ASSERT_EQUAL((size_t)0, ps.getUnusedPeers().size());
+  CPPUNIT_ASSERT_EQUAL(peer2->getIPAddress(), usedPeer->getIPAddress());
+  CPPUNIT_ASSERT_EQUAL(peer3.get(), usedPeer.get());
+}
+
 void DefaultPeerStorageTest::testIsPeerAvailable()
 {
   DefaultPeerStorage ps;