From 47c1b93dc96cfeb0992e518bdf2efc9d0ff4c241 Mon Sep 17 00:00:00 2001
From: Alexander Bondarenko <486682+dpwiz@users.noreply.github.com>
Date: Thu, 18 Apr 2024 17:22:18 +0300
Subject: [PATCH] fix encoding

---
 src/Simplex/Messaging/Encoding/String.hs |  8 ++++---
 src/Simplex/Messaging/Protocol.hs        | 29 +++++++++++++++++++-----
 2 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/src/Simplex/Messaging/Encoding/String.hs b/src/Simplex/Messaging/Encoding/String.hs
index fcefdc73d..e72c21c27 100644
--- a/src/Simplex/Messaging/Encoding/String.hs
+++ b/src/Simplex/Messaging/Encoding/String.hs
@@ -32,6 +32,7 @@ import qualified Data.List.NonEmpty as L
 import Data.Set (Set)
 import qualified Data.Set as S
 import Data.Text (Text)
+import qualified Data.Text as T
 import Data.Text.Encoding (decodeLatin1, encodeUtf8)
 import Data.Time.Clock (UTCTime)
 import Data.Time.Clock.System (SystemTime (..))
@@ -39,7 +40,7 @@ import Data.Time.Format.ISO8601
 import Data.Word (Word16, Word32)
 import Simplex.Messaging.Encoding
 import Simplex.Messaging.Parsers (parseAll)
-import Simplex.Messaging.Util ((<$?>))
+import Simplex.Messaging.Util (safeDecodeUtf8, (<$?>))
 
 class TextEncoding a where
   textEncode :: a -> Text
@@ -75,9 +76,10 @@ instance StrEncoding Str where
   strEncode = unStr
   strP = Str <$> A.takeTill (== ' ') <* optional A.space
 
+-- inherited from ByteString, the parser only allows non-empty strings
 instance StrEncoding String where
-  strEncode = strEncode . B.pack
-  strP = B.unpack <$> strP
+  strEncode = strEncode . encodeUtf8 . T.pack
+  strP = T.unpack . safeDecodeUtf8 <$> strP
 
 instance ToJSON Str where
   toJSON (Str s) = strToJSON s
diff --git a/src/Simplex/Messaging/Protocol.hs b/src/Simplex/Messaging/Protocol.hs
index c548fb053..fafd8a340 100644
--- a/src/Simplex/Messaging/Protocol.hs
+++ b/src/Simplex/Messaging/Protocol.hs
@@ -195,6 +195,8 @@ import Data.List.NonEmpty (NonEmpty (..))
 import qualified Data.List.NonEmpty as L
 import Data.Maybe (isJust, isNothing)
 import Data.String
+import qualified Data.Text as T
+import Data.Text.Encoding (encodeUtf8)
 import Data.Time.Clock.System (SystemTime (..))
 import Data.Type.Equality
 import Data.Word (Word16)
@@ -210,7 +212,7 @@ import Simplex.Messaging.Parsers
 import Simplex.Messaging.ServiceScheme
 import Simplex.Messaging.Transport
 import Simplex.Messaging.Transport.Client (TransportHost, TransportHosts (..))
-import Simplex.Messaging.Util (bshow, eitherToMaybe, (<$?>))
+import Simplex.Messaging.Util (bshow, eitherToMaybe, safeDecodeUtf8, (<$?>))
 import Simplex.Messaging.Version
 import Simplex.Messaging.Version.Internal
 
@@ -1167,11 +1169,11 @@ data ErrorType
 instance StrEncoding ErrorType where
   strEncode = \case
     CMD e -> "CMD " <> bshow e
-    PROXY e -> "PROXY " <> bshow e
+    PROXY e -> "PROXY " <> strEncode e
     e -> bshow e
   strP =
     "CMD " *> (CMD <$> parseRead1)
-      <|> "PROXY " *> (PROXY <$> parseRead1)
+      <|> "PROXY " *> (PROXY <$> strP)
       <|> parseRead1
 
 -- | SMP command error type.
@@ -1190,7 +1192,6 @@ data CommandError
     NO_ENTITY
   deriving (Eq, Read, Show)
 
--- TODO keep error params
 data ProxyError
   = -- | Correctly parsed SMP server ERR response.
     -- This error is forwarded to the agent client as `ERR SMP err`.
@@ -1475,7 +1476,7 @@ instance Encoding ProxyError where
   smpEncode e = case e of
     PROTOCOL et -> "PROTOCOL " <> smpEncode et
     RESPONSE et -> "RESPONSE " <> smpEncode et
-    UNEXPECTED s -> "UNEXPECTED " <> smpEncode (B.pack s)
+    UNEXPECTED s -> "UNEXPECTED " <> smpEncode (encodeUtf8 $ T.pack s)
     TIMEOUT -> "TIMEOUT"
     NETWORK -> "NETWORK"
     BAD_HOST -> "BAD_HOST"
@@ -1485,7 +1486,7 @@ instance Encoding ProxyError where
     A.takeTill (== ' ') >>= \case
       "PROTOCOL" -> PROTOCOL <$> _smpP
       "RESPONSE" -> RESPONSE <$> _smpP
-      "UNEXPECTED" -> UNEXPECTED . B.unpack <$> _smpP
+      "UNEXPECTED" -> UNEXPECTED . (T.unpack . safeDecodeUtf8) <$> _smpP
       "TIMEOUT" -> pure TIMEOUT
       "NETWORK" -> pure NETWORK
       "BAD_HOST" -> pure BAD_HOST
@@ -1493,6 +1494,22 @@ instance Encoding ProxyError where
       "TRANSPORT" -> TRANSPORT <$> (A.space *> transportErrorP)
       _ -> fail "bad command error type"
 
+instance StrEncoding ProxyError where
+  strEncode = \case
+    PROTOCOL et -> "PROTOCOL " <> strEncode et
+    RESPONSE et -> "RESPONSE " <> strEncode et
+    UNEXPECTED "" -> "UNEXPECTED" -- Arbitrary instance generates empty strings which String instance can't handle
+    UNEXPECTED s -> "UNEXPECTED " <> strEncode s
+    TRANSPORT t -> "TRANSPORT " <> serializeTransportError t
+    e -> bshow e
+  strP =
+    "PROTOCOL " *> (PROTOCOL <$> strP)
+      <|> "RESPONSE " *> (RESPONSE <$> strP)
+      <|> "UNEXPECTED " *> (UNEXPECTED <$> strP)
+      <|> "UNEXPECTED" $> UNEXPECTED ""
+      <|> "TRANSPORT " *> (TRANSPORT <$> transportErrorP)
+      <|> parseRead1
+
 -- | Send signed SMP transmission to TCP transport.
 tPut :: Transport c => THandle v c p -> NonEmpty (Either TransportError SentRawTransmission) -> IO [Either TransportError ()]
 tPut th@THandle {params} = fmap concat . mapM tPutBatch . batchTransmissions (batch params) (blockSize params)