From 857988ceb7c9d73926c07bb1522ce86a1669f4c5 Mon Sep 17 00:00:00 2001 From: Sam Truzjan Date: Mon, 23 Dec 2013 03:28:54 +0400 Subject: Add spec for Message module --- tests/Network/KRPC/MessageSpec.hs | 71 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 71 insertions(+) create mode 100644 tests/Network/KRPC/MessageSpec.hs (limited to 'tests/Network/KRPC') diff --git a/tests/Network/KRPC/MessageSpec.hs b/tests/Network/KRPC/MessageSpec.hs new file mode 100644 index 00000000..7aca4489 --- /dev/null +++ b/tests/Network/KRPC/MessageSpec.hs @@ -0,0 +1,71 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} +module Network.KRPC.MessageSpec (spec) where +import Control.Applicative +import Data.ByteString.Lazy as BL +import Test.Hspec +import Test.QuickCheck +import Test.QuickCheck.Instances () + +import Data.BEncode as BE +import Network.KRPC.Message + +instance Arbitrary ErrorCode where + arbitrary = arbitraryBoundedEnum + +instance Arbitrary KError where + arbitrary = KError <$> arbitrary <*> arbitrary <*> arbitrary + +instance Arbitrary KQuery where + arbitrary = KQuery <$> pure (BInteger 0) <*> arbitrary <*> arbitrary + +instance Arbitrary KResponse where + arbitrary = KResponse <$> pure (BList []) <*> arbitrary + +instance Arbitrary KMessage where + arbitrary = frequency + [ (1, Q <$> arbitrary) + , (1, R <$> arbitrary) + , (1, E <$> arbitrary) + ] + +spec :: Spec +spec = do + describe "error message" $ do + it "properly bencoded (iso)" $ property $ \ ke -> + BE.decode (BL.toStrict (BE.encode ke)) `shouldBe` Right (ke :: KError) + + it "properly bencoded" $ do + BE.decode "d1:eli201e23:A Generic Error Ocurrede1:t2:aa1:y1:ee" + `shouldBe` Right (KError GenericError "A Generic Error Ocurred" "aa") + + BE.decode "d1:eli202e22:A Server Error Ocurrede1:t2:bb1:y1:ee" + `shouldBe` Right (KError ServerError "A Server Error Ocurred" "bb") + + BE.decode "d1:eli203e24:A Protocol Error Ocurrede1:t2:cc1:y1:ee" + `shouldBe` Right (KError ProtocolError "A Protocol Error Ocurred" "cc") + + BE.decode "d1:eli204e30:Attempt to call unknown methode1:t2:dd1:y1:ee" + `shouldBe` Right + (KError MethodUnknown "Attempt to call unknown method" "dd") + + describe "query message" $ do + it "properly bencoded (iso)" $ property $ \ kq -> + BE.decode (BL.toStrict (BE.encode kq)) `shouldBe` Right (kq :: KQuery) + + it "properly bencoded" $ do + BE.decode "d1:ale1:q4:ping1:t2:aa1:y1:qe" `shouldBe` + Right (KQuery (BList []) "ping" "aa") + + + describe "response message" $ do + it "properly bencoded (iso)" $ property $ \ kr -> + BE.decode (BL.toStrict (BE.encode kr)) `shouldBe` Right (kr :: KResponse) + + it "properly bencoded" $ do + BE.decode "d1:rle1:t2:aa1:y1:re" `shouldBe` + Right (KResponse (BList []) "aa") + + describe "generic message" $ do + it "properly bencoded (iso)" $ property $ \ km -> + BE.decode (BL.toStrict (BE.encode km)) `shouldBe` Right (km :: KMessage) \ No newline at end of file -- cgit v1.2.3 From 46b6ba10202b73ba413d18bd21a284e3897c12b0 Mon Sep 17 00:00:00 2001 From: Sam Truzjan Date: Tue, 24 Dec 2013 23:50:23 +0400 Subject: Update tests --- krpc.cabal | 54 +++++++-------------------- src/Network/KRPC.hs | 4 +- src/Network/KRPC/Manager.hs | 12 ++++++ tests/Client.hs | 80 ---------------------------------------- tests/Network/KRPC/MethodSpec.hs | 52 ++++++++++++++++++++++++++ tests/Network/KRPCSpec.hs | 33 +++++++++++++++++ tests/Server.hs | 20 ---------- tests/Shared.hs | 39 -------------------- 8 files changed, 114 insertions(+), 180 deletions(-) delete mode 100644 tests/Client.hs create mode 100644 tests/Network/KRPC/MethodSpec.hs create mode 100644 tests/Network/KRPCSpec.hs delete mode 100644 tests/Server.hs delete mode 100644 tests/Shared.hs (limited to 'tests/Network/KRPC') diff --git a/krpc.cabal b/krpc.cabal index fb7b01fe..c4c0ae10 100644 --- a/krpc.cabal +++ b/krpc.cabal @@ -64,8 +64,13 @@ test-suite spec default-language: Haskell2010 hs-source-dirs: tests main-is: Spec.hs + other-modules: Network.KRPCSpec + Network.KRPC.MethodSpec + Network.KRPC.MessageSpec build-depends: base == 4.* , bytestring + , network + , mtl , hspec , QuickCheck @@ -74,55 +79,24 @@ test-suite spec , bencoding , krpc ---test-suite test-client --- type: exitcode-stdio-1.0 --- default-language: Haskell2010 --- hs-source-dirs: tests --- main-is: Client.hs --- other-modules: Shared --- build-depends: base == 4.* --- , bytestring --- , process --- , filepath --- --- , bencoding --- , krpc --- , network --- --- , HUnit --- , test-framework --- , test-framework-hunit - - ---executable test-server +--executable bench-server -- default-language: Haskell2010 --- hs-source-dirs: tests +-- hs-source-dirs: bench -- main-is: Server.hs --- other-modules: Shared -- build-depends: base == 4.* -- , bytestring --- , bencoding -- , krpc -- , network +-- ghc-options: -fforce-recomp ---executable bench-server +--benchmark bench-client +-- type: exitcode-stdio-1.0 -- default-language: Haskell2010 -- hs-source-dirs: bench --- main-is: Server.hs --- build-depends: base == 4.* +-- main-is: Main.hs +-- build-depends: base == 4.* -- , bytestring +-- , criterion -- , krpc -- , network --- ghc-options: -fforce-recomp - -benchmark bench-client - type: exitcode-stdio-1.0 - default-language: Haskell2010 - hs-source-dirs: bench - main-is: Main.hs - build-depends: base == 4.* - , bytestring - , criterion - , krpc - , network - ghc-options: -O2 -fforce-recomp \ No newline at end of file +-- ghc-options: -O2 -fforce-recomp \ No newline at end of file diff --git a/src/Network/KRPC.hs b/src/Network/KRPC.hs index e10fcb58..10d2eb55 100644 --- a/src/Network/KRPC.hs +++ b/src/Network/KRPC.hs @@ -94,7 +94,6 @@ module Network.KRPC -- * RPC , Handler , handler - , listen , query -- * Manager @@ -102,9 +101,12 @@ module Network.KRPC , Manager , newManager , closeManager + , withManager + , listen -- * Exceptions , KError (..) + , ErrorCode (..) ) where import Network.KRPC.Message diff --git a/src/Network/KRPC/Manager.hs b/src/Network/KRPC/Manager.hs index 304f43f2..9d8688d3 100644 --- a/src/Network/KRPC/Manager.hs +++ b/src/Network/KRPC/Manager.hs @@ -9,6 +9,8 @@ module Network.KRPC.Manager , Manager , newManager , closeManager + , withManager + , query , Handler @@ -102,6 +104,9 @@ closeManager Manager {..} = do -- TODO unblock calls close sock +withManager :: SockAddr -> [Handler h] -> (Manager h -> IO a) -> IO a +withManager addr hs = bracket (newManager addr hs) closeManager + sendMessage :: MonadIO m => BEncode a => Socket -> SockAddr -> a -> m () sendMessage sock addr a = do liftIO $ sendManyTo sock (BL.toChunks (BE.encode a)) addr @@ -136,6 +141,11 @@ queryResponse ares = do Right r -> pure r Left e -> throwIO $ decodeError e respId +-- | +-- +-- This function will throw exception if quered node respond with +-- @error@ message or timeout expires. +-- query :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m b query addr params = do Manager {..} <- getManager @@ -161,6 +171,8 @@ query addr params = do -- Handlers -----------------------------------------------------------------------} +-- | Any thrown exception will be supressed and send over wire back to +-- the quering node. handler :: forall h a b. (KRPC a b, Monad h) => (SockAddr -> a -> h b) -> Handler h handler body = (name, wrapper) diff --git a/tests/Client.hs b/tests/Client.hs deleted file mode 100644 index 2b49bd82..00000000 --- a/tests/Client.hs +++ /dev/null @@ -1,80 +0,0 @@ -{-# LANGUAGE OverloadedStrings #-} -module Main (main) where - -import Control.Concurrent -import Control.Exception -import qualified Data.ByteString as B -import Data.BEncode as BE -import Data.BEncode.BDict as BE -import System.Process -import System.FilePath - -import Test.HUnit hiding (Test) -import Test.Framework -import Test.Framework.Providers.HUnit - -import Network.KRPC -import Network.Socket -import Shared - - -addr :: SockAddr -addr = SockAddrInet 6000 0 - -withServ :: FilePath -> IO () -> IO () -withServ serv_path = bracket up terminateProcess . const - where - up = do - (_, _, _, h) <- createProcess (proc serv_path []) - threadDelay 1000000 - return h - -main :: IO () -main = do - let serv_path = "dist" "build" "test-server" "test-server" - withServ serv_path $ - defaultMain tests - - -(==?) :: (Eq a, Show a) => a -> IO a -> Assertion -expected ==? action = do - actual <- action - expected @=? actual - -tests :: [Test] -tests = - [ testCase "unit" $ - () ==? call addr unitM () - - , testCase "echo int" $ - 1234 ==? call addr echoM 1234 - - , testCase "reverse 1..100" $ - reverse [1..100] ==? call addr reverseM [1..100] - - , testCase "reverse empty list" $ - reverse [] ==? call addr reverseM [] - - , testCase "reverse singleton list" $ - reverse [1] ==? call addr reverseM [1] - - , testCase "swap pair" $ - (1, 0) ==? call addr swapM (0, 1) - - , testCase "shift triple" $ - ([2..10], (), 1) ==? call addr shiftR ((), 1, [2..10]) - - , testCase "echo bytestring" $ - let bs = B.replicate 400 0 in - bs ==? call addr echoBytes bs - - , testCase "raw method" $ - BInteger 10 ==? call addr rawM (BInteger 10) - - , testCase "raw dict" $ - let dict = BDict $ BE.fromAscList - [ ("some_int", BInteger 100) - , ("some_list", BList [BInteger 10]) - ] - in dict ==? call addr rawDictM dict - ] diff --git a/tests/Network/KRPC/MethodSpec.hs b/tests/Network/KRPC/MethodSpec.hs new file mode 100644 index 00000000..c1c58282 --- /dev/null +++ b/tests/Network/KRPC/MethodSpec.hs @@ -0,0 +1,52 @@ +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} +module Network.KRPC.MethodSpec where +import Control.Applicative +import Data.BEncode +import Data.ByteString as BS +import Data.Typeable +import Network.KRPC +import Test.Hspec + + +data Ping = Ping + deriving (Show, Eq, Typeable) + +instance BEncode Ping where + toBEncode Ping = toBEncode () + fromBEncode b = Ping <$ (fromBEncode b :: Result ()) + +instance KRPC Ping Ping + +ping :: Monad h => Handler h +ping = handler $ \ _ Ping -> return Ping + +newtype Echo a = Echo a + deriving (Show, Eq, BEncode, Typeable) + +echo :: Monad h => Handler h +echo = handler $ \ _ (Echo a) -> return (Echo (a :: ByteString)) + +instance (Typeable a, BEncode a) => KRPC (Echo a) (Echo a) + +spec :: Spec +spec = do + describe "ping method" $ do + it "name is ping" $ do + (method :: Method Ping Ping) `shouldBe` "ping" + + it "has pretty Show instance" $ do + show (method :: Method Ping Ping) `shouldBe` "ping :: Ping -> Ping" + + describe "echo method" $ do + it "is overloadable" $ do + (method :: Method (Echo Int ) (Echo Int )) `shouldBe` "echo int" + (method :: Method (Echo Bool) (Echo Bool)) `shouldBe` "echo bool" + + it "has pretty Show instance" $ do + show (method :: Method (Echo Int) (Echo Int)) + `shouldBe` "echo int :: Echo Int -> Echo Int" \ No newline at end of file diff --git a/tests/Network/KRPCSpec.hs b/tests/Network/KRPCSpec.hs new file mode 100644 index 00000000..27148682 --- /dev/null +++ b/tests/Network/KRPCSpec.hs @@ -0,0 +1,33 @@ +{-# LANGUAGE OverloadedStrings #-} +module Network.KRPCSpec (spec) where +import Control.Monad.Reader +import Network.Socket (SockAddr (..)) +import Network.KRPC +import Network.KRPC.MethodSpec hiding (spec) +import Test.Hspec + +servAddr :: SockAddr +servAddr = SockAddrInet 6000 (256 * 256 * 256 + 127) + +handlers :: [Handler IO] +handlers = + [ handler $ \ _ Ping -> return Ping + , handler $ \ _ (Echo a) -> return (Echo (a :: Bool)) + , handler $ \ _ (Echo a) -> return (Echo (a :: Int)) + ] + +spec :: Spec +spec = do + describe "query" $ do + it "run handlers" $ do + let int = 0xabcd :: Int + (withManager servAddr handlers $ runReaderT $ do + listen + query servAddr (Echo int)) + `shouldReturn` Echo int + + it "throw timeout exception" $ do + (withManager servAddr handlers $ runReaderT $ do + query servAddr (Echo (0xabcd :: Int)) + ) + `shouldThrow` (== KError GenericError "timeout expired" "0") diff --git a/tests/Server.hs b/tests/Server.hs deleted file mode 100644 index b4b34891..00000000 --- a/tests/Server.hs +++ /dev/null @@ -1,20 +0,0 @@ -{-# LANGUAGE IncoherentInstances #-} -module Main (main) where - -import Data.BEncode -import Network.KRPC -import Network.Socket -import Shared - - -main :: IO () -main = server (SockAddrInet 6000 0) - [ unitM ==> return - , echoM ==> return - , echoBytes ==> return - , swapM ==> \(a, b) -> return (b, a) - , reverseM ==> return . reverse - , shiftR ==> \(a, b, c) -> return (c, a, b) - , rawM ==> return - , rawDictM ==> return - ] diff --git a/tests/Shared.hs b/tests/Shared.hs deleted file mode 100644 index 16547644..00000000 --- a/tests/Shared.hs +++ /dev/null @@ -1,39 +0,0 @@ -{-# LANGUAGE OverloadedStrings #-} -module Shared - ( echoM - , echoBytes - , unitM - , swapM - , reverseM - , shiftR - , rawM - , rawDictM - ) where - -import Data.ByteString (ByteString) -import Data.BEncode -import Network.KRPC - -unitM :: Method () () -unitM = method "unit" [] [] - -echoM :: Method Int Int -echoM = method "echo" ["x"] ["x"] - -echoBytes :: Method ByteString ByteString -echoBytes = method "echoBytes" ["x"] ["x"] - -reverseM :: Method [Int] [Int] -reverseM = method "reverse" ["xs"] ["ys"] - -swapM :: Method (Int, Int) (Int, Int) -swapM = method "swap" ["x", "y"] ["b", "a"] - -shiftR :: Method ((), Int, [Int]) ([Int], (), Int) -shiftR = method "shiftR" ["x", "y", "z"] ["a", "b", "c"] - -rawM :: Method BValue BValue -rawM = method "rawM" [""] [""] - -rawDictM :: Method BValue BValue -rawDictM = method "m" [] [] \ No newline at end of file -- cgit v1.2.3 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). --- krpc.cabal | 14 +++++++-- src/Network/KRPC.hs | 4 ++- src/Network/KRPC/Manager.hs | 61 ++++++++++++++++++++++++--------------- src/Network/KRPC/Message.hs | 45 ++++++++++++++++++++++++++--- src/Network/KRPC/Method.hs | 3 +- tests/Network/KRPC/MessageSpec.hs | 7 +++-- 6 files changed, 99 insertions(+), 35 deletions(-) (limited to 'tests/Network/KRPC') diff --git a/krpc.cabal b/krpc.cabal index c565bd2a..66c08ccb 100644 --- a/krpc.cabal +++ b/krpc.cabal @@ -34,6 +34,11 @@ source-repository this branch: master tag: v0.6.1.0 +flag builder + description: Use older bytestring package and bytestring-builder. + default: False + + library default-language: Haskell2010 default-extensions: PatternGuards @@ -44,7 +49,6 @@ library Network.KRPC.Method Network.KRPC.Manager build-depends: base == 4.* - , bytestring >= 0.10 , text >= 0.11 , data-default-class , lifted-base >= 0.1.1 @@ -54,7 +58,13 @@ library , monad-logger >= 0.3 , bencoding >= 0.4.3 , network >= 2.3 + , cereal , containers + if flag(builder) + build-depends: bytestring >= 0.9, bytestring-builder + else + build-depends: bytestring >= 0.10 + if impl(ghc < 7.6) build-depends: ghc-prim ghc-options: -Wall @@ -89,4 +99,4 @@ benchmark bench , monad-logger , criterion , krpc - ghc-options: -O2 -fforce-recomp \ No newline at end of file + ghc-options: -O2 -fforce-recomp 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 " :: " <> diff --git a/tests/Network/KRPC/MessageSpec.hs b/tests/Network/KRPC/MessageSpec.hs index 7aca4489..498ef679 100644 --- a/tests/Network/KRPC/MessageSpec.hs +++ b/tests/Network/KRPC/MessageSpec.hs @@ -20,7 +20,8 @@ instance Arbitrary KQuery where arbitrary = KQuery <$> pure (BInteger 0) <*> arbitrary <*> arbitrary instance Arbitrary KResponse where - arbitrary = KResponse <$> pure (BList []) <*> arbitrary + -- TODO: Abitrary instance for ReflectedIP + arbitrary = KResponse <$> pure (BList []) <*> arbitrary <*> pure Nothing instance Arbitrary KMessage where arbitrary = frequency @@ -64,8 +65,8 @@ spec = do it "properly bencoded" $ do BE.decode "d1:rle1:t2:aa1:y1:re" `shouldBe` - Right (KResponse (BList []) "aa") + Right (KResponse (BList []) "aa" Nothing) describe "generic message" $ do it "properly bencoded (iso)" $ property $ \ km -> - BE.decode (BL.toStrict (BE.encode km)) `shouldBe` Right (km :: KMessage) \ No newline at end of file + BE.decode (BL.toStrict (BE.encode km)) `shouldBe` Right (km :: KMessage) -- cgit v1.2.3