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/Manager.hs | 61 ++++++++++++++++++++++++++++----------------- 1 file changed, 38 insertions(+), 23 deletions(-) (limited to 'src/Network/KRPC/Manager.hs') 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 -- cgit v1.2.3