1 files changed, 74 insertions, 20 deletions

diff --git a/src/Network/DHT/Mainline.hs b/src/Network/DHT/Mainline.hs
index d7aed430..2b7db3c7 100644
--- a/src/Network/DHT/Mainline.hs
+++ b/src/Network/DHT/Mainline.hs
@@ -1,15 +1,23 @@
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE FlexibleInstances          #-}
 {-# LANGUAGE DeriveDataTypeable         #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 module Network.DHT.Mainline where
 
+import Data.Digest.CRC32C
+import Control.Applicative
+import Data.Maybe
+import Data.Monoid
+import Data.Word
+import Data.IP
 import Data.BEncode                    as BE
 import Data.Bits
 import Data.ByteString                 (ByteString)
 import Data.ByteString.Base16          as Base16
 import qualified Data.ByteString.Char8 as Char8
+import qualified Data.ByteString       as BS
 import qualified Data.ByteString.Lazy  as L
 import Data.Default
 import Data.LargeWord
@@ -17,28 +25,16 @@ import Data.Serialize                  as S
 import Data.String
 import Data.Typeable
 import Network.KRPC.Message            as KRPC
-import qualified Network.RPC           as RPC (NodeId)
-         ;import Network.RPC           as RPC hiding (NodeId)
+import Network.RPC as RPC
 import Text.PrettyPrint                as PP hiding ((<>))
 import Text.PrettyPrint.HughesPJClass  hiding (($$), (<>))
 
--- | Each node has a globally unique identifier known as the \"node
--- ID.\"
---
---   Normally, /this/ node id should be saved between invocations
---   of the client software.
-newtype NodeId = NodeId Word160
-  deriving (Show, Eq, Ord, Typeable, Bits, FiniteBits)
+nodeIdSize = finiteBitSize (undefined :: NodeId KMessageOf) `div` 8
 
-instance BEncode NodeId where
+instance BEncode (NodeId KMessageOf) where
     toBEncode (NodeId w) = toBEncode $ S.encode w
     fromBEncode bval = fromBEncode bval >>= S.decode
 
--- | NodeId size in bytes.
-nodeIdSize :: Int
-nodeIdSize = 20
-
-
 -- instance BEncode NodeId where TODO
 
 -- TODO: put this somewhere appropriate
@@ -46,14 +42,14 @@ instance (Serialize a, Serialize b) => Serialize (LargeKey a b) where
    put (LargeKey lo hi) = put hi >> put lo
    get = flip LargeKey <$> get <*> get
 
-instance Serialize NodeId where
+instance Serialize (NodeId KMessageOf) where
   get = NodeId <$> get
   {-# INLINE get #-}
   put (NodeId bs) = put bs
   {-# INLINE put #-}
 
 -- | ASCII encoded.
-instance IsString NodeId where
+instance IsString (NodeId KMessageOf) where
   fromString str
     | length str == nodeIdSize     = NodeId (either error id $ S.decode (fromString str :: ByteString))
     | length str == 2 * nodeIdSize = NodeId (either error id $ S.decode (fst $ Base16.decode $ fromString str))
@@ -61,16 +57,74 @@ instance IsString NodeId where
   {-# INLINE fromString #-}
 
 -- | Meaningless node id, for testing purposes only.
-instance Default NodeId where
+instance Default (NodeId KMessageOf) where
   def = NodeId 0
 
 -- | base16 encoded.
-instance Pretty NodeId where
+instance Pretty (NodeId KMessageOf) where
   pPrint (NodeId nid) = PP.text $ Char8.unpack $ Base16.encode $ S.encode nid
 
+-- | KRPC 'compact list' compatible encoding: contact information for
+-- nodes is encoded as a 26-byte string. Also known as "Compact node
+-- info" the 20-byte Node ID in network byte order has the compact
+-- IP-address/port info concatenated to the end.
+instance Serialize a => Serialize (NodeInfo KMessageOf a ()) where
+  get = (\a b -> NodeInfo a b ()) <$> get <*> get
+  put NodeInfo {..} = put nodeId >> put nodeAddr
+
+instance Pretty ip => Pretty (NodeInfo KMessageOf ip ()) where
+  pPrint NodeInfo {..} = pPrint nodeId <> "@(" <> pPrint nodeAddr <> ")"
+
+instance Pretty ip => Pretty [NodeInfo KMessageOf ip ()] where
+  pPrint = PP.vcat . PP.punctuate "," . map pPrint
+
+
+-- | 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
+
+
+
 instance Envelope KMessageOf where
     type TransactionID KMessageOf = KRPC.TransactionId
-    type NodeId KMessageOf = Network.DHT.Mainline.NodeId
+
+    -- | Each node has a globally unique identifier known as the \"node
+    -- ID.\"
+    --
+    --   Normally, /this/ node id should be saved between invocations
+    --   of the client software.
+    newtype NodeId KMessageOf = NodeId Word160
+      deriving (Show, Eq, Ord, Typeable, Bits, FiniteBits)
 
     envelopePayload (Q q) = queryArgs q
     envelopePayload (R r) = respVals r