Method to recover thisNodeId from stored table.

3 files changed, 37 insertions, 1 deletions

diff --git a/src/Network/BitTorrent/Address.hs b/src/Network/BitTorrent/Address.hs
index 3d2616cc..7ef837db 100644
--- a/src/Network/BitTorrent/Address.hs
+++ b/src/Network/BitTorrent/Address.hs
@@ -56,7 +56,8 @@ module Network.BitTorrent.Address
 
          -- * Node
          -- ** Id
-       , NodeId
+       , NodeId(..)
+       , nodeIdSize
        , testIdBit
        , NodeDistance
        , distance
diff --git a/src/Network/BitTorrent/DHT.hs b/src/Network/BitTorrent/DHT.hs
index a5b4612f..1c62a0e0 100644
--- a/src/Network/BitTorrent/DHT.hs
+++ b/src/Network/BitTorrent/DHT.hs
@@ -129,6 +129,8 @@ dht opts addr action = do
 --
 --   * "router.utorrent.com" since it is just an alias to
 --   "router.bittorrent.com".
+--      XXX: ignoring this advise as it resolves to a different
+--      ip address for me.
 
 -- | List of bootstrap nodes maintained by different bittorrent
 -- software authors.
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