1 files changed, 17 insertions, 14 deletions

diff --git a/src/Network/BitTorrent/DHT/Routing.hs b/src/Network/BitTorrent/DHT/Routing.hs
index 5c9788dc..cf4a4de3 100644
--- a/src/Network/BitTorrent/DHT/Routing.hs
+++ b/src/Network/BitTorrent/DHT/Routing.hs
@@ -442,16 +442,16 @@ size = L.sum . shape
 depth :: Table ip -> BucketCount
 depth = L.length . shape
 
-lookupBucket :: NodeId -> Table ip -> Maybe (Bucket ip)
-lookupBucket nid = go 0
+lookupBucket :: NodeId -> Table ip -> [Bucket ip]
+lookupBucket nid = go 0 []
   where
-    go i (Zero table bucket)
-      |  testIdBit nid i  = pure bucket
-      |     otherwise     = go (succ i) table
-    go i (One  bucket table)
-      |  testIdBit nid i  = go (succ i) table
-      |     otherwise     = pure bucket
-    go _ (Tip _ _ bucket) = pure bucket
+    go i bs (Zero table bucket)
+      |  testIdBit nid i     = bucket : toBucketList table ++ bs
+      |     otherwise        = go (succ i) (bucket:bs) table
+    go i bs (One  bucket table)
+      |  testIdBit nid i     = go (succ i) (bucket:bs) table
+      |     otherwise        = bucket : toBucketList table ++ bs
+    go _ bs (Tip _ _ bucket) = bucket : bs
 
 compatibleNodeId :: Table ip -> IO NodeId
 compatibleNodeId tbl = genBucketSample prefix br
@@ -504,11 +504,14 @@ instance TableKey InfoHash where
 -- | Get a list of /K/ closest nodes using XOR metric. Used in
 -- 'find_node' and 'get_peers' queries.
 kclosest :: Eq ip => TableKey a => K -> a -> Table ip -> [NodeInfo ip]
-kclosest k (toNodeId -> nid)
-  = L.take k . rank nodeId nid
-  . L.map PSQ.key . PSQ.toList . fromMaybe PSQ.empty
-  . fmap bktNodes
-  . lookupBucket nid
+kclosest k (toNodeId -> nid) tbl = take k $ rank nodeId nid (L.concat bucket)
+                                          ++ rank nodeId nid (L.concat everyone)
+ where
+    (bucket,everyone) =
+        L.splitAt 1
+        . L.map (L.map PSQ.key . PSQ.toList . bktNodes)
+        . lookupBucket nid
+        $ tbl
 
 {-----------------------------------------------------------------------
 --  Routing