From 73bbaf36bb189a4e09c3e547e1bf82aad6c8322a Mon Sep 17 00:00:00 2001
From: MudDev <trademud@gmail.com>
Date: Fri, 24 Oct 2025 08:11:30 -0600
Subject: [PATCH] Fix KeyError in aiobtdht by applying a monkey patch to the
 Bucket.add method

---
 src/DHT/DHTServer.py    |  4 ++++
 src/DHT/aiobtdht_fix.py | 52 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+)
 create mode 100644 src/DHT/aiobtdht_fix.py

diff --git a/src/DHT/DHTServer.py b/src/DHT/DHTServer.py
index cf2ff0002..c02b19c89 100644
--- a/src/DHT/DHTServer.py
+++ b/src/DHT/DHTServer.py
@@ -7,6 +7,10 @@
 
 import asyncio_gevent
 
+# Apply fix for aiobtdht KeyError bug before importing DHT
+from .aiobtdht_fix import patch_aiobtdht
+patch_aiobtdht()
+
 from aiobtdht import DHT
 from aioudp import UDPServer
 
diff --git a/src/DHT/aiobtdht_fix.py b/src/DHT/aiobtdht_fix.py
new file mode 100644
index 000000000..489df944d
--- /dev/null
+++ b/src/DHT/aiobtdht_fix.py
@@ -0,0 +1,52 @@
+"""
+Monkey patch for aiobtdht to fix KeyError bug in routing_table/bucket.py
+
+The bug: When removing nodes with negative rate, the code tries to pop a tuple (node, stat)
+from the _nodes dictionary, but the dictionary keys are just node objects.
+
+The fix: Extract the node object from the tuple before popping.
+"""
+
+def patch_aiobtdht():
+    """Apply the fix to aiobtdht.routing_table.bucket.Bucket.add method"""
+    try:
+        from aiobtdht.routing_table.bucket import Bucket
+        from aiobtdht.routing_table.node_stat import NodeStat
+        
+        # Store the original add method
+        original_add = Bucket.add
+        
+        def fixed_add(self, node):
+            """Fixed version of Bucket.add that properly unpacks the tuple when deleting nodes"""
+            if not self.id_in_range(node.id):
+                raise IndexError("Node id not in bucket range")
+
+            if node in self._nodes:
+                self._nodes[node].renew()
+                return True
+            elif len(self._nodes) < self._max_capacity:
+                self._nodes[node] = NodeStat()
+                return True
+            else:
+                can_delete = list(filter(lambda it: it[1].rate < 0, self._enum_nodes()))
+                if can_delete:
+                    # FIX: Unpack the tuple (node, stat) and only pop the node
+                    for node_to_delete, _ in can_delete:
+                        self._nodes.pop(node_to_delete)
+
+                    return self.add(node)
+                else:
+                    return False
+        
+        # Replace the method
+        Bucket.add = fixed_add
+        return True
+        
+    except Exception as e:
+        print(f"Warning: Failed to patch aiobtdht: {e}")
+        return False
+
+
+# Apply the patch when this module is imported
+patch_aiobtdht()
+