{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiWayIf #-}
{-# language OverloadedStrings #-}
{-# language ScopedTypeVariables #-}
module Web.Scotty.Body (
  newBodyInfo,
  cloneBodyInfo

  , getFormParamsAndFilesAction
  , getBodyAction
  , getBodyChunkAction
  -- wai-extra
  , W.RequestParseException(..)
  ) where

import           Control.Concurrent.MVar
import           Control.Monad.IO.Class
import Control.Monad.Trans.Resource (InternalState)
import Data.Bifunctor (first, bimap)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as B
import qualified Data.ByteString.Lazy.Char8 as BL
import qualified GHC.Exception as E (throw)
import           Network.Wai (Request(..), getRequestBodyChunk)
import qualified Network.Wai.Handler.Warp as Warp (InvalidRequest(..))
import qualified Network.Wai.Parse as W (File, Param, getRequestBodyType, tempFileBackEnd, RequestBodyType(..), sinkRequestBodyEx, RequestParseException(..), ParseRequestBodyOptions)
-- import UnliftIO (MonadUnliftIO(..))
import UnliftIO.Exception (Handler(..), catches, throwIO)

import           Web.Scotty.Internal.Types (BodyInfo(..), BodyChunkBuffer(..), BodyPartiallyStreamed(..), RouteOptions(..), File, ScottyException(..), Param)
import           Web.Scotty.Util (readRequestBody, decodeUtf8Lenient)


-- | 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.
--
-- NB : catches exceptions from 'warp' and 'wai-extra' and wraps them into 'ScottyException'
getFormParamsAndFilesAction ::
  InternalState
  -> W.ParseRequestBodyOptions
  -> Request -- ^ only used for its body type
  -> BodyInfo -- ^ the request body contents are read from here
  -> RouteOptions
  -> IO ([Param], [File FilePath])
getFormParamsAndFilesAction :: InternalState
-> ParseRequestBodyOptions
-> Request
-> BodyInfo
-> RouteOptions
-> IO ([Param], [File FilePath])
getFormParamsAndFilesAction InternalState
istate ParseRequestBodyOptions
prbo Request
req BodyInfo
bodyInfo RouteOptions
opts = do
  let
    bs2t :: ByteString -> Text
bs2t = ByteString -> Text
decodeUtf8Lenient
    convertBoth :: (ByteString, ByteString) -> Param
convertBoth = (ByteString -> Text)
-> (ByteString -> Text) -> (ByteString, ByteString) -> Param
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap ByteString -> Text
bs2t ByteString -> Text
bs2t
    convertKey :: (ByteString, c) -> (Text, c)
convertKey = (ByteString -> Text) -> (ByteString, c) -> (Text, c)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ByteString -> Text
bs2t
  ByteString
bs <- BodyInfo -> RouteOptions -> IO ByteString
getBodyAction BodyInfo
bodyInfo RouteOptions
opts
  let
    wholeBody :: [ByteString]
wholeBody = ByteString -> [ByteString]
BL.toChunks ByteString
bs
  ([(ByteString, ByteString)]
formparams, [File FilePath]
fs) <- InternalState
-> ParseRequestBodyOptions
-> [ByteString]
-> Maybe RequestBodyType
-> IO ([(ByteString, ByteString)], [File FilePath])
forall (m :: * -> *).
MonadIO m =>
InternalState
-> ParseRequestBodyOptions
-> [ByteString]
-> Maybe RequestBodyType
-> m ([(ByteString, ByteString)], [File FilePath])
parseRequestBodyExBS InternalState
istate ParseRequestBodyOptions
prbo [ByteString]
wholeBody (Request -> Maybe RequestBodyType
W.getRequestBodyType Request
req) IO ([(ByteString, ByteString)], [File FilePath])
-> [Handler IO ([(ByteString, ByteString)], [File FilePath])]
-> IO ([(ByteString, ByteString)], [File FilePath])
forall (m :: * -> *) a.
MonadUnliftIO m =>
m a -> [Handler m a] -> m a
`catches` [Handler IO ([(ByteString, ByteString)], [File FilePath])]
forall (m :: * -> *) a. MonadIO m => [Handler m a]
handleWaiParseSafeExceptions
  ([Param], [File FilePath]) -> IO ([Param], [File FilePath])
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ((ByteString, ByteString) -> Param
convertBoth ((ByteString, ByteString) -> Param)
-> [(ByteString, ByteString)] -> [Param]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(ByteString, ByteString)]
formparams, File FilePath -> File FilePath
forall {c}. (ByteString, c) -> (Text, c)
convertKey (File FilePath -> File FilePath)
-> [File FilePath] -> [File FilePath]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [File FilePath]
fs)

-- | Wrap exceptions from upstream libraries into 'ScottyException'
handleWaiParseSafeExceptions :: MonadIO m => [Handler m a]
handleWaiParseSafeExceptions :: forall (m :: * -> *) a. MonadIO m => [Handler m a]
handleWaiParseSafeExceptions = [Handler m a
forall {a}. Handler m a
h1, Handler m a
forall {a}. Handler m a
h2]
  where
    h1 :: Handler m a
h1 = (RequestParseException -> m a) -> Handler m a
forall (m :: * -> *) a e. Exception e => (e -> m a) -> Handler m a
Handler (\ (RequestParseException
e :: W.RequestParseException ) -> ScottyException -> m a
forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO (ScottyException -> m a) -> ScottyException -> m a
forall a b. (a -> b) -> a -> b
$ RequestParseException -> ScottyException
WaiRequestParseException RequestParseException
e)
    h2 :: Handler m a
h2 = (InvalidRequest -> m a) -> Handler m a
forall (m :: * -> *) a e. Exception e => (e -> m a) -> Handler m a
Handler (\(InvalidRequest
e :: Warp.InvalidRequest) -> ScottyException -> m a
forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO (ScottyException -> m a) -> ScottyException -> m a
forall a b. (a -> b) -> a -> b
$ InvalidRequest -> ScottyException
WarpRequestException InvalidRequest
e)

-- | Adapted 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.
parseRequestBodyExBS :: MonadIO m =>
                        InternalState
                     -> W.ParseRequestBodyOptions
                     -> [B.ByteString]
                     -> Maybe W.RequestBodyType
                     -> m ([W.Param], [W.File FilePath])
parseRequestBodyExBS :: forall (m :: * -> *).
MonadIO m =>
InternalState
-> ParseRequestBodyOptions
-> [ByteString]
-> Maybe RequestBodyType
-> m ([(ByteString, ByteString)], [File FilePath])
parseRequestBodyExBS InternalState
istate ParseRequestBodyOptions
o [ByteString]
bl Maybe RequestBodyType
rty =
    case Maybe RequestBodyType
rty of
        Maybe RequestBodyType
Nothing -> ([(ByteString, ByteString)], [File FilePath])
-> m ([(ByteString, ByteString)], [File FilePath])
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 ([(ByteString, ByteString)], [File FilePath])
-> m ([(ByteString, ByteString)], [File FilePath])
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ([(ByteString, ByteString)], [File FilePath])
 -> m ([(ByteString, ByteString)], [File FilePath]))
-> IO ([(ByteString, ByteString)], [File FilePath])
-> m ([(ByteString, ByteString)], [File FilePath])
forall a b. (a -> b) -> a -> b
$ ParseRequestBodyOptions
-> BackEnd FilePath
-> RequestBodyType
-> IO ByteString
-> IO ([(ByteString, ByteString)], [File FilePath])
forall y.
ParseRequestBodyOptions
-> BackEnd y
-> RequestBodyType
-> IO ByteString
-> IO ([(ByteString, ByteString)], [File y])
W.sinkRequestBodyEx ParseRequestBodyOptions
o (InternalState -> BackEnd FilePath
forall ignored1 ignored2.
InternalState
-> ignored1 -> ignored2 -> IO ByteString -> IO FilePath
W.tempFileBackEnd InternalState
istate) RequestBodyType
rbt IO ByteString
provider


-- | 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.
--
-- throw 'ScottyException' if request body too big
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 -> Bool
B.null ByteString
newChunk) ([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))