From 82691067f8e1dbaa5445b1c724641bfb84a44d26 Mon Sep 17 00:00:00 2001
From: Maddiaa0 <47148561+Maddiaa0@users.noreply.github.com>
Date: Fri, 20 Dec 2024 16:51:50 +0000
Subject: [PATCH] test for dropping bad scoring peers

---
 .../p2p/src/services/peer_manager.test.ts     | 36 +++++++++++++++++++
 yarn-project/p2p/src/services/peer_manager.ts |  3 +-
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/yarn-project/p2p/src/services/peer_manager.test.ts b/yarn-project/p2p/src/services/peer_manager.test.ts
index 27991c8626b9..2feba167ffd5 100644
--- a/yarn-project/p2p/src/services/peer_manager.test.ts
+++ b/yarn-project/p2p/src/services/peer_manager.test.ts
@@ -264,6 +264,42 @@ describe('PeerManager', () => {
       expect(peer2?.status).toBe('cached'); // in dial queue
     });
 
+    it('should disconnect from unhealthy peers during heartbeat', async () => {
+      // Create two peers with different states
+      const bannedPeerId = await createSecp256k1PeerId();
+      const disconnectPeerId = await createSecp256k1PeerId();
+      const healthyPeerId = await createSecp256k1PeerId();
+
+      // Mock the connections to return our test peers
+      mockLibP2PNode.getConnections.mockReturnValue([
+        { remotePeer: bannedPeerId },
+        { remotePeer: disconnectPeerId },
+        { remotePeer: healthyPeerId },
+      ]);
+
+      // Set the peer scores to trigger different states
+      peerManager.penalizePeer(bannedPeerId, PeerErrorSeverity.LowToleranceError); // Will set score below -100
+      peerManager.penalizePeer(bannedPeerId, PeerErrorSeverity.LowToleranceError); // Additional penalty to ensure banned state
+
+      peerManager.penalizePeer(disconnectPeerId, PeerErrorSeverity.LowToleranceError); // Will set score between -100 and -50
+      peerManager.penalizePeer(disconnectPeerId, PeerErrorSeverity.HighToleranceError);
+
+      // Trigger heartbeat which should call pruneUnhealthyPeers
+      peerManager.heartbeat();
+
+      await sleep(100);
+
+      // Verify that hangUp was called for both unhealthy peers
+      expect(mockLibP2PNode.hangUp).toHaveBeenCalledWith(bannedPeerId);
+      expect(mockLibP2PNode.hangUp).toHaveBeenCalledWith(disconnectPeerId);
+
+      // Verify that hangUp was not called for the healthy peer
+      expect(mockLibP2PNode.hangUp).not.toHaveBeenCalledWith(healthyPeerId);
+
+      // Verify hangUp was called exactly twice (once for each unhealthy peer)
+      expect(mockLibP2PNode.hangUp).toHaveBeenCalledTimes(2);
+    });
+
     it('should properly clean up peers on stop', async () => {
       const enr = await createMockENR();
       await discoveredPeerCallback(enr);
diff --git a/yarn-project/p2p/src/services/peer_manager.ts b/yarn-project/p2p/src/services/peer_manager.ts
index 07e0e3ae4534..d6133bfe9e42 100644
--- a/yarn-project/p2p/src/services/peer_manager.ts
+++ b/yarn-project/p2p/src/services/peer_manager.ts
@@ -225,7 +225,7 @@ export class PeerManager extends WithTracer {
     const connectedHealthyPeers: Connection[] = [];
 
     for (const peer of connections) {
-      const score = this.peerScoring.getScore(peer.remotePeer.toString());
+      const score = this.peerScoring.getScoreState(peer.remotePeer.toString());
       switch (score) {
         // TODO: add goodbye and give reasons
         case PeerScoreState.Banned:
@@ -242,6 +242,7 @@ export class PeerManager extends WithTracer {
 
   // TODO: send a goodbye with a reason to the peer
   private async disconnectPeer(peer: PeerId) {
+    this.logger.debug(`Disconnecting peer ${peer.toString()}`);
     await this.libP2PNode.hangUp(peer);
   }