1 files changed, 69 insertions, 2 deletions

diff --git a/src/Network/BitTorrent/Address.hs b/src/Network/BitTorrent/Address.hs
index c0c0a3f8..3d2616cc 100644
--- a/src/Network/BitTorrent/Address.hs
+++ b/src/Network/BitTorrent/Address.hs
@@ -58,9 +58,12 @@ module Network.BitTorrent.Address
          -- ** Id
        , NodeId
        , testIdBit
-       , genNodeId
        , NodeDistance
        , distance
+       , genNodeId
+       , bucketRange
+       , genBucketSample
+       , bep42
 
          -- ** Info
        , NodeAddr (..)
@@ -118,6 +121,7 @@ import Text.PrettyPrint.HughesPJClass hiding ((<>),($$))
 import System.Locale    (defaultTimeLocale)
 #endif
 import System.Entropy
+import Data.Digest.CRC32C
 
 -- import Paths_bittorrent (version)
 
@@ -662,7 +666,7 @@ testIdBit :: NodeId -> Word -> Bool
 testIdBit (NodeId bs) i
   | fromIntegral i < nodeIdSize * 8
   , (q, r) <- quotRem (fromIntegral i) 8
-  = testBit (BS.index bs q) r
+  = testBit (BS.index bs q) (7 - r)
   |     otherwise      = False
 {-# INLINE testIdBit #-}
 
@@ -690,6 +694,47 @@ instance Pretty NodeDistance where
 distance :: NodeId -> NodeId -> NodeDistance
 distance (NodeId a) (NodeId b) = NodeDistance (BS.pack (BS.zipWith xor a b))
 
+-- | Accepts a depth/index of a bucket and whether or not it is the last one,
+-- yields:
+--
+--  count of leading bytes to be copied from your node id.
+--
+--  mask to clear the extra bits of the last copied byte
+--
+--  mask to toggle the last copied bit if it is not the last bucket
+--
+-- Normally this is used with 'genBucketSample' to obtain a random id suitable
+-- for refreshing a particular bucket.
+bucketRange :: Int -> Bool -> (Int, Word8, Word8)
+bucketRange depth is_last = (q,m,b)
+ where
+     (q,r) = divMod ((if is_last then (+7) else (+8)) depth) 8
+     m = 2^(7-r) - 1
+     b = if is_last then 0 else 2^(7-r)
+
+-- | Generate a random 'NodeId' within a range suitable for a bucket.  To
+-- obtain a sample for bucket number /index/ where /is_last/ indicates if this
+-- is for the current deepest bucket in our routing table:
+--
+-- > sample <- genBucketSample nid (bucketRange index is_last)
+genBucketSample :: NodeId -> (Int,Word8,Word8) -> IO NodeId
+genBucketSample n qmb = genBucketSample' getEntropy n qmb
+
+-- | Generalizion of 'genBucketSample' that accepts a byte generator
+-- function to use instead of the system entropy.
+genBucketSample' :: Applicative m =>
+    (Int -> m ByteString) -> NodeId -> (Int,Word8,Word8) -> m NodeId
+genBucketSample' gen (NodeId self) (q,m,b)
+    | q <= 0           =  NodeId <$> gen nodeIdSize
+    | q >= nodeIdSize  =  pure (NodeId self)
+    | otherwise        =  NodeId . build <$> gen (nodeIdSize - q + 1)
+ where
+    build tl = BS.init hd <> BS.cons (h .|. t) (BS.tail tl)
+     where
+        hd = BS.take q self
+        h = xor b (complement m .&. BS.last hd)
+        t = m .&. BS.head tl
+
 ------------------------------------------------------------------------
 
 data NodeAddr a = NodeAddr
@@ -1176,3 +1221,25 @@ fingerprint pid = either (const def) id $ runGet getCI (getPeerId pid)
     getShadowVersion = do
       str <- BC.unpack <$> getByteString 5
       return $ Version (catMaybes $ L.map decodeShadowVerNr str) []
+
+
+-- | Modifies a purely random 'NodeId' to one that is related to a given
+-- routable address in accordance with BEP 42.
+bep42 :: Address a => a -> NodeId -> Maybe NodeId
+bep42 addr (NodeId r)
+ | Just ip <- fmap S.encode (fromAddr addr :: Maybe IPv4)
+              <|> fmap S.encode (fromAddr addr :: Maybe IPv6)
+    = genBucketSample' retr (NodeId $ crc $ masked ip) (3,0x07,0)
+ | otherwise
+    = Nothing
+ where
+    ip4mask = "\x03\x0f\x3f\xff" :: ByteString
+    ip6mask = "\x01\x03\x07\x0f\x1f\x3f\x7f\xff" :: ByteString
+    rbyte = BS.last r
+    retr n = pure $ BS.drop (BS.length r - n) r
+    crc = S.encode . crc32c . BS.pack
+    masked ip = case BS.zipWith (.&.) msk ip of
+                    (b:bs) -> (b .|. shiftL (rbyte .&. 7) 5) : bs
+                    bs     -> bs
+        where msk | BS.length ip == 4  =  ip4mask
+                  | otherwise          =  ip6mask