    {-# LANGUAGE FlexibleContexts, FlexibleInstances, RecordWildCards,
                 OverloadedStrings, MultiWayIf #-}
    module Web.Scotty.Body (
      newBodyInfo,
      cloneBodyInfo

      , getFormParamsAndFilesAction
      , getBodyAction
      , getBodyChunkAction
      ) where

    import           Control.Concurrent.MVar
    import           Control.Monad.IO.Class
    import qualified Data.ByteString as BS
    import qualified Data.ByteString.Char8 as B
    import qualified Data.ByteString.Lazy.Char8 as BL
    import           Data.Maybe
    import qualified GHC.Exception as E (throw)
    import           Network.Wai (Request(..), getRequestBodyChunk)
    import qualified Network.Wai.Parse as W (File, Param, getRequestBodyType, BackEnd, lbsBackEnd, sinkRequestBody)
    import           Web.Scotty.Action (Param)
    import           Web.Scotty.Internal.Types (BodyInfo(..), BodyChunkBuffer(..), BodyPartiallyStreamed(..), RouteOptions(..))
    import           Web.Scotty.Util (readRequestBody, strictByteStringToLazyText)

    -- | Make a new BodyInfo with readProgress at 0 and an empty BodyChunkBuffer.
    newBodyInfo :: (MonadIO m) => Request -> m BodyInfo
    newBodyInfo :: forall (m :: * -> *). MonadIO m => Request -> m BodyInfo
    newBodyInfo Request
    req = IO BodyInfo -> m BodyInfo
    forall a. IO a -> m a
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
    liftIO (IO BodyInfo -> m BodyInfo) -> IO BodyInfo -> m BodyInfo
    forall a b. (a -> b) -> a -> b
    $ do
      MVar Int
    readProgress <- Int -> IO (MVar Int)
    forall a. a -> IO (MVar a)
    newMVar Int
    0
      MVar BodyChunkBuffer
    chunkBuffer <- BodyChunkBuffer -> IO (MVar BodyChunkBuffer)
    forall a. a -> IO (MVar a)
    newMVar (Bool -> [ByteString] -> BodyChunkBuffer
    BodyChunkBuffer Bool
    False [])
      BodyInfo -> IO BodyInfo
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (BodyInfo -> IO BodyInfo) -> BodyInfo -> IO BodyInfo
    forall a b. (a -> b) -> a -> b
    $ MVar Int -> MVar BodyChunkBuffer -> IO ByteString -> BodyInfo
    BodyInfo MVar Int
    readProgress MVar BodyChunkBuffer
    chunkBuffer (Request -> IO ByteString
    getRequestBodyChunk Request
    req)

    -- | Make a copy of a BodyInfo, sharing the previous BodyChunkBuffer but with the
    -- readProgress MVar reset to 0.
    cloneBodyInfo :: (MonadIO m) => BodyInfo -> m BodyInfo
    cloneBodyInfo :: forall (m :: * -> *). MonadIO m => BodyInfo -> m BodyInfo
    cloneBodyInfo (BodyInfo MVar Int
    _ MVar BodyChunkBuffer
    chunkBufferVar IO ByteString
    getChunk) = IO BodyInfo -> m BodyInfo
    forall a. IO a -> m a
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
    liftIO (IO BodyInfo -> m BodyInfo) -> IO BodyInfo -> m BodyInfo
    forall a b. (a -> b) -> a -> b
    $ do
      MVar Int
    cleanReadProgressVar <- Int -> IO (MVar Int)
    forall a. a -> IO (MVar a)
    newMVar Int
    0
      BodyInfo -> IO BodyInfo
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (BodyInfo -> IO BodyInfo) -> BodyInfo -> IO BodyInfo
    forall a b. (a -> b) -> a -> b
    $ MVar Int -> MVar BodyChunkBuffer -> IO ByteString -> BodyInfo
    BodyInfo MVar Int
    cleanReadProgressVar MVar BodyChunkBuffer
    chunkBufferVar IO ByteString
    getChunk

    -- | Get the form params and files from the request. Requires reading the whole body.
    getFormParamsAndFilesAction :: Request -> BodyInfo -> RouteOptions -> IO ([Param], [W.File BL.ByteString])
    getFormParamsAndFilesAction :: Request
    -> BodyInfo -> RouteOptions -> IO ([Param], [File ByteString])
    getFormParamsAndFilesAction Request
    req BodyInfo
    bodyInfo RouteOptions
    opts = do
      let shouldParseBody :: Bool
    shouldParseBody = Maybe RequestBodyType -> Bool
    forall a. Maybe a -> Bool
    isJust (Maybe RequestBodyType -> Bool) -> Maybe RequestBodyType -> Bool
    forall a b. (a -> b) -> a -> b
    $ Request -> Maybe RequestBodyType
    W.getRequestBodyType Request
    req

      if Bool
    shouldParseBody
        then
        do
          ByteString
    bs <- BodyInfo -> RouteOptions -> IO ByteString
    getBodyAction BodyInfo
    bodyInfo RouteOptions
    opts
          let wholeBody :: [ByteString]
    wholeBody = ByteString -> [ByteString]
    BL.toChunks ByteString
    bs
          ([Param]
    formparams, [File ByteString]
    fs) <- [ByteString]
    -> BackEnd ByteString -> Request -> IO ([Param], [File ByteString])
    forall (m :: * -> *) y.
    MonadIO m =>
    [ByteString] -> BackEnd y -> Request -> m ([Param], [File y])
    parseRequestBody [ByteString]
    wholeBody BackEnd ByteString
    forall (m :: * -> *) ignored1 ignored2.
    Monad m =>
    ignored1 -> ignored2 -> m ByteString -> m ByteString
    W.lbsBackEnd Request
    req -- NB this loads the whole body into memory
          let convert :: Param -> Param
    convert (ByteString
    k, ByteString
    v) = (ByteString -> Text
    strictByteStringToLazyText ByteString
    k, ByteString -> Text
    strictByteStringToLazyText ByteString
    v)
          ([Param], [File ByteString]) -> IO ([Param], [File ByteString])
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Param -> Param
    convert (Param -> Param) -> [Param] -> [Param]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> [Param]
    formparams, [File ByteString]
    fs)
        else
        ([Param], [File ByteString]) -> IO ([Param], [File ByteString])
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ([], [])

    -- | Retrieve the entire body, using the cached chunks in the BodyInfo and reading any other
    -- chunks if they still exist.
    -- Mimic the previous behavior by throwing BodyPartiallyStreamed if the user has already
    -- started reading the body by chunks.
    getBodyAction :: BodyInfo -> RouteOptions -> IO (BL.ByteString)
    getBodyAction :: BodyInfo -> RouteOptions -> IO ByteString
    getBodyAction (BodyInfo MVar Int
    readProgress MVar BodyChunkBuffer
    chunkBufferVar IO ByteString
    getChunk) RouteOptions
    opts =
      MVar Int -> (Int -> IO (Int, ByteString)) -> IO ByteString
    forall a b. MVar a -> (a -> IO (a, b)) -> IO b
    modifyMVar MVar Int
    readProgress ((Int -> IO (Int, ByteString)) -> IO ByteString)
    -> (Int -> IO (Int, ByteString)) -> IO ByteString
    forall a b. (a -> b) -> a -> b
    $ \Int
    index ->
        MVar BodyChunkBuffer
    -> (BodyChunkBuffer -> IO (BodyChunkBuffer, (Int, ByteString)))
    -> IO (Int, ByteString)
    forall a b. MVar a -> (a -> IO (a, b)) -> IO b
    modifyMVar MVar BodyChunkBuffer
    chunkBufferVar ((BodyChunkBuffer -> IO (BodyChunkBuffer, (Int, ByteString)))
     -> IO (Int, ByteString))
    -> (BodyChunkBuffer -> IO (BodyChunkBuffer, (Int, ByteString)))
    -> IO (Int, ByteString)
    forall a b. (a -> b) -> a -> b
    $ \bcb :: BodyChunkBuffer
    bcb@(BodyChunkBuffer Bool
    hasFinished [ByteString]
    chunks) -> do
          if | Int
    index Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    > Int
    0 -> BodyPartiallyStreamed -> IO (BodyChunkBuffer, (Int, ByteString))
    forall a e. Exception e => e -> a
    E.throw BodyPartiallyStreamed
    BodyPartiallyStreamed
             | Bool
    hasFinished -> (BodyChunkBuffer, (Int, ByteString))
    -> IO (BodyChunkBuffer, (Int, ByteString))
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (BodyChunkBuffer
    bcb, (Int
    index, [ByteString] -> ByteString
    BL.fromChunks [ByteString]
    chunks))
             | Bool
    otherwise -> do
                 [ByteString]
    newChunks <- IO ByteString
    -> ([ByteString] -> IO [ByteString])
    -> Maybe Int
    -> IO [ByteString]
    readRequestBody IO ByteString
    getChunk [ByteString] -> IO [ByteString]
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (RouteOptions -> Maybe Int
    maxRequestBodySize RouteOptions
    opts)
                 (BodyChunkBuffer, (Int, ByteString))
    -> IO (BodyChunkBuffer, (Int, ByteString))
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ((BodyChunkBuffer, (Int, ByteString))
     -> IO (BodyChunkBuffer, (Int, ByteString)))
    -> (BodyChunkBuffer, (Int, ByteString))
    -> IO (BodyChunkBuffer, (Int, ByteString))
    forall a b. (a -> b) -> a -> b
    $ (Bool -> [ByteString] -> BodyChunkBuffer
    BodyChunkBuffer Bool
    True ([ByteString]
    chunks [ByteString] -> [ByteString] -> [ByteString]
    forall a. [a] -> [a] -> [a]
    ++ [ByteString]
    newChunks), (Int
    index, [ByteString] -> ByteString
    BL.fromChunks ([ByteString]
    chunks [ByteString] -> [ByteString] -> [ByteString]
    forall a. [a] -> [a] -> [a]
    ++ [ByteString]
    newChunks)))

    -- | Retrieve a chunk from the body at the index stored in the readProgress MVar.
    -- Serve the chunk from the cached array if it's already present; otherwise read another
    -- chunk from WAI and advance the index.
    getBodyChunkAction :: BodyInfo -> IO BS.ByteString
    getBodyChunkAction :: BodyInfo -> IO ByteString
    getBodyChunkAction (BodyInfo MVar Int
    readProgress MVar BodyChunkBuffer
    chunkBufferVar IO ByteString
    getChunk) =
      MVar Int -> (Int -> IO (Int, ByteString)) -> IO ByteString
    forall a b. MVar a -> (a -> IO (a, b)) -> IO b
    modifyMVar MVar Int
    readProgress ((Int -> IO (Int, ByteString)) -> IO ByteString)
    -> (Int -> IO (Int, ByteString)) -> IO ByteString
    forall a b. (a -> b) -> a -> b
    $ \Int
    index ->
        MVar BodyChunkBuffer
    -> (BodyChunkBuffer -> IO (BodyChunkBuffer, (Int, ByteString)))
    -> IO (Int, ByteString)
    forall a b. MVar a -> (a -> IO (a, b)) -> IO b
    modifyMVar MVar BodyChunkBuffer
    chunkBufferVar ((BodyChunkBuffer -> IO (BodyChunkBuffer, (Int, ByteString)))
     -> IO (Int, ByteString))
    -> (BodyChunkBuffer -> IO (BodyChunkBuffer, (Int, ByteString)))
    -> IO (Int, ByteString)
    forall a b. (a -> b) -> a -> b
    $ \bcb :: BodyChunkBuffer
    bcb@(BodyChunkBuffer Bool
    hasFinished [ByteString]
    chunks) -> do
          if | Int
    index Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < [ByteString] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [ByteString]
    chunks -> (BodyChunkBuffer, (Int, ByteString))
    -> IO (BodyChunkBuffer, (Int, ByteString))
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (BodyChunkBuffer
    bcb, (Int
    index Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    1, [ByteString]
    chunks [ByteString] -> Int -> ByteString
    forall a. HasCallStack => [a] -> Int -> a
    !! Int
    index))
             | Bool
    hasFinished -> (BodyChunkBuffer, (Int, ByteString))
    -> IO (BodyChunkBuffer, (Int, ByteString))
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (BodyChunkBuffer
    bcb, (Int
    index, ByteString
    forall a. Monoid a => a
    mempty))
             | Bool
    otherwise -> do
                 ByteString
    newChunk <- IO ByteString
    getChunk
                 (BodyChunkBuffer, (Int, ByteString))
    -> IO (BodyChunkBuffer, (Int, ByteString))
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Bool -> [ByteString] -> BodyChunkBuffer
    BodyChunkBuffer (ByteString
    newChunk ByteString -> ByteString -> Bool
    forall a. Eq a => a -> a -> Bool
    == ByteString
    forall a. Monoid a => a
    mempty) ([ByteString]
    chunks [ByteString] -> [ByteString] -> [ByteString]
    forall a. [a] -> [a] -> [a]
    ++ [ByteString
    newChunk]), (Int
    index Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    1, ByteString
    newChunk))

    -- Stolen from wai-extra's Network.Wai.Parse, modified to accept body as list of Bytestrings.
    -- Reason: WAI's requestBody is an IO action that returns the body as chunks. Once read,
    -- they can't be read again. We read them into a lazy Bytestring, so Scotty user can get
    -- the raw body, even if they also want to call wai-extra's parsing routines.
    parseRequestBody :: MonadIO m
                     => [B.ByteString]
                     -> W.BackEnd y
                     -> Request
                     -> m ([W.Param], [W.File y])
    parseRequestBody :: forall (m :: * -> *) y.
    MonadIO m =>
    [ByteString] -> BackEnd y -> Request -> m ([Param], [File y])
    parseRequestBody [ByteString]
    bl BackEnd y
    s Request
    r =
        case Request -> Maybe RequestBodyType
    W.getRequestBodyType Request
    r of
            Maybe RequestBodyType
    Nothing -> ([Param], [File y]) -> m ([Param], [File y])
    forall a. a -> m a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ([], [])
            Just RequestBodyType
    rbt -> do
                MVar [ByteString]
    mvar <- IO (MVar [ByteString]) -> m (MVar [ByteString])
    forall a. IO a -> m a
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
    liftIO (IO (MVar [ByteString]) -> m (MVar [ByteString]))
    -> IO (MVar [ByteString]) -> m (MVar [ByteString])
    forall a b. (a -> b) -> a -> b
    $ [ByteString] -> IO (MVar [ByteString])
    forall a. a -> IO (MVar a)
    newMVar [ByteString]
    bl -- MVar is a bit of a hack so we don't have to inline
                                            -- large portions of Network.Wai.Parse
                let provider :: IO ByteString
    provider = MVar [ByteString]
    -> ([ByteString] -> IO ([ByteString], ByteString)) -> IO ByteString
    forall a b. MVar a -> (a -> IO (a, b)) -> IO b
    modifyMVar MVar [ByteString]
    mvar (([ByteString] -> IO ([ByteString], ByteString)) -> IO ByteString)
    -> ([ByteString] -> IO ([ByteString], ByteString)) -> IO ByteString
    forall a b. (a -> b) -> a -> b
    $ \[ByteString]
    bsold -> case [ByteString]
    bsold of
                                                    []     -> ([ByteString], ByteString) -> IO ([ByteString], ByteString)
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ([], ByteString
    B.empty)
                                                    (ByteString
    b:[ByteString]
    bs) -> ([ByteString], ByteString) -> IO ([ByteString], ByteString)
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ([ByteString]
    bs, ByteString
    b)
                IO ([Param], [File y]) -> m ([Param], [File y])
    forall a. IO a -> m a
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
    liftIO (IO ([Param], [File y]) -> m ([Param], [File y]))
    -> IO ([Param], [File y]) -> m ([Param], [File y])
    forall a b. (a -> b) -> a -> b
    $ BackEnd y
    -> RequestBodyType -> IO ByteString -> IO ([Param], [File y])
    forall y.
    BackEnd y
    -> RequestBodyType -> IO ByteString -> IO ([Param], [File y])
    W.sinkRequestBody BackEnd y
    s RequestBodyType
    rbt IO ByteString
    provider