3 files changed, 284 insertions, 3 deletions

diff --git a/src/Network/KRPC/Manager.hs b/src/Network/KRPC/Manager.hs
new file mode 100644
index 00000000..9aa1bea7
--- /dev/null
+++ b/src/Network/KRPC/Manager.hs
@@ -0,0 +1,179 @@
+{-# LANGUAGE FlexibleContexts    #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+module Network.KRPC.Manager
+       ( MonadKRPC (..)
+       , newManager
+       , query
+       , handler
+       ) where
+
+import Control.Applicative
+import Control.Arrow
+import Control.Concurrent
+--import Control.Exception hiding (Handler)
+import Control.Exception.Lifted as Lifted hiding (Handler)
+import Control.Monad
+import Control.Monad.Trans.Control
+import Control.Monad.IO.Class
+import Data.BEncode as BE
+import Data.ByteString.Char8 as BC
+import Data.ByteString.Lazy as BL
+import Data.IORef
+import Data.List as L
+import Data.Map as M
+import Network.KRPC.Message
+import Network.KRPC.Method
+import Network.Socket
+import Network.Socket.ByteString as BS
+
+
+type KResult = Either KError KResponse
+
+type TransactionCounter = IORef Int
+type CallId             = (TransactionId, SockAddr)
+type CallRes            = MVar KResult
+type PendingCalls       = IORef (Map CallId CallRes)
+
+type HandlerBody m = SockAddr -> BValue -> m (BE.Result BValue)
+type Handler     m = (MethodName, HandlerBody m)
+
+data Manager m = Manager
+  { sock               :: !Socket
+  , transactionCounter :: {-# UNPACK #-} !TransactionCounter
+  , pendingCalls       :: {-# UNPACK #-} !PendingCalls
+  , handlers           :: [Handler m]
+  }
+
+class (MonadBaseControl IO m, MonadIO m) => MonadKRPC m where
+  getManager :: m (Manager a)
+
+sockAddrFamily :: SockAddr -> Family
+sockAddrFamily (SockAddrInet  _ _    ) = AF_INET
+sockAddrFamily (SockAddrInet6 _ _ _ _) = AF_INET6
+sockAddrFamily (SockAddrUnix  _      ) = AF_UNIX
+
+seedTransaction :: Int
+seedTransaction = 0
+
+newManager :: SockAddr -> IO (Manager a)
+newManager servAddr = do
+    sock  <- bindServ
+    tran  <- newIORef seedTransaction
+    calls <- newIORef M.empty
+    return $ Manager sock tran calls []
+  where
+    bindServ = do
+      let family = sockAddrFamily servAddr
+      sock <- socket family Datagram defaultProtocol
+      when (family == AF_INET6) $ do
+        setSocketOption sock IPv6Only 0
+      bindSocket sock servAddr
+      return sock
+
+sendMessage :: MonadIO m => BEncode a => Socket -> SockAddr -> a -> m ()
+sendMessage sock addr a =
+  liftIO $ sendManyTo sock (BL.toChunks (BE.encode a)) addr
+
+{-----------------------------------------------------------------------
+--  Client
+-----------------------------------------------------------------------}
+
+genTransactionId :: TransactionCounter -> IO TransactionId
+genTransactionId ref = do
+  cur <- atomicModifyIORef' ref $ \ cur -> (succ cur, cur)
+  return $ BC.pack (show cur)
+
+registerQuery :: CallId -> PendingCalls -> IO CallRes
+registerQuery cid ref = do
+  ares <- newEmptyMVar
+  atomicModifyIORef' ref $ \ m -> (M.insert cid ares m, ())
+  return ares
+
+unregisterQuery :: CallId -> PendingCalls -> IO ()
+unregisterQuery cid ref = do
+  atomicModifyIORef' ref $ \ m -> (M.delete cid m, ())
+
+queryResponse :: BEncode a => CallRes -> IO a
+queryResponse ares = do
+  res  <- readMVar ares
+  case res of
+    Left  e -> throwIO e
+    Right (KResponse {..}) ->
+      case fromBEncode respVals of
+        Left e  -> throwIO (KError ProtocolError (BC.pack e) respId)
+        Right a -> return a
+
+query :: forall m a b. (MonadKRPC m, KRPC a b) => SockAddr -> a -> m b
+query addr params = do
+  Manager {..} <- getManager
+  liftIO $ do
+    tid <- genTransactionId transactionCounter
+    let Method name = method :: Method a b
+    let q = KQuery (toBEncode params) name tid
+    ares <- registerQuery (tid, addr) pendingCalls
+    sendMessage sock addr q
+      `onException` unregisterQuery (tid, addr) pendingCalls
+    queryResponse ares
+
+{-----------------------------------------------------------------------
+--  Handlers
+-----------------------------------------------------------------------}
+
+handler :: forall m a b. (KRPC a b, MonadKRPC m)
+        => (SockAddr -> a -> m b) -> Handler m
+handler body = (name, wrapper)
+  where
+    Method name = method :: Method a b
+    wrapper addr args =
+      case fromBEncode args of
+        Left  e -> return $ Left e
+        Right a -> (Right . toBEncode) <$> body addr a
+
+runHandler :: MonadKRPC m => HandlerBody m -> SockAddr -> KQuery -> m KResult
+runHandler handler addr KQuery {..} = wrapper `Lifted.catch` failback
+  where
+    wrapper = ((`decodeError` queryId) +++ (`KResponse` queryId))
+           <$> handler addr queryArgs
+    failback e = return $ Left $ serverError e queryId
+
+dispatchHandler :: MonadKRPC m => KQuery -> SockAddr -> m KResult
+dispatchHandler q @ KQuery {..} addr = do
+  Manager {..} <- getManager
+  case L.lookup queryMethod handlers of
+    Nothing      -> return $ Left $ unknownMethod queryMethod queryId
+    Just handler -> runHandler handler addr q
+
+{-----------------------------------------------------------------------
+--  Listener
+-----------------------------------------------------------------------}
+
+handleQuery :: MonadKRPC m => KQuery -> SockAddr -> m ()
+handleQuery q addr = do
+  Manager {..} <- getManager
+  res <- dispatchHandler q addr
+  sendMessage sock addr $ either toBEncode toBEncode res
+
+handleResponse :: MonadKRPC m => KResult -> SockAddr -> m ()
+handleResponse result addr = do
+  Manager {..} <- getManager
+  mcall <- undefined (addr, respId) pendingCalls
+  case mcall of
+    Nothing   -> return ()
+    Just ares -> liftIO $ putMVar ares result
+
+handleMessage :: MonadKRPC m => KMessage -> SockAddr -> m ()
+handleMessage (Q q) = handleQuery    q
+handleMessage (R r) = handleResponse (Right r)
+handleMessage (E e) = handleResponse (Left  e)
+
+maxMsgSize :: Int
+maxMsgSize = 64 * 1024
+
+listener :: MonadKRPC m => m ()
+listener = do
+  Manager {..} <- getManager
+  forever $ do
+    (bs, addr) <- liftIO $ BS.recvFrom sock maxMsgSize
+    case BE.decode bs of
+      Left  e -> liftIO $ sendMessage sock addr  $ unknownMessage e
+      Right m -> handleMessage m addr
diff --git a/src/Network/KRPC/Message.hs b/src/Network/KRPC/Message.hs
index a70c2ea9..3bbfb1db 100644
--- a/src/Network/KRPC/Message.hs
+++ b/src/Network/KRPC/Message.hs
@@ -19,10 +19,17 @@
 {-# LANGUAGE DefaultSignatures      #-}
 {-# LANGUAGE DeriveDataTypeable     #-}
 module Network.KRPC.Message
-       ( -- * Error
-         ErrorCode (..)
+       ( -- * Transaction
+         TransactionId
+       , unknownTransaction
+
+         -- * Error
+       , ErrorCode (..)
        , KError(..)
        , serverError
+       , decodeError
+       , unknownMethod
+       , unknownMessage
 
          -- * Query
        , KQuery(..)
@@ -30,6 +37,9 @@ module Network.KRPC.Message
 
          -- * Response
        , KResponse(..)
+
+         -- * Message
+       , KMessage (..)
        ) where
 
 import Control.Applicative
@@ -47,6 +57,9 @@ import Data.Typeable
 -- are enough as they cover 2^16 outstanding queries.
 type TransactionId = ByteString
 
+unknownTransaction :: TransactionId
+unknownTransaction = ""
+
 {-----------------------------------------------------------------------
 -- Error messages
 -----------------------------------------------------------------------}
@@ -120,6 +133,15 @@ instance Exception KError
 serverError :: SomeException -> TransactionId -> KError
 serverError e = KError ServerError (BC.pack (show e))
 
+decodeError :: String -> TransactionId -> KError
+decodeError msg = KError ProtocolError (BC.pack msg)
+
+unknownMethod :: MethodName -> TransactionId -> KError
+unknownMethod = KError MethodUnknown
+
+unknownMessage :: String -> KError
+unknownMessage msg = KError ProtocolError (BC.pack msg) ""
+
 {-----------------------------------------------------------------------
 -- Query messages
 -----------------------------------------------------------------------}
@@ -183,4 +205,23 @@ instance BEncode KResponse where
   fromBEncode = fromDict $ do
     lookAhead $ match "y" (BString "r")
     KResponse <$>! "r" <*>! "t"
-  {-# INLINE fromBEncode #-}
\ No newline at end of file
+  {-# INLINE fromBEncode #-}
+
+{-----------------------------------------------------------------------
+-- Summed messages
+-----------------------------------------------------------------------}
+
+data KMessage
+  = Q KQuery
+  | R KResponse
+  | E KError
+
+instance BEncode KMessage where
+  toBEncode (Q q) = toBEncode q
+  toBEncode (R r) = toBEncode r
+  toBEncode (E e) = toBEncode e
+
+  fromBEncode b =
+        Q <$> fromBEncode b
+    <|> R <$> fromBEncode b
+    <|> E <$> fromBEncode b
diff --git a/src/Network/KRPC/Method.hs b/src/Network/KRPC/Method.hs
new file mode 100644
index 00000000..54aa8ef0
--- /dev/null
+++ b/src/Network/KRPC/Method.hs
@@ -0,0 +1,61 @@
+{-# LANGUAGE RankNTypes                 #-}
+{-# LANGUAGE MultiParamTypeClasses      #-}
+{-# LANGUAGE FunctionalDependencies     #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+module Network.KRPC.Method
+       ( Method (..)
+       , KRPC (..)
+       ) where
+
+import Data.BEncode (BEncode)
+import Data.Monoid
+import Data.String
+import Data.Typeable
+import Network.KRPC.Message
+
+
+-- | Method datatype used to describe name, parameters and return
+--   values of procedure. Client use a method to /invoke/, server
+--   /implements/ the method to make the actual work.
+--
+--   We use the following fantom types to ensure type-safiety:
+--
+--     * param: Type of method parameters. Ordinary Tuple type used
+--     to specify more than one parameter, so for example @Method
+--     (Int, Int) result@ will take two arguments.
+--
+--     * result: Type of return value of the method. Similarly,
+--     tuple used to specify more than one return value, so for
+--     exsample @Method (Foo, Bar) (Bar, Foo)@ will take two arguments
+--     and return two values.
+--
+newtype Method param result = Method MethodName
+  deriving (Eq, Ord, IsString, BEncode)
+
+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 (Method name) =
+    shows name <>
+    showString " :: " <>
+    shows paramsTy <>
+    showString " -> " <>
+    shows valuesTy
+  where
+    impossible = error "KRPC.showsMethod: impossible"
+    paramsTy = typeOf (impossible :: a)
+    valuesTy = typeOf (impossible :: b)
+
+-- | Example:
+--   @
+--   data Ping = Ping Text deriving BEncode
+--   data Pong = Pong Text deriving BEncode
+--
+--   instance KRPC Ping Pong where
+--     method = "ping"
+--   @
+class (BEncode req, BEncode resp) => KRPC req resp | req -> resp where
+  method :: Method req resp