Make PeerAddr polimorphic in its address type

5 files changed, 137 insertions, 47 deletions

diff --git a/src/Network/BitTorrent/Client/Swarm.hs b/src/Network/BitTorrent/Client/Swarm.hs
index a9dca048..1901905c 100644
--- a/src/Network/BitTorrent/Client/Swarm.hs
+++ b/src/Network/BitTorrent/Client/Swarm.hs
@@ -43,7 +43,7 @@ getAnnounceQuery Swarm {..} = AnnounceQuery
   , reqEvent    = Nothing
   }
 
-askPeers :: Swarm -> IO [PeerAddr]
+askPeers :: Swarm -> IO [PeerAddr IP]
 askPeers s @ Swarm {..} = do
   AnnounceInfo {..} <- RPC.announce (getAnnounceQuery s) trackerConn
   return (getPeerList respPeers)
diff --git a/src/Network/BitTorrent/Core/PeerAddr.hs b/src/Network/BitTorrent/Core/PeerAddr.hs
index 94510bba..1da4c81a 100644
--- a/src/Network/BitTorrent/Core/PeerAddr.hs
+++ b/src/Network/BitTorrent/Core/PeerAddr.hs
@@ -12,12 +12,18 @@
 {-# LANGUAGE StandaloneDeriving         #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE DeriveDataTypeable         #-}
+{-# LANGUAGE FlexibleInstances          #-}
+{-# LANGUAGE DeriveFunctor              #-}
 {-# OPTIONS -fno-warn-orphans           #-} -- for PortNumber instances
 module Network.BitTorrent.Core.PeerAddr
        ( -- * Peer address
          PeerAddr(..)
        , defaultPorts
        , peerSockAddr
+       , mergeIPLists
+       , splitIPList
+       , IP, IPv4, IPv6 --re-export Data.IP constructors
+       , IPAddress ()
        ) where
 
 import Control.Applicative
@@ -38,10 +44,14 @@ import Data.String
 import Data.Typeable
 import Data.Word
 import Data.IP
+import Data.Maybe
+import Data.Foldable
+import Data.Either
 import Network.Socket
 import Text.PrettyPrint
 import Text.PrettyPrint.Class
 import Text.Read (readMaybe)
+import qualified Text.ParserCombinators.ReadP as RP
 import System.IO.Unsafe
 
 import Data.Torrent.JSON
@@ -61,32 +71,74 @@ instance Serialize PortNumber where
   put = putWord16be . fromIntegral
   {-# INLINE put #-}
 
+class (Show i, Read i) => IPAddress i where
+    showIp :: i -> String
+    showIp = show
+
+    readIp :: String -> i
+    readIp = read
+
+    toHostAddr :: i -> Either HostAddress HostAddress6
+
+instance IPAddress IPv4 where
+    toHostAddr = Left . toHostAddress
+
+instance IPAddress IPv6 where
+    toHostAddr = Right . toHostAddress6
+
+instance IPAddress IP where
+    toHostAddr (IPv4 ip) = toHostAddr ip
+    toHostAddr (IPv6 ip) = toHostAddr ip
+
+
+deriving instance Typeable IP
+deriving instance Typeable IPv4
+deriving instance Typeable IPv6
+
+ipToBEncode ip = BString $ BS8.pack $ showIp ip
+ipFromBEncode (BString ip) = return $ readIp $ BS8.unpack ip
+
+instance BEncode IP where
+    toBEncode = ipToBEncode
+    fromBEncode = ipFromBEncode
+
+instance BEncode IPv4 where
+    toBEncode = ipToBEncode
+    fromBEncode = ipFromBEncode
+
+instance BEncode IPv6 where
+    toBEncode = ipToBEncode
+    fromBEncode = ipFromBEncode
+
+instance Serialize IPv4 where
+    put ip = put $ toHostAddress ip
+    get = fromHostAddress <$> get
+
+instance Serialize IPv6 where
+    put ip = put $ toHostAddress6 ip
+    get = fromHostAddress6 <$> get
+
 -- TODO check semantic of ord and eq instances
 -- TODO use SockAddr instead of peerIP and peerPort
 
 -- | Peer address info normally extracted from peer list or peer
 -- compact list encoding.
-data PeerAddr = PeerAddr
+data PeerAddr a = PeerAddr
   { peerId   :: !(Maybe PeerId)
-  , peerIP   :: {-# UNPACK #-} !IP
+  , peerAddr :: a
   , peerPort :: {-# UNPACK #-} !PortNumber
-  } deriving (Show, Eq, Typeable)
-
-instance BEncode IP where
-    toBEncode ip = toBEncode $ BS8.pack $ show ip
-    fromBEncode (BString ip) = return $ fromString $ BS8.unpack ip
+  } deriving (Show, Eq, Typeable, Functor)
 
 peer_id_key, peer_ip_key, peer_port_key :: BKey
 peer_id_key   = "peer id"
 peer_ip_key   = "ip"
 peer_port_key = "port"
 
--- FIXME do we need to byteswap peerIP in bencode instance?
 -- | The tracker's 'announce response' compatible encoding.
-instance BEncode PeerAddr where
+instance (Typeable a, BEncode a) => BEncode (PeerAddr a) where
   toBEncode PeerAddr {..} = toDict $
        peer_id_key   .=? peerId
-    .: peer_ip_key   .=! BS8.pack (show peerIP)
+    .: peer_ip_key   .=! peerAddr
     .: peer_port_key .=! peerPort
     .: endDict
 
@@ -95,19 +147,32 @@ instance BEncode PeerAddr where
              <*>! peer_ip_key
              <*>! peer_port_key
 
+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'.
 --
 --   For more info see: <http://www.bittorrent.org/beps/bep_0023.html>
 --
-instance Serialize PeerAddr where -- TODO do it properly
-  put PeerAddr {..} = (putWord32host $ toHostAddress $ ipv4 peerIP) >> put peerPort
-  {-# INLINE put #-}
-  get = PeerAddr Nothing <$> (IPv4 . fromHostAddress <$> getWord32host) <*> get
-  {-# INLINE get #-}
+-- TODO: test byte order
+instance (Serialize a) => Serialize (PeerAddr a) where
+    put PeerAddr {..} =
+        put peerAddr >> put peerPort
+    get =
+        PeerAddr Nothing <$> get <*> get
 
 -- | @127.0.0.1:6881@
-instance Default PeerAddr where
+instance Default (PeerAddr IPv4) where
   def = "127.0.0.1:6881"
 
 -- inet_addr is pure; so it is safe to throw IO
@@ -122,30 +187,49 @@ unsafeCatchIO m = unsafePerformIO $
 --
 --   @peerPort \"127.0.0.1:6881\" == 6881@
 --
-instance IsString PeerAddr where
-  fromString str -- TODO IPv6
+instance IsString (PeerAddr IPv4) where
+  fromString str
     | [hostAddrStr, portStr] <- splitWhen (== ':') str
-    , Just hostAddr <- read hostAddrStr
+    , hostAddr <- read hostAddrStr
     , Just portNum  <- toEnum <$> readMaybe portStr
                 = PeerAddr Nothing hostAddr portNum
-    | otherwise = error $ "fromString: unable to parse PeerAddr: " ++ str
+    | otherwise = error $ "fromString: unable to parse (PeerAddr IPv4): " ++ str
+
+readsIPv6_port :: String -> [((IPv6, PortNumber), String)]
+readsIPv6_port = RP.readP_to_S $ do
+  ip <- RP.char '[' *> (RP.readS_to_P reads) <* RP.char ']'
+  RP.char ':'
+  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
 
 -- | fingerprint + "at" + dotted.host.inet.addr:port
-instance Pretty PeerAddr where
+-- TODO: instances for IPv6, HostName
+instance Pretty (PeerAddr IP) where
   pretty p @ PeerAddr {..}
     | Just pid <- peerId = pretty (fingerprint pid) <+> "at" <+> paddr
     |     otherwise      = paddr
     where
-      paddr = text (show (peerSockAddr p))
+      paddr = text (show peerAddr ++ ":" ++ show peerPort)
 
 -- | Ports typically reserved for bittorrent P2P listener.
 defaultPorts :: [PortNumber]
 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.
 --
-peerSockAddr :: PeerAddr -> SockAddr
+peerSockAddr :: (IPAddress i) => PeerAddr i -> SockAddr
 peerSockAddr PeerAddr {..}
-    | IPv4 v4 <- peerIP = SockAddrInet peerPort (toHostAddress v4)
-    | IPv6 v6 <- peerIP = SockAddrInet6 peerPort 0 (toHostAddress6 v6) 0
+    | Left hAddr <- toHostAddr peerAddr =
+        SockAddrInet peerPort hAddr
+    | Right hAddr <- toHostAddr peerAddr =
+        SockAddrInet6 peerPort 0 hAddr 0
diff --git a/src/Network/BitTorrent/Exchange/Assembler.hs b/src/Network/BitTorrent/Exchange/Assembler.hs
index 5dc7c5ca..aa009f49 100644
--- a/src/Network/BitTorrent/Exchange/Assembler.hs
+++ b/src/Network/BitTorrent/Exchange/Assembler.hs
@@ -65,6 +65,7 @@ import Data.IntMap.Strict as IM
 import Data.List as L
 import Data.Map as M
 import Data.Maybe
+import Data.IP
 
 import Data.Torrent.Piece
 import Network.BitTorrent.Core
@@ -79,7 +80,7 @@ type PieceMap = IntMap
 
 data Assembler = Assembler
   { -- | A set of blocks that have been 'Request'ed but not yet acked.
-    _inflight :: Map PeerAddr (PieceMap [BlockRange])
+    _inflight :: Map (PeerAddr IP) (PieceMap [BlockRange])
 
     -- | A set of blocks that but not yet assembled.
   , _pending  :: PieceMap Bucket
@@ -114,7 +115,7 @@ allowPiece pix a @ Assembler {..}  = over pending (IM.insert pix bkt) a
   where
     bkt = B.empty (piPieceLength info)
 
-allowedSet :: PeerAddr -> Assembler -> [BlockIx]
+allowedSet :: (PeerAddr IP) -> Assembler -> [BlockIx]
 allowedSet = undefined
 
 --inflight :: PeerAddr -> BlockIx -> Assembler -> Assembler
@@ -123,7 +124,7 @@ allowedSet = undefined
 --  You should check if a returned by peer block is actually have
 -- been requested and in-flight. This is needed to avoid "I send
 -- random corrupted block" attacks.
-insert :: PeerAddr -> Block a -> Assembler -> Assembler
+insert :: (PeerAddr IP) -> Block a -> Assembler -> Assembler
 insert = undefined
 
 {-
@@ -156,4 +157,4 @@ inserta :: Block a
         -> (PieceMap [Block a], Maybe (Piece a))
 inserta = undefined
 
--}
\ No newline at end of file
+-}
diff --git a/src/Network/BitTorrent/Exchange/Wire.hs b/src/Network/BitTorrent/Exchange/Wire.hs
index ae9babb3..27b4be12 100644
--- a/src/Network/BitTorrent/Exchange/Wire.hs
+++ b/src/Network/BitTorrent/Exchange/Wire.hs
@@ -505,7 +505,7 @@ initiateHandshake sock hs = do
   recvHandshake sock
 
 -- | Tries to connect to peer using reasonable default parameters.
-connectToPeer :: PeerAddr -> IO Socket
+connectToPeer :: (IPAddress i) => PeerAddr i -> IO Socket
 connectToPeer p = do
   sock <- socket AF_INET Stream Network.Socket.defaultProtocol
   connect sock (peerSockAddr p)
@@ -628,7 +628,7 @@ reconnect = undefined
 --
 -- This function can throw 'WireFailure' exception.
 --
-connectWire :: Handshake -> PeerAddr -> ExtendedCaps -> Wire () -> IO ()
+connectWire :: (IPAddress i) => Handshake -> PeerAddr i -> ExtendedCaps -> Wire () -> IO ()
 connectWire hs addr extCaps wire =
   bracket (connectToPeer addr) close $ \ sock -> do
     hs' <- initiateHandshake sock hs
@@ -673,7 +673,7 @@ connectWire hs addr extCaps wire =
 --
 --   This function can throw 'WireFailure' exception.
 --
-acceptWire :: Socket -> PeerAddr -> Wire () -> IO ()
+acceptWire :: (IPAddress i) => Socket -> PeerAddr i -> 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 6249cdc4..95b9c7ca 100644
--- a/src/Network/BitTorrent/Tracker/Message.hs
+++ b/src/Network/BitTorrent/Tracker/Message.hs
@@ -22,6 +22,8 @@
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE TemplateHaskell            #-}
 {-# LANGUAGE DeriveDataTypeable         #-}
+{-# LANGUAGE DeriveFunctor              #-}
+{-# LANGUAGE ScopedTypeVariables        #-}
 {-# OPTIONS -fno-warn-orphans           #-}
 module Network.BitTorrent.Tracker.Message
        ( -- * Announce
@@ -83,6 +85,7 @@ import Data.Text (Text)
 import Data.Text.Encoding
 import Data.Typeable
 import Data.Word
+import Data.IP
 import Network
 import Network.HTTP.Types.QueryLike
 import Network.HTTP.Types.URI hiding (urlEncode)
@@ -431,18 +434,18 @@ renderAnnounceRequest = queryToSimpleQuery . toQuery
 --
 --   For more info see: <http://www.bittorrent.org/beps/bep_0023.html>
 --
-data PeerList
-  =        PeerList { getPeerList :: [PeerAddr] }
-  | CompactPeerList { getPeerList :: [PeerAddr] }
-    deriving (Show, Eq, Typeable)
+data PeerList a
+  =        PeerList { getPeerList :: [PeerAddr a] }
+  | CompactPeerList { getPeerList :: [PeerAddr a] }
+    deriving (Show, Eq, Typeable, Functor)
 
-putCompactPeerList :: S.Putter [PeerAddr]
+putCompactPeerList :: (Serialize a) => S.Putter [PeerAddr a]
 putCompactPeerList = mapM_ put
 
-getCompactPeerList :: S.Get [PeerAddr]
+getCompactPeerList :: (Serialize a) => S.Get [PeerAddr a]
 getCompactPeerList = many get
 
-instance BEncode PeerList where
+instance (Typeable a, BEncode a, Serialize a) => BEncode (PeerList a) where
   toBEncode (PeerList        xs) = toBEncode xs
   toBEncode (CompactPeerList xs) = toBEncode $ runPut (putCompactPeerList xs)
 
@@ -473,7 +476,7 @@ data AnnounceInfo =
      , respMinInterval :: !(Maybe Int)
 
        -- | Peers that must be contacted.
-     , respPeers       :: !PeerList
+     , respPeers       :: !(PeerList IP)
 
        -- | Human readable warning.
      , respWarning     :: !(Maybe Text)
@@ -490,19 +493,21 @@ instance BEncode AnnounceInfo where
     .: "incomplete"      .=? respIncomplete
     .: "interval"        .=! respInterval
     .: "min interval"    .=? respMinInterval
-    .: "peers"           .=! respPeers
+    .: "peers"           .=! peers
+    .: "peers6"          .=! peers6
     .: "warning message" .=? respWarning
     .: endDict
+        where (peers,peers6) = splitIPList $ getPeerList respPeers
 
   fromBEncode (BDict d)
     | Just t <- BE.lookup "failure reason" d = Failure <$> fromBEncode t
-    | otherwise = (`fromDict` (BDict d)) $ do
-      AnnounceInfo
+    | otherwise = (`fromDict` (BDict d)) $
+       AnnounceInfo
         <$>? "complete"
         <*>? "incomplete"
         <*>! "interval"
         <*>? "min interval"
-        <*>! "peers"
+        <*>  (PeerList <$> (mergeIPLists <$>! "peers" <*>? "peers6"))
         <*>? "warning message"
   fromBEncode _ = decodingError "Announce info"
 
@@ -513,13 +518,13 @@ instance Serialize AnnounceInfo where
     putWord32be $ fromIntegral respInterval
     putWord32be $ fromIntegral $ fromMaybe 0 respIncomplete
     putWord32be $ fromIntegral $ fromMaybe 0 respComplete
-    forM_ (getPeerList respPeers) put
+    forM_ (fmap ipv4 <$> getPeerList respPeers) put
 
   get = do
     interval <- getWord32be
     leechers <- getWord32be
     seeders  <- getWord32be
-    peers    <- many get
+    peers    <- many $ fmap IPv4 <$> get
 
     return $ AnnounceInfo {
         respWarning     = Nothing