Handle reflected IP addresses (see bep 42).

1 files changed, 41 insertions, 4 deletions

diff --git a/src/Network/KRPC/Message.hs b/src/Network/KRPC/Message.hs
index ebf5573e..6f4ae620 100644
--- a/src/Network/KRPC/Message.hs
+++ b/src/Network/KRPC/Message.hs
@@ -35,17 +35,22 @@ module Network.KRPC.Message
 
          -- * Response
        , KResponse(..)
+       , ReflectedIP(..)
 
          -- * Message
        , KMessage (..)
        ) where
 
 import Control.Applicative
+import Control.Arrow
 import Control.Exception.Lifted as Lifted
 import Data.BEncode as BE
 import Data.ByteString as B
 import Data.ByteString.Char8 as BC
+import qualified Data.Serialize as S
+import Data.Word
 import Data.Typeable
+import Network.Socket (SockAddr (..),PortNumber,HostAddress)
 
 
 -- | This transaction ID is generated by the querying node and is
@@ -188,6 +193,35 @@ instance BEncode KQuery where
     KQuery <$>! "a" <*>! "q" <*>! "t"
   {-# INLINE fromBEncode #-}
 
+newtype ReflectedIP = ReflectedIP SockAddr
+  deriving (Eq, Ord, Show)
+
+instance BEncode ReflectedIP where
+  toBEncode (ReflectedIP addr) = BString (encodeAddr addr)
+  fromBEncode (BString bs) = ReflectedIP <$> decodeAddr bs
+  fromBEncode _            = Left "ReflectedIP should be a bencoded string"
+
+port16 :: Word16 -> PortNumber
+port16 = fromIntegral
+
+decodeAddr :: ByteString -> Either String SockAddr
+decodeAddr bs | B.length bs == 6
+    = ( \(a,p) -> SockAddrInet <$> fmap port16 p <*> a )
+      $ (S.runGet S.getWord32host *** S.decode )
+      $ B.splitAt 4 bs
+decodeAddr bs | B.length bs == 18
+    = ( \(a,p) -> flip SockAddrInet6 0 <$> fmap port16 p <*> a <*> pure 0 )
+      $ (S.decode *** S.decode )
+      $ B.splitAt 16 bs
+decodeAddr _ = Left "incorrectly sized address and port"
+
+encodeAddr :: SockAddr -> ByteString
+encodeAddr (SockAddrInet port addr)
+    = S.runPut (S.putWord32host addr >> S.put (fromIntegral port :: Word16))
+encodeAddr (SockAddrInet6 port _ addr _)
+    = S.runPut (S.put addr >> S.put (fromIntegral port :: Word16))
+encodeAddr _ = B.empty
+
 {-----------------------------------------------------------------------
 -- Response messages
 -----------------------------------------------------------------------}
@@ -206,7 +240,8 @@ instance BEncode KQuery where
 data KResponse = KResponse
   { respVals :: BValue         -- ^ 'BDict' containing return values;
   , respId   :: TransactionId  -- ^ match to the corresponding 'queryId'.
-  } deriving (Show, Read, Eq, Ord, Typeable)
+  , respIP   :: Maybe ReflectedIP
+  } deriving (Show, Eq, Ord, Typeable)
 
 -- | Responses, or KRPC message dictionaries with a \"y\" value of
 -- \"r\", contain one additional key \"r\". The value of \"r\" is a
@@ -218,7 +253,8 @@ data KResponse = KResponse
 --
 instance BEncode KResponse where
   toBEncode KResponse {..} = toDict $
-       "r" .=! respVals
+       "ip" .=? respIP
+    .: "r" .=! respVals
     .: "t" .=! respId
     .: "y" .=! ("r" :: ByteString)
     .: endDict
@@ -226,7 +262,8 @@ instance BEncode KResponse where
 
   fromBEncode = fromDict $ do
     lookAhead $ match "y" (BString "r")
-    KResponse <$>! "r" <*>! "t"
+    addr <- optional (field (req "ip"))
+    (\r t -> KResponse r t addr) <$>! "r" <*>! "t"
   {-# INLINE fromBEncode #-}
 
 {-----------------------------------------------------------------------
@@ -249,4 +286,4 @@ instance BEncode KMessage where
         Q <$> fromBEncode b
     <|> R <$> fromBEncode b
     <|> E <$> fromBEncode b
-    <|> decodingError "KMessage: unknown message or message tag"
\ No newline at end of file
+    <|> decodingError "KMessage: unknown message or message tag"