Parse IPv4 addresses even when using variadic IP type.

1 files changed, 10 insertions, 3 deletions

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"