summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--krpc.cabal14
-rw-r--r--src/Network/KRPC.hs4
-rw-r--r--src/Network/KRPC/Manager.hs61
-rw-r--r--src/Network/KRPC/Message.hs45
-rw-r--r--src/Network/KRPC/Method.hs3
-rw-r--r--tests/Network/KRPC/MessageSpec.hs7
6 files changed, 99 insertions, 35 deletions
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
34 branch: master 34 branch: master
35 tag: v0.6.1.0 35 tag: v0.6.1.0
36 36
37flag builder
38 description: Use older bytestring package and bytestring-builder.
39 default: False
40
41
37library 42library
38 default-language: Haskell2010 43 default-language: Haskell2010
39 default-extensions: PatternGuards 44 default-extensions: PatternGuards
@@ -44,7 +49,6 @@ library
44 Network.KRPC.Method 49 Network.KRPC.Method
45 Network.KRPC.Manager 50 Network.KRPC.Manager
46 build-depends: base == 4.* 51 build-depends: base == 4.*
47 , bytestring >= 0.10
48 , text >= 0.11 52 , text >= 0.11
49 , data-default-class 53 , data-default-class
50 , lifted-base >= 0.1.1 54 , lifted-base >= 0.1.1
@@ -54,7 +58,13 @@ library
54 , monad-logger >= 0.3 58 , monad-logger >= 0.3
55 , bencoding >= 0.4.3 59 , bencoding >= 0.4.3
56 , network >= 2.3 60 , network >= 2.3
61 , cereal
57 , containers 62 , containers
63 if flag(builder)
64 build-depends: bytestring >= 0.9, bytestring-builder
65 else
66 build-depends: bytestring >= 0.10
67
58 if impl(ghc < 7.6) 68 if impl(ghc < 7.6)
59 build-depends: ghc-prim 69 build-depends: ghc-prim
60 ghc-options: -Wall 70 ghc-options: -Wall
@@ -89,4 +99,4 @@ benchmark bench
89 , monad-logger 99 , monad-logger
90 , criterion 100 , criterion
91 , krpc 101 , krpc
92 ghc-options: -O2 -fforce-recomp \ No newline at end of file 102 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
59 -- ** Query 59 -- ** Query
60 , QueryFailure (..) 60 , QueryFailure (..)
61 , query 61 , query
62 , query'
63 , queryRaw
62 , getQueryCount 64 , getQueryCount
63 65
64 -- ** Handler 66 -- ** Handler
@@ -86,4 +88,4 @@ import Data.Default.Class
86import Network.KRPC.Message 88import Network.KRPC.Message
87import Network.KRPC.Method 89import Network.KRPC.Method
88import Network.KRPC.Manager 90import Network.KRPC.Manager
89import Network.Socket (SockAddr (..)) \ No newline at end of file 91import 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
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
49import Control.Monad.Reader 51import Control.Monad.Reader
50import Control.Monad.Trans.Control 52import Control.Monad.Trans.Control
51import Data.BEncode as BE 53import Data.BEncode as BE
54import Data.BEncode.Internal as BE
52import Data.ByteString as BS 55import Data.ByteString as BS
53import Data.ByteString.Char8 as BC 56import Data.ByteString.Char8 as BC
54import Data.ByteString.Lazy as BL 57import Data.ByteString.Lazy as BL
@@ -118,7 +121,7 @@ type KResult = Either KError KResponse
118 121
119type TransactionCounter = IORef Int 122type TransactionCounter = IORef Int
120type CallId = (TransactionId, SockAddr) 123type CallId = (TransactionId, SockAddr)
121type CallRes = MVar KResult 124type CallRes = MVar (BValue, KResult)
122type PendingCalls = IORef (Map CallId CallRes) 125type PendingCalls = IORef (Map CallId CallRes)
123 126
124type HandlerBody h = SockAddr -> BValue -> h (BE.Result BValue) 127type HandlerBody h = SockAddr -> BValue -> h (BE.Result BValue)
@@ -163,6 +166,7 @@ sockAddrFamily :: SockAddr -> Family
163sockAddrFamily (SockAddrInet _ _ ) = AF_INET 166sockAddrFamily (SockAddrInet _ _ ) = AF_INET
164sockAddrFamily (SockAddrInet6 _ _ _ _) = AF_INET6 167sockAddrFamily (SockAddrInet6 _ _ _ _) = AF_INET6
165sockAddrFamily (SockAddrUnix _ ) = AF_UNIX 168sockAddrFamily (SockAddrUnix _ ) = AF_UNIX
169sockAddrFamily (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
264queryResponse :: BEncode a => CallRes -> IO a
265queryResponse 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)
275sendQuery :: BEncode a => Socket -> SockAddr -> a -> IO () 270sendQuery :: 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--
286query :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m b 281query :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m b
287query addr params = do 282query addr params = queryK addr params (\_ x _ -> x)
283
284-- | Like 'query' but possibly returns your externally routable IP address.
285query' :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m (b, Maybe ReflectedIP)
286query' 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.
291queryRaw :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m (b, BValue)
292queryRaw addr params = queryK addr params (\raw x _ -> (x,raw))
293
294queryK :: forall h m a b x. (MonadKRPC h m, KRPC a b) =>
295 SockAddr -> a -> (BValue -> b -> Maybe ReflectedIP -> x) -> m x
296queryK 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
422handleResponse :: MonadKRPC h m => KResult -> SockAddr -> m () 437handleResponse :: MonadKRPC h m => BValue -> KResult -> SockAddr -> m ()
423handleResponse result addr = do 438handleResponse 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
432handleMessage :: MonadKRPC h m => KMessage -> SockAddr -> m () 447handleMessage :: MonadKRPC h m => BValue -> KMessage -> SockAddr -> m ()
433handleMessage (Q q) = handleQuery q 448handleMessage _ (Q q) = handleQuery q
434handleMessage (R r) = handleResponse (Right r) 449handleMessage raw (R r) = handleResponse raw (Right r)
435handleMessage (E e) = handleResponse (Left e) 450handleMessage raw (E e) = handleResponse raw (Left e)
436 451
437listener :: MonadKRPC h m => m () 452listener :: MonadKRPC h m => m ()
438listener = do 453listener = 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
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
35 35
36 -- * Response 36 -- * Response
37 , KResponse(..) 37 , KResponse(..)
38 , ReflectedIP(..)
38 39
39 -- * Message 40 -- * Message
40 , KMessage (..) 41 , KMessage (..)
41 ) where 42 ) where
42 43
43import Control.Applicative 44import Control.Applicative
45import Control.Arrow
44import Control.Exception.Lifted as Lifted 46import Control.Exception.Lifted as Lifted
45import Data.BEncode as BE 47import Data.BEncode as BE
46import Data.ByteString as B 48import Data.ByteString as B
47import Data.ByteString.Char8 as BC 49import Data.ByteString.Char8 as BC
50import qualified Data.Serialize as S
51import Data.Word
48import Data.Typeable 52import Data.Typeable
53import Network.Socket (SockAddr (..),PortNumber,HostAddress)
49 54
50 55
51-- | This transaction ID is generated by the querying node and is 56-- | This transaction ID is generated by the querying node and is
@@ -188,6 +193,35 @@ instance BEncode KQuery where
188 KQuery <$>! "a" <*>! "q" <*>! "t" 193 KQuery <$>! "a" <*>! "q" <*>! "t"
189 {-# INLINE fromBEncode #-} 194 {-# INLINE fromBEncode #-}
190 195
196newtype ReflectedIP = ReflectedIP SockAddr
197 deriving (Eq, Ord, Show)
198
199instance BEncode ReflectedIP where
200 toBEncode (ReflectedIP addr) = BString (encodeAddr addr)
201 fromBEncode (BString bs) = ReflectedIP <$> decodeAddr bs
202 fromBEncode _ = Left "ReflectedIP should be a bencoded string"
203
204port16 :: Word16 -> PortNumber
205port16 = fromIntegral
206
207decodeAddr :: ByteString -> Either String SockAddr
208decodeAddr bs | B.length bs == 6
209 = ( \(a,p) -> SockAddrInet <$> fmap port16 p <*> a )
210 $ (S.runGet S.getWord32host *** S.decode )
211 $ B.splitAt 4 bs
212decodeAddr bs | B.length bs == 18
213 = ( \(a,p) -> flip SockAddrInet6 0 <$> fmap port16 p <*> a <*> pure 0 )
214 $ (S.decode *** S.decode )
215 $ B.splitAt 16 bs
216decodeAddr _ = Left "incorrectly sized address and port"
217
218encodeAddr :: SockAddr -> ByteString
219encodeAddr (SockAddrInet port addr)
220 = S.runPut (S.putWord32host addr >> S.put (fromIntegral port :: Word16))
221encodeAddr (SockAddrInet6 port _ addr _)
222 = S.runPut (S.put addr >> S.put (fromIntegral port :: Word16))
223encodeAddr _ = B.empty
224
191{----------------------------------------------------------------------- 225{-----------------------------------------------------------------------
192-- Response messages 226-- Response messages
193-----------------------------------------------------------------------} 227-----------------------------------------------------------------------}
@@ -206,7 +240,8 @@ instance BEncode KQuery where
206data KResponse = KResponse 240data KResponse = KResponse
207 { respVals :: BValue -- ^ 'BDict' containing return values; 241 { respVals :: BValue -- ^ 'BDict' containing return values;
208 , respId :: TransactionId -- ^ match to the corresponding 'queryId'. 242 , respId :: TransactionId -- ^ match to the corresponding 'queryId'.
209 } deriving (Show, Read, Eq, Ord, Typeable) 243 , respIP :: Maybe ReflectedIP
244 } deriving (Show, Eq, Ord, Typeable)
210 245
211-- | Responses, or KRPC message dictionaries with a \"y\" value of 246-- | Responses, or KRPC message dictionaries with a \"y\" value of
212-- \"r\", contain one additional key \"r\". The value of \"r\" is a 247-- \"r\", contain one additional key \"r\". The value of \"r\" is a
@@ -218,7 +253,8 @@ data KResponse = KResponse
218-- 253--
219instance BEncode KResponse where 254instance BEncode KResponse where
220 toBEncode KResponse {..} = toDict $ 255 toBEncode KResponse {..} = toDict $
221 "r" .=! respVals 256 "ip" .=? respIP
257 .: "r" .=! respVals
222 .: "t" .=! respId 258 .: "t" .=! respId
223 .: "y" .=! ("r" :: ByteString) 259 .: "y" .=! ("r" :: ByteString)
224 .: endDict 260 .: endDict
@@ -226,7 +262,8 @@ instance BEncode KResponse where
226 262
227 fromBEncode = fromDict $ do 263 fromBEncode = fromDict $ do
228 lookAhead $ match "y" (BString "r") 264 lookAhead $ match "y" (BString "r")
229 KResponse <$>! "r" <*>! "t" 265 addr <- optional (field (req "ip"))
266 (\r t -> KResponse r t addr) <$>! "r" <*>! "t"
230 {-# INLINE fromBEncode #-} 267 {-# INLINE fromBEncode #-}
231 268
232{----------------------------------------------------------------------- 269{-----------------------------------------------------------------------
@@ -249,4 +286,4 @@ instance BEncode KMessage where
249 Q <$> fromBEncode b 286 Q <$> fromBEncode b
250 <|> R <$> fromBEncode b 287 <|> R <$> fromBEncode b
251 <|> E <$> fromBEncode b 288 <|> E <$> fromBEncode b
252 <|> decodingError "KMessage: unknown message or message tag" \ No newline at end of file 289 <|> 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 }
47instance (Typeable a, Typeable b) => Show (Method a b) where 47instance (Typeable a, Typeable b) => Show (Method a b) where
48 showsPrec _ = showsMethod 48 showsPrec _ = showsMethod
49 49
50showsMethod :: forall a. forall b. Typeable a => Typeable b 50showsMethod :: forall a b. ( Typeable a , Typeable b ) => Method a b -> ShowS
51 => Method a b -> ShowS
52showsMethod (Method name) = 51showsMethod (Method name) =
53 showString (BC.unpack name) <> 52 showString (BC.unpack name) <>
54 showString " :: " <> 53 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
20 arbitrary = KQuery <$> pure (BInteger 0) <*> arbitrary <*> arbitrary 20 arbitrary = KQuery <$> pure (BInteger 0) <*> arbitrary <*> arbitrary
21 21
22instance Arbitrary KResponse where 22instance Arbitrary KResponse where
23 arbitrary = KResponse <$> pure (BList []) <*> arbitrary 23 -- TODO: Abitrary instance for ReflectedIP
24 arbitrary = KResponse <$> pure (BList []) <*> arbitrary <*> pure Nothing
24 25
25instance Arbitrary KMessage where 26instance Arbitrary KMessage where
26 arbitrary = frequency 27 arbitrary = frequency
@@ -64,8 +65,8 @@ spec = do
64 65
65 it "properly bencoded" $ do 66 it "properly bencoded" $ do
66 BE.decode "d1:rle1:t2:aa1:y1:re" `shouldBe` 67 BE.decode "d1:rle1:t2:aa1:y1:re" `shouldBe`
67 Right (KResponse (BList []) "aa") 68 Right (KResponse (BList []) "aa" Nothing)
68 69
69 describe "generic message" $ do 70 describe "generic message" $ do
70 it "properly bencoded (iso)" $ property $ \ km -> 71 it "properly bencoded (iso)" $ property $ \ km ->
71 BE.decode (BL.toStrict (BE.encode km)) `shouldBe` Right (km :: KMessage) \ No newline at end of file 72 BE.decode (BL.toStrict (BE.encode km)) `shouldBe` Right (km :: KMessage)
generated by cgit v1.2.3 (git 2.39.1) at 2024-11-11 18:11:23 +0000