Merge pull request #110 from input-output-hk/ETCM-443-fix-discovery-k…

…ey-not-found ETCM-443: Fix the KTable and the kademlia ID mapping getting out of sync
input-output-hk · Dec 3, 2020 · d5a2123 · d5a2123
2 parents dc8142a + b00832b
commit d5a2123
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 1 deletion.
diff --git a/scalanet/discovery/src/io/iohk/scalanet/discovery/ethereum/v4/DiscoveryService.scala b/scalanet/discovery/src/io/iohk/scalanet/discovery/ethereum/v4/DiscoveryService.scala
@@ -188,8 +188,13 @@ object DiscoveryService {
       )
     }
 
+    /** Update the timestamp of the peer in the K-table, if it's still part of it. */
     def withTouch(peer: Peer[A]): State[A] =
-      copy(kBuckets = kBuckets.touch(peer.kademliaId))
+      if (kBuckets.contains(peer.kademliaId))
+        copy(kBuckets = kBuckets.touch(peer.kademliaId))
+      else
+        // Not adding because `kademliaIdToNodeId` and `nodeMap` may no longer have this peer.
+        this
 
     def clearBondingResults(peer: Peer[A]): State[A] =
       copy(bondingResultsMap = bondingResultsMap - peer)

diff --git a/scalanet/discovery/ut/src/io/iohk/scalanet/discovery/ethereum/v4/DiscoveryServiceSpec.scala b/scalanet/discovery/ut/src/io/iohk/scalanet/discovery/ethereum/v4/DiscoveryServiceSpec.scala
@@ -1003,6 +1003,30 @@ class DiscoveryServiceSpec extends AsyncFlatSpec with Matchers {
       }
     }
   }
+
+  behavior of "withTouch"
+
+  it should "not touch a peer not already in the k-table" in test {
+    new Fixture {
+      override val test = for {
+        state0 <- stateRef.get
+        state1 = state0.withTouch(remotePeer)
+        state2 = state1.withEnrAndAddress(remotePeer, remoteENR, remoteNode.address)
+        _ <- Task.sleep(1.milli) // If we're too quick then TimeSet will assign the same timestamp.
+        state3 = state2.withTouch(remotePeer)
+      } yield {
+        val peerId = remotePeer.kademliaId
+
+        state0.kBuckets.contains(peerId) shouldBe false
+        state1.kBuckets.contains(peerId) shouldBe false
+        state2.kBuckets.contains(peerId) shouldBe true
+
+        val (_, bucket2) = state2.kBuckets.getBucket(peerId)
+        val (_, bucket3) = state3.kBuckets.getBucket(peerId)
+        bucket2.timestamps should not equal bucket3.timestamps
+      }
+    }
+  }
 }
 
 object DiscoveryServiceSpec {

diff --git a/scalanet/discovery/ut/src/io/iohk/scalanet/kademlia/KBucketsSpec.scala b/scalanet/discovery/ut/src/io/iohk/scalanet/kademlia/KBucketsSpec.scala
@@ -36,6 +36,10 @@ class KBucketsSpec extends FlatSpec {
     kb.add(v).contains(v) shouldBe true
   }
 
+  they should "retrieve any node added via touch" in forAll(genBitVector()) { v =>
+    kb.touch(v).contains(v) shouldBe true
+  }
+
   they should "not retrieve any node removed via remove" in forAll(genBitVector()) { v =>
     kb.add(v).remove(v).contains(v) shouldBe false
   }