1 files changed, 64 insertions, 45 deletions

diff --git a/src/Network/BitTorrent/Core/PeerAddr.hs b/src/Network/BitTorrent/Core/PeerAddr.hs
index 60ada54d..3c3e98c5 100644
--- a/src/Network/BitTorrent/Core/PeerAddr.hs
+++ b/src/Network/BitTorrent/Core/PeerAddr.hs
@@ -14,6 +14,7 @@
 {-# LANGUAGE DeriveDataTypeable         #-}
 {-# LANGUAGE FlexibleInstances          #-}
 {-# LANGUAGE DeriveFunctor              #-}
+{-# LANGUAGE ViewPatterns               #-}
 {-# OPTIONS -fno-warn-orphans           #-} -- for PortNumber instances
 module Network.BitTorrent.Core.PeerAddr
        ( -- * Peer address
@@ -24,25 +25,26 @@ module Network.BitTorrent.Core.PeerAddr
          -- * IP
        , mergeIPLists
        , splitIPList
-       , IP, IPv4, IPv6 --re-export Data.IP constructors
        , IPAddress ()
        ) where
 
 import Control.Applicative
+import Control.Monad
 import Data.Aeson (ToJSON, FromJSON)
 import Data.BEncode   as BS
 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
 import Data.Serialize as S
 import Data.String
 import Data.Typeable
-import Data.IP
-import Data.Foldable
-import Data.Either
+import Data.Word
 import Network.Socket
 import Text.PrettyPrint
 import Text.PrettyPrint.Class
@@ -52,12 +54,22 @@ import qualified Text.ParserCombinators.ReadP as RP
 import Network.BitTorrent.Core.PeerId
 
 
+{-----------------------------------------------------------------------
+--  Port number
+-----------------------------------------------------------------------}
+
 deriving instance ToJSON PortNumber
 deriving instance FromJSON PortNumber
 
 instance BEncode PortNumber where
-  toBEncode = toBEncode . fromEnum
-  fromBEncode b = toEnum <$> fromBEncode b
+  toBEncode   = toBEncode    .  fromEnum
+  fromBEncode = fromBEncode >=> portNumber
+    where
+      portNumber :: Integer -> BS.Result PortNumber
+      portNumber n
+        | 0 <= n && n <= fromIntegral (maxBound :: Word16)
+        = pure $ fromIntegral n
+        | otherwise = decodingError $ "PortNumber: " ++ show n
 
 instance Serialize PortNumber where
   get = fromIntegral <$> getWord16be
@@ -65,57 +77,59 @@ instance Serialize PortNumber where
   put = putWord16be . fromIntegral
   {-# INLINE put #-}
 
-class (Show i, Read i) => IPAddress i where
-    showIp :: i -> String
-    showIp = show
+{-----------------------------------------------------------------------
+--  IP addr
+-----------------------------------------------------------------------}
 
-    readIp :: String -> i
-    readIp = read
-
-    toHostAddr :: i -> Either HostAddress HostAddress6
+class IPAddress i where
+  toHostAddr :: i -> Either HostAddress HostAddress6
 
 instance IPAddress IPv4 where
-    toHostAddr = Left . toHostAddress
+  toHostAddr = Left . toHostAddress
 
 instance IPAddress IPv6 where
-    toHostAddr = Right . toHostAddress6
+  toHostAddr = Right . toHostAddress6
 
 instance IPAddress IP where
-    toHostAddr (IPv4 ip) = toHostAddr ip
-    toHostAddr (IPv6 ip) = toHostAddr ip
-
+  toHostAddr (IPv4 ip) = toHostAddr ip
+  toHostAddr (IPv6 ip) = toHostAddr ip
 
 deriving instance Typeable IP
 deriving instance Typeable IPv4
 deriving instance Typeable IPv6
 
-ipToBEncode :: IPAddress i => i -> BValue
-ipToBEncode ip = BString $ BS8.pack $ showIp ip
+ipToBEncode :: Show i => i -> BValue
+ipToBEncode ip = BString $ BS8.pack $ show ip
 
-ipFromBEncode :: Monad m => IPAddress a => BValue -> m a
-ipFromBEncode (BString ip) = return $ readIp $ BS8.unpack ip
-ipFromBEncode _            = fail "ipFromBEncode"
+ipFromBEncode :: Read a => BValue -> BS.Result a
+ipFromBEncode (BString (BS8.unpack -> ipStr))
+  | Just ip <- readMaybe (ipStr) = pure ip
+  |         otherwise            = decodingError $ "IP: " ++ ipStr
+ipFromBEncode _    = decodingError $ "IP: addr should be a bstring"
 
 instance BEncode IP where
-    toBEncode = ipToBEncode
-    fromBEncode = ipFromBEncode
+  toBEncode   = ipToBEncode
+  fromBEncode = ipFromBEncode
 
 instance BEncode IPv4 where
-    toBEncode = ipToBEncode
-    fromBEncode = ipFromBEncode
+  toBEncode   = ipToBEncode
+  fromBEncode = ipFromBEncode
 
 instance BEncode IPv6 where
-    toBEncode = ipToBEncode
-    fromBEncode = ipFromBEncode
+  toBEncode   = ipToBEncode
+  fromBEncode = ipFromBEncode
 
 instance Serialize IPv4 where
-    put ip = put $ toHostAddress ip
-    get = fromHostAddress <$> get
+    put = putWord32host    .  toHostAddress
+    get = fromHostAddress <$> getWord32host
 
 instance Serialize IPv6 where
     put ip = put $ toHostAddress6 ip
     get = fromHostAddress6 <$> get
 
+{-----------------------------------------------------------------------
+--  Peer addr
+-----------------------------------------------------------------------}
 -- TODO check semantic of ord and eq instances
 
 -- | Peer address info normally extracted from peer list or peer
@@ -126,23 +140,25 @@ data PeerAddr a = PeerAddr
   , peerPort :: {-# UNPACK #-} !PortNumber
   } deriving (Show, Eq, Typeable, Functor)
 
-peer_id_key, peer_ip_key, peer_port_key :: BKey
-peer_id_key   = "peer id"
+peer_ip_key, peer_id_key, peer_port_key :: BKey
 peer_ip_key   = "ip"
+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
   toBEncode PeerAddr {..} = toDict $
-       peer_id_key   .=? peerId
-    .: peer_ip_key   .=! peerAddr
+       peer_ip_key   .=! peerAddr
+    .: peer_id_key   .=? peerId
     .: peer_port_key .=! peerPort
     .: endDict
 
   fromBEncode = fromDict $ do
-    PeerAddr <$>? peer_id_key
+    peerAddr <$>? peer_id_key
              <*>! peer_ip_key
              <*>! peer_port_key
+    where
+      peerAddr ip pid port = PeerAddr ip pid port
 
 mergeIPLists :: [PeerAddr IPv4] -> Maybe [PeerAddr IPv6] -> [PeerAddr IP]
 mergeIPLists v4 v6 = (fmap IPv4 `L.map` v4)
@@ -162,10 +178,8 @@ splitIPList xs = partitionEithers $ toEither <$> xs
 --
 -- TODO: test byte order
 instance (Serialize a) => Serialize (PeerAddr a) where
-    put PeerAddr {..} =
-        put peerAddr >> put peerPort
-    get =
-        PeerAddr Nothing <$> get <*> get
+  put PeerAddr {..} = put peerAddr >> put peerPort
+  get = PeerAddr Nothing <$> get <*> get
 
 -- | @127.0.0.1:6881@
 instance Default (PeerAddr IPv4) where
@@ -178,7 +192,7 @@ instance Default (PeerAddr IPv4) where
 instance IsString (PeerAddr IPv4) where
   fromString str
     | [hostAddrStr, portStr] <- splitWhen (== ':') str
-    , hostAddr <- read hostAddrStr
+    , Just hostAddr <- readMaybe hostAddrStr
     , Just portNum  <- toEnum <$> readMaybe portStr
                 = PeerAddr Nothing hostAddr portNum
     | otherwise = error $ "fromString: unable to parse (PeerAddr IPv4): " ++ str
@@ -196,6 +210,11 @@ instance IsString (PeerAddr IPv6) where
         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 (PeerAddr IP) where
@@ -215,8 +234,8 @@ _resolvePeerAddr = undefined
 -- | Convert peer info from tracker response to socket address.  Used
 --   for establish connection between peers.
 --
-peerSockAddr :: (IPAddress i) => PeerAddr i -> SockAddr
+peerSockAddr :: PeerAddr IP -> SockAddr
 peerSockAddr PeerAddr {..} =
-  case toHostAddr peerAddr of
-    Left  host4 -> SockAddrInet  peerPort   host4
-    Right host6 -> SockAddrInet6 peerPort 0 host6 0
+  case peerAddr of
+    IPv4 ipv4 -> SockAddrInet  peerPort   (toHostAddress  ipv4)
+    IPv6 ipv6 -> SockAddrInet6 peerPort 0 (toHostAddress6 ipv6) 0