Parse IPv4 addresses even when using variadic IP type.

2 files changed, 15 insertions, 6 deletions

diff --git a/src/Network/BitTorrent/Address.hs b/src/Network/BitTorrent/Address.hs
index 8dbe1583..c0c0a3f8 100644
--- a/src/Network/BitTorrent/Address.hs
+++ b/src/Network/BitTorrent/Address.hs
@@ -18,6 +18,8 @@
 {-# LANGUAGE MultiParamTypeClasses      #-}
 {-# LANGUAGE DeriveDataTypeable         #-}
 {-# LANGUAGE DeriveFunctor              #-}
+{-# LANGUAGE DeriveFoldable             #-}
+{-# LANGUAGE DeriveTraversable          #-}
 {-# LANGUAGE TemplateHaskell            #-}
 {-# OPTIONS -fno-warn-orphans           #-}
 module Network.BitTorrent.Address
@@ -443,7 +445,7 @@ instance Serialize IP where
       case n of
         4  -> IPv4 <$> get
         16 -> IPv6 <$> get
-        _ -> fail "Wrong number of bytes remaining to parse IP"
+        _ -> fail (show n ++ " is the wrong number of remaining bytes to parse IP")
 
 instance Serialize IPv4 where
     put = putWord32host    .  toHostAddress
@@ -693,7 +695,7 @@ distance (NodeId a) (NodeId b) = NodeDistance (BS.pack (BS.zipWith xor a b))
 data NodeAddr a = NodeAddr
   { nodeHost ::                !a
   , nodePort :: {-# UNPACK #-} !PortNumber
-  } deriving (Eq, Typeable, Functor)
+  } deriving (Eq, Typeable, Functor, Foldable, Traversable)
 
 instance Show a => Show (NodeAddr a) where
   showsPrec i NodeAddr {..}
@@ -745,7 +747,7 @@ fromPeerAddr PeerAddr {..} = NodeAddr
 data NodeInfo a = NodeInfo
   { nodeId   :: !NodeId
   , nodeAddr :: !(NodeAddr a)
-  } deriving (Show, Eq, Functor)
+  } deriving (Show, Eq, Functor, Foldable, Traversable)
 
 instance Eq a => Ord (NodeInfo a) where
   compare = comparing nodeId
diff --git a/src/Network/BitTorrent/DHT/Message.hs b/src/Network/BitTorrent/DHT/Message.hs
index 145141ee..008145de 100644
--- a/src/Network/BitTorrent/DHT/Message.hs
+++ b/src/Network/BitTorrent/DHT/Message.hs
@@ -59,6 +59,7 @@
 {-# LANGUAGE FlexibleInstances  #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE UndecidableInstances  #-}
+{-# LANGUAGE ScopedTypeVariables   #-}
 module Network.BitTorrent.DHT.Message
        ( -- * Envelopes
          Query (..)
@@ -184,21 +185,27 @@ newtype NodeFound ip = NodeFound [NodeInfo ip]
 nodes_key :: BKey
 nodes_key = "nodes"
 
+-- Convert IPv4 address.  Useful for using variadic IP type.
+from4 :: forall s. Address s => NodeInfo IPv4 -> Either String (NodeInfo s)
+from4 n = maybe (Left "Error converting IPv4") Right
+                    $ traverse (fromAddr :: IPv4 -> Maybe s) n
+
 binary :: Serialize a => BKey -> BE.Get [a]
 binary k = field (req k) >>= either (fail . format) return .
     runGet (many get)
   where
     format str = "fail to deserialize " ++ show k ++ " field: " ++ str
 
-instance (Typeable ip, Serialize ip) => BEncode (NodeFound ip) where
+instance (Typeable ip, Address ip) => BEncode (NodeFound ip) where
   toBEncode (NodeFound ns) = toDict $
        nodes_key .=! runPut (mapM_ put ns)
     .: endDict
 
-  fromBEncode = fromDict $ NodeFound <$> binary nodes_key
+  -- TODO: handle IPv6 by reading the "nodes6" key (see bep 32)
+  fromBEncode bval = NodeFound <$> (traverse from4 =<< fromDict (binary nodes_key) bval)
 
 -- | \"q\" == \"find_node\"
-instance (Serialize ip, Typeable ip)
+instance (Address ip, Typeable ip)
       => KRPC (Query FindNode) (Response (NodeFound ip)) where
   method = "find_node"