1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
|
{-# 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 req = liftIO $ do
readProgress <- newMVar 0
chunkBuffer <- newMVar (BodyChunkBuffer False [])
return $ BodyInfo readProgress chunkBuffer (getRequestBodyChunk 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 (BodyInfo _ chunkBufferVar getChunk) = liftIO $ do
cleanReadProgressVar <- newMVar 0
return $ BodyInfo cleanReadProgressVar chunkBufferVar 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 req bodyInfo opts = do
let shouldParseBody = isJust $ W.getRequestBodyType req
if shouldParseBody
then
do
bs <- getBodyAction bodyInfo opts
let wholeBody = BL.toChunks bs
(formparams, fs) <- parseRequestBody wholeBody W.lbsBackEnd req -- NB this loads the whole body into memory
let convert (k, v) = (strictByteStringToLazyText k, strictByteStringToLazyText v)
return (convert <$> formparams, fs)
else
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 readProgress chunkBufferVar getChunk) opts =
modifyMVar readProgress $ \index ->
modifyMVar chunkBufferVar $ \bcb@(BodyChunkBuffer hasFinished chunks) -> do
if | index > 0 -> E.throw BodyPartiallyStreamed
| hasFinished -> return (bcb, (index, BL.fromChunks chunks))
| otherwise -> do
newChunks <- readRequestBody getChunk return (maxRequestBodySize opts)
return $ (BodyChunkBuffer True (chunks ++ newChunks), (index, BL.fromChunks (chunks ++ 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 readProgress chunkBufferVar getChunk) =
modifyMVar readProgress $ \index ->
modifyMVar chunkBufferVar $ \bcb@(BodyChunkBuffer hasFinished chunks) -> do
if | index < length chunks -> return (bcb, (index + 1, chunks !! index))
| hasFinished -> return (bcb, (index, mempty))
| otherwise -> do
newChunk <- getChunk
return (BodyChunkBuffer (newChunk == mempty) (chunks ++ [newChunk]), (index + 1, 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 bl s r =
case W.getRequestBodyType r of
Nothing -> return ([], [])
Just rbt -> do
mvar <- liftIO $ newMVar bl -- MVar is a bit of a hack so we don't have to inline
-- large portions of Network.Wai.Parse
let provider = modifyMVar mvar $ \bsold -> case bsold of
[] -> return ([], B.empty)
(b:bs) -> return (bs, b)
liftIO $ W.sinkRequestBody s rbt provider
|
