Method to recover thisNodeId from stored table.

1 files changed, 33 insertions, 0 deletions

diff --git a/src/Network/BitTorrent/DHT/Routing.hs b/src/Network/BitTorrent/DHT/Routing.hs
index 14aec612..a4da8445 100644
--- a/src/Network/BitTorrent/DHT/Routing.hs
+++ b/src/Network/BitTorrent/DHT/Routing.hs
@@ -37,6 +37,7 @@ module Network.BitTorrent.DHT.Routing
        , shape
        , Network.BitTorrent.DHT.Routing.size
        , Network.BitTorrent.DHT.Routing.depth
+       , compatibleNodeId
 
          -- * Lookup
        , K
@@ -77,6 +78,8 @@ import Data.Word
 import GHC.Generics
 import Text.PrettyPrint as PP hiding ((<>))
 import Text.PrettyPrint.HughesPJClass (pPrint,Pretty)
+import qualified Data.ByteString as BS
+import Data.Bits
 
 import Data.Torrent
 import Network.BitTorrent.Address
@@ -442,6 +445,36 @@ lookupBucket nid = go 0
       |     otherwise     = pure bucket
     go _ (Tip _ _ bucket) = pure bucket
 
+compatibleNodeId :: Table ip -> IO NodeId
+compatibleNodeId tbl = genBucketSample prefix br
+ where
+    br = bucketRange (L.length (shape tbl) - 1) True
+    bs = BS.pack $ take nodeIdSize $ tablePrefix tbl ++ repeat 0
+    prefix = NodeId bs
+
+tablePrefix :: Table ip -> [Word8]
+tablePrefix = map (packByte . take 8 . (++repeat False))
+              . chunksOf 8
+              . tableBits
+ where
+    packByte = foldl1' (.|.) . zipWith bitmask [7,6 .. 0]
+    bitmask ix True = bit ix
+    bitmask _  _    = 0
+
+tableBits :: Table ip -> [Bool]
+tableBits (One _ tbl)  = True : tableBits tbl
+tableBits (Zero tbl _) = False : tableBits tbl
+tableBits (Tip _ _ _)  = []
+
+chunksOf :: Int -> [e] -> [[e]]
+chunksOf i ls = map (take i) (build (splitter ls)) where
+  splitter :: [e] -> ([e] -> a -> a) -> a -> a
+  splitter [] _ n = n
+  splitter l c n  = l `c` splitter (drop i l) c n
+
+build :: ((a -> [a] -> [a]) -> [a] -> [a]) -> [a]
+build g = g (:) []
+
 -- | Count of closest nodes in find_node request.
 type K = Int