From 5d0791e6ed2e500c08e7dadda39a254c8340cef5 Mon Sep 17 00:00:00 2001 From: joe Date: Tue, 17 Jan 2017 18:42:09 -0500 Subject: Handle reflected IP addresses (see bep 42). --- src/Network/KRPC.hs | 4 ++- src/Network/KRPC/Manager.hs | 61 ++++++++++++++++++++++++++++----------------- src/Network/KRPC/Message.hs | 45 ++++++++++++++++++++++++++++++--- src/Network/KRPC/Method.hs | 3 +-- 4 files changed, 83 insertions(+), 30 deletions(-) (limited to 'src/Network') diff --git a/src/Network/KRPC.hs b/src/Network/KRPC.hs index b15927cf..d185fb4c 100644 --- a/src/Network/KRPC.hs +++ b/src/Network/KRPC.hs @@ -59,6 +59,8 @@ module Network.KRPC -- ** Query , QueryFailure (..) , query + , query' + , queryRaw , getQueryCount -- ** Handler @@ -86,4 +88,4 @@ import Data.Default.Class import Network.KRPC.Message import Network.KRPC.Method import Network.KRPC.Manager -import Network.Socket (SockAddr (..)) \ No newline at end of file +import Network.Socket (SockAddr (..)) 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 " :: " <> -- cgit v1.2.3