diff options
author | joe <joe@jerkface.net> | 2017-01-17 18:42:09 -0500 |
---|---|---|
committer | joe <joe@jerkface.net> | 2017-01-17 18:42:09 -0500 |
commit | 5d0791e6ed2e500c08e7dadda39a254c8340cef5 (patch) | |
tree | 1232e01ea7452473941e488af01b98bc90202554 /src/Network/KRPC/Manager.hs | |
parent | 5c54f6570a27e1509ddf048a91bd69c05052f2f1 (diff) |
Handle reflected IP addresses (see bep 42).
Diffstat (limited to 'src/Network/KRPC/Manager.hs')
-rw-r--r-- | src/Network/KRPC/Manager.hs | 61 |
1 files changed, 38 insertions, 23 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 | |||
30 | -- * Queries | 30 | -- * Queries |
31 | , QueryFailure (..) | 31 | , QueryFailure (..) |
32 | , query | 32 | , query |
33 | , query' | ||
34 | , queryRaw | ||
33 | , getQueryCount | 35 | , getQueryCount |
34 | 36 | ||
35 | -- * Handlers | 37 | -- * Handlers |
@@ -49,6 +51,7 @@ import Control.Monad.Logger | |||
49 | import Control.Monad.Reader | 51 | import Control.Monad.Reader |
50 | import Control.Monad.Trans.Control | 52 | import Control.Monad.Trans.Control |
51 | import Data.BEncode as BE | 53 | import Data.BEncode as BE |
54 | import Data.BEncode.Internal as BE | ||
52 | import Data.ByteString as BS | 55 | import Data.ByteString as BS |
53 | import Data.ByteString.Char8 as BC | 56 | import Data.ByteString.Char8 as BC |
54 | import Data.ByteString.Lazy as BL | 57 | import Data.ByteString.Lazy as BL |
@@ -118,7 +121,7 @@ type KResult = Either KError KResponse | |||
118 | 121 | ||
119 | type TransactionCounter = IORef Int | 122 | type TransactionCounter = IORef Int |
120 | type CallId = (TransactionId, SockAddr) | 123 | type CallId = (TransactionId, SockAddr) |
121 | type CallRes = MVar KResult | 124 | type CallRes = MVar (BValue, KResult) |
122 | type PendingCalls = IORef (Map CallId CallRes) | 125 | type PendingCalls = IORef (Map CallId CallRes) |
123 | 126 | ||
124 | type HandlerBody h = SockAddr -> BValue -> h (BE.Result BValue) | 127 | type HandlerBody h = SockAddr -> BValue -> h (BE.Result BValue) |
@@ -163,6 +166,7 @@ sockAddrFamily :: SockAddr -> Family | |||
163 | sockAddrFamily (SockAddrInet _ _ ) = AF_INET | 166 | sockAddrFamily (SockAddrInet _ _ ) = AF_INET |
164 | sockAddrFamily (SockAddrInet6 _ _ _ _) = AF_INET6 | 167 | sockAddrFamily (SockAddrInet6 _ _ _ _) = AF_INET6 |
165 | sockAddrFamily (SockAddrUnix _ ) = AF_UNIX | 168 | sockAddrFamily (SockAddrUnix _ ) = AF_UNIX |
169 | sockAddrFamily (SockAddrCan _ ) = AF_CAN | ||
166 | 170 | ||
167 | -- | Bind socket to the specified address. To enable query handling | 171 | -- | Bind socket to the specified address. To enable query handling |
168 | -- run 'listen'. | 172 | -- run 'listen'. |
@@ -261,15 +265,6 @@ unregisterQuery cid ref = do | |||
261 | atomicModifyIORef' ref $ swap . | 265 | atomicModifyIORef' ref $ swap . |
262 | M.updateLookupWithKey (const (const Nothing)) cid | 266 | M.updateLookupWithKey (const (const Nothing)) cid |
263 | 267 | ||
264 | queryResponse :: BEncode a => CallRes -> IO a | ||
265 | queryResponse ares = do | ||
266 | res <- readMVar ares | ||
267 | case res of | ||
268 | Left (KError c m _) -> throwIO $ QueryFailed c (T.decodeUtf8 m) | ||
269 | Right (KResponse {..}) -> | ||
270 | case fromBEncode respVals of | ||
271 | Right r -> pure r | ||
272 | Left e -> throwIO $ QueryFailed ProtocolError (T.pack e) | ||
273 | 268 | ||
274 | -- (sendmsg EINVAL) | 269 | -- (sendmsg EINVAL) |
275 | sendQuery :: BEncode a => Socket -> SockAddr -> a -> IO () | 270 | sendQuery :: BEncode a => Socket -> SockAddr -> a -> IO () |
@@ -284,7 +279,21 @@ sendQuery sock addr q = handle sockError $ sendMessage sock addr q | |||
284 | -- respond with @error@ message or the query timeout expires. | 279 | -- respond with @error@ message or the query timeout expires. |
285 | -- | 280 | -- |
286 | query :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m b | 281 | query :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m b |
287 | query addr params = do | 282 | query addr params = queryK addr params (\_ x _ -> x) |
283 | |||
284 | -- | Like 'query' but possibly returns your externally routable IP address. | ||
285 | query' :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m (b, Maybe ReflectedIP) | ||
286 | query' addr params = queryK addr params (const (,)) | ||
287 | |||
288 | -- | Enqueue a query, but give us the complete BEncoded content sent by the | ||
289 | -- remote Node. This is useful for handling extensions that this library does | ||
290 | -- not otherwise support. | ||
291 | queryRaw :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m (b, BValue) | ||
292 | queryRaw addr params = queryK addr params (\raw x _ -> (x,raw)) | ||
293 | |||
294 | queryK :: forall h m a b x. (MonadKRPC h m, KRPC a b) => | ||
295 | SockAddr -> a -> (BValue -> b -> Maybe ReflectedIP -> x) -> m x | ||
296 | queryK addr params kont = do | ||
288 | Manager {..} <- getManager | 297 | Manager {..} <- getManager |
289 | tid <- liftIO $ genTransactionId transactionCounter | 298 | tid <- liftIO $ genTransactionId transactionCounter |
290 | let queryMethod = method :: Method a b | 299 | let queryMethod = method :: Method a b |
@@ -299,7 +308,13 @@ query addr params = do | |||
299 | `onException` unregisterQuery (tid, addr) pendingCalls | 308 | `onException` unregisterQuery (tid, addr) pendingCalls |
300 | 309 | ||
301 | timeout (optQueryTimeout options * 10 ^ (6 :: Int)) $ do | 310 | timeout (optQueryTimeout options * 10 ^ (6 :: Int)) $ do |
302 | queryResponse ares | 311 | (raw,res) <- readMVar ares |
312 | case res of | ||
313 | Left (KError c m _) -> throwIO $ QueryFailed c (T.decodeUtf8 m) | ||
314 | Right (KResponse {..}) -> | ||
315 | case fromBEncode respVals of | ||
316 | Right r -> pure $ kont raw r respIP | ||
317 | Left e -> throwIO $ QueryFailed ProtocolError (T.pack e) | ||
303 | 318 | ||
304 | case mres of | 319 | case mres of |
305 | Just res -> do | 320 | Just res -> do |
@@ -378,7 +393,7 @@ runHandler h addr KQuery {..} = Lifted.catches wrapper failbacks | |||
378 | 393 | ||
379 | Right a -> do | 394 | Right a -> do |
380 | $(logDebugS) "handler.success" signature | 395 | $(logDebugS) "handler.success" signature |
381 | return $ Right $ KResponse a queryId | 396 | return $ Right $ KResponse a queryId (Just $ ReflectedIP addr) |
382 | 397 | ||
383 | failbacks = | 398 | failbacks = |
384 | [ E.Handler $ \ (e :: HandlerFailure) -> do | 399 | [ E.Handler $ \ (e :: HandlerFailure) -> do |
@@ -419,20 +434,20 @@ handleQuery q addr = void $ fork $ do | |||
419 | res <- dispatchHandler q addr | 434 | res <- dispatchHandler q addr |
420 | sendMessage sock addr $ either toBEncode toBEncode res | 435 | sendMessage sock addr $ either toBEncode toBEncode res |
421 | 436 | ||
422 | handleResponse :: MonadKRPC h m => KResult -> SockAddr -> m () | 437 | handleResponse :: MonadKRPC h m => BValue -> KResult -> SockAddr -> m () |
423 | handleResponse result addr = do | 438 | handleResponse raw result addr = do |
424 | Manager {..} <- getManager | 439 | Manager {..} <- getManager |
425 | liftIO $ do | 440 | liftIO $ do |
426 | let resultId = either errorId respId result | 441 | let resultId = either errorId respId result |
427 | mcall <- unregisterQuery (resultId, addr) pendingCalls | 442 | mcall <- unregisterQuery (resultId, addr) pendingCalls |
428 | case mcall of | 443 | case mcall of |
429 | Nothing -> return () | 444 | Nothing -> return () |
430 | Just ares -> putMVar ares result | 445 | Just ares -> putMVar ares (raw,result) |
431 | 446 | ||
432 | handleMessage :: MonadKRPC h m => KMessage -> SockAddr -> m () | 447 | handleMessage :: MonadKRPC h m => BValue -> KMessage -> SockAddr -> m () |
433 | handleMessage (Q q) = handleQuery q | 448 | handleMessage _ (Q q) = handleQuery q |
434 | handleMessage (R r) = handleResponse (Right r) | 449 | handleMessage raw (R r) = handleResponse raw (Right r) |
435 | handleMessage (E e) = handleResponse (Left e) | 450 | handleMessage raw (E e) = handleResponse raw (Left e) |
436 | 451 | ||
437 | listener :: MonadKRPC h m => m () | 452 | listener :: MonadKRPC h m => m () |
438 | listener = do | 453 | listener = do |
@@ -441,10 +456,10 @@ listener = do | |||
441 | (bs, addr) <- liftIO $ do | 456 | (bs, addr) <- liftIO $ do |
442 | handle exceptions $ BS.recvFrom sock (optMaxMsgSize options) | 457 | handle exceptions $ BS.recvFrom sock (optMaxMsgSize options) |
443 | 458 | ||
444 | case BE.decode bs of | 459 | case BE.parse bs >>= \r -> (,) r <$> BE.decode bs of |
445 | -- TODO ignore unknown messages at all? | 460 | -- TODO ignore unknown messages at all? |
446 | Left e -> liftIO $ sendMessage sock addr $ unknownMessage e | 461 | Left e -> liftIO $ sendMessage sock addr $ unknownMessage e |
447 | Right m -> handleMessage m addr | 462 | Right (raw,m) -> handleMessage raw m addr |
448 | where | 463 | where |
449 | exceptions :: IOError -> IO (BS.ByteString, SockAddr) | 464 | exceptions :: IOError -> IO (BS.ByteString, SockAddr) |
450 | exceptions e | 465 | exceptions e |