Don't tell peers about themselves in kademlia

Add test about find peer responses
Peergos · Nov 4, 2023 · 0993997 · 0993997
1 parent 22cc9be
commit 0993997
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 5 deletions.
diff --git a/src/main/java/org/peergos/protocol/dht/KademliaEngine.java b/src/main/java/org/peergos/protocol/dht/KademliaEngine.java
@@ -60,11 +60,11 @@ public void setAddressBook(AddressBook addrs) {
         this.addressBook = addrs;
     }
 
-    public synchronized void addOutgoingConnection(PeerId peer, Multiaddr addr) {
+    public synchronized void addOutgoingConnection(PeerId peer) {
         router.touch(Instant.now(), new Node(Id.create(Hash.sha256(peer.getBytes()), 256), peer.toString()));
     }
 
-    public synchronized void addIncomingConnection(PeerId peer, Multiaddr addr) {
+    public synchronized void addIncomingConnection(PeerId peer) {
         router.touch(Instant.now(), new Node(Id.create(Hash.sha256(peer.getBytes()), 256), peer.toString()));
     }
 
@@ -152,8 +152,11 @@ public void receiveRequest(Dht.Message msg, PeerId source, Stream stream) {
             }
             case FIND_NODE: {
                 Dht.Message.Builder builder = msg.toBuilder();
-                builder = builder.addAllCloserPeers(getKClosestPeers(msg.getKey().toByteArray())
+                Multihash sourcePeer = Multihash.deserialize(source.getBytes());
+                byte[] target = msg.getKey().toByteArray();
+                builder = builder.addAllCloserPeers(getKClosestPeers(target)
                         .stream()
+                        .filter(p -> ! p.peerId.equals(sourcePeer)) // don't tell a peer about themselves
                         .map(PeerAddresses::toProtobuf)
                         .collect(Collectors.toList()));
                 Dht.Message reply = builder.build();

diff --git a/src/main/java/org/peergos/protocol/dht/KademliaProtocol.java b/src/main/java/org/peergos/protocol/dht/KademliaProtocol.java
@@ -31,7 +31,7 @@ public KademliaProtocol(KademliaEngine engine) {
     @NotNull
     @Override
     protected CompletableFuture<KademliaController> onStartInitiator(@NotNull Stream stream) {
-        engine.addOutgoingConnection(stream.remotePeerId(), stream.getConnection().remoteAddress());
+        engine.addOutgoingConnection(stream.remotePeerId());
         ReplyHandler handler = new ReplyHandler(stream, initiatorSentBytes, initiatorReceivedBytes);
         stream.pushHandler(handler);
         return CompletableFuture.completedFuture(handler);
@@ -40,7 +40,7 @@ protected CompletableFuture<KademliaController> onStartInitiator(@NotNull Stream
     @NotNull
     @Override
     protected CompletableFuture<KademliaController> onStartResponder(@NotNull Stream stream) {
-        engine.addIncomingConnection(stream.remotePeerId(), stream.getConnection().remoteAddress());
+        engine.addIncomingConnection(stream.remotePeerId());
         IncomingRequestHandler handler = new IncomingRequestHandler(engine);
         stream.pushHandler(handler);
         return CompletableFuture.completedFuture(handler);

diff --git a/src/test/java/org/peergos/KademliaTest.java b/src/test/java/org/peergos/KademliaTest.java
@@ -4,6 +4,7 @@
 import io.ipfs.multihash.Multihash;
 import io.libp2p.core.*;
 import io.libp2p.core.crypto.*;
+import io.libp2p.core.multiformats.*;
 import io.libp2p.crypto.keys.*;
 import org.junit.*;
 import org.peergos.blockstore.*;
@@ -122,4 +123,25 @@ public void ipnsBenchmark() throws Exception {
             node2.stop();
         }
     }
+
+    @Test
+    public void kademliaFindNodeLimitTest() {
+        PeerId us = new HostBuilder().generateIdentity().getPeerId();
+        KademliaEngine kad = new KademliaEngine(Multihash.fromBase58(us.toBase58()),
+                new RamProviderStore(1000), new RamRecordStore(), new RamBlockstore());
+        RamAddressBook addrs = new RamAddressBook();
+        kad.setAddressBook(addrs);
+        for (int i=0; i < 1000; i++) {
+            PeerId peer = new HostBuilder().generateIdentity().getPeerId();
+            for (int j=0; j < 100; j++) {
+                kad.addIncomingConnection(peer);
+                addrs.addAddrs(peer, 0, new Multiaddr[]{new Multiaddr("/ip4/127.0.0.1/tcp/4001/p2p/" + peer.toBase58())});
+            }
+        }
+        List<PeerAddresses> closest = kad.getKClosestPeers(new byte[32]);
+        Assert.assertTrue(closest.size() <= 20);
+        for (PeerAddresses addr : closest) {
+            Assert.assertTrue(addr.addresses.size() == 1);
+        }
+    }
 }