Properly escape characters in JSON strings.

Author: pranaysashank

src/Streamly/Internal/Data/Json/Stream.hs

@@ -31,6 +31,7 @@ import Streamly.Internal.Data.Array (Array)
 import Streamly.Internal.Data.Fold.Types (Fold(..))
 import Streamly.Internal.Data.Strict (Tuple' (..))
 import qualified Streamly.Internal.Data.Parser as PR
+import qualified Streamly.Internal.Data.Parser.ParserK.Types as K
 import qualified Streamly.Internal.Data.Parser.ParserD as P
 import qualified Streamly.Internal.Data.Array as A
 import qualified Streamly.Internal.Data.Fold as IFL
@@ -67,6 +68,22 @@ instance Enum a => Hashable (A.Array a) where
     hashWithSalt salt arr = fromIntegral $ runIdentity $
         IUF.fold A.read (IFL.rollingHashWithSalt $ fromIntegral salt) arr
 
+{-# INLINE jsonEscapes #-}
+jsonEscapes :: Word8 -> Maybe Word8
+jsonEscapes 34  = Just 34 -- \" -> "
+jsonEscapes 92  = Just 92 -- \\ -> \
+jsonEscapes 47  = Just 47 -- \/ -> /
+jsonEscapes 98  = Just 8  -- \b -> BS
+jsonEscapes 102 = Just 12 -- \f -> FF
+jsonEscapes 110 = Just 10 -- \n -> LF
+jsonEscapes 114 = Just 13 -- \r -> CR
+jsonEscapes 116 = Just 9  -- \t -> TAB
+jsonEscapes _   = Nothing
+
+{-# INLINE escapeFoldUtf8With #-}
+escapeFoldUtf8With :: Monad m => Fold m Char container -> Fold m Word8 container
+escapeFoldUtf8With = Uni.escapeFoldUtf8With 92 jsonEscapes
+
 type JsonString = Array Char
 
 type JsonArray = Array Value
@@ -168,13 +185,13 @@ parseJsonString = do
     w <- P.peek
     case w of
         DOUBLE_QUOTE -> skip 1 >> return s
-        BACKSLASH -> (fmap (s <>) escapeParseJsonString) <* skip 1
+        BACKSLASH -> fmap (s <>) escapeParseJsonString
         _ -> do
             P.die $ [(chr . fromIntegral) w] ++ " : String without end."
 
 {-# INLINE escapeParseJsonString #-}
 escapeParseJsonString :: MonadCatch m => Parser m Word8 JsonString
-escapeParseJsonString = P.scan startState go (Uni.foldUtf8With A.unsafeWrite)
+escapeParseJsonString = P.scan startState go (escapeFoldUtf8With A.unsafeWrite)
   where
     startState = False
     go s a
@@ -236,15 +253,15 @@ parseJsonArray = do
 {-# INLINE parseJsonEOF #-}
 parseJsonEOF :: MonadCatch m => PR.Parser m Word8 Value
 parseJsonEOF =
-    P.toParserK $ do
+    K.toParserK $ do
         v <- parseJsonValue
         skipSpace
         P.eof
         return v
 
 {-# INLINE parseJson #-}
 parseJson :: MonadCatch m => PR.Parser m Word8 Value
-parseJson = P.toParserK $ parseJsonValue
+parseJson = K.toParserK $ parseJsonValue
 
 {-
 

src/Streamly/Internal/Data/Parser.hs

@@ -293,7 +293,7 @@ dieM = K.toParserK . D.dieM
 
 {-# INLINE peekMaybe #-}
 peekMaybe :: MonadCatch m => Parser m a (Maybe a)
-peekMaybe = D.toParserK D.peekMaybe
+peekMaybe = K.toParserK D.peekMaybe
 
 -------------------------------------------------------------------------------
 -- Failing Parsers
@@ -520,7 +520,7 @@ sliceSepByP _cond = undefined -- K.toParserK . D.sliceSepByP cond
 --
 {-# INLINE scan #-}
 scan :: MonadCatch m => s -> (s -> a -> Maybe s) -> Fold m a b -> Parser m a b
-scan s f fl = D.toParserK $ D.scan s f fl
+scan s f fl = K.toParserK $ D.scan s f fl
 
 -- | @sepBy fl p sep@ collects zero or more stream elements separated by @sep@.
 --
@@ -538,7 +538,7 @@ sepBy :: MonadCatch m
       -> Parser m a b
       -> Parser m a sep
       -> Parser m a c
-sepBy fl pa = D.toParserK . D.sepBy fl (D.fromParserK pa) . D.fromParserK
+sepBy fl pa = K.toParserK . D.sepBy fl (K.fromParserK pa) . K.fromParserK
 
 -- Note: Keep this consistent with S.splitOn. In fact we should eliminate
 -- S.splitOn in favor of the parser.

src/Streamly/Internal/Unicode/Stream.hs

@@ -45,6 +45,7 @@ module Streamly.Internal.Unicode.Stream
     , decodeUtf8ArraysD
     , decodeUtf8ArraysLenientD
     , foldUtf8With
+    , escapeFoldUtf8With
 
     -- * Transformation
     , stripStart
@@ -813,9 +814,13 @@ oDecode table state codep byte = do
 data Fp m s = FreshPoint !CodePoint !DecodeState (m s)
 
 {-# INLINE_NORMAL foldUtf8WithE #-}
-foldUtf8WithE :: Monad m => CodingFailureMode -> Fold m Char container -> Fold m Word8 container
+foldUtf8WithE
+    :: Monad m
+    => CodingFailureMode
+    -> Fold m Char container
+    -> Fold m Word8 container
 foldUtf8WithE cfm (FL.Fold arrayStep arrayInit arrayExtract) =
-    let A.Array p _ _ = utf8d
+    let A.Array p _ = utf8d
         !ptr = (unsafeForeignPtrToPtr p)
     in FL.Fold (step' ptr) (return $ FreshPoint 0 0 arrayInit) extract
   where
@@ -877,6 +882,111 @@ foldUtf8WithE cfm (FL.Fold arrayStep arrayInit arrayExtract) =
         aSt <- arrayState
         if statePtr /= 0 then inputUnderflow aSt else arrayExtract aSt
 
+data EscapeState m s
+    = EscapeState (m s)
+    | NotEscapeState !CodePoint !DecodeState (m s)
+
+{-# INLINE_NORMAL escapeFoldUtf8WithE #-}
+escapeFoldUtf8WithE
+    :: Monad m
+    => CodingFailureMode
+    -> Word8
+    -> (Word8 -> Maybe Word8)
+    -> Fold m Char container
+    -> Fold m Word8 container
+escapeFoldUtf8WithE cfm escape trans (FL.Fold arrayStep arrayInit arrayExtract) =
+    let A.Array p _ = utf8d
+        !ptr = (unsafeForeignPtrToPtr p)
+    in FL.Fold (step' ptr) (return $ NotEscapeState 0 0 arrayInit) extract
+  where
+    {-# INLINE transliterateOrError #-}
+    transliterateOrError e arrayState =
+        case cfm of
+            ErrorOnCodingFailure -> error e
+            TransliterateCodingFailure ->
+                return $ NotEscapeState 0 0 (arrayStep arrayState replacementChar)
+
+    {-# INLINE transliterateOrError1 #-}
+    transliterateOrError1 table statePtr codepointPtr e arrayState x =
+        case cfm of
+            ErrorOnCodingFailure -> error e
+            TransliterateCodingFailure -> do
+                aS <- arrayState
+                aSt <- arrayStep aS replacementChar
+                if x <= 0x7f
+                    then return $ NotEscapeState 0 0 (arrayStep aSt (unsafeChr (fromIntegral x)))
+                else
+                    let (Tuple' sv cp) = oDecode table statePtr codepointPtr x
+                    in
+                    case sv of
+                        12 -> transliterateOrError e aSt
+                        0 -> return $ NotEscapeState cp sv (arrayStep aSt (unsafeChr cp))
+                        _ -> return $ NotEscapeState cp sv (return aSt)
+
+    step' _ (NotEscapeState _ _ arrayState) c
+        | c == escape = return $ EscapeState arrayState
+
+    step' table (NotEscapeState codepointPtr statePtr arrayState) x =
+        if statePtr == 0 && x <= 0x7f
+        then do
+            aSt <- arrayState
+            return $ NotEscapeState 0 0 (arrayStep aSt (unsafeChr (fromIntegral x)))
+        else
+            let (Tuple' sv cp) = oDecode table statePtr codepointPtr x
+            in
+            case sv of
+                12 ->
+                    transliterateOrError1
+                        table
+                        statePtr
+                        codepointPtr
+                        "Streamly.Streams.StreamD.escapeFoldUtf8With: Invalid UTF8 codepoint encountered"
+                        arrayState
+                        x
+                0 -> do
+                    aSt <- arrayState
+                    return $ NotEscapeState cp sv (arrayStep aSt (unsafeChr cp))
+                _ -> return $ NotEscapeState cp sv arrayState
+
+    step' _ (EscapeState arrayState) c = do
+        aSt <- arrayState
+        case trans c of
+            Just x ->
+                return $
+                NotEscapeState 0 0 (arrayStep aSt (unsafeChr (fromIntegral x)))
+            Nothing ->
+                transliterateOrError
+                    "Streamly.Streams.StreamD.escapeFoldUtf8With: Invalid UTF8 codepoint encountered"
+                    aSt
+
+    {-# INLINE inputUnderflow #-}
+    inputUnderflow arrayState =
+        case cfm of
+            ErrorOnCodingFailure ->
+                error "Streamly.Internal.Data.Stream.StreamD.escapeFoldUtf8With: Input Underflow"
+            TransliterateCodingFailure -> do
+                aSt <- arrayStep arrayState replacementChar
+                arrayExtract aSt
+
+    extract (NotEscapeState _ statePtr arrayState) = do
+        aSt <- arrayState
+        if statePtr /= 0
+            then inputUnderflow aSt
+            else arrayExtract aSt
+
+    extract (EscapeState arrayState) = do
+        aSt <- arrayState
+        inputUnderflow aSt
+
 {-# INLINE foldUtf8With #-}
 foldUtf8With :: Monad m => Fold m Char container -> Fold m Word8 container
 foldUtf8With = foldUtf8WithE ErrorOnCodingFailure
+
+{-# INLINE escapeFoldUtf8With #-}
+escapeFoldUtf8With ::
+       Monad m
+    => Word8
+    -> (Word8 -> Maybe Word8)
+    -> Fold m Char container
+    -> Fold m Word8 container
+escapeFoldUtf8With = escapeFoldUtf8WithE ErrorOnCodingFailure