Initial commit.

3 files changed, 415 insertions, 0 deletions

diff --git a/src/Remote/KRPC.hs b/src/Remote/KRPC.hs
new file mode 100644
index 00000000..a6318ccd
--- /dev/null
+++ b/src/Remote/KRPC.hs
@@ -0,0 +1,132 @@
+-- |
+--   Copyright   :  (c) Sam T. 2013
+--   License     :  MIT
+--   Maintainer  :  pxqr.sta@gmail.com
+--   Stability   :  experimental
+--   Portability :  portable
+--
+--   This module provides remote procedure call.
+--
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE FlexibleContexts, DeriveDataTypeable #-}
+{-# LANGUAGE ExplicitForAll, KindSignatures #-}
+module Remote.KRPC
+       ( module Remote.KRPC.Method, RemoteAddr
+
+         -- * Client
+       , call, async, await
+
+         -- * Server
+       , handler, server
+       ) where
+
+import Control.Exception
+import Control.Monad
+import Control.Monad.Trans.Control
+import Control.Monad.IO.Class
+import Data.BEncode
+import Data.List as L
+import Data.Map  as M
+import Data.Text as T
+import Data.Typeable
+import Network
+
+import Remote.KRPC.Protocol
+import Remote.KRPC.Method
+
+
+data RPCException = RPCException KError
+                  deriving (Show, Eq, Typeable)
+
+instance Exception RPCException
+
+
+type RemoteAddr = KRemoteAddr
+
+
+queryCall :: BEncodable param
+          => KRemote -> KRemoteAddr
+          -> Method remote param result -> param -> IO ()
+queryCall sock addr m arg = sendMessage q addr sock
+  where
+    q = kquery (L.head (methodName m)) [(L.head (methodParams m), toBEncode arg)]
+
+getResult :: BEncodable result
+          => KRemote -> KRemoteAddr
+          -> Method remote param result -> IO result
+getResult sock addr m = do
+  resp <- recvResponse addr sock
+  case resp of
+    Left e -> throw (RPCException e)
+    Right (KResponse dict) -> do
+      let valName = L.head (methodVals m)
+      case M.lookup valName dict of
+        Just val | Right res <- fromBEncode val -> return res
+        Nothing     -> throw (RPCException (ProtocolError msg))
+          where
+            msg = T.concat
+              [ "Unable to find return value: ", T.pack (show valName), "\n"
+              , "in response: ", T.pack (show dict)
+              ]
+
+-- TODO async call
+-- | Makes remote procedure call. Throws RPCException if server
+-- returns error or decode error occurred.
+--
+call :: (MonadBaseControl IO host, MonadIO host)
+     => (BEncodable param, BEncodable result)
+     => RemoteAddr
+     -> Method remote param result
+     -> param
+     -> host result
+call addr m arg = liftIO $ withRemote $ \sock -> do
+  queryCall sock addr m arg
+  getResult sock addr m
+
+
+newtype Async result = Async { waitResult :: IO result }
+
+async :: MonadIO host
+      => (BEncodable param, BEncodable result)
+      => RemoteAddr
+      -> Method remote param result
+      -> param
+      -> host (Async result)
+async addr m arg = do
+  liftIO $ withRemote $ \sock ->
+     queryCall sock addr m arg
+  return $ Async $ withRemote $ \sock ->
+     getResult sock addr m
+
+await :: MonadIO host => Async result -> host result
+await = liftIO . waitResult
+
+-- TODO better name
+type MHandler remote = Method remote BEncode (Result BEncode)
+
+handler :: forall (remote :: * -> *) (param :: *) (result :: *).
+           (BEncodable param, BEncodable result)
+        => Monad remote
+        => Method remote param result
+        -> Method remote BEncode (Result BEncode)
+handler m = m { methodBody = \x -> do
+                  case fromBEncode x of
+                    Right a -> liftM (Right . toBEncode) (methodBody m a)
+                    Left e  -> return (Left e)
+              }
+
+-- TODO: allow forkIO
+server :: (MonadBaseControl IO remote, MonadIO remote)
+       => PortNumber
+       -> [MHandler remote]
+       -> remote ()
+server servport ms = remoteServer servport $ \_ q -> do
+  let name = queryMethod q
+  let args = queryArgs q
+  let m = L.head ms
+  res <- methodBody m (snd (L.head (M.toList args)))
+  case res of
+   Left  r -> return (Left (ProtocolError (T.pack r)))
+   Right r -> do
+     let retName = L.head (methodVals m)
+     return (Right (kresponse [(retName, r)]))
diff --git a/src/Remote/KRPC/Method.hs b/src/Remote/KRPC/Method.hs
new file mode 100644
index 00000000..f4b0bb9a
--- /dev/null
+++ b/src/Remote/KRPC/Method.hs
@@ -0,0 +1,76 @@
+{-# LANGUAGE OverloadedStrings #-}
+module Remote.KRPC.Method
+       ( Method(methodName, methodParams, methodVals, methodBody)
+
+         -- * Construction
+       , method
+
+         -- * Predefined methods
+       , idM, composeM, concatM
+       ) where
+
+import Prelude hiding ((.), id)
+import Control.Category
+import Control.Monad
+
+import Remote.KRPC.Protocol
+
+-- | The
+--
+--   * argument: type of method parameter
+--
+--   * remote: A monad used by server-side.
+--
+--   * result: type of return value of the method.
+--
+data Method remote param result = Method {
+    -- | Name used in query and
+    methodName   :: [MethodName]
+
+    -- | Description of each method parameter in right to left order.
+  , methodParams :: [ParamName]
+
+    -- | Description of each method return value in right to left order.
+  , methodVals   :: [ValName]
+
+    -- | Description of method body.
+  , methodBody   :: param -> remote result
+  }
+
+instance Monad remote => Category (Method remote) where
+  id  = idM
+  (.) = composeM
+
+-- | Remote identity function. Could be used for echo servers for example.
+--
+--   idM = method "id" ["x"] ["y"] return
+--
+idM :: Monad m => Method m a a
+idM = method "id" ["x"] ["y"] return
+
+-- | Pipelining of two or more methods.
+--
+--   NOTE: composed methods will work only with this implementation of
+--   KRPC, so both server and client should use this implementation,
+--   otherwise you more likely get the 'ProtocolError'.
+--
+composeM :: Monad m => Method m b c -> Method m a b -> Method m a c
+composeM g h = Method (methodName g ++ methodName h)
+                      (methodParams h)
+                      (methodVals g)
+                      (methodBody h >=> methodBody g)
+
+-- | Concat list of list. Could be used for performance tests.
+--
+--   concatM = method "concat" ["xxs"] ["xs"] $ return . Prelude.concat
+--
+concatM :: Monad m => Method m [[a]] [a]
+concatM = method "concat" ["xxs"] ["xs"] $ return . Prelude.concat
+
+
+method :: MethodName
+       -> [ParamName]
+       -> [ValName]
+       -> (param -> remote result)
+       -> Method remote param result
+method name = Method [name]
diff --git a/src/Remote/KRPC/Protocol.hs b/src/Remote/KRPC/Protocol.hs
new file mode 100644
index 00000000..0aa7e100
--- /dev/null
+++ b/src/Remote/KRPC/Protocol.hs
@@ -0,0 +1,207 @@
+-- |
+--   Copyright   :  (c) Sam T. 2013
+--   License     :  MIT
+--   Maintainer  :  pxqr.sta@gmail.com
+--   Stability   :  experimental
+--   Portability :  portable
+--
+--   This module provides straightforward implementation of KRPC
+--   protocol. In many situations Network.KRPC should be prefered
+--   since it gives more safe, convenient and high level api.
+--
+--   > See http://www.bittorrent.org/beps/bep_0005.html#krpc-protocol
+--
+{-# LANGUAGE OverloadedStrings, FlexibleContexts #-}
+module Remote.KRPC.Protocol
+       (
+         -- * Error
+         KError(..), errorCode, mkKError
+
+         -- * Query
+       , KQuery(..), MethodName, ParamName, kquery
+
+         -- * Response
+       , KResponse(..), ValName, kresponse
+       , sendMessage, recvResponse
+
+         -- * Remote
+       , KRemote, KRemoteAddr, withRemote, remoteServer
+
+         -- * Re-exports
+       , encode, encoded, decode, decoded, toBEncode, fromBEncode
+       ) where
+
+import Control.Applicative
+import Control.Exception.Lifted
+import Control.Monad
+import Control.Monad.IO.Class
+import Control.Monad.Trans.Control
+import Data.BEncode
+import Data.ByteString as B
+import qualified Data.ByteString.Lazy as LB
+import Data.Map as M
+import Data.Text as T
+import Network.Socket hiding (recvFrom)
+import Network.Socket.ByteString
+
+
+data KError
+  = GenericError { errorMessage :: Text }
+  | ServerError  { errorMessage :: Text }
+    -- | Malformed packet, invalid arguments or bad token.
+  | ProtocolError { errorMessage :: Text }
+  | MethodUnknown { errorMessage :: Text }
+   deriving (Show, Read, Eq, Ord)
+
+instance BEncodable KError where
+  toBEncode e = fromAssocs
+    [ "y" --> ("e" :: ByteString)
+    , "e" --> (errorCode e, errorMessage e)
+    ]
+
+  fromBEncode (BDict d)
+    |  M.lookup "y" d == Just (BString "e") =
+       uncurry mkKError <$> d >-- "e"
+
+  fromBEncode _ = decodingError "KError"
+
+type ErrorCode = Int
+
+errorCode :: KError -> ErrorCode
+errorCode (GenericError _)  = 201
+errorCode (ServerError _)   = 202
+errorCode (ProtocolError _) = 203
+errorCode (MethodUnknown _) = 204
+
+mkKError :: ErrorCode -> Text -> KError
+mkKError 201 = GenericError
+mkKError 202 = ServerError
+mkKError 203 = ProtocolError
+mkKError 204 = MethodUnknown
+mkKError _   = GenericError
+
+
+
+type MethodName = ByteString
+type ParamName  = ByteString
+
+data KQuery = KQuery {
+    queryMethod :: MethodName
+  , queryArgs   :: Map ParamName BEncode
+  } deriving (Show, Read, Eq, Ord)
+
+instance BEncodable KQuery where
+  toBEncode (KQuery m args) = fromAssocs
+    [ "y" --> ("q" :: ByteString)
+    , "q" --> m
+    , "a" --> BDict args
+    ]
+
+  fromBEncode (BDict d)
+    | M.lookup "y" d == Just (BString "q") =
+      KQuery <$> d >-- "q"
+             <*> d >-- "a"
+
+  fromBEncode _ = decodingError "KQuery"
+
+kquery :: MethodName -> [(ParamName, BEncode)] -> KQuery
+kquery name args = KQuery name (M.fromList args)
+
+
+
+
+type ValName = ByteString
+
+newtype KResponse = KResponse (Map ValName BEncode)
+                    deriving (Show, Read, Eq, Ord)
+
+instance BEncodable KResponse where
+  toBEncode (KResponse vals) = fromAssocs
+    [ "y" --> ("r" :: ByteString)
+    , "r" --> vals
+    ]
+
+
+  fromBEncode (BDict d)
+    | M.lookup "y" d == Just (BString "r") =
+      KResponse <$> d >-- "r"
+
+  fromBEncode _ = decodingError "KDict"
+
+kresponse :: [(ValName, BEncode)] -> KResponse
+kresponse = KResponse . M.fromList
+
+
+type KRemoteAddr = (HostAddress, PortNumber)
+
+remoteAddr :: KRemoteAddr -> SockAddr
+remoteAddr = SockAddrInet <$> snd <*> fst
+
+type KRemote = Socket
+
+withRemote :: (MonadBaseControl IO m, MonadIO m) => (KRemote -> m a) -> m a
+withRemote = bracket (liftIO (socket AF_INET Datagram defaultProtocol))
+                     (liftIO .  sClose)
+
+maxMsgSize :: Int
+maxMsgSize = 16 * 1024
+
+-- TODO eliminate toStrict
+sendMessage :: BEncodable msg => msg -> KRemoteAddr -> KRemote -> IO ()
+sendMessage msg (host, port) sock =
+  sendAllTo sock (LB.toStrict (encoded msg)) (SockAddrInet port host)
+
+recvResponse :: KRemoteAddr -> KRemote -> IO (Either KError KResponse)
+recvResponse addr sock = do
+  connect sock (remoteAddr addr)
+  (raw, _) <- recvFrom sock maxMsgSize
+  return $ case decoded raw of
+    Right resp -> Right resp
+    Left decE -> Left $ case decoded raw of
+      Right kerror -> kerror
+      _ -> ProtocolError (T.pack decE)
+
+remoteServer :: (MonadBaseControl IO remote, MonadIO remote)
+             => PortNumber
+             -> (KRemoteAddr -> KQuery -> remote (Either KError KResponse))
+             -> remote ()
+remoteServer servport action = bracket (liftIO bind) (liftIO . sClose) loop
+  where
+    bind = do
+     sock <- socket AF_INET Datagram defaultProtocol
+     bindSocket sock (SockAddrInet servport iNADDR_ANY)
+     return sock
+
+    loop sock = forever $ do
+      (bs, addr) <- liftIO $ recvFrom sock maxMsgSize
+
+      case addr of
+        SockAddrInet port host ->
+          case decoded bs of
+            Right query -> do
+              res <- action (host, port) query
+              case res of
+                Right resp -> liftIO $ sendMessage resp (host, port) sock
+                Left err   -> liftIO $ sendMessage err  (host, port) sock
+
+            Left decodeE -> liftIO $ sendMessage rpcE (host, port) sock
+              where
+                rpcE = ProtocolError $ T.concat
+                  ["Unable to decode query: ", T.pack (show bs), "\n"
+                  ,"Specifically: ", T.pack decodeE
+                  ]
+        _ -> return ()
+
+
+
+-- TODO to bencodable
+instance (BEncodable a, BEncodable b) => BEncodable (a, b) where
+  {-# SPECIALIZE instance (BEncodable a, BEncodable b) => BEncodable (a, b) #-}
+  toBEncode (a, b) = BList [toBEncode a, toBEncode b]
+  {-# INLINE toBEncode #-}
+
+  fromBEncode be = case fromBEncode be of
+    Right [a, b] -> (,) <$> fromBEncode a <*> fromBEncode b
+    Right _      -> decodingError "Unable to decode a pair."
+    Left  e      -> Left e
+  {-# INLINE fromBEncode #-}
\ No newline at end of file