1 files changed, 73 insertions, 45 deletions

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