1 files changed, 17 insertions, 26 deletions

diff --git a/src/Network/KRPC/Manager.hs b/src/Network/KRPC/Manager.hs
index e7f0563b..b1e93101 100644
--- a/src/Network/KRPC/Manager.hs
+++ b/src/Network/KRPC/Manager.hs
@@ -76,6 +76,7 @@ import Data.Text as T
 import Data.Text.Encoding as T
 import Data.Tuple
 import Data.Typeable
+import Network.RPC
 import Network.KRPC.Message
 import Network.KRPC.Method
 import Network.Socket hiding (listen)
@@ -136,11 +137,11 @@ type CallId             = (TransactionId, SockAddr)
 type CallRes            = MVar (KQueryArgs, KResult) -- (raw response, decoded response)
 type PendingCalls       = IORef (Map CallId CallRes)
 
-type HandlerBody h = SockAddr -> KQueryArgs -> h (Either String KQueryArgs)
+type HandlerBody h msg v = SockAddr -> msg v -> h (Either String v)
 
 -- | Handler is a function which will be invoked then some /remote/
 -- node querying /this/ node.
-type Handler     h = (MethodName, HandlerBody h)
+type Handler     h msg v = (MethodName, HandlerBody h msg v)
 
 -- | Keep track pending queries made by /this/ node and handle queries
 -- made by /remote/ nodes.
@@ -150,7 +151,7 @@ data Manager h = Manager
   , listenerThread     :: !(MVar ThreadId)
   , transactionCounter :: {-# UNPACK #-} !TransactionCounter
   , pendingCalls       :: {-# UNPACK #-} !PendingCalls
-  , handlers           :: [Handler h]
+  , handlers           :: [Handler h KMessageOf BValue]
   }
 
 -- | A monad which can perform or handle queries.
@@ -182,10 +183,10 @@ sockAddrFamily (SockAddrCan _        ) = AF_CAN
 
 -- | Bind socket to the specified address. To enable query handling
 -- run 'listen'.
-newManager :: Options        -- ^ various protocol options;
-           -> SockAddr       -- ^ address to listen on;
-           -> [Handler h]    -- ^ handlers to run on incoming queries.
-           -> IO (Manager h) -- ^ new rpc manager.
+newManager :: Options                       -- ^ various protocol options;
+           -> SockAddr                      -- ^ address to listen on;
+           -> [Handler h KMessageOf BValue] -- ^ handlers to run on incoming queries.
+           -> IO (Manager h)                -- ^ new rpc manager.
 newManager opts @ Options {..} servAddr handlers = do
     validateOptions opts
     sock  <- bindServ
@@ -217,7 +218,7 @@ isActive Manager {..} = liftIO $ isBound sock
 
 -- | Normally you should use Control.Monad.Trans.Resource.allocate
 -- function.
-withManager :: Options -> SockAddr -> [Handler h]
+withManager :: Options -> SockAddr -> [Handler h KMessageOf BValue]
             -> (Manager h -> IO a) -> IO a
 withManager opts addr hs = bracket (newManager opts addr hs) closeManager
 
@@ -408,35 +409,25 @@ prettyQF e = T.encodeUtf8 $ "handler fail while performing query: "
 --   If the handler make some 'query' normally it /should/ handle
 --   corresponding 'QueryFailure's.
 --
-handler :: forall h a b. (KRPC a b, Monad h)
-        => (SockAddr -> a -> h b) -> Handler h
-handler body = (name, wrapper)
+handler :: forall h a b msg. (KRPC a b, Applicative h, Functor msg)
+        => Messaging msg TransactionId (Envelope a b) -> (SockAddr -> a -> h b) -> Handler h msg (Envelope a b)
+handler msging body = (name, wrapper)
   where
     Method name = method :: Method a b
     wrapper addr args =
-#ifdef VERSION_bencoding
-      case fromBEncode args of
-#else
-      case S.decode args of
-#endif
-        Left  e -> return $ Left e
-        Right a -> do
-          r <- body addr a
-#ifdef VERSION_bencoding
-          return $ Right $ toBEncode r
-#else
-          return $ Right $ S.encode r
-#endif
+      case unseal $ messagePayload msging args of
+        Left  e -> pure $ Left e
+        Right a -> Right . seal <$> body addr a
 
 runHandler :: MonadKRPC h m
-           => HandlerBody h -> SockAddr -> KQuery -> m KResult
+           => HandlerBody h KMessageOf BValue -> SockAddr -> KQuery -> m KResult
 runHandler h addr m = Lifted.catches wrapper failbacks
   where
     signature = querySignature (queryMethod m) (queryId m) addr
 
     wrapper = do
       $(logDebugS) "handler.quered" signature
-      result <- liftHandler (h addr (queryArgs m))
+      result <- liftHandler (h addr (Q m))
 
       case result of
         Left msg -> do