diff options
Diffstat (limited to 'src/Network/KRPC/Protocol.hs')
-rw-r--r-- | src/Network/KRPC/Protocol.hs | 249 |
1 files changed, 249 insertions, 0 deletions
diff --git a/src/Network/KRPC/Protocol.hs b/src/Network/KRPC/Protocol.hs new file mode 100644 index 00000000..d28fdbeb --- /dev/null +++ b/src/Network/KRPC/Protocol.hs | |||
@@ -0,0 +1,249 @@ | |||
1 | -- | | ||
2 | -- Copyright : (c) Sam Truzjan 2013 | ||
3 | -- License : BSD3 | ||
4 | -- Maintainer : pxqr.sta@gmail.com | ||
5 | -- Stability : experimental | ||
6 | -- Portability : portable | ||
7 | -- | ||
8 | -- This module provides straightforward implementation of KRPC | ||
9 | -- protocol. In many situations 'Network.KRPC' should be prefered | ||
10 | -- since it gives more safe, convenient and high level api. | ||
11 | -- | ||
12 | -- > See http://www.bittorrent.org/beps/bep_0005.html#krpc-protocol | ||
13 | -- | ||
14 | {-# LANGUAGE OverloadedStrings #-} | ||
15 | {-# LANGUAGE FlexibleContexts #-} | ||
16 | {-# LANGUAGE TypeSynonymInstances #-} | ||
17 | {-# LANGUAGE MultiParamTypeClasses #-} | ||
18 | {-# LANGUAGE FunctionalDependencies #-} | ||
19 | {-# LANGUAGE DefaultSignatures #-} | ||
20 | module Remote.KRPC.Protocol | ||
21 | ( -- * Error | ||
22 | KError(..), ErrorCode, errorCode, mkKError | ||
23 | |||
24 | -- * Query | ||
25 | , KQuery(queryMethod, queryArgs), MethodName, ParamName, kquery | ||
26 | |||
27 | -- * Response | ||
28 | , KResponse(respVals), ValName, kresponse | ||
29 | |||
30 | , sendMessage, recvResponse | ||
31 | |||
32 | -- * Remote | ||
33 | , KRemote, KRemoteAddr, withRemote, remoteServer | ||
34 | |||
35 | -- * Re-exports | ||
36 | , encode, encoded, decode, decoded, toBEncode, fromBEncode | ||
37 | ) where | ||
38 | |||
39 | import Control.Applicative | ||
40 | import Control.Exception.Lifted as Lifted | ||
41 | import Control.Monad | ||
42 | import Control.Monad.IO.Class | ||
43 | import Control.Monad.Trans.Control | ||
44 | |||
45 | import Data.BEncode | ||
46 | import Data.ByteString as B | ||
47 | import Data.ByteString.Char8 as BC | ||
48 | import qualified Data.ByteString.Lazy as LB | ||
49 | import Data.Map as M | ||
50 | |||
51 | import Network.Socket hiding (recvFrom) | ||
52 | import Network.Socket.ByteString | ||
53 | |||
54 | |||
55 | -- | Errors used to signal that some error occurred while processing a | ||
56 | -- procedure call. Error may be send only from server to client but | ||
57 | -- not in the opposite direction. | ||
58 | -- | ||
59 | -- Errors are encoded as bencoded dictionary: | ||
60 | -- | ||
61 | -- > { "y" : "e", "e" : [<error_code>, <human_readable_error_reason>] } | ||
62 | -- | ||
63 | data KError | ||
64 | -- | Some error doesn't fit in any other category. | ||
65 | = GenericError { errorMessage :: ByteString } | ||
66 | |||
67 | -- | Occur when server fail to process procedure call. | ||
68 | | ServerError { errorMessage :: ByteString } | ||
69 | |||
70 | -- | Malformed packet, invalid arguments or bad token. | ||
71 | | ProtocolError { errorMessage :: ByteString } | ||
72 | |||
73 | -- | Occur when client trying to call method server don't know. | ||
74 | | MethodUnknown { errorMessage :: ByteString } | ||
75 | deriving (Show, Read, Eq, Ord) | ||
76 | |||
77 | instance BEncode KError where | ||
78 | {-# SPECIALIZE instance BEncode KError #-} | ||
79 | {-# INLINE toBEncode #-} | ||
80 | toBEncode e = fromAscAssocs -- WARN: keep keys sorted | ||
81 | [ "e" --> (errorCode e, errorMessage e) | ||
82 | , "y" --> ("e" :: ByteString) | ||
83 | ] | ||
84 | |||
85 | {-# INLINE fromBEncode #-} | ||
86 | fromBEncode (BDict d) | ||
87 | | M.lookup "y" d == Just (BString "e") | ||
88 | = uncurry mkKError <$> d >-- "e" | ||
89 | |||
90 | fromBEncode _ = decodingError "KError" | ||
91 | |||
92 | type ErrorCode = Int | ||
93 | |||
94 | errorCode :: KError -> ErrorCode | ||
95 | errorCode (GenericError _) = 201 | ||
96 | errorCode (ServerError _) = 202 | ||
97 | errorCode (ProtocolError _) = 203 | ||
98 | errorCode (MethodUnknown _) = 204 | ||
99 | {-# INLINE errorCode #-} | ||
100 | |||
101 | mkKError :: ErrorCode -> ByteString -> KError | ||
102 | mkKError 201 = GenericError | ||
103 | mkKError 202 = ServerError | ||
104 | mkKError 203 = ProtocolError | ||
105 | mkKError 204 = MethodUnknown | ||
106 | mkKError _ = GenericError | ||
107 | {-# INLINE mkKError #-} | ||
108 | |||
109 | serverError :: SomeException -> KError | ||
110 | serverError = ServerError . BC.pack . show | ||
111 | |||
112 | -- TODO Asc everywhere | ||
113 | |||
114 | |||
115 | type MethodName = ByteString | ||
116 | type ParamName = ByteString | ||
117 | |||
118 | -- | Query used to signal that caller want to make procedure call to | ||
119 | -- callee and pass arguments in. Therefore query may be only sent from | ||
120 | -- client to server but not in the opposite direction. | ||
121 | -- | ||
122 | -- Queries are encoded as bencoded dictionary: | ||
123 | -- | ||
124 | -- > { "y" : "q", "q" : "<method_name>", "a" : [<arg1>, <arg2>, ...] } | ||
125 | -- | ||
126 | data KQuery = KQuery { | ||
127 | queryMethod :: MethodName | ||
128 | , queryArgs :: Map ParamName BValue | ||
129 | } deriving (Show, Read, Eq, Ord) | ||
130 | |||
131 | instance BEncode KQuery where | ||
132 | {-# SPECIALIZE instance BEncode KQuery #-} | ||
133 | {-# INLINE toBEncode #-} | ||
134 | toBEncode (KQuery m args) = fromAscAssocs -- WARN: keep keys sorted | ||
135 | [ "a" --> BDict args | ||
136 | , "q" --> m | ||
137 | , "y" --> ("q" :: ByteString) | ||
138 | ] | ||
139 | |||
140 | {-# INLINE fromBEncode #-} | ||
141 | fromBEncode (BDict d) | ||
142 | | M.lookup "y" d == Just (BString "q") = | ||
143 | KQuery <$> d >-- "q" | ||
144 | <*> d >-- "a" | ||
145 | |||
146 | fromBEncode _ = decodingError "KQuery" | ||
147 | |||
148 | kquery :: MethodName -> [(ParamName, BValue)] -> KQuery | ||
149 | kquery name args = KQuery name (M.fromList args) | ||
150 | {-# INLINE kquery #-} | ||
151 | |||
152 | |||
153 | |||
154 | |||
155 | type ValName = ByteString | ||
156 | |||
157 | -- | KResponse used to signal that callee successufully process a | ||
158 | -- procedure call and to return values from procedure. KResponse should | ||
159 | -- not be sent if error occurred during RPC. Thus KResponse may be only | ||
160 | -- sent from server to client. | ||
161 | -- | ||
162 | -- Responses are encoded as bencoded dictionary: | ||
163 | -- | ||
164 | -- > { "y" : "r", "r" : [<val1>, <val2>, ...] } | ||
165 | -- | ||
166 | newtype KResponse = KResponse { respVals :: BDict } | ||
167 | deriving (Show, Read, Eq, Ord) | ||
168 | |||
169 | instance BEncode KResponse where | ||
170 | {-# INLINE toBEncode #-} | ||
171 | toBEncode (KResponse vals) = fromAscAssocs -- WARN: keep keys sorted | ||
172 | [ "r" --> vals | ||
173 | , "y" --> ("r" :: ByteString) | ||
174 | ] | ||
175 | |||
176 | {-# INLINE fromBEncode #-} | ||
177 | fromBEncode (BDict d) | ||
178 | | M.lookup "y" d == Just (BString "r") = | ||
179 | KResponse <$> d >-- "r" | ||
180 | |||
181 | fromBEncode _ = decodingError "KDict" | ||
182 | |||
183 | |||
184 | kresponse :: [(ValName, BValue)] -> KResponse | ||
185 | kresponse = KResponse . M.fromList | ||
186 | {-# INLINE kresponse #-} | ||
187 | |||
188 | |||
189 | |||
190 | type KRemoteAddr = (HostAddress, PortNumber) | ||
191 | |||
192 | type KRemote = Socket | ||
193 | |||
194 | withRemote :: (MonadBaseControl IO m, MonadIO m) => (KRemote -> m a) -> m a | ||
195 | withRemote = bracket (liftIO (socket AF_INET Datagram defaultProtocol)) | ||
196 | (liftIO . sClose) | ||
197 | {-# SPECIALIZE withRemote :: (KRemote -> IO a) -> IO a #-} | ||
198 | |||
199 | |||
200 | maxMsgSize :: Int | ||
201 | {-# INLINE maxMsgSize #-} | ||
202 | -- release | ||
203 | --maxMsgSize = 512 -- size of payload of one udp packet | ||
204 | -- bench | ||
205 | maxMsgSize = 64 * 1024 -- max udp size | ||
206 | |||
207 | |||
208 | -- TODO eliminate toStrict | ||
209 | sendMessage :: BEncode msg => msg -> KRemoteAddr -> KRemote -> IO () | ||
210 | sendMessage msg (host, port) sock = | ||
211 | sendAllTo sock (LB.toStrict (encoded msg)) (SockAddrInet port host) | ||
212 | {-# INLINE sendMessage #-} | ||
213 | |||
214 | recvResponse :: KRemote -> IO (Either KError KResponse) | ||
215 | recvResponse sock = do | ||
216 | (raw, _) <- recvFrom sock maxMsgSize | ||
217 | return $ case decoded raw of | ||
218 | Right resp -> Right resp | ||
219 | Left decE -> Left $ case decoded raw of | ||
220 | Right kerror -> kerror | ||
221 | _ -> ProtocolError (BC.pack decE) | ||
222 | |||
223 | -- | Run server using a given port. Method invocation should be done manually. | ||
224 | remoteServer :: (MonadBaseControl IO remote, MonadIO remote) | ||
225 | => PortNumber -- ^ Port number to listen. | ||
226 | -> (KRemoteAddr -> KQuery -> remote (Either KError KResponse)) | ||
227 | -- ^ Handler. | ||
228 | -> remote () | ||
229 | remoteServer servport action = bracket (liftIO bindServ) (liftIO . sClose) loop | ||
230 | where | ||
231 | bindServ = do | ||
232 | sock <- socket AF_INET Datagram defaultProtocol | ||
233 | bindSocket sock (SockAddrInet servport iNADDR_ANY) | ||
234 | return sock | ||
235 | |||
236 | loop sock = forever $ do | ||
237 | (bs, addr) <- liftIO $ recvFrom sock maxMsgSize | ||
238 | case addr of | ||
239 | SockAddrInet port host -> do | ||
240 | let kaddr = (host, port) | ||
241 | reply <- handleMsg bs kaddr | ||
242 | liftIO $ sendMessage reply kaddr sock | ||
243 | _ -> return () | ||
244 | |||
245 | where | ||
246 | handleMsg bs addr = case decoded bs of | ||
247 | Right query -> (either toBEncode toBEncode <$> action addr query) | ||
248 | `Lifted.catch` (return . toBEncode . serverError) | ||
249 | Left decodeE -> return $ toBEncode (ProtocolError (BC.pack decodeE)) | ||