String.hs

{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

{-# OPTIONS_GHC -fprof-auto #-}
{-# LANGUAGE ApplicativeDo #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NondecreasingIndentation #-}
module Thrift.Protocol.JSON.String
  ( parseJSONString
  ) where

import Control.Exception
import Control.Monad
import Thrift.Binary.Parser
import Data.Bits
import qualified Data.ByteString as ByteString
import Data.ByteString.Internal (ByteString(..))
import Data.ByteString.Unsafe
import Data.Char
import Data.Text (Text)
import Data.Text.Encoding
import Data.Word
import Data.Word8
import Foreign.ForeignPtr
import Foreign.Ptr
import Foreign.Storable
import System.IO.Unsafe

import Util.FFI

parseJSONString :: Parser Text
parseJSONString = do
  skipSpaces *> word8 _quotedbl
  rawBytes <- peekBS
  case unescape rawBytes of
    Left (StringError err) -> fail err
    Right (UnescapeResult bytes inLen) -> do
      skipN (inLen + 1) -- Skip over the string and closing "
      case decodeUtf8' bytes of
        Left err -> fail $ show err
        Right txt -> return txt

-- NOTE: The cpp2 implementation only accepts escaped "unicode" characters of
-- the form "\u00XX" where the X's are hex digits, so we will do the same here.
-- This means that all \u characters fit into a single byte which simplifies the
-- unescaping logic signifigantly

-- Datatype where all fields are strict so that we can avoid calling `evaluate`
data UnescapeResult = UnescapeResult !ByteString {-# UNPACK #-} !Int

-- Low-level unescaping logic. Works on the underlying byte array
unescape :: ByteString -> Either StringError UnescapeResult
unescape rawBytes = unsafeDupablePerformIO $ try $
  unsafeUseAsCStringLen rawBytes $ \(charBuf, fullLen) -> do
  let input = castPtr charBuf :: Ptr Word8

  -- Find the lengths of the input and output
  (inLen, outLen) <- unescapedLength input fullLen

  -- If the input has the same length as the output, then there are no escape
  -- characters, so we can just return the ByteString as is
  if inLen == outLen
    then return $! UnescapeResult (ByteString.take outLen rawBytes) inLen
    else do

  -- Allocate a buffer for the output
  outBuf <- mallocForeignPtrArray outLen
  unsafeWithForeignPtr outBuf $ \output -> do
    let
      -- Unescape the input character by character.
      -- i is the input index and j is the output index
      go !i !j
        | j >= outLen = return ()
        | otherwise = do
            byte <- peekByteOff input i
            if
              -- Case 1: we have reached an escape character
              -- NOTE: no bounds checks are needed because we already checked
              -- bounds in unescapeLength
              | byte == _backslash -> do
                  escChr <- peekByteOff input $ i + 1
                  if
                    -- Case 1.1: We have reached a 4-hex-digit unicode point
                    | escChr == _u -> do
                        c <- parseUnicodeChar =<<
                             peekByteOff (castPtr input) (i + 2)
                        pokeByteOff output j c
                        go (i + 6) (j + 1)
                    -- Case 1.2: Byte is a valid escape char
                    | Just unesc <- unescapeChar escChr -> do
                        pokeByteOff output j unesc
                        go (i + 2) (j + 1)
                    -- Case 1.3: Invalid escape character
                    | otherwise ->
                        throwIO $ StringError $
                        "parseJSONString: Not a valid escape character: " ++
                        [chr $ fromIntegral byte]
              -- Case 2: Not an escape character
              | otherwise -> do
                  pokeByteOff output j byte
                  go (i + 1) (j + 1)
    go 0 0
  return $! UnescapeResult (PS outBuf 0 outLen) inLen

{-# INLINE unescapedLength #-}
-- Get the escaped and unescaped lengths of the input
unescapedLength :: Ptr Word8 -> Int -> IO (Int, Int)
unescapedLength buf len = go 0 0
  where
    -- In this function i is the index in the input and j is the output length
    go !i !j
      | i >= len = stringError
      | otherwise = do
          byte <- peekByteOff buf i
          if | byte == _quotedbl  -> return (i, j)
             | byte == _backslash ->
                 if i + 1 < len
                 then do
                   nextByte <- peekByteOff buf (i + 1)
                   if | nextByte == _u -> go (i + 6) (j + 1) -- skip over \uXXXX
                      | otherwise      -> go (i + 2) (j + 1) -- skip over \X
                 else stringError
             | otherwise -> go (i + 1) (j + 1)
    stringError =
      throwIO $ StringError "parseJSONString: string does not terminate"

{-# INLINE unescapeChar #-}
unescapeChar :: Word8 -> Maybe Word8
unescapeChar = \case
  0x22 {- " -} -> Just _quotedbl
  0x27 {- ' -} -> Just _quotesingle
  0x2F {- / -} -> Just _slash
  0x5C {- \ -} -> Just _backslash
  0x30 {- 0 -} -> Just _0
  0x61 {- a -} -> Just $ fromIntegral $ ord '\a'
  0x62 {- b -} -> Just $ fromIntegral $ ord '\b'
  0x66 {- f -} -> Just $ fromIntegral $ ord '\f'
  0x6E {- n -} -> Just _lf -- linefeed
  0x72 {- r -} -> Just _cr -- carriage return
  0x74 {- t -} -> Just _tab
  0x76 {- v -} -> Just _vt -- vertical tab
  _ -> Nothing

{-# INLINE parseUnicodeChar #-}
parseUnicodeChar :: Word32 -> IO Word8
parseUnicodeChar !word = do
  -- Must have the form '\u00XX' (see note above).
  -- 0x30 is the acsii character '0'
  unless (word .&. 0x0000FFFF == 0x3030) unicodeError
  d0 <- hexdigit $ fromIntegral $ (word `shiftR` 16) .&. 0x000000FF
  d1 <- hexdigit $ fromIntegral $ (word `shiftR` 24) .&. 0x000000FF
  return $ d0 `shiftL` 4 .|. d1
  where
    unicodeError = throwIO $ StringError "Improperly formatted unicode string"
    hexdigit :: Word8 -> IO Word8
    hexdigit n
      | n `inRange` (_0, _9) = pure $ n - _0
      | n `inRange` (_a, _f) = pure $ n - _a + 10
      | n `inRange` (_A, _F) = pure $ n - _A + 10
      | otherwise = unicodeError

inRange :: Word8 -> (Word8, Word8) -> Bool
inRange n (x, y) = n >= x && n <= y

newtype StringError = StringError String
  deriving Show

instance Exception StringError