1 files changed, 130 insertions, 162 deletions

diff --git a/src/Network/DatagramServer.hs b/src/Network/DatagramServer.hs
index 21300108..e1bf91c5 100644
--- a/src/Network/DatagramServer.hs
+++ b/src/Network/DatagramServer.hs
@@ -60,6 +60,7 @@
 {-# LANGUAGE FunctionalDependencies #-}
 {-# LANGUAGE DeriveDataTypeable     #-}
 {-# LANGUAGE TemplateHaskell        #-}
+{-# LANGUAGE KindSignatures #-}
 module Network.DatagramServer
        ( -- * Methods
          Method
@@ -88,6 +89,7 @@ module Network.DatagramServer
        , withManager
        , isActive
        , listen
+       , Protocol(..)
 
          -- * Re-exports
        , ErrorCode (..)
@@ -96,7 +98,6 @@ module Network.DatagramServer
 
 import Data.Default.Class
 import Network.DatagramServer.Mainline
-import Network.KRPC.Method
 import Network.Socket (SockAddr (..))
 
 import Control.Applicative
@@ -192,42 +193,38 @@ validateOptions Options {..}
 --  Options
 -----------------------------------------------------------------------}
 
-type KResult = Either KError KMessage -- Response
+type KResult msg raw = Either (KError (TransactionID msg)) (msg raw)-- Response
 
 type TransactionCounter = IORef Int
-type CallId             = (TransactionId, SockAddr)
-type CallRes            = MVar (KQueryArgs, KResult) -- (raw response, decoded response)
-type PendingCalls       = IORef (Map CallId CallRes)
+type CallId msg           = (TransactionID msg, SockAddr)
+type CallRes msg raw      = MVar (raw, KResult msg raw) -- (raw response, decoded response)
+type PendingCalls msg raw = IORef (Map (CallId msg) (CallRes msg raw))
 
 type HandlerBody h msg v = SockAddr -> msg v -> h (Either String (msg v))
 
 -- | Handler is a function which will be invoked then some /remote/
 -- node querying /this/ node.
-type Handler     h msg v = (MethodName, HandlerBody h msg v)
+type Handler     h msg v = (QueryMethod msg, HandlerBody h msg v)
 
 -- | Keep track pending queries made by /this/ node and handle queries
 -- made by /remote/ nodes.
-data Manager h = Manager
+data Manager h raw msg = Manager
   { sock               :: !Socket
   , options            :: !Options
   , listenerThread     :: !(MVar ThreadId)
   , transactionCounter :: {-# UNPACK #-} !TransactionCounter
-  , pendingCalls       :: {-# UNPACK #-} !PendingCalls
-#ifdef VERSION_bencoding
-  , handlers           :: [Handler h KMessageOf BValue]
-#else
-  , handlers           :: [Handler h KMessageOf BC.ByteString]
-#endif
+  , pendingCalls       :: {-# UNPACK #-} !(PendingCalls msg raw)
+  , handlers           :: [Handler h msg raw]
   }
 
 -- | A monad which can perform or handle queries.
 class (MonadBaseControl IO m, MonadLogger m, MonadIO m)
-    => MonadKRPC h m | m -> h where
+    => MonadKRPC h m raw msg | m -> h, m -> raw, m -> msg  where
 
   -- | Ask for manager.
-  getManager :: m (Manager h)
+  getManager :: m (Manager h raw msg)
 
-  default getManager :: MonadReader (Manager h) m => m (Manager h)
+  default getManager :: MonadReader (Manager h raw msg) m => m (Manager h raw msg)
   getManager = ask
 
   -- | Can be used to add logging for instance.
@@ -237,7 +234,7 @@ class (MonadBaseControl IO m, MonadLogger m, MonadIO m)
   liftHandler = id
 
 instance (MonadBaseControl IO h, MonadLogger h, MonadIO h)
-      => MonadKRPC h (ReaderT (Manager h) h) where
+      => MonadKRPC h (ReaderT (Manager h raw msg) h) raw msg where
 
   liftHandler = lift
 
@@ -251,12 +248,8 @@ sockAddrFamily (SockAddrCan _        ) = AF_CAN
 -- run 'listen'.
 newManager :: Options                       -- ^ various protocol options;
            -> SockAddr                      -- ^ address to listen on;
-#ifdef VERSION_bencoding
-           -> [Handler h KMessageOf BValue] -- ^ handlers to run on incoming queries.
-#else
-           -> [Handler h KMessageOf BC.ByteString] -- ^ handlers to run on incoming queries.
-#endif
-           -> IO (Manager h)                -- ^ new rpc manager.
+           -> [Handler h msg raw] -- ^ handlers to run on incoming queries.
+           -> IO (Manager h raw msg)                -- ^ new rpc manager.
 newManager opts @ Options {..} servAddr handlers = do
     validateOptions opts
     sock  <- bindServ
@@ -274,7 +267,7 @@ newManager opts @ Options {..} servAddr handlers = do
       return sock
 
 -- | Unblock all pending calls and close socket.
-closeManager :: Manager m -> IO ()
+closeManager :: Manager m raw msg -> IO ()
 closeManager Manager {..} = do
   maybe (return ()) killThread =<< tryTakeMVar listenerThread
   -- TODO unblock calls
@@ -282,18 +275,14 @@ closeManager Manager {..} = do
 
 -- | Check if the manager is still active. Manager becomes active
 -- until 'closeManager' called.
-isActive :: Manager m -> IO Bool
+isActive :: Manager m raw msg -> IO Bool
 isActive Manager {..} = liftIO $ isBound sock
 {-# INLINE isActive #-}
 
 -- | Normally you should use Control.Monad.Trans.Resource.allocate
 -- function.
-#ifdef VERSION_bencoding
-withManager :: Options -> SockAddr -> [Handler h KMessageOf BValue]
-#else
-withManager :: Options -> SockAddr -> [Handler h KMessageOf BC.ByteString]
-#endif
-            -> (Manager h -> IO a) -> IO a
+withManager :: Options -> SockAddr -> [Handler h msg raw]
+            -> (Manager h raw msg -> IO a) -> IO a
 withManager opts addr hs = bracket (newManager opts addr hs) closeManager
 
 {-----------------------------------------------------------------------
@@ -301,15 +290,12 @@ withManager opts addr hs = bracket (newManager opts addr hs) closeManager
 -----------------------------------------------------------------------}
 
 -- TODO prettify log messages
-querySignature :: MethodName -> TransactionId -> SockAddr -> Text
+querySignature :: ( Show      ( QueryMethod   msg )
+                  , Serialize ( TransactionID msg ) )
+                  => QueryMethod msg -> TransactionID msg -> SockAddr -> Text
 querySignature name transaction addr = T.concat
-#ifdef VERSION_bencoding
-  [  "&", T.decodeUtf8 name
-  , " #", T.decodeUtf8 (Base16.encode transaction) -- T.decodeUtf8 transaction
-#else
   [  "&", T.pack (show name)
   , " #", T.decodeUtf8 (Base16.encode $ S.encode transaction)
-#endif
   , " @", T.pack (show addr)
   ]
 
@@ -332,23 +318,19 @@ sendMessage :: MonadIO m => Socket -> SockAddr -> BC.ByteString -> m ()
 sendMessage sock addr a = do
   liftIO $ sendManyTo sock [a] addr
 
-genTransactionId :: TransactionCounter -> IO TransactionId
+genTransactionId :: Envelope msg => TransactionCounter -> IO (TransactionID msg)
 genTransactionId ref = do
   cur <- atomicModifyIORef' ref $ \ cur -> (succ cur, cur)
-#ifdef VERSION_bencoding
-  return $ BC.pack (show cur)
-#else
-  return $ either (error "failed to create TransactionId") id $ S.decode $ BC.pack (L.take 24 $ show cur ++ L.repeat ' ')
-#endif
+  uniqueTransactionId cur
 
 -- | How many times 'query' call have been performed.
-getQueryCount :: MonadKRPC h m => m Int
+getQueryCount :: MonadKRPC h m raw msg => m Int
 getQueryCount = do
   Manager {..} <- getManager
   curTrans <- liftIO $ readIORef transactionCounter
   return $ curTrans - optSeedTransaction options
 
-registerQuery :: CallId -> PendingCalls -> IO CallRes
+registerQuery :: Ord (TransactionID msg) => CallId msg -> PendingCalls msg raw -> IO (CallRes msg raw)
 registerQuery cid ref = do
   ares <- newEmptyMVar
   atomicModifyIORef' ref $ \ m -> (M.insert cid ares m, ())
@@ -356,7 +338,7 @@ registerQuery cid ref = do
 
 -- simultaneous M.lookup and M.delete guarantees that we never get two
 -- or more responses to the same query
-unregisterQuery :: CallId -> PendingCalls -> IO (Maybe CallRes)
+unregisterQuery :: Ord (TransactionID msg) => CallId msg -> PendingCalls msg raw -> IO (Maybe (CallRes msg raw))
 unregisterQuery cid ref = do
   atomicModifyIORef' ref $ swap .
     M.updateLookupWithKey (const (const Nothing)) cid
@@ -374,35 +356,37 @@ sendQuery sock addr q = handle sockError $ sendMessage sock addr q
 --  This function should throw 'QueryFailure' exception if quered node
 --  respond with @error@ message or the query timeout expires.
 --
-query :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m b
-query addr params = queryK addr params (\_ x _ -> x)
+query :: forall h m a b raw msg. (SerializableTo raw b, Show (QueryMethod msg), Ord (TransactionID msg), Serialize (TransactionID msg), SerializableTo raw a, MonadKRPC h m raw msg, WireFormat raw msg, KRPC a b) => QueryMethod msg -> SockAddr -> a -> m b
+query meth addr params = queryK meth addr params (\_ x _ -> x)
 
 -- | Like 'query' but possibly returns your externally routable IP address.
-query' :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m (b, Maybe ReflectedIP)
-query' addr params = queryK addr params (const (,))
+query' :: forall h m a b raw msg. (SerializableTo raw b, Show (QueryMethod msg), Ord (TransactionID msg), Serialize (TransactionID msg), SerializableTo raw a, MonadKRPC h m raw msg, WireFormat raw msg, KRPC a b) => QueryMethod msg -> SockAddr -> a -> m (b, Maybe ReflectedIP)
+query' meth addr params = queryK meth addr params (const (,))
 
 -- | Enqueue a query, but give us the complete BEncoded content sent by the
 -- remote Node.  This is useful for handling extensions that this library does
 -- not otherwise support.
-queryRaw :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m (b, KQueryArgs)
-queryRaw addr params = queryK addr params (\raw x _ -> (x,raw))
+queryRaw :: forall h m a b raw msg. (SerializableTo raw b, Show (QueryMethod msg), Ord (TransactionID msg), Serialize (TransactionID msg), SerializableTo raw a, MonadKRPC h m raw msg, WireFormat raw msg, KRPC a b) => QueryMethod msg -> SockAddr -> a -> m (b, raw)
+queryRaw meth addr params = queryK meth addr params (\raw x _ -> (x,raw))
 
-queryK :: forall h m a b x. (MonadKRPC h m, KRPC a b) =>
-    SockAddr -> a -> (KQueryArgs -> b -> Maybe ReflectedIP -> x) -> m x
-queryK addr params kont = do
+queryK :: forall h m a b x raw msg. (SerializableTo raw b, Show (QueryMethod msg), Ord (TransactionID msg), Serialize (TransactionID msg), SerializableTo raw a, MonadKRPC h m raw msg, WireFormat raw msg, KRPC a b) =>
+    QueryMethod msg -> SockAddr -> a -> (raw -> b -> Maybe ReflectedIP -> x) -> m x
+queryK meth addr params kont = do
   Manager {..} <- getManager
   tid <- liftIO $ genTransactionId transactionCounter
-  let queryMethod = method :: Method a b
-  let signature = querySignature (methodName queryMethod) tid addr
+  -- let queryMethod = method :: Method a b
+  let signature = querySignature meth tid addr
   $(logDebugS) "query.sending" signature
 
   mres <- liftIO $ do
     ares <- registerQuery (tid, addr) pendingCalls
 
+    let cli = error "TODO TOX client node id"
+        ctx = error "TODO TOX ToxCipherContext or () for Mainline"
+    q <- buildQuery cli addr meth tid params
+    let qb = encodePayload (q :: msg a) :: msg raw
+        qbs = encodeHeaders ctx qb
 #ifdef VERSION_bencoding
-    let q = Q (KQuery (toBEncode params) (methodName queryMethod) tid)
-        qb = encodePayload q :: KMessage
-        qbs = encodeHeaders () qb :: BC.ByteString
 #else
     let q = Tox.Message (methodName queryMethod) cli tid params
         cli = error "TODO TOX client node id"
@@ -416,18 +400,13 @@ queryK addr params kont = do
     timeout (optQueryTimeout options * 10 ^ (6 :: Int)) $ do
       (raw,res) <- readMVar ares -- MVar (KQueryArgs, KResult)
       case res of
-#ifdef VERSION_bencoding
         Left  (KError   c m _) -> throwIO $ QueryFailed c (T.decodeUtf8 m)
-        Right (R (KResponse {..})) ->
-          case fromBEncode respVals of
-            Right r -> pure $ kont raw r respIP
-#else
-        Left  _                  -> throwIO $ QueryFailed GenericError "TODO: TOX ERROR"
-        Right (Tox.Message {..}) ->
-          case S.decode msgPayload of
-            Right r -> pure $ kont raw r Nothing
-#endif
-            Left  e -> throwIO $ QueryFailed ProtocolError (T.pack e)
+        Right m -> case decodePayload m of
+                    Right r -> case envelopeClass (r :: msg b) of
+                                Response reflectedAddr -> pure $ kont raw (envelopePayload r) reflectedAddr
+                                Error (KError   c m _) -> throwIO $ QueryFailed c (T.decodeUtf8 m) -- XXX neccessary?
+                                Query _                -> throwIO $ QueryFailed ProtocolError "BUG!! UNREACHABLE"
+                    Left e  -> throwIO $ QueryFailed ProtocolError (T.pack e)
 
   case mres of
     Just res -> do
@@ -477,43 +456,33 @@ 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 msg raw. (KRPC a b, Applicative h, Functor msg, WireFormat raw msg, SerializableTo raw a, SerializableTo raw b)
-        => (SockAddr -> a -> h b) -> Handler h msg raw
-handler body = (name, wrapper)
+handler :: forall h a b msg raw. (Applicative h, Functor msg, WireFormat raw msg, SerializableTo raw a, SerializableTo raw b)
+        => QueryMethod msg -> (SockAddr -> a -> h b) -> Handler h msg raw
+handler name body = (name, wrapper)
   where
-    Method name = method :: Method a b
     wrapper :: SockAddr -> msg raw -> h (Either String (msg raw))
     wrapper addr args =
       case decodePayload args of
         Left  e -> pure $ Left e
         Right a -> Right . encodePayload . buildReply (error "self node-id") addr args <$> body addr (envelopePayload a)
 
-runHandler :: MonadKRPC h m
-#ifdef VERSION_bencoding
-           => HandlerBody h KMessageOf BValue -> SockAddr -> KQuery -> m KResult
-#else
-           => HandlerBody h KMessageOf BC.ByteString -> SockAddr -> KQuery -> m KResult
-#endif
-runHandler h addr m = Lifted.catches wrapper failbacks
+runHandler :: ( MonadKRPC h m raw msg
+              , Envelope msg
+              , Show (QueryMethod msg)
+              , Serialize (TransactionID msg))
+           => QueryMethod msg -> HandlerBody h msg raw -> SockAddr -> msg raw -> m (KResult msg raw)
+runHandler meth h addr m = Lifted.catches wrapper failbacks
   where
-    signature = querySignature (queryMethod m) (queryId m) addr
+    signature = querySignature meth (envelopeTransaction m) addr
 
     wrapper = do
       $(logDebugS) "handler.quered" signature
-#ifdef VERSION_bencoding
-      result <- liftHandler (h addr (Q m))
-#else
       result <- liftHandler (h addr m)
-#endif
 
       case result of
         Left msg -> do
           $(logDebugS) "handler.bad_query" $ signature <> " !" <> T.pack msg
-#ifdef VERSION_bencoding
-          return $ Left $ KError ProtocolError (BC.pack msg) (queryId m)
-#else
-          return $ Left $ decodeError "TODO TOX ProtocolError" (queryId m)
-#endif
+          return $ Left $ KError ProtocolError (BC.pack msg) (envelopeTransaction m)
 
         Right a -> do -- KQueryArgs
           $(logDebugS) "handler.success" signature
@@ -522,41 +491,30 @@ runHandler h addr m = Lifted.catches wrapper failbacks
     failbacks =
       [ E.Handler $ \ (e :: HandlerFailure) -> do
           $(logDebugS) "handler.failed" signature
-#ifdef VERSION_bencoding
-          return $ Left $ KError ProtocolError (prettyHF e) (queryId m)
-#else
-          return $ Left $ decodeError "TODO TOX ProtocolError 2" (queryId m)
-#endif
+          return $ Left $ KError ProtocolError (prettyHF e) (envelopeTransaction m)
 
 
         -- may happen if handler makes query and fail
       , E.Handler $ \ (e :: QueryFailure) -> do
-#ifdef VERSION_bencoding
-          return $ Left $ KError ServerError (prettyQF e) (queryId m)
-#else
-          return $ Left $ decodeError "TODO TOX ServerError" (queryId m)
-#endif
+          return $ Left $ KError ServerError (prettyQF e) (envelopeTransaction m)
 
         -- since handler thread exit after sendMessage we can safely
         -- suppress async exception here
       , E.Handler $ \ (e :: SomeException) -> do
-#ifdef VERSION_bencoding
-          return $ Left $ KError GenericError (BC.pack (show e)) (queryId m)
-#else
-          return $ Left $ decodeError "TODO TOX GenericError" (queryId m)
-#endif
+          return $ Left $ KError GenericError (BC.pack (show e)) (envelopeTransaction m)
       ]
 
-dispatchHandler :: MonadKRPC h m => KQuery -> SockAddr -> m KResult
-dispatchHandler q addr = do
+dispatchHandler :: ( MonadKRPC h m raw msg
+                   , Eq (QueryMethod msg)
+                   , Show (QueryMethod msg)
+                   , Serialize (TransactionID msg)
+                   , Envelope msg
+                   ) => QueryMethod msg -> msg raw -> SockAddr -> m (KResult msg raw)
+dispatchHandler meth q addr = do
   Manager {..} <- getManager
-  case L.lookup (queryMethod q) handlers of
-#ifdef VERSION_bencoding
-    Nothing -> return $ Left $ KError MethodUnknown (queryMethod q) (queryId q)
-#else
-    Nothing -> return $ Left $ decodeError "TODO TOX Error MethodUnknown" (queryId q)
-#endif
-    Just h  -> runHandler h addr q
+  case L.lookup meth handlers of
+    Nothing -> return $ Left $ KError MethodUnknown ("Unknown method " <> BC.pack (show meth)) (envelopeTransaction q)
+    Just h  -> runHandler meth h addr q
 
 {-----------------------------------------------------------------------
 --  Listener
@@ -569,71 +527,75 @@ dispatchHandler q addr = do
 -- peer B fork too many threads
 -- ... space leak
 --
-handleQuery :: MonadKRPC h m => KQueryArgs -> KQuery -> SockAddr -> m ()
-handleQuery raw q addr = void $ fork $ do
+handleQuery :: ( MonadKRPC h m raw msg
+               , WireFormat raw msg
+               , Eq (QueryMethod msg)
+               , Show (QueryMethod msg)
+               , Serialize (TransactionID msg)
+               ) => QueryMethod msg -> raw -> msg raw -> SockAddr -> m ()
+handleQuery meth raw q addr = void $ fork $ do
   myThreadId >>= liftIO . flip labelThread "KRPC.handleQuery"
   Manager {..} <- getManager
-  res <- dispatchHandler q addr
+  res <- dispatchHandler meth q addr
 #ifdef VERSION_bencoding
-  let  res' = either E id res
-       resbe = either toBEncode toBEncode res
-  $(logOther "q") $ T.unlines
-    [ either (const "<unicode-fail>") id $ T.decodeUtf8' (BL.toStrict $ showBEncode raw)
-    , "==>"
-    , either (const "<unicode-fail>") id $ T.decodeUtf8' (BL.toStrict $ showBEncode resbe)
-    ]
-  sendMessage sock addr $ encodeHeaders () res'
+  let  res' = either buildError Just res
+       ctx = error "TODO TOX ToxCipherContext 2 or () for Mainline"
+       resbs = fmap (encodeHeaders ctx) res' :: Maybe BS.ByteString
+-- TODO: Generalize this debug print.
+--     resbe = either toBEncode toBEncode res
+-- $(logOther "q") $ T.unlines
+--   [ either (const "<unicode-fail>") id $ T.decodeUtf8' (BL.toStrict $ showBEncode raw)
+--   , "==>"
+--   , either (const "<unicode-fail>") id $ T.decodeUtf8' (BL.toStrict $ showBEncode resbe)
+--   ]
+  maybe (return ()) (sendMessage sock addr) resbs
 #else
   -- Errors not sent for Tox.
   let ctx = error "TODO TOX ToxCipherContext 2"
   either (const $ return ()) (sendMessage sock addr . encodeHeaders ctx) res
 #endif
 
-handleResponse :: MonadKRPC h m => KQueryArgs -> KResult -> SockAddr -> m ()
+handleResponse :: ( MonadKRPC h m raw msg
+                  , Ord (TransactionID msg)
+                  , Envelope msg
+                  ) => raw -> KResult msg raw -> SockAddr -> m ()
 handleResponse raw result addr = do
   Manager {..} <- getManager
   liftIO $ do
-#ifdef VERSION_bencoding
     let resultId = either errorId envelopeTransaction result
-#else
-    let resultId = either Tox.msgNonce Tox.msgNonce result
-#endif
     mcall <- unregisterQuery (resultId, addr) pendingCalls
     case mcall of
       Nothing   -> return ()
       Just ares -> putMVar ares (raw,result)
 
-#ifdef VERSION_bencoding
-handleMessage :: MonadKRPC h m => KQueryArgs -> KMessage -> SockAddr -> m ()
-handleMessage raw (Q q) = handleQuery    raw q
-handleMessage raw (R r) = handleResponse raw (Right (R r))
-handleMessage raw (E e) = handleResponse raw (Left  e)
-#else
-handleMessage :: MonadKRPC h m => KQueryArgs -> Tox.Message BC.ByteString -> SockAddr -> m ()
-handleMessage raw q | Tox.isQuery q    = handleQuery    raw q
-handleMessage raw r | Tox.isResponse r = handleResponse raw (Right r)
-handleMessage raw e | Tox.isError e    = handleResponse raw (Left  e)
-#endif
-
-listener :: MonadKRPC h m => m ()
-listener = do
+data Protocol raw (msg :: * -> *) = Protocol { rawProxy :: !(Proxy raw)
+                                             , msgProxy :: !(Proxy msg)
+                                             }
+
+listener :: forall h m raw msg.
+            ( MonadKRPC h m raw msg
+            , WireFormat raw msg
+            , Ord (TransactionID msg)
+            , Eq (QueryMethod msg)
+            , Show (QueryMethod msg)
+            , Serialize (TransactionID msg)
+            ) => Protocol raw msg -> m ()
+listener p = do
   Manager {..} <- getManager
   fix $ \again -> do
-    let ctx = error "TODO TOX ToxCipherContext 3"
+    let ctx = error "TODO TOX ToxCipherContext or () for Mainline"
     (bs, addr) <- liftIO $ do
       handle exceptions $ BS.recvFrom sock (optMaxMsgSize options)
-#ifdef VERSION_bencoding
-    case BE.parse bs >>= \r -> (,) r <$> BE.decode bs of
-#else
-    case return bs >>= \r -> (,) r <$> decodeHeaders ctx bs of
-#endif
-      -- TODO ignore unknown messages at all?
-#ifdef VERSION_bencoding
-      Left  e       -> liftIO $ sendMessage sock addr $ encodeHeaders () (E (unknownMessage e) :: KMessage)
-#else
-      Left  _       -> return () -- TODO TOX send unknownMessage error
-#endif
-      Right (raw,m) -> handleMessage raw m addr
+    case parsePacket (msgProxy p) bs >>= \r -> (,) r <$> decodeHeaders ctx r of
+      Left  e       -> -- XXX: Send parse failure message?
+                       --      liftIO $ sendMessage sock addr $ encodeHeaders ctx (unknownMessage e)
+                       return () -- Without transaction id, error message isn't very useful.
+      Right (raw,m) ->
+        case envelopeClass m of
+            Query meth  -> handleQuery meth (raw::raw) m         addr
+            Response _  -> handleResponse   raw        (Right m) addr
+            Error e     -> handleResponse   raw        (Left  e) addr
+
     again
   where
     exceptions :: IOError -> IO (BS.ByteString, SockAddr)
@@ -644,11 +606,17 @@ listener = do
 
 -- | Should be run before any 'query', otherwise they will never
 -- succeed.
-listen :: MonadKRPC h m => m ()
-listen = do
+listen :: ( MonadKRPC h m raw msg
+          , WireFormat raw msg
+          , Ord (TransactionID msg)
+          , Eq (QueryMethod msg)
+          , Show (QueryMethod msg)
+          , Serialize (TransactionID msg)
+          ) => Protocol raw msg -> m ()
+listen p = do
   Manager {..} <- getManager
   tid <- fork $ do
     myThreadId >>= liftIO . flip labelThread "KRPC.listen"
-    listener `Lifted.finally`
+    listener p `Lifted.finally`
       liftIO (takeMVar listenerThread)
   liftIO $ putMVar listenerThread tid