1 files changed, 49 insertions, 79 deletions

diff --git a/src/Network/BitTorrent/Address.hs b/src/Network/BitTorrent/Address.hs
index 2132f8f9..560ac1ef 100644
--- a/src/Network/BitTorrent/Address.hs
+++ b/src/Network/BitTorrent/Address.hs
@@ -11,6 +11,7 @@
 --
 {-# LANGUAGE CPP                        #-}
 {-# LANGUAGE FlexibleInstances          #-}
+{-# LANGUAGE FlexibleContexts           #-}
 {-# LANGUAGE RecordWildCards            #-}
 {-# LANGUAGE StandaloneDeriving         #-}
 {-# LANGUAGE ViewPatterns               #-}
@@ -59,11 +60,8 @@ module Network.BitTorrent.Address
          -- * Node
          -- ** Id
        , NodeId
-       , asNodeId
        , nodeIdSize
        , testIdBit
-       , NodeDistance
-       , distance
        , genNodeId
        , bucketRange
        , genBucketSample
@@ -73,6 +71,8 @@ module Network.BitTorrent.Address
          -- ** Info
        , NodeAddr (..)
        , NodeInfo (..)
+       , mapAddress
+       , traverseAddress
        , rank
 
          -- * Fingerprint
@@ -98,7 +98,6 @@ import Data.BEncode.BDict (BKey)
 import Data.Bits
 import qualified Data.ByteString as BS
 import qualified Data.ByteString.Internal as BS
-import Data.ByteString.Base16 as Base16
 import Data.ByteString.Char8 as BC
 import Data.ByteString.Char8 as BS8
 import qualified Data.ByteString.Lazy as BL
@@ -130,6 +129,9 @@ import System.Locale    (defaultTimeLocale)
 #endif
 import System.Entropy
 import Data.Digest.CRC32C
+import qualified Network.RPC as RPC
+import Network.KRPC.Message (KMessageOf)
+import Network.DHT.Mainline
 
 -- import Paths_bittorrent (version)
 
@@ -646,48 +648,10 @@ peerSocket socketType pa = do
 --   in the DHT to get the location of peers to download from using
 --   the BitTorrent protocol.
 
--- TODO more compact representation ('ShortByteString's?)
+-- asNodeId :: ByteString -> NodeId
+-- asNodeId bs = NodeId $ BS.take nodeIdSize bs
 
--- | 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 ByteString
-  deriving (Show, Eq, Ord, Typeable
-#ifdef VERSION_bencoding
-           , BEncode
-#endif
-           )
-
-
-nodeIdSize :: Int
-nodeIdSize = 20
-
-asNodeId :: ByteString -> NodeId
-asNodeId bs = NodeId $ BS.take nodeIdSize bs
-
--- | Meaningless node id, for testing purposes only.
-instance Default NodeId where
-  def = NodeId (BS.replicate nodeIdSize 0)
-
-instance Serialize NodeId where
-  get = NodeId <$> getByteString nodeIdSize
-  {-# INLINE get #-}
-  put (NodeId bs) = putByteString bs
-  {-# INLINE put #-}
-
--- | ASCII encoded.
-instance IsString NodeId where
-  fromString str
-    | L.length str == nodeIdSize     = NodeId (fromString str)
-    | L.length str == 2 * nodeIdSize = NodeId (fst $ Base16.decode $ fromString str)
-    |           otherwise = error "fromString: invalid NodeId length"
-  {-# INLINE fromString #-}
-
--- | base16 encoded.
-instance Pretty NodeId where
-  pPrint (NodeId nid) = PP.text $ BC.unpack $ Base16.encode nid
+{-
 
 -- | Test if the nth bit is set.
 testIdBit :: NodeId -> Word -> Bool
@@ -696,6 +660,10 @@ testIdBit (NodeId bs) i
   , (q, r) <- quotRem (fromIntegral i) 8
   = testBit (BS.index bs q) (7 - r)
   |     otherwise      = False
+-}
+
+testIdBit :: FiniteBits bs => bs -> Word -> Bool
+testIdBit bs i = testBit bs (fromIntegral (finiteBitSize bs - fromIntegral i))
 {-# INLINE testIdBit #-}
 
 -- TODO WARN is the 'system' random suitable for this?
@@ -703,25 +671,10 @@ testIdBit (NodeId bs) i
 --   Distribution of ID's should be as uniform as possible.
 --
 genNodeId :: IO NodeId
-genNodeId = NodeId <$> getEntropy nodeIdSize
+genNodeId = NodeId . either error id . S.decode <$> getEntropy nodeIdSize
 
 ------------------------------------------------------------------------
 
--- | In Kademlia, the distance metric is XOR and the result is
--- interpreted as an unsigned integer.
-newtype NodeDistance = NodeDistance BS.ByteString
-  deriving (Eq, Ord)
-
-instance Pretty NodeDistance where
-  pPrint (NodeDistance bs) = text $ BC.unpack (Base16.encode bs)
-
-instance Show NodeDistance where
-  show (NodeDistance bs) = BC.unpack (Base16.encode bs)
-
--- | distance(A,B) = |A xor B| Smaller values are closer.
-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:
 --
@@ -753,13 +706,13 @@ genBucketSample n qmb = genBucketSample' getEntropy n qmb
 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 <= 0           =  NodeId . either error id . S.decode <$> gen nodeIdSize
     | q >= nodeIdSize  =  pure (NodeId self)
-    | otherwise        =  NodeId . build <$> gen (nodeIdSize - q + 1)
+    | otherwise        =  NodeId . either error id . S.decode .  build <$> gen (nodeIdSize - q + 1)
  where
     build tl = BS.init hd <> BS.cons (h .|. t) (BS.tail tl)
      where
-        hd = BS.take q self
+        hd = BS.take q $ S.encode self
         h = xor b (complement m .&. BS.last hd)
         t = m .&. BS.head tl
 
@@ -819,32 +772,46 @@ fromPeerAddr PeerAddr {..} = NodeAddr
 
 ------------------------------------------------------------------------
 
-data NodeInfo a = NodeInfo
-  { nodeId   :: !NodeId
-  , nodeAddr :: !(NodeAddr a)
-  } deriving (Show, Eq, Functor, Foldable, Traversable)
+data NodeInfo dht addr u = NodeInfo
+  { nodeId   :: !(RPC.NodeId dht)
+  , nodeAddr :: !(NodeAddr addr)
+  , nodeAnnotation :: u
+  } deriving (Functor, Foldable, Traversable)
+
+deriving instance ( Show (RPC.NodeId dht)
+                  , Show addr
+                  , Show u ) => Show (NodeInfo dht addr u)
 
+mapAddress :: (addr -> b) -> NodeInfo dht addr u -> NodeInfo dht b u
+mapAddress f ni = ni { nodeAddr = fmap f (nodeAddr ni) }
 
-instance Eq a => Ord (NodeInfo a) where
+traverseAddress :: Applicative f => (addr -> f b) -> NodeInfo dht addr u -> f (NodeInfo dht b u)
+traverseAddress f ni = fmap (\addr -> ni { nodeAddr = addr }) $ traverse f (nodeAddr ni)
+
+-- Warning: Eq and Ord only look at the nodeId field.
+instance Eq (RPC.NodeId dht) => Eq (NodeInfo dht a u) where
+    a == b  = (nodeId a == nodeId b)
+
+instance Ord (RPC.NodeId dht) => Ord (NodeInfo dht a u) where
   compare = comparing nodeId
 
 -- | 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 a) where
-  get = NodeInfo <$> get <*> get
+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 ip) where
+instance Pretty ip => Pretty (NodeInfo KMessageOf ip ()) where
   pPrint NodeInfo {..} = pPrint nodeId <> "@(" <> pPrint nodeAddr <> ")"
 
-instance Pretty ip => Pretty [NodeInfo ip] where
+instance Pretty ip => Pretty [NodeInfo KMessageOf ip ()] where
   pPrint = PP.vcat . PP.punctuate "," . L.map pPrint
 
 -- | Order by closeness: nearest nodes first.
 rank :: (x -> NodeId) -> NodeId -> [x] -> [x]
-rank f nid = L.sortBy (comparing (distance nid . f))
+rank f nid = L.sortBy (comparing (RPC.distance nid . f))
 
 {-----------------------------------------------------------------------
 -- Fingerprint
@@ -1259,8 +1226,11 @@ bep42s addr (NodeId r) = mapMaybe (bep42 addr) rs
  where
     rs = L.map (NodeId . change3bits r) [0..7]
 
-change3bits :: ByteString -> Word8 -> ByteString
-change3bits bs n = BS.snoc (BS.init bs) (BS.last bs .&. 0xF8 .|. n)
+-- 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.
@@ -1274,9 +1244,9 @@ bep42 addr (NodeId r)
  where
     ip4mask = "\x03\x0f\x3f\xff" :: ByteString
     ip6mask = "\x01\x03\x07\x0f\x1f\x3f\x7f\xff" :: ByteString
-    nbhood_select = BS.last r .&. 7
-    retr n = pure $ BS.drop (BS.length r - n) r
-    crc = S.encode . crc32c . BS.pack
+    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