Handle reflected IP addresses (see bep 42).

3 files changed, 80 insertions, 29 deletions

diff --git a/src/Network/KRPC/Manager.hs b/src/Network/KRPC/Manager.hs
index 4436a9ba..9477d23c 100644
--- a/src/Network/KRPC/Manager.hs
+++ b/src/Network/KRPC/Manager.hs
@@ -30,6 +30,8 @@ module Network.KRPC.Manager
          -- * Queries
        , QueryFailure (..)
        , query
+       , query'
+       , queryRaw
        , getQueryCount
 
          -- * Handlers
@@ -49,6 +51,7 @@ import Control.Monad.Logger
 import Control.Monad.Reader
 import Control.Monad.Trans.Control
 import Data.BEncode as BE
+import Data.BEncode.Internal as BE
 import Data.ByteString as BS
 import Data.ByteString.Char8 as BC
 import Data.ByteString.Lazy as BL
@@ -118,7 +121,7 @@ type KResult = Either KError KResponse
 
 type TransactionCounter = IORef Int
 type CallId             = (TransactionId, SockAddr)
-type CallRes            = MVar KResult
+type CallRes            = MVar (BValue, KResult)
 type PendingCalls       = IORef (Map CallId CallRes)
 
 type HandlerBody h = SockAddr -> BValue -> h (BE.Result BValue)
@@ -163,6 +166,7 @@ sockAddrFamily :: SockAddr -> Family
 sockAddrFamily (SockAddrInet  _ _    ) = AF_INET
 sockAddrFamily (SockAddrInet6 _ _ _ _) = AF_INET6
 sockAddrFamily (SockAddrUnix  _      ) = AF_UNIX
+sockAddrFamily (SockAddrCan _        ) = AF_CAN
 
 -- | Bind socket to the specified address. To enable query handling
 -- run 'listen'.
@@ -261,15 +265,6 @@ unregisterQuery cid ref = do
   atomicModifyIORef' ref $ swap .
     M.updateLookupWithKey (const (const Nothing)) cid
 
-queryResponse :: BEncode a => CallRes -> IO a
-queryResponse ares = do
-  res <- readMVar ares
-  case res of
-    Left  (KError   c m _) -> throwIO $ QueryFailed c (T.decodeUtf8 m)
-    Right (KResponse {..}) ->
-      case fromBEncode respVals of
-        Right r -> pure r
-        Left  e -> throwIO $ QueryFailed ProtocolError (T.pack e)
 
 -- (sendmsg EINVAL)
 sendQuery :: BEncode a => Socket -> SockAddr -> a -> IO ()
@@ -284,7 +279,21 @@ sendQuery sock addr q = handle sockError $ sendMessage sock addr q
 --  respond with @error@ message or the query timeout expires.
 --
 query :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m b
-query addr params = do
+query addr params = queryK addr params (\_ x _ -> x)
+
+-- | Like 'query' but possibly returns your externally routable IP address.
+query' :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m (b, Maybe ReflectedIP)
+query' addr params = queryK addr params (const (,))
+
+-- | Enqueue a query, but give us the complete BEncoded content sent by the
+-- remote Node.  This is useful for handling extensions that this library does
+-- not otherwise support.
+queryRaw :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m (b, BValue)
+queryRaw addr params = queryK addr params (\raw x _ -> (x,raw))
+
+queryK :: forall h m a b x. (MonadKRPC h m, KRPC a b) =>
+    SockAddr -> a -> (BValue -> b -> Maybe ReflectedIP -> x) -> m x
+queryK addr params kont = do
   Manager {..} <- getManager
   tid <- liftIO $ genTransactionId transactionCounter
   let queryMethod = method :: Method a b
@@ -299,7 +308,13 @@ query addr params = do
       `onException` unregisterQuery (tid, addr) pendingCalls
 
     timeout (optQueryTimeout options * 10 ^ (6 :: Int)) $ do
-      queryResponse ares
+      (raw,res) <- readMVar ares
+      case res of
+        Left  (KError   c m _) -> throwIO $ QueryFailed c (T.decodeUtf8 m)
+        Right (KResponse {..}) ->
+          case fromBEncode respVals of
+            Right r -> pure $ kont raw r respIP
+            Left  e -> throwIO $ QueryFailed ProtocolError (T.pack e)
 
   case mres of
     Just res -> do
@@ -378,7 +393,7 @@ runHandler h addr KQuery {..} = Lifted.catches wrapper failbacks
 
         Right a -> do
           $(logDebugS) "handler.success" signature
-          return $ Right $ KResponse a queryId
+          return $ Right $ KResponse a queryId (Just $ ReflectedIP addr)
 
     failbacks =
       [ E.Handler $ \ (e :: HandlerFailure) -> do
@@ -419,20 +434,20 @@ handleQuery q addr = void $ fork $ do
   res <- dispatchHandler q addr
   sendMessage sock addr $ either toBEncode toBEncode res
 
-handleResponse :: MonadKRPC h m => KResult -> SockAddr -> m ()
-handleResponse result addr = do
+handleResponse :: MonadKRPC h m => BValue -> KResult -> SockAddr -> m ()
+handleResponse raw result addr = do
   Manager {..} <- getManager
   liftIO $ do
     let resultId = either errorId respId  result
     mcall <- unregisterQuery (resultId, addr) pendingCalls
     case mcall of
       Nothing   -> return ()
-      Just ares -> putMVar ares result
+      Just ares -> putMVar ares (raw,result)
 
-handleMessage :: MonadKRPC h m => KMessage -> SockAddr -> m ()
-handleMessage (Q q) = handleQuery    q
-handleMessage (R r) = handleResponse (Right r)
-handleMessage (E e) = handleResponse (Left  e)
+handleMessage :: MonadKRPC h m => BValue -> KMessage -> SockAddr -> m ()
+handleMessage _   (Q q) = handleQuery    q
+handleMessage raw (R r) = handleResponse raw (Right r)
+handleMessage raw (E e) = handleResponse raw (Left  e)
 
 listener :: MonadKRPC h m => m ()
 listener = do
@@ -441,10 +456,10 @@ listener = do
     (bs, addr) <- liftIO $ do
       handle exceptions $ BS.recvFrom sock (optMaxMsgSize options)
 
-    case BE.decode bs of
+    case BE.parse bs >>= \r -> (,) r <$> BE.decode bs of
       -- TODO ignore unknown messages at all?
-      Left  e -> liftIO $ sendMessage sock addr $ unknownMessage e
-      Right m -> handleMessage m addr
+      Left  e       -> liftIO $ sendMessage sock addr $ unknownMessage e
+      Right (raw,m) -> handleMessage raw m addr
   where
     exceptions :: IOError -> IO (BS.ByteString, SockAddr)
     exceptions e
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"
diff --git a/src/Network/KRPC/Method.hs b/src/Network/KRPC/Method.hs
index ea9da958..916b38a8 100644
--- a/src/Network/KRPC/Method.hs
+++ b/src/Network/KRPC/Method.hs
@@ -47,8 +47,7 @@ newtype Method param result = Method { methodName :: MethodName }
 instance (Typeable a, Typeable b) => Show (Method a b) where
   showsPrec _ = showsMethod
 
-showsMethod :: forall a. forall b. Typeable a => Typeable b
-            => Method a b -> ShowS
+showsMethod :: forall a b. ( Typeable a , Typeable b ) => Method a b -> ShowS
 showsMethod (Method name) =
     showString (BC.unpack name) <>
     showString " :: " <>