3 files changed, 57 insertions, 33 deletions

diff --git a/src/Network/KRPC.hs b/src/Network/KRPC.hs
index 96971803..69a4efca 100644
--- a/src/Network/KRPC.hs
+++ b/src/Network/KRPC.hs
@@ -56,13 +56,15 @@ module Network.KRPC
        , KRPC (..)
 
          -- * RPC
-       , Handler
-       , handler
-
          -- ** Query
        , QueryFailure (..)
        , query
 
+         -- ** Handler
+       , HandlerFailure (..)
+       , Handler
+       , handler
+
          -- * Manager
        , MonadKRPC (..)
        , Options (..)
@@ -73,11 +75,8 @@ module Network.KRPC
        , withManager
        , listen
 
-         -- * Exceptions
-       , KError (..)
+         -- * Re-expor
        , ErrorCode (..)
-
-         -- * Re-export
        , SockAddr (..)
        ) where
 
diff --git a/src/Network/KRPC/Manager.hs b/src/Network/KRPC/Manager.hs
index 6799277f..222b961a 100644
--- a/src/Network/KRPC/Manager.hs
+++ b/src/Network/KRPC/Manager.hs
@@ -31,6 +31,7 @@ module Network.KRPC.Manager
        , query
 
          -- * Handlers
+       , HandlerFailure (..)
        , Handler
        , handler
        ) where
@@ -39,7 +40,8 @@ import Control.Applicative
 import Control.Concurrent
 import Control.Concurrent.Lifted (fork)
 import Control.Exception hiding (Handler)
-import Control.Exception.Lifted as Lifted (catch, finally)
+import qualified Control.Exception.Lifted as E (Handler (..))
+import Control.Exception.Lifted as Lifted (catches, finally)
 import Control.Monad
 import Control.Monad.Logger
 import Control.Monad.Reader
@@ -288,9 +290,38 @@ query addr params = do
 {-----------------------------------------------------------------------
 --  Handlers
 -----------------------------------------------------------------------}
+-- we already throw:
+--
+--   * ErrorCode(MethodUnknown) in the 'dispatchHandler';
+--
+--   * ErrorCode(ServerError) in the 'runHandler'; (those can be
+--   async exception too)
+--
+--   * ErrorCode(GenericError) on
+
+-- | Used to signal protocol errors.
+data HandlerFailure
+  = BadAddress            -- ^ for e.g.: node calls herself;
+  | InvalidParameter Text -- ^ for e.g.: bad session token.
+    deriving (Show, Eq, Typeable)
+
+instance Exception HandlerFailure
+
+prettyHF :: HandlerFailure -> BS.ByteString
+prettyHF  BadAddress               = T.encodeUtf8 "bad address"
+prettyHF (InvalidParameter reason) = T.encodeUtf8 $
+  "invalid parameter: " <> reason
+
+prettyQF :: QueryFailure -> BS.ByteString
+prettyQF e = T.encodeUtf8 $ "handler fail while performing query: "
+  <> T.pack (show e)
 
 -- | Make handler from handler function. Any thrown exception will be
 -- supressed and send over the wire back to the querying node.
+--
+--   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)
@@ -305,7 +336,7 @@ handler body = (name, wrapper)
 
 runHandler :: MonadKRPC h m
            => HandlerBody h -> SockAddr -> KQuery -> m KResult
-runHandler h addr KQuery {..} = wrapper `Lifted.catch` failback
+runHandler h addr KQuery {..} = Lifted.catches wrapper failbacks
   where
     signature = querySignature queryMethod queryId addr
 
@@ -315,22 +346,33 @@ runHandler h addr KQuery {..} = wrapper `Lifted.catch` failback
 
       case result of
         Left msg -> do
-          $(logDebugS) "handler.failed" $ signature <> " !" <> T.pack msg
-          return $ Left  $ decodeError msg queryId
+          $(logDebugS) "handler.bad_query" $ signature <> " !" <> T.pack msg
+          return $ Left $ KError ProtocolError (BC.pack msg) queryId
 
         Right a -> do
           $(logDebugS) "handler.success" signature
-          return $ Right $ a `KResponse`   queryId
+          return $ Right $ KResponse a queryId
+
+    failbacks =
+      [ E.Handler $ \ (e :: HandlerFailure) -> do
+          $(logDebugS) "handler.failed" signature
+          return $ Left $ KError ProtocolError (prettyHF e) queryId
+
+        -- may happen if handler makes query and fail
+      , E.Handler $ \ (e :: QueryFailure) -> do
+          return $ Left $ KError ServerError (prettyQF e) queryId
 
-    failback e = do
-      $(logDebugS) "handler.errored" signature
-      return $ Left $ serverError e queryId
+        -- since handler thread exit after sendMessage we can safely
+        -- suppress async exception here
+      , E.Handler $ \ (e :: SomeException) -> do
+          return $ Left $ KError GenericError (BC.pack (show e)) queryId
+      ]
 
 dispatchHandler :: MonadKRPC h m => KQuery -> SockAddr -> m KResult
 dispatchHandler q @ KQuery {..} addr = do
   Manager {..} <- getManager
   case L.lookup queryMethod handlers of
-    Nothing -> return $ Left $ unknownMethod queryMethod queryId
+    Nothing -> return $ Left $ KError MethodUnknown queryMethod queryId
     Just h  -> runHandler h addr q
 
 {-----------------------------------------------------------------------
diff --git a/src/Network/KRPC/Message.hs b/src/Network/KRPC/Message.hs
index d6279f11..96945843 100644
--- a/src/Network/KRPC/Message.hs
+++ b/src/Network/KRPC/Message.hs
@@ -26,11 +26,8 @@ module Network.KRPC.Message
          -- * Error
        , ErrorCode (..)
        , KError(..)
-       , serverError
        , decodeError
-       , unknownMethod
        , unknownMessage
-       , timeoutExpired
 
          -- * Query
        , KQuery(..)
@@ -143,29 +140,15 @@ instance BEncode KError where
 
 instance Exception KError
 
--- | Happen when some query handler fail.
-serverError :: SomeException -> TransactionId -> KError
-serverError e = KError ServerError (BC.pack (show e))
-
 -- | Received 'queryArgs' or 'respVals' can not be decoded.
 decodeError :: String -> TransactionId -> KError
 decodeError msg = KError ProtocolError (BC.pack msg)
 
--- | If /remote/ node send query /this/ node doesn't know about then
--- this error message should be sent in response.
-unknownMethod :: MethodName -> TransactionId -> KError
-unknownMethod = KError MethodUnknown
-
 -- | A remote node has send some 'KMessage' this node is unable to
 -- decode.
 unknownMessage :: String -> KError
 unknownMessage msg = KError ProtocolError (BC.pack msg) unknownTransaction
 
--- | A /remote/ node is not responding to the /our/ request the for
--- specified period of time.
-timeoutExpired :: TransactionId -> KError
-timeoutExpired = KError GenericError "timeout expired"
-
 {-----------------------------------------------------------------------
 -- Query messages
 -----------------------------------------------------------------------}