3 files changed, 80 insertions, 30 deletions

diff --git a/src/Network/BitTorrent/Core/PeerAddr.hs b/src/Network/BitTorrent/Core/PeerAddr.hs
index 6c6056c9..88239d0b 100644
--- a/src/Network/BitTorrent/Core/PeerAddr.hs
+++ b/src/Network/BitTorrent/Core/PeerAddr.hs
@@ -21,11 +21,6 @@ module Network.BitTorrent.Core.PeerAddr
          PeerAddr(..)
        , defaultPorts
        , peerSockAddr
-
-         -- * IP
-       , mergeIPLists
-       , splitIPList
-       , IPAddress ()
        ) where
 
 import Control.Applicative
@@ -36,8 +31,6 @@ import Data.BEncode.BDict (BKey)
 import Data.ByteString.Char8 as BS8
 import Data.Char
 import Data.Default
-import Data.Either
-import Data.Foldable
 import Data.IP
 import Data.List      as L
 import Data.List.Split
@@ -86,13 +79,16 @@ class IPAddress i where
 
 instance IPAddress IPv4 where
   toHostAddr = Left . toHostAddress
+  {-# INLINE toHostAddr #-}
 
 instance IPAddress IPv6 where
   toHostAddr = Right . toHostAddress6
+  {-# INLINE toHostAddr #-}
 
 instance IPAddress IP where
   toHostAddr (IPv4 ip) = toHostAddr ip
   toHostAddr (IPv6 ip) = toHostAddr ip
+  {-# INLINE toHostAddr #-}
 
 deriving instance Typeable IP
 deriving instance Typeable IPv4
@@ -100,6 +96,7 @@ deriving instance Typeable IPv6
 
 ipToBEncode :: Show i => i -> BValue
 ipToBEncode ip = BString $ BS8.pack $ show ip
+{-# INLINE ipToBEncode #-}
 
 ipFromBEncode :: Read a => BValue -> BS.Result a
 ipFromBEncode (BString (BS8.unpack -> ipStr))
@@ -107,17 +104,25 @@ ipFromBEncode (BString (BS8.unpack -> ipStr))
   |         otherwise            = decodingError $ "IP: " ++ ipStr
 ipFromBEncode _    = decodingError $ "IP: addr should be a bstring"
 
+instance Ord IP where
+
 instance BEncode IP where
   toBEncode   = ipToBEncode
+  {-# INLINE toBEncode #-}
   fromBEncode = ipFromBEncode
+  {-# INLINE fromBEncode #-}
 
 instance BEncode IPv4 where
   toBEncode   = ipToBEncode
+  {-# INLINE toBEncode #-}
   fromBEncode = ipFromBEncode
+  {-# INLINE fromBEncode #-}
 
 instance BEncode IPv6 where
   toBEncode   = ipToBEncode
+  {-# INLINE toBEncode #-}
   fromBEncode = ipFromBEncode
+  {-# INLINE fromBEncode #-}
 
 instance Serialize IPv4 where
     put = putWord32host    .  toHostAddress
@@ -136,9 +141,14 @@ instance Serialize IPv6 where
 -- compact list encoding.
 data PeerAddr a = PeerAddr
   { peerId   :: !(Maybe PeerId)
+
+    -- | This is usually 'IPv4', 'IPv6', 'IP' or unresolved
+    -- 'HostName'.
   , peerHost :: !a
+
+    -- | The port the peer listenning for incoming P2P sessions.
   , peerPort :: {-# UNPACK #-} !PortNumber
-  } deriving (Show, Eq, Typeable, Functor)
+  } deriving (Show, Eq, Ord, Typeable, Functor)
 
 peer_ip_key, peer_id_key, peer_port_key :: BKey
 peer_ip_key   = "ip"
@@ -160,17 +170,6 @@ instance (Typeable a, BEncode a) => BEncode (PeerAddr a) where
     where
       peerAddr = flip PeerAddr
 
-mergeIPLists :: [PeerAddr IPv4] -> Maybe [PeerAddr IPv6] -> [PeerAddr IP]
-mergeIPLists v4 v6 = (fmap IPv4 `L.map` v4)
-                  ++ (fmap IPv6 `L.map` Data.Foldable.concat v6)
-
-splitIPList :: [PeerAddr IP] -> ([PeerAddr IPv4],[PeerAddr IPv6])
-splitIPList xs = partitionEithers $ toEither <$> xs
-    where
-      toEither :: PeerAddr IP -> Either (PeerAddr IPv4) (PeerAddr IPv6)
-      toEither pa@(PeerAddr _ (IPv4 _) _) = Left  (ipv4 <$> pa)
-      toEither pa@(PeerAddr _ (IPv6 _) _) = Right (ipv6 <$> pa)
-
 -- | The tracker's 'compact peer list' compatible encoding. The
 -- 'peerId' is always 'Nothing'.
 --
@@ -231,8 +230,8 @@ defaultPorts =  [6881..6889]
 _resolvePeerAddr :: (IPAddress i) => PeerAddr HostName -> PeerAddr i
 _resolvePeerAddr = undefined
 
--- | Convert peer info from tracker response to socket address.  Used
---   for establish connection between peers.
+-- | Convert peer info from tracker or DHT announce query response to
+-- socket address. Usually used to intiate connection between peers.
 --
 peerSockAddr :: PeerAddr IP -> SockAddr
 peerSockAddr PeerAddr {..} =
diff --git a/src/Network/BitTorrent/Exchange/Wire.hs b/src/Network/BitTorrent/Exchange/Wire.hs
index 2a7d2aeb..5f7b0ebe 100644
--- a/src/Network/BitTorrent/Exchange/Wire.hs
+++ b/src/Network/BitTorrent/Exchange/Wire.hs
@@ -673,7 +673,7 @@ connectWire hs addr extCaps wire =
 --
 --   This function can throw 'WireFailure' exception.
 --
-acceptWire :: (IPAddress i) => Socket -> PeerAddr i -> Wire () -> IO ()
+acceptWire :: Socket -> PeerAddr IP -> Wire () -> IO ()
 acceptWire sock peerAddr wire = do
   bracket (return sock) close $ \ _ -> do
     error "acceptWire: not implemented"
diff --git a/src/Network/BitTorrent/Tracker/Message.hs b/src/Network/BitTorrent/Tracker/Message.hs
index d0be1c36..da46628b 100644
--- a/src/Network/BitTorrent/Tracker/Message.hs
+++ b/src/Network/BitTorrent/Tracker/Message.hs
@@ -70,7 +70,6 @@ module Network.BitTorrent.Tracker.Message
 
 import Control.Applicative
 import Control.Monad
-import Data.Aeson (ToJSON(..), FromJSON(..))
 import Data.Aeson.TH
 import Data.BEncode as BE hiding (Result)
 import Data.BEncode.BDict as BE
@@ -79,8 +78,10 @@ import Data.ByteString.Char8 as BC
 import Data.Char as Char
 import Data.Convertible
 import Data.Default
+import Data.Either
 import Data.List as L
 import Data.Maybe
+import Data.Monoid
 import Data.Serialize as S hiding (Result)
 import Data.String
 import Data.Text (Text)
@@ -193,9 +194,9 @@ instance Serialize AnnounceQuery where
     put           reqPeerId
     put           reqProgress
     putEvent      reqEvent
-    putWord32be $ fromMaybe 0 reqIP
-    putWord32be $ 0 -- TODO what the fuck is "key"?
-    putWord32be $ fromIntegral $ fromMaybe (-1) reqNumWant
+    putWord32host $ fromMaybe 0 reqIP
+    putWord32be   $ 0 -- TODO what the fuck is "key"?
+    putWord32be   $ fromIntegral $ fromMaybe (-1) reqNumWant
 
     put           reqPort
 
@@ -415,7 +416,9 @@ data AnnounceRequest = AnnounceRequest
   } deriving (Show, Eq, Typeable)
 
 instance QueryLike AnnounceRequest where
-  toQuery AnnounceRequest{..} = toQuery announceAdvises ++ toQuery announceQuery
+  toQuery AnnounceRequest{..} =
+    toQuery announceAdvises <>
+    toQuery announceQuery
 
 -- | Parse announce request from query string.
 parseAnnounceRequest :: SimpleQuery -> ParseResult AnnounceRequest
@@ -441,6 +444,11 @@ data PeerList ip
   | CompactPeerList [PeerAddr ip]
     deriving (Show, Eq, Typeable, Functor)
 
+-- | The empty non-compact peer list.
+instance Default (PeerList IP) where
+  def = PeerList []
+  {-# INLINE def #-}
+
 getPeerList :: PeerList IP -> [PeerAddr IP]
 getPeerList (PeerList        xs) = xs
 getPeerList (CompactPeerList xs) = xs
@@ -482,6 +490,17 @@ data AnnounceInfo =
      , respWarning     :: !(Maybe Text)
      } deriving (Show, Eq, Typeable)
 
+-- | Empty peer list with default reannounce interval.
+instance Default AnnounceInfo where
+  def = AnnounceInfo
+    { respComplete    = Nothing
+    , respIncomplete  = Nothing
+    , respInterval    = defaultReannounceInterval
+    , respMinInterval = Nothing
+    , respPeers       = def
+    , respWarning     = Nothing
+    }
+
 -- | HTTP tracker protocol compatible encoding.
 instance BEncode AnnounceInfo where
   toBEncode (Failure t)        = toDict $
@@ -494,10 +513,24 @@ instance BEncode AnnounceInfo where
     .: "interval"        .=! respInterval
     .: "min interval"    .=? respMinInterval
     .: "peers"           .=! peers
-    .: "peers6"          .=! peers6
+    .: "peers6"          .=? peers6
     .: "warning message" .=? respWarning
     .: endDict
-        where (peers,peers6) = splitIPList $ getPeerList respPeers
+    where
+      (peers, peers6) = prttn respPeers
+
+      prttn :: PeerList IP -> (PeerList IPv4, Maybe (PeerList IPv6))
+      prttn (PeerList        xs) = (PeerList xs, Nothing)
+      prttn (CompactPeerList xs) = mk $ partitionEithers $ toEither <$> xs
+        where
+          mk (v4s, v6s)
+            | L.null v6s = (CompactPeerList v4s, Nothing)
+            | otherwise  = (CompactPeerList v4s, Just (CompactPeerList v6s))
+
+          toEither :: PeerAddr IP -> Either (PeerAddr IPv4) (PeerAddr IPv6)
+          toEither PeerAddr {..} = case peerHost of
+            IPv4 ipv4 -> Left  $ PeerAddr peerId ipv4 peerPort
+            IPv6 ipv6 -> Right $ PeerAddr peerId ipv6 peerPort
 
   fromBEncode (BDict d)
     | Just t <- BE.lookup "failure reason" d = Failure <$> fromBEncode t
@@ -507,8 +540,26 @@ instance BEncode AnnounceInfo where
         <*>? "incomplete"
         <*>! "interval"
         <*>? "min interval"
-        <*>  (PeerList <$> (mergeIPLists <$>! "peers" <*>? "peers6"))
+        <*>  (uncurry merge =<< (,) <$>! "peers" <*>? "peers6")
         <*>? "warning message"
+    where
+      merge :: PeerList IPv4 -> Maybe (PeerList IPv6) -> BE.Get (PeerList IP)
+      merge (PeerList ips)          Nothing  = pure (PeerList ips)
+      merge (PeerList _  )          (Just _)
+        = fail "PeerList: non-compact peer list provided, \
+                         \but the `peers6' field present"
+
+      merge (CompactPeerList ipv4s) Nothing
+        = pure $ CompactPeerList (fmap IPv4 <$> ipv4s)
+
+      merge (CompactPeerList _    ) (Just (PeerList _))
+        = fail "PeerList: the `peers6' field value \
+                         \should contain *compact* peer list"
+
+      merge (CompactPeerList ipv4s) (Just (CompactPeerList ipv6s))
+        = pure $ CompactPeerList $
+                 (fmap IPv4 <$> ipv4s) <> (fmap IPv6 <$> ipv6s)
+
   fromBEncode _ = decodingError "Announce info"
 
 -- | UDP tracker protocol compatible encoding.