3 files changed, 93 insertions, 4 deletions

diff --git a/krpc.cabal b/krpc.cabal
index 66c08ccb..452f1132 100644
--- a/krpc.cabal
+++ b/krpc.cabal
@@ -38,6 +38,9 @@ flag builder
   description: Use older bytestring package and bytestring-builder.
   default: False
 
+flag aeson
+  description: Use aeson for pretty-printing bencoded data.
+  default: True
 
 library
   default-language:    Haskell2010
@@ -48,6 +51,7 @@ library
                        Network.KRPC.Message
                        Network.KRPC.Method
                        Network.KRPC.Manager
+                       Data.BEncode.Pretty
   build-depends:       base          == 4.*
                      , text          >= 0.11
                      , data-default-class
@@ -60,6 +64,9 @@ library
                      , network       >= 2.3
                      , cereal
                      , containers
+  if flag(aeson)
+    build-depends: aeson, aeson-pretty, unordered-containers, vector
+    ghc-options:   -DBENCODE_AESON
   if flag(builder)
     build-depends: bytestring >= 0.9, bytestring-builder
   else
diff --git a/src/Data/BEncode/Pretty.hs b/src/Data/BEncode/Pretty.hs
new file mode 100644
index 00000000..7b0d46a0
--- /dev/null
+++ b/src/Data/BEncode/Pretty.hs
@@ -0,0 +1,75 @@
+{-# LANGUAGE CPP                #-}
+module Data.BEncode.Pretty where -- (showBEncode) where
+
+import Data.BEncode.Types
+import qualified Data.ByteString as BS
+import qualified Data.ByteString.Lazy as BL
+import qualified Data.ByteString.Lazy.Char8 as BL8
+import Data.Text (Text)
+import qualified Data.Text as T
+#ifdef BENCODE_AESON
+import Data.BEncode.BDict hiding (map)
+import Data.Aeson.Types hiding (parse)
+import Data.Aeson.Encode.Pretty
+import qualified Data.HashMap.Strict as HashMap
+import qualified Data.Vector as Vector
+import Data.Foldable as Foldable
+import Data.Text.Encoding
+import Text.Printf
+#endif
+
+#ifdef BENCODE_AESON
+
+unhex :: Text -> BS.ByteString
+unhex t = BS.pack $ map unhex1 [0 .. BS.length nibs `div` 2]
+ where
+    nibs = encodeUtf8 t
+    unhex1 i = unnib (BS.index nibs (i * 2)) * 0x10
+                + unnib (BS.index nibs (i * 2 + 1))
+    unnib a | a <= 0x39 = a - 0x30
+            | otherwise = a - (0x41 - 10)
+
+hex :: BS.ByteString -> Text
+hex bs = T.concat $ map (T.pack . printf "%02X") $ BS.unpack bs
+
+quote_chr :: Char
+quote_chr = ' '
+
+quote :: Text -> Text
+quote t = quote_chr `T.cons` t `T.snoc` quote_chr
+
+
+instance ToJSON BValue where
+    toJSON (BInteger x) = Number $ fromIntegral x
+    toJSON (BString s) = String $ either (const $ hex s) quote $ decodeUtf8' s
+    toJSON (BList xs) = Array $ Vector.fromList $ map toJSON xs
+    toJSON (BDict d) = toJSON d
+
+instance ToJSON a => ToJSON (BDictMap a) where
+    toJSON d = Object $ HashMap.fromList $ map convert $ toAscList d
+       where
+        convert (k,v) = (decodeUtf8 k,toJSON v)
+
+instance FromJSON BValue where
+    parseJSON (Number x) = pure $ BInteger (truncate x)
+    parseJSON (Bool x)   = pure $ BInteger $ if x then 1 else 0
+    parseJSON (String s)
+        | T.head s==quote_chr = pure $ BString $ encodeUtf8 (T.takeWhile (/=quote_chr) $ T.drop 1 s)
+        | otherwise           = pure $ BString $ unhex s
+    parseJSON (Array v)  = BList <$> traverse parseJSON (Foldable.toList v)
+    parseJSON (Object d) = BDict <$> parseJSON (Object d)
+    parseJSON (Null)     = pure $ BDict Nil
+
+instance FromJSON v => FromJSON (BDictMap v) where
+    parseJSON (Object d) = fromAscList <$> traverse convert (HashMap.toList d)
+        where
+            convert (k,v) = (,) (encodeUtf8 k) <$> parseJSON v
+    parseJSON _ = fail "Not a BDict"
+#endif
+
+showBEncode :: BValue -> BL.ByteString
+#ifdef BENCODE_AESON
+showBEncode b = encodePretty $ toJSON b
+#else
+showBEncode b = BL8.pack (show b)
+#endif
diff --git a/src/Network/KRPC/Manager.hs b/src/Network/KRPC/Manager.hs
index 9477d23c..c90c92f9 100644
--- a/src/Network/KRPC/Manager.hs
+++ b/src/Network/KRPC/Manager.hs
@@ -52,6 +52,7 @@ import Control.Monad.Reader
 import Control.Monad.Trans.Control
 import Data.BEncode as BE
 import Data.BEncode.Internal as BE
+import Data.BEncode.Pretty (showBEncode)
 import Data.ByteString as BS
 import Data.ByteString.Char8 as BC
 import Data.ByteString.Lazy as BL
@@ -428,11 +429,17 @@ dispatchHandler q @ KQuery {..} addr = do
 -- peer B fork too many threads
 -- ... space leak
 --
-handleQuery :: MonadKRPC h m => KQuery -> SockAddr -> m ()
-handleQuery q addr = void $ fork $ do
+handleQuery :: MonadKRPC h m => BValue -> KQuery -> SockAddr -> m ()
+handleQuery raw q addr = void $ fork $ do
   Manager {..} <- getManager
   res <- dispatchHandler q addr
-  sendMessage sock addr $ either toBEncode toBEncode res
+  let 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 resbe
 
 handleResponse :: MonadKRPC h m => BValue -> KResult -> SockAddr -> m ()
 handleResponse raw result addr = do
@@ -445,7 +452,7 @@ handleResponse raw result addr = do
       Just ares -> putMVar ares (raw,result)
 
 handleMessage :: MonadKRPC h m => BValue -> KMessage -> SockAddr -> m ()
-handleMessage _   (Q q) = handleQuery    q
+handleMessage raw (Q q) = handleQuery    raw q
 handleMessage raw (R r) = handleResponse raw (Right r)
 handleMessage raw (E e) = handleResponse raw (Left  e)