From d4a253f6004c15723b0033606fd591c2a1a5364e Mon Sep 17 00:00:00 2001
From: Joshua Karp <joshua.karp@matrix.ai>
Date: Thu, 11 Nov 2021 14:53:03 +1100
Subject: [PATCH] Fixing an update to a node (as opposed to a newly added node)
 on a full bucket

---
 src/nodes/NodeGraph.ts        | 21 ++++++++-----
 src/nodes/errors.ts           |  5 +++
 tests/nodes/NodeGraph.test.ts | 57 +++++++++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+), 7 deletions(-)

diff --git a/src/nodes/NodeGraph.ts b/src/nodes/NodeGraph.ts
index f59154477c..cf57a1d9f6 100644
--- a/src/nodes/NodeGraph.ts
+++ b/src/nodes/NodeGraph.ts
@@ -209,17 +209,24 @@ class NodeGraph {
     if (bucket == null) {
       bucket = {};
     }
-    const bucketEntries = Object.entries(bucket);
-    if (bucketEntries.length === this.maxNodesPerBucket) {
-      const leastActive = bucketEntries.reduce((prev, curr) => {
-        return prev[1].lastUpdated < curr[1].lastUpdated ? prev : curr;
-      });
-      delete bucket[leastActive[0]];
-    }
     bucket[nodeId] = {
       address: nodeAddress,
       lastUpdated: new Date(),
     };
+    // Perform the check on size after we add/update the node. If it's an update,
+    // then we don't need to perform the deletion.
+    let bucketEntries = Object.entries(bucket);
+    if (bucketEntries.length > this.maxNodesPerBucket) {
+      const leastActive = bucketEntries.reduce((prev, curr) => {
+        return new Date(prev[1].lastUpdated) < new Date(curr[1].lastUpdated) ? prev : curr;
+      });
+      delete bucket[leastActive[0]];
+      bucketEntries = Object.entries(bucket);
+      // For safety, make sure that the bucket is actually at maxNodesPerBucket
+      if (bucketEntries.length != this.maxNodesPerBucket) {
+        throw new nodesErrors.ErrorNodeGraphOversizedBucket();
+      }
+    }
     return [
       {
         type: 'put',
diff --git a/src/nodes/errors.ts b/src/nodes/errors.ts
index b12e5da619..8447bb5e25 100644
--- a/src/nodes/errors.ts
+++ b/src/nodes/errors.ts
@@ -19,6 +19,10 @@ class ErrorNodeGraphEmptyDatabase extends ErrorNodes {}
 
 class ErrorNodeGraphInvalidBucketIndex extends ErrorNodes {}
 
+class ErrorNodeGraphOversizedBucket extends ErrorNodes {
+  description: 'Bucket invalidly contains more nodes than capacity'
+}
+
 class ErrorNodeConnectionNotStarted extends ErrorNodes {}
 
 class ErrorNodeConnectionDestroyed extends ErrorNodes {}
@@ -53,6 +57,7 @@ export {
   ErrorNodeGraphSelfConnect,
   ErrorNodeGraphEmptyDatabase,
   ErrorNodeGraphInvalidBucketIndex,
+  ErrorNodeGraphOversizedBucket,
   ErrorNodeConnectionNotStarted,
   ErrorNodeConnectionDestroyed,
   ErrorNodeConnectionTimeout,
diff --git a/tests/nodes/NodeGraph.test.ts b/tests/nodes/NodeGraph.test.ts
index 80dace9a48..8624326c5b 100644
--- a/tests/nodes/NodeGraph.test.ts
+++ b/tests/nodes/NodeGraph.test.ts
@@ -388,6 +388,63 @@ describe('NodeGraph', () => {
       fail('Bucket undefined');
     }
   });
+  test('enforces k-bucket size, retaining all nodes if adding a pre-existing node', async () => {
+    // Add k nodes to the database (importantly, they all go into the same bucket)
+    let currNodeId = nodesTestUtils.generateNodeIdForBucket(nodeId, 59);
+    // Keep a record of the first node ID that we added
+    const firstNodeId = currNodeId;
+    for (let i = 1; i <= nodeGraph.maxNodesPerBucket; i++) {
+      // Add the current node ID
+      const nodeAddress = {
+        ip: (i + '.' + i + '.' + i + '.' + i) as Host,
+        port: i as Port,
+      };
+      await nodeGraph.setNode(currNodeId, nodeAddress);
+      // Increment the current node ID - skip for the last one to keep currNodeId
+      // as the last added node ID
+      if (i != nodeGraph.maxNodesPerBucket) {
+        const incrementedNodeId = nodesTestUtils.incrementNodeId(currNodeId);
+        currNodeId = incrementedNodeId;
+      }
+    }
+    // All of these nodes are in bucket 59
+    const originalBucket = await nodeGraph.getBucket(59);
+    if (originalBucket) {
+      expect(Object.keys(originalBucket).length).toBe(
+        nodeGraph.maxNodesPerBucket,
+      );
+    } else {
+      // Should be unreachable
+      fail('Bucket undefined');
+    }
+
+    // If we tried to re-add the first node, it would simply remove the original
+    // first node, as this is the "least active".
+    // We instead want to check that we don't mistakenly delete a node if we're
+    // updating an existing one.
+    // So, re-add the last node
+    const newLastAddress: NodeAddress = {
+      ip: '30.30.30.30' as Host,
+      port: 30 as Port,
+    };
+    await nodeGraph.setNode(currNodeId, newLastAddress);
+
+    const finalBucket = await nodeGraph.getBucket(59);
+    if (finalBucket) {
+      // We should still have a full bucket
+      expect(Object.keys(finalBucket).length).toEqual(
+        nodeGraph.maxNodesPerBucket,
+      );
+      // Ensure that this new node is in the bucket
+      expect(finalBucket[currNodeId]).toEqual({
+        address: newLastAddress,
+        lastUpdated: expect.any(Date),
+      });
+    } else {
+      // Should be unreachable
+      fail('Bucket undefined');
+    }
+  });
   test('retrieves all buckets (in expected lexicographic order)', async () => {
     // Bucket 0 is expected to never have any nodes (as nodeId XOR 0 = nodeId)
     // Bucket 1 (minimum):