2 files changed, 53 insertions, 76 deletions

diff --git a/src/Network/Address.hs b/src/Network/Address.hs
index d34b0431..9ecd89a3 100644
--- a/src/Network/Address.hs
+++ b/src/Network/Address.hs
@@ -71,7 +71,6 @@ module Network.Address
        , NodeInfo (..)
        , mapAddress
        , traverseAddress
-       , rank
 
          -- * Fingerprint
          -- $fingerprint
@@ -83,7 +82,7 @@ module Network.Address
          -- * Utils
        , libUserAgent
        , sockAddrPort
-       , withPort
+       , setPort
        , getBindAddress
        ) where
 
@@ -148,17 +147,11 @@ sockAddrPort (SockAddrInet6 p _ _ _) = Just p
 sockAddrPort _                       = Nothing
 {-# INLINE sockAddrPort #-}
 
-withPort :: SockAddr -> PortNumber -> SockAddr
-withPort (SockAddrInet  _ ip           ) port = SockAddrInet port ip
-withPort (SockAddrInet6 _ flow ip scope) port = SockAddrInet6 port flow ip scope
-withPort addr                            _    = addr
-{-# INLINE withPort #-}
-
 instance Address a => Address (NodeAddr a) where
   toSockAddr NodeAddr {..} = setPort nodePort $ toSockAddr nodeHost
   fromSockAddr sa = NodeAddr <$> fromSockAddr sa <*> sockAddrPort sa
 
-instance Address a => Address (PeerAddr a) where
+instance Address PeerAddr where
   toSockAddr PeerAddr {..} = setPort peerPort $ toSockAddr peerHost
   fromSockAddr sa = PeerAddr Nothing <$> fromSockAddr sa <*> sockAddrPort sa
 
@@ -412,23 +405,18 @@ instance BEncode IPv6 where
   {-# INLINE fromBEncode #-}
 #endif
 
-{-----------------------------------------------------------------------
---  Peer addr
------------------------------------------------------------------------}
--- TODO check semantic of ord and eq instances
-
 -- | Peer address info normally extracted from peer list or peer
 -- compact list encoding.
-data PeerAddr a = PeerAddr
+data PeerAddr = PeerAddr
   { peerId   :: !(Maybe PeerId)
 
     -- | This is usually 'IPv4', 'IPv6', 'IP' or unresolved
     -- 'HostName'.
-  , peerHost :: !a
+  , peerHost :: !IP
 
     -- | The port the peer listenning for incoming P2P sessions.
   , peerPort :: {-# UNPACK #-} !PortNumber
-  } deriving (Show, Eq, Ord, Typeable, Functor)
+  } deriving (Show, Eq, Ord, Typeable)
 
 #ifdef VERSION_bencoding
 peer_ip_key, peer_id_key, peer_port_key :: BKey
@@ -437,7 +425,7 @@ peer_id_key   = "peer id"
 peer_port_key = "port"
 
 -- | The tracker's 'announce response' compatible encoding.
-instance (Typeable a, BEncode a) => BEncode (PeerAddr a) where
+instance BEncode PeerAddr where
   toBEncode PeerAddr {..} = toDict $
        peer_ip_key   .=! peerHost
     .: peer_id_key   .=? peerId
@@ -457,37 +445,39 @@ instance (Typeable a, BEncode a) => BEncode (PeerAddr a) where
 --
 --   For more info see: <http://www.bittorrent.org/beps/bep_0023.html>
 --
--- TODO: test byte order
-instance (Serialize a) => Serialize (PeerAddr a) where
+-- WARNING: Input must be exactly 6 or 18 bytes so that we can identify IP version.
+--
+instance Serialize PeerAddr where
   put PeerAddr {..} = put peerHost >> put peerPort
-  get = PeerAddr Nothing <$> get <*> get
+  get = do
+    cnt <- remaining
+    PeerAddr Nothing <$> isolate (cnt - 2) get <*> get
 
 -- | @127.0.0.1:6881@
-instance Default (PeerAddr IPv4) where
+instance Default PeerAddr where
   def = "127.0.0.1:6881"
 
--- | @127.0.0.1:6881@
-instance Default (PeerAddr IP) where
-  def = IPv4 <$> def
-
 -- | Example:
 --
 --   @peerPort \"127.0.0.1:6881\" == 6881@
 --
-instance IsString (PeerAddr IPv4) where
+instance IsString PeerAddr where
   fromString str
     | [hostAddrStr, portStr] <- splitWhen (== ':') str
     , Just hostAddr <- readMaybe hostAddrStr
     , Just portNum  <- toEnum <$> readMaybe portStr
-                = PeerAddr Nothing hostAddr portNum
-    | otherwise = error $ "fromString: unable to parse (PeerAddr IPv4): " ++ str
+                = PeerAddr Nothing (IPv4 hostAddr) portNum
+    | [((ip,port),"")] <- readsIPv6_port str =
+        PeerAddr Nothing (IPv6 ip) port
+    | otherwise = error $ "fromString: unable to parse IP: " ++ str
 
-instance Read (PeerAddr IPv4) where
+instance Read PeerAddr where
   readsPrec i = RP.readP_to_S $ do
-    ipv4 <- RP.readS_to_P (readsPrec i)
+    ip <-     IPv4 <$> ( RP.readS_to_P (readsPrec i)                         )
+          <|> IPv6 <$> ( RP.char '[' *> (RP.readS_to_P reads) <* RP.char ']' )
     _    <- RP.char ':'
     port <- toEnum <$> RP.readS_to_P (readsPrec i)
-    return $ PeerAddr Nothing ipv4 port
+    return $ PeerAddr Nothing ip port
 
 readsIPv6_port :: String -> [((IPv6, PortNumber), String)]
 readsIPv6_port = RP.readP_to_S $ do
@@ -496,27 +486,16 @@ readsIPv6_port = RP.readP_to_S $ do
   port <- toEnum <$> read <$> (RP.many1 $ RP.satisfy isDigit) <* RP.eof
   return (ip,port)
 
-instance IsString (PeerAddr IPv6) where
-  fromString str
-    | [((ip,port),"")] <- readsIPv6_port str =
-        PeerAddr Nothing ip port
-    | otherwise = error $ "fromString: unable to parse (PeerAddr IPv6): " ++ str
-
-instance IsString (PeerAddr IP) where
-  fromString str
-    | '[' `L.elem` str = IPv6 <$> fromString str
-    |      otherwise   = IPv4 <$> fromString str
 
 -- | fingerprint + "at" + dotted.host.inet.addr:port
--- TODO: instances for IPv6, HostName
-instance Pretty a => Pretty (PeerAddr a) where
+instance Pretty PeerAddr where
   pPrint PeerAddr {..}
     | Just pid <- peerId = pPrint (fingerprint pid) <+> "at" <+> paddr
     |     otherwise      = paddr
     where
       paddr = pPrint peerHost <> ":" <> text (show peerPort)
 
-instance Hashable a => Hashable (PeerAddr a) where
+instance Hashable PeerAddr where
   hashWithSalt s PeerAddr {..} =
     s `hashWithSalt` peerId `hashWithSalt` peerHost `hashWithSalt` peerPort
 
@@ -524,10 +503,7 @@ instance Hashable a => Hashable (PeerAddr a) where
 defaultPorts :: [PortNumber]
 defaultPorts =  [6881..6889]
 
-_resolvePeerAddr :: (IPAddress i) => PeerAddr HostName -> PeerAddr i
-_resolvePeerAddr = undefined
-
-_peerSockAddr :: PeerAddr IP -> (Family, SockAddr)
+_peerSockAddr :: PeerAddr -> (Family, SockAddr)
 _peerSockAddr PeerAddr {..} =
     case peerHost of
           IPv4 ipv4 ->
@@ -535,11 +511,11 @@ _peerSockAddr PeerAddr {..} =
           IPv6 ipv6 ->
               (AF_INET6, SockAddrInet6 peerPort 0 (toHostAddress6 ipv6) 0)
 
-peerSockAddr :: PeerAddr IP -> SockAddr
+peerSockAddr :: PeerAddr -> SockAddr
 peerSockAddr = snd . _peerSockAddr
 
 -- | Create a socket connected to the address specified in a peerAddr
-peerSocket :: SocketType -> PeerAddr IP -> IO Socket
+peerSocket :: SocketType -> PeerAddr -> IO Socket
 peerSocket socketType pa = do
     let (family, addr) = _peerSockAddr pa
     sock <- socket family socketType defaultProtocol
@@ -607,7 +583,7 @@ instance BEncode a => BEncode (NodeAddr a) where
   {-# INLINE fromBEncode #-}
 #endif
 
-fromPeerAddr :: PeerAddr a -> NodeAddr a
+fromPeerAddr :: PeerAddr -> NodeAddr IP
 fromPeerAddr PeerAddr {..} = NodeAddr
   { nodeHost = peerHost
   , nodePort = peerPort
@@ -615,12 +591,6 @@ fromPeerAddr PeerAddr {..} = NodeAddr
 
 ------------------------------------------------------------------------
 
--- | Order by closeness: nearest nodes first.
-rank :: ( Ord (NodeId dht)
-        , Bits (NodeId dht)
-        ) => (x -> NodeId dht) -> NodeId dht -> [x] -> [x]
-rank f nid = L.sortBy (comparing (RPC.distance nid . f))
-
 {-----------------------------------------------------------------------
 -- Fingerprint
 -----------------------------------------------------------------------}
diff --git a/src/Network/BitTorrent/DHT/ContactInfo.hs b/src/Network/BitTorrent/DHT/ContactInfo.hs
index 3db3d5a8..172306a1 100644
--- a/src/Network/BitTorrent/DHT/ContactInfo.hs
+++ b/src/Network/BitTorrent/DHT/ContactInfo.hs
@@ -1,6 +1,7 @@
 {-# LANGUAGE BangPatterns #-}
 module Network.BitTorrent.DHT.ContactInfo
        ( PeerStore
+       , PeerAddr(..)
        , Network.BitTorrent.DHT.ContactInfo.lookup
        , Network.BitTorrent.DHT.ContactInfo.freshPeers
        , Network.BitTorrent.DHT.ContactInfo.insertPeer
@@ -36,17 +37,17 @@ import Network.Address
 -- --  PeerSet
 -- -----------------------------------------------------------------------}
 --
--- type PeerSet a = [(PeerAddr a, NodeInfo a, Timestamp)]
+-- type PeerSet a = [(PeerAddr, NodeInfo a, Timestamp)]
 --
 -- -- compare PSQueue vs Ordered list
 --
--- takeNewest :: PeerSet a -> [PeerAddr a]
+-- takeNewest :: PeerSet a -> [PeerAddr]
 -- takeNewest = undefined
 --
 -- dropOld :: Timestamp -> PeerSet a -> PeerSet a
 -- dropOld = undefined
 --
--- insert :: PeerAddr a -> Timestamp -> PeerSet a -> PeerSet a
+-- insert :: PeerAddr -> Timestamp -> PeerSet a -> PeerSet a
 -- insert = undefined
 --
 -- type Mask = Int
@@ -116,16 +117,19 @@ import Network.Address
 
 -- | Storage used to keep track a set of known peers in client,
 -- tracker or DHT sessions.
-newtype PeerStore ip = PeerStore (HashMap InfoHash (SwarmData ip))
+newtype PeerStore = PeerStore (HashMap InfoHash SwarmData)
 
 type Timestamp = POSIXTime
 
-data SwarmData ip = SwarmData
-    { peers :: !(PSQ (PeerAddr ip) Timestamp)
+data SwarmData = SwarmData
+    { peers :: !(PSQ PeerAddr Timestamp)
     , name :: !(Maybe ByteString)
     }
 
-
+-- | This wrapper will serialize an ip address with a '4' or '6' prefix byte
+-- to indicate whether it is IPv4 or IPv6.
+--
+-- Note: it does not serialize port numbers.
 newtype SerializeAddress a = SerializeAddress { unserializeAddress :: a }
 
 instance Address a => Serialize (SerializeAddress a) where
@@ -145,7 +149,8 @@ instance Address a => Serialize (SerializeAddress a) where
         | otherwise = return $ error "cannot serialize non-IP SerializeAddress"
 
 
-instance (Ord ip, Address ip) => Serialize (SwarmData ip) where
+{- XXX: What happened to the ports?
+instance Serialize SwarmData where
     get = flip SwarmData <$> get
                          <*> ( PSQ.fromList . L.map parseAddr <$> get )
      where
@@ -157,17 +162,17 @@ instance (Ord ip, Address ip) => Serialize (SwarmData ip) where
         put $ L.map (\(addr :-> _) -> (SerializeAddress <$> addr))
                 -- XXX: should we serialize the timestamp?
             $ PSQ.toList peers
+-}
 
-
-knownSwarms :: PeerStore ip -> [ (InfoHash, Int, Maybe ByteString) ]
+knownSwarms :: PeerStore -> [ (InfoHash, Int, Maybe ByteString) ]
 knownSwarms (PeerStore m) = L.map (\(ih,SwarmData q n) -> (ih,PSQ.size q,n)) $ HM.toList m
 
-swarmSingleton :: Ord ip => PeerAddr ip -> SwarmData ip
+swarmSingleton :: PeerAddr -> SwarmData
 swarmSingleton a = SwarmData
     { peers = PSQ.singleton a 0
     , name = Nothing }
 
-swarmInsert :: Ord ip => SwarmData ip -> SwarmData ip -> SwarmData ip
+swarmInsert :: SwarmData -> SwarmData -> SwarmData
 swarmInsert old new = SwarmData
     { peers = L.foldl' (\q (a :-> t) -> PSQ.insertWith (\p _ -> p) a t q) (peers old) (PSQ.toList $ peers new)
     , name = name new <|> name old -- TODO: decodeUtf8' check
@@ -176,12 +181,12 @@ swarmInsert old new = SwarmData
 isSwarmOccupied SwarmData{..} = not $ PSQ.null peers
 
 -- | Empty store.
-instance Default (PeerStore a) where
+instance Default (PeerStore) where
   def = PeerStore HM.empty
   {-# INLINE def #-}
 
 -- | Monoid under union operation.
-instance Ord a => Monoid (PeerStore a) where
+instance Monoid PeerStore where
   mempty  = def
   {-# INLINE mempty #-}
 
@@ -189,20 +194,22 @@ instance Ord a => Monoid (PeerStore a) where
     PeerStore (HM.unionWith swarmInsert a b)
   {-# INLINE mappend #-}
 
+{-
 -- | Can be used to store peers between invocations of the client
 -- software.
-instance (Ord a, Address a) => Serialize (PeerStore a) where
+instance Serialize PeerStore where
   get = PeerStore . HM.fromList <$> get
   put (PeerStore m) = put (L.filter (isSwarmOccupied . snd) $ HM.toList m)
+-}
 
 -- | Returns all peers associated with a given info hash.
-lookup :: Ord a => InfoHash -> PeerStore a -> [PeerAddr a]
+lookup :: InfoHash -> PeerStore -> [PeerAddr]
 lookup ih (PeerStore m) = maybe [] (PSQ.keys . peers) $ HM.lookup ih m
 
 batchSize = 64
 
 -- | Used in 'get_peers' DHT queries.
-freshPeers :: Ord a => InfoHash -> Timestamp -> PeerStore a -> ([PeerAddr a], PeerStore a)
+freshPeers :: InfoHash -> Timestamp -> PeerStore -> ([PeerAddr], PeerStore)
 freshPeers ih tm (PeerStore m) = fromMaybe ([],PeerStore m) $ do
     swarm <- HM.lookup ih m
     let ps0 = take batchSize $ unfoldr (incomp minView) (peers swarm)
@@ -219,7 +226,7 @@ incomp !f !x = do
    pure $! ( (result,x'), x' )
 
 -- | Used in 'announce_peer' DHT queries.
-insertPeer :: Ord a => InfoHash -> Maybe ByteString -> PeerAddr a -> PeerStore a -> PeerStore a
+insertPeer :: InfoHash -> Maybe ByteString -> PeerAddr -> PeerStore -> PeerStore
 insertPeer !ih !name !a !(PeerStore m) = seq a' $ PeerStore (HM.insertWith swarmInsert ih a' m)
  where
     a' = SwarmData { peers = PSQ.singleton a 0