1 files changed, 43 insertions, 3 deletions

diff --git a/src/Network/BitTorrent/DHT/Message.hs b/src/Network/BitTorrent/DHT/Message.hs
index 04cc0d41..e21e0e70 100644
--- a/src/Network/BitTorrent/DHT/Message.hs
+++ b/src/Network/BitTorrent/DHT/Message.hs
@@ -74,6 +74,8 @@ module Network.BitTorrent.DHT.Message
          -- ** find_node
        , FindNode (..)
        , NodeFound (..)
+       , bep42s
+       -- , bep42
 
 
 #ifdef VERSION_bencoding
@@ -92,7 +94,7 @@ import Control.Applicative
 import Data.Bool
 #ifdef VERSION_bencoding
 import Data.BEncode as BE
-import Data.BEncode.BDict as BDict
+import Data.BEncode.BDict as BDict hiding (map)
 #else
 import qualified Network.DatagramServer.Tox as Tox
 import Network.DatagramServer.Tox (NodeId)
@@ -101,21 +103,24 @@ import Control.Monad
 #endif
 import Network.KRPC.Method
 import Network.Address hiding (NodeId)
+import Data.Bits
 import Data.ByteString (ByteString)
+import qualified Data.ByteString as BS
+import Data.Digest.CRC32C
 import Data.List as L
 import Data.Monoid
 import Data.Serialize as S
 import Data.Typeable
+import Data.Word
 import Network
 import Network.DatagramServer
-import Network.DatagramServer.Mainline (KMessageOf)
+import Network.DatagramServer.Mainline
 import Data.Maybe
 
 import Data.Torrent (InfoHash)
 import Network.BitTorrent.DHT.Token
 #ifdef VERSION_bencoding
 import Network.DatagramServer ()
-import Network.DHT.Mainline ()
 #endif
 import Network.DatagramServer.Types hiding (Query,Response)
 
@@ -467,3 +472,38 @@ instance KRPC (Query Announce) (Response Announced) where
 
 -- endif VERSION_bencoding
 #endif
+
+-- | Yields all 8 DHT neighborhoods available to you given a particular ip
+-- address.
+bep42s :: Address a => a -> NodeId KMessageOf -> [NodeId KMessageOf]
+bep42s addr (NodeId r) = mapMaybe (bep42 addr) rs
+ where
+    rs = map (NodeId . change3bits r) [0..7]
+
+-- change3bits :: ByteString -> Word8 -> ByteString
+-- change3bits bs n = BS.snoc (BS.init bs) (BS.last bs .&. 0xF8 .|. n)
+
+change3bits :: (Num b, Bits b) => b -> b -> b
+change3bits bs n = (bs .&. complement 7) .|. n
+
+-- | 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 KMessageOf -> Maybe (NodeId KMessageOf)
+bep42 addr (NodeId r)
+ | Just ip <- fmap S.encode (fromAddr addr :: Maybe IPv4)
+              <|> fmap S.encode (fromAddr addr :: Maybe IPv6)
+    = genBucketSample' retr (NodeId $ crc $ applyMask ip) (3,0x07,0)
+ | otherwise
+    = Nothing
+ where
+    ip4mask = "\x03\x0f\x3f\xff" :: ByteString
+    ip6mask = "\x01\x03\x07\x0f\x1f\x3f\x7f\xff" :: ByteString
+    nbhood_select = (fromIntegral r :: Word8) .&. 7
+    retr n = pure $ BS.drop (nodeIdSize - n) $ S.encode r
+    crc = flip shiftL (finiteBitSize (NodeId undefined) - 32) . fromIntegral . crc32c . BS.pack
+    applyMask ip = case BS.zipWith (.&.) msk ip of
+                    (b:bs) -> (b .|. shiftL nbhood_select 5) : bs
+                    bs     -> bs
+        where msk | BS.length ip == 4  =  ip4mask
+                  | otherwise          =  ip6mask
+