2 files changed, 87 insertions, 56 deletions

diff --git a/src/Network/KRPC.hs b/src/Network/KRPC.hs
index 2c3a1b48..a96d8da9 100644
--- a/src/Network/KRPC.hs
+++ b/src/Network/KRPC.hs
@@ -190,7 +190,7 @@ recvResponse sock = do
     Right resp -> Right resp
     Left decE -> Left $ case decode raw of
       Right kerror -> kerror
-      _ -> ProtocolError (BC.pack decE)
+      _ -> KError ProtocolError (BC.pack decE) undefined
 
 withRemote :: (MonadBaseControl IO m, MonadIO m) => (Socket -> m a) -> m a
 withRemote = bracket (liftIO (socket AF_INET6 Datagram defaultProtocol))
@@ -199,8 +199,10 @@ withRemote = bracket (liftIO (socket AF_INET6 Datagram defaultProtocol))
 
 getResult :: BEncode result => Socket -> IO result
 getResult sock = do
-  resp <- either throw (return . respVals) =<< recvResponse sock
-  either (throw . ProtocolError . BC.pack) return $ fromBEncode resp
+  KResponse {..} <- either throw return =<< recvResponse sock
+  case fromBEncode respVals of
+    Left msg -> throw $ KError ProtocolError (BC.pack msg) respId
+    Right r  -> return r
 
 -- | Makes remote procedure call. Throws RPCException on any error
 -- occurred.
@@ -233,10 +235,10 @@ handler body = (name, newbody)
     {-# INLINE newbody #-}
     newbody addr KQuery {..} =
       case fromBEncode queryArgs of
-        Left  e -> return (Left (ProtocolError (BC.pack e)))
+        Left  e -> return $ Left $ KError ProtocolError (BC.pack e) queryId
         Right a -> do
           r <- body addr a
-          return (Right (KResponse (toBEncode r) queryId))
+          return $ Right $ KResponse (toBEncode r) queryId
 
 sockAddrFamily :: SockAddr -> Family
 sockAddrFamily (SockAddrInet  _ _    ) = AF_INET
@@ -265,8 +267,9 @@ remoteServer servAddr action = bracket (liftIO bindServ) (liftIO . sClose) loop
       where
         handleMsg bs addr = case decode bs of
           Right query -> (either toBEncode toBEncode <$> action addr query)
-                        `Lifted.catch` (return . toBEncode . serverError)
-          Left decodeE   -> return $ toBEncode (ProtocolError (BC.pack decodeE))
+              `Lifted.catch` (return . toBEncode . (`serverError` undefined ))
+          Left decodeE   -> return $ toBEncode $
+             KError ProtocolError (BC.pack decodeE) undefined
 
 -- | Run RPC server on specified port by using list of handlers.
 --   Server will dispatch procedure specified by callee, but note that
@@ -277,7 +280,7 @@ server :: (MonadBaseControl IO remote, MonadIO remote)
        -> [MethodHandler remote]  -- ^ Method table.
        -> remote ()
 server servAddr handlers = do
-    remoteServer servAddr $ \addr q -> do
-      case L.lookup  (queryMethod q) handlers of
-        Nothing -> return $ Left $ MethodUnknown (queryMethod q)
-        Just  m -> m addr q
+  remoteServer servAddr $ \addr q @ KQuery {..} -> do
+    case L.lookup  queryMethod handlers of
+      Nothing -> return $ Left $ KError MethodUnknown queryMethod queryId
+      Just  m -> m addr q
diff --git a/src/Network/KRPC/Message.hs b/src/Network/KRPC/Message.hs
index 1a004c64..a70c2ea9 100644
--- a/src/Network/KRPC/Message.hs
+++ b/src/Network/KRPC/Message.hs
@@ -20,7 +20,8 @@
 {-# LANGUAGE DeriveDataTypeable     #-}
 module Network.KRPC.Message
        ( -- * Error
-         KError(..)
+         ErrorCode (..)
+       , KError(..)
        , serverError
 
          -- * Query
@@ -31,13 +32,60 @@ module Network.KRPC.Message
        , KResponse(..)
        ) where
 
+import Control.Applicative
 import Control.Exception.Lifted as Lifted
 import Data.BEncode as BE
-import Data.BEncode.BDict as BE
 import Data.ByteString as B
 import Data.ByteString.Char8 as BC
 import Data.Typeable
 
+
+-- | This transaction ID is generated by the querying node and is
+-- echoed in the response, so responses may be correlated with
+-- multiple queries to the same node. The transaction ID should be
+-- encoded as a short string of binary numbers, typically 2 characters
+-- are enough as they cover 2^16 outstanding queries.
+type TransactionId = ByteString
+
+{-----------------------------------------------------------------------
+-- Error messages
+-----------------------------------------------------------------------}
+
+data ErrorCode
+    -- | Some error doesn't fit in any other category.
+  = GenericError
+
+    -- | Occur when server fail to process procedure call.
+  | ServerError
+
+    -- | Malformed packet, invalid arguments or bad token.
+  | ProtocolError
+
+    -- | Occur when client trying to call method server don't know.
+  | MethodUnknown
+    deriving (Show, Read, Eq, Ord, Bounded, Typeable)
+
+instance Enum ErrorCode where
+  fromEnum GenericError  = 201
+  fromEnum ServerError   = 202
+  fromEnum ProtocolError = 203
+  fromEnum MethodUnknown = 204
+  {-# INLINE fromEnum #-}
+
+  toEnum 201 = GenericError
+  toEnum 202 = ServerError
+  toEnum 203 = ProtocolError
+  toEnum 204 = MethodUnknown
+  toEnum _   = GenericError
+  {-# INLINE toEnum #-}
+
+instance BEncode ErrorCode where
+  toBEncode = toBEncode . fromEnum
+  {-# INLINE toBEncode #-}
+
+  fromBEncode b = toEnum <$> fromBEncode b
+  {-# INLINE fromBEncode #-}
+
 -- | Errors used to signal that some error occurred while processing a
 -- procedure call. Error may be send only from server to client but
 -- not in the opposite direction.
@@ -46,62 +94,38 @@ import Data.Typeable
 --
 --   > { "y" : "e", "e" : [<error_code>, <human_readable_error_reason>] }
 --
-data KError
-    -- | Some error doesn't fit in any other category.
-  = GenericError { errorMessage :: !ByteString }
-
-    -- | Occur when server fail to process procedure call.
-  | ServerError  { errorMessage :: !ByteString }
-
-    -- | Malformed packet, invalid arguments or bad token.
-  | ProtocolError { errorMessage :: !ByteString }
-
-    -- | Occur when client trying to call method server don't know.
-  | MethodUnknown { errorMessage :: !ByteString }
-   deriving (Show, Read, Eq, Ord, Typeable)
+data KError = KError
+  { errorCode    :: !ErrorCode
+  , errorMessage :: !ByteString
+  , errorId      :: !TransactionId
+  } deriving (Show, Read, Eq, Ord, Typeable)
 
 instance BEncode KError where
-  {-# SPECIALIZE instance BEncode KError #-}
-  {-# INLINE toBEncode #-}
-  toBEncode e = toDict $
-       "e" .=! (errorCode e, errorMessage e)
+
+  toBEncode KError {..} = toDict $
+       "e" .=! (errorCode, errorMessage)
+    .: "t" .=! errorId
     .: "y" .=! ("e" :: ByteString)
     .: endDict
+  {-# INLINE toBEncode #-}
 
+  fromBEncode = fromDict $ do
+    lookAhead $ match "y" (BString "e")
+    (code, msg) <- field (req "e")
+    KError code msg <$>! "t"
   {-# INLINE fromBEncode #-}
-  fromBEncode be @ (BDict d)
-    | BE.lookup "y" d == Just (BString "e")
-    = (`fromDict` be) $ do
-      uncurry mkKError <$>! "e"
-
-  fromBEncode _ = decodingError "KError"
 
 instance Exception KError
 
-type ErrorCode = Int
-
-errorCode :: KError -> ErrorCode
-errorCode (GenericError _)  = 201
-errorCode (ServerError _)   = 202
-errorCode (ProtocolError _) = 203
-errorCode (MethodUnknown _) = 204
-{-# INLINE errorCode #-}
+serverError :: SomeException -> TransactionId -> KError
+serverError e = KError ServerError (BC.pack (show e))
 
-mkKError :: ErrorCode -> ByteString -> KError
-mkKError 201 = GenericError
-mkKError 202 = ServerError
-mkKError 203 = ProtocolError
-mkKError 204 = MethodUnknown
-mkKError _   = GenericError
-{-# INLINE mkKError #-}
-
-serverError :: SomeException -> KError
-serverError = ServerError . BC.pack . show
+{-----------------------------------------------------------------------
+-- Query messages
+-----------------------------------------------------------------------}
 
 type MethodName = ByteString
 
-type TransactionId = ByteString
-
 -- | Query used to signal that caller want to make procedure call to
 -- callee and pass arguments in. Therefore query may be only sent from
 -- client to server but not in the opposite direction.
@@ -130,6 +154,10 @@ instance BEncode KQuery where
     KQuery <$>! "a" <*>! "q" <*>! "t"
   {-# INLINE fromBEncode #-}
 
+{-----------------------------------------------------------------------
+-- Response messages
+-----------------------------------------------------------------------}
+
 -- | KResponse used to signal that callee successufully process a
 -- procedure call and to return values from procedure. KResponse should
 -- not be sent if error occurred during RPC. Thus KResponse may be only