1 files changed, 173 insertions, 17 deletions

diff --git a/src/Network/Address.hs b/src/Network/Address.hs
index 9ecd89a3..8715a82d 100644
--- a/src/Network/Address.hs
+++ b/src/Network/Address.hs
@@ -59,18 +59,13 @@ module Network.Address
        , peerSocket
 
          -- * Node
+       , NodeAddr (..)
+
          -- ** Id
-       , NodeId
        , testIdBit
-       , genNodeId
        , bucketRange
        , genBucketSample
-
-         -- ** Info
-       , NodeAddr (..)
-       , NodeInfo (..)
-       , mapAddress
-       , traverseAddress
+       , genBucketSample'
 
          -- * Fingerprint
          -- $fingerprint
@@ -126,7 +121,6 @@ import Text.PrettyPrint.HughesPJClass hiding ((<>),($$))
 import System.Locale    (defaultTimeLocale)
 #endif
 import System.Entropy
-import Network.DatagramServer.Types as RPC
 
 -- import Paths_bittorrent (version)
 
@@ -147,9 +141,82 @@ sockAddrPort (SockAddrInet6 p _ _ _) = Just p
 sockAddrPort _                       = Nothing
 {-# INLINE sockAddrPort #-}
 
-instance Address a => Address (NodeAddr a) where
-  toSockAddr NodeAddr {..} = setPort nodePort $ toSockAddr nodeHost
-  fromSockAddr sa = NodeAddr <$> fromSockAddr sa <*> sockAddrPort sa
+class (Eq a, Serialize a, Typeable a, Hashable a, Pretty a)
+    => Address a where
+  toSockAddr   :: a        -> SockAddr
+  fromSockAddr :: SockAddr -> Maybe a
+
+fromAddr :: (Address a, Address b) => a -> Maybe b
+fromAddr = fromSockAddr . toSockAddr
+
+-- | Note that port is zeroed.
+instance Address IPv4 where
+  toSockAddr = SockAddrInet 0 . toHostAddress
+  fromSockAddr (SockAddrInet _ h) = Just (fromHostAddress h)
+  fromSockAddr  _                 = Nothing
+
+-- | Note that port is zeroed.
+instance Address IPv6 where
+  toSockAddr h = SockAddrInet6 0 0 (toHostAddress6 h) 0
+  fromSockAddr (SockAddrInet6 _ _ h _) = Just (fromHostAddress6 h)
+  fromSockAddr  _                      = Nothing
+
+-- | Note that port is zeroed.
+instance Address IP where
+  toSockAddr (IPv4 h) = toSockAddr h
+  toSockAddr (IPv6 h) = toSockAddr h
+  fromSockAddr sa =
+        IPv4 <$> fromSockAddr sa
+    <|> IPv6 <$> fromSockAddr sa
+
+data NodeAddr a = NodeAddr
+  { nodeHost ::                !a
+  , nodePort :: {-# UNPACK #-} !PortNumber
+  } deriving (Eq, Ord, Typeable, Functor, Foldable, Traversable)
+
+instance Show a => Show (NodeAddr a) where
+  showsPrec i NodeAddr {..}
+    = showsPrec i nodeHost <> showString ":" <> showsPrec i nodePort
+
+instance Read (NodeAddr IPv4) where
+  readsPrec i = RP.readP_to_S $ do
+    ipv4 <- RP.readS_to_P (readsPrec i)
+    _    <- RP.char ':'
+    port <- toEnum <$> RP.readS_to_P (readsPrec i)
+    return $ NodeAddr ipv4 port
+
+-- | @127.0.0.1:6882@
+instance Default (NodeAddr IPv4) where
+  def = "127.0.0.1:6882"
+
+-- | KRPC compatible encoding.
+instance Serialize a => Serialize (NodeAddr a) where
+  get = NodeAddr <$> get <*> get
+  {-# INLINE get #-}
+  put NodeAddr {..} = put nodeHost >> put nodePort
+  {-# INLINE put #-}
+
+-- | Example:
+--
+--   @nodePort \"127.0.0.1:6881\" == 6881@
+--
+instance IsString (NodeAddr IPv4) where
+  fromString str
+    | [hostAddrStr, portStr] <- splitWhen (== ':') str
+    , Just hostAddr <- readMaybe hostAddrStr
+    , Just portNum  <- toEnum <$> readMaybe portStr
+                = NodeAddr hostAddr portNum
+    | otherwise = error $ "fromString: unable to parse (NodeAddr IPv4): " ++ str
+
+
+instance Hashable a => Hashable (NodeAddr a) where
+  hashWithSalt s NodeAddr {..} = hashWithSalt s (nodeHost, nodePort)
+  {-# INLINE hashWithSalt #-}
+
+instance Pretty ip => Pretty (NodeAddr ip) where
+  pPrint NodeAddr {..} = pPrint nodeHost <> ":" <> pPrint nodePort
+
+
 
 instance Address PeerAddr where
   toSockAddr PeerAddr {..} = setPort peerPort $ toSockAddr peerHost
@@ -552,6 +619,37 @@ testIdBit :: FiniteBits bs => bs -> Word -> Bool
 testIdBit bs i = testBit bs (fromIntegral (finiteBitSize bs - 1 - fromIntegral i))
 {-# INLINE testIdBit #-}
 
+-- | 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 :: ( FiniteBits nid
+                   , Serialize nid
+                   ) => nid -> (Int,Word8,Word8) -> IO nid
+genBucketSample n qmb = genBucketSample' getEntropy n qmb
+
+-- | Generalizion of 'genBucketSample' that accepts a byte generator
+-- function to use instead of the system entropy.
+genBucketSample' :: forall m dht nid.
+                    ( Applicative m
+                    , FiniteBits nid
+                    , Serialize nid
+                    ) =>
+    (Int -> m ByteString) -> nid -> (Int,Word8,Word8) -> m nid
+genBucketSample' gen self (q,m,b)
+    | q <= 0           =  either error id . S.decode <$> gen nodeIdSize
+    | q >= nodeIdSize  =  pure self
+    | otherwise        =  either error id . S.decode .  build <$> gen (nodeIdSize - q + 1)
+ where
+    nodeIdSize = finiteBitSize (undefined :: nid) `div` 8
+    build tl = BS.init hd <> BS.cons (h .|. t) (BS.tail tl)
+     where
+        hd = BS.take q $ S.encode self
+        h = xor b (complement m .&. BS.last hd)
+        t = m .&. BS.head tl
+
+
 ------------------------------------------------------------------------
 
 -- | Accepts a depth/index of a bucket and whether or not it is the last one,
@@ -583,11 +681,69 @@ instance BEncode a => BEncode (NodeAddr a) where
   {-# INLINE fromBEncode #-}
 #endif
 
-fromPeerAddr :: PeerAddr -> NodeAddr IP
-fromPeerAddr PeerAddr {..} = NodeAddr
-  { nodeHost = peerHost
-  , nodePort = peerPort
-  }
+
+instance Hashable PortNumber where
+  hashWithSalt s = hashWithSalt s . fromEnum
+  {-# INLINE hashWithSalt #-}
+
+instance Pretty PortNumber where
+  pPrint = PP.int . fromEnum
+  {-# INLINE pPrint #-}
+
+instance Serialize PortNumber where
+  get = fromIntegral <$> getWord16be
+  {-# INLINE get #-}
+  put = putWord16be . fromIntegral
+  {-# INLINE put #-}
+
+instance Pretty IPv4 where
+  pPrint = PP.text . show
+  {-# INLINE pPrint #-}
+
+instance Pretty IPv6 where
+  pPrint = PP.text . show
+  {-# INLINE pPrint #-}
+
+instance Pretty IP where
+  pPrint = PP.text . show
+  {-# INLINE pPrint #-}
+
+
+-- | When 'get'ing an IP it must be 'isolate'd to the appropriate
+-- number of bytes since we have no other way of telling which
+-- address type we are trying to parse
+instance Serialize IP where
+    put (IPv4 ip) = put ip
+    put (IPv6 ip) = put ip
+
+    get = do
+      n <- remaining
+      case n of
+        4  -> IPv4 <$> get
+        16 -> IPv6 <$> get
+        _ -> fail (show n ++ " is the wrong number of remaining bytes to parse IP")
+
+instance Serialize IPv4 where
+    put = putWord32host    .  toHostAddress
+    get = fromHostAddress <$> getWord32host
+
+instance Serialize IPv6 where
+    put ip = put $ toHostAddress6 ip
+    get = fromHostAddress6 <$> get
+
+
+instance Hashable IPv4 where
+  hashWithSalt = hashUsing toHostAddress
+  {-# INLINE hashWithSalt #-}
+
+instance Hashable IPv6 where
+  hashWithSalt s a = hashWithSalt s (toHostAddress6 a)
+
+instance Hashable IP where
+  hashWithSalt s (IPv4 h) = hashWithSalt s h
+  hashWithSalt s (IPv6 h) = hashWithSalt s h
+
+
 
 ------------------------------------------------------------------------