Request body is preserved across 'next' calls #147 (#308)

* add test case

* add Web.Scotty.Body module

* test case fixed but other test fails :D

* move request max size test case under ScottyT and cleanup imports

* add counterexample test to setMaxRequestBodySize

---------

Co-authored-by: Marco Zocca <[email protected]>

Web/Scotty/Action.hs

@@ -47,7 +47,7 @@ import           Blaze.ByteString.Builder   (fromLazyByteString)
 
 import qualified Control.Exception          as E
 import           Control.Monad              (liftM, when)
-import           Control.Monad.Error.Class
+import           Control.Monad.Error.Class (throwError, catchError)
 import           Control.Monad.IO.Class     (MonadIO(..))
 import           Control.Monad.Reader       (MonadReader(..), ReaderT(..))
 import qualified Control.Monad.State.Strict as MS

Web/Scotty/Body.hs

@@ -0,0 +1,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           GHC.Exception
+import           Network.Wai (Request(..), getRequestBodyChunk)
+import qualified Network.Wai.Parse as W (File, Param, getRequestBodyType, BackEnd, lbsBackEnd, sinkRequestBody)
+import           Web.Scotty.Action
+import           Web.Scotty.Internal.Types (BodyInfo(..), BodyChunkBuffer(..), BodyPartiallyStreamed(..), RouteOptions(..))
+import           Web.Scotty.Util
+
+-- | 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 -> 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

Web/Scotty/Internal/Types.hs

@@ -12,6 +12,7 @@ module Web.Scotty.Internal.Types where
 import           Blaze.ByteString.Builder (Builder)
 
 import           Control.Applicative
+import Control.Concurrent.MVar
 import           Control.Exception (Exception)
 import qualified Control.Exception as E
 import qualified Control.Monad as Monad
@@ -28,7 +29,7 @@ import           Control.Monad.Trans.Control (MonadBaseControl, StM, liftBaseWit
 import           Control.Monad.Trans.Except
 
 import qualified Data.ByteString as BS
-import           Data.ByteString.Lazy.Char8 (ByteString)
+import qualified Data.ByteString.Lazy.Char8 as LBS8 (ByteString)
 import           Data.Default.Class (Default, def)
 import           Data.String (IsString(..))
 import           Data.Text.Lazy (Text, pack)
@@ -69,10 +70,21 @@ type Kilobytes = Int
 type Middleware m = Application m -> Application m
 type Application m = Request -> m Response
 
+------------------ Scotty Request Body --------------------
+
+data BodyChunkBuffer = BodyChunkBuffer { hasFinishedReadingChunks :: Bool
+                                       , chunksReadSoFar :: [BS.ByteString] }
+
+data BodyInfo = BodyInfo { bodyInfoReadProgress :: MVar Int
+                         , bodyInfoChunkBuffer :: MVar BodyChunkBuffer
+                         , bodyInfoDirectChunkRead :: IO BS.ByteString
+                         }
+
 --------------- Scotty Applications -----------------
+
 data ScottyState e m =
     ScottyState { middlewares :: [Wai.Middleware]
-                , routes :: [Middleware m]
+                , routes :: [BodyInfo -> Middleware m]
                 , handler :: ErrorHandler e m
                 , routeOptions :: RouteOptions
                 }
@@ -83,7 +95,7 @@ instance Default (ScottyState e m) where
 addMiddleware :: Wai.Middleware -> ScottyState e m -> ScottyState e m
 addMiddleware m s@(ScottyState {middlewares = ms}) = s { middlewares = m:ms }
 
-addRoute :: Middleware m -> ScottyState e m -> ScottyState e m
+addRoute :: (BodyInfo -> Middleware m) -> ScottyState e m -> ScottyState e m
 addRoute r s@(ScottyState {routes = rs}) = s { routes = r:rs }
 
 addHandler :: ErrorHandler e m -> ScottyState e m -> ScottyState e m
@@ -131,19 +143,19 @@ instance Exception ScottyException
 ------------------ Scotty Actions -------------------
 type Param = (Text, Text)
 
-type File = (Text, FileInfo ByteString)
+type File = (Text, FileInfo LBS8.ByteString)
 
 data ActionEnv = Env { getReq       :: Request
                      , getCaptureParams :: [Param]
                      , getFormParams    :: [Param]
                      , getQueryParams :: [Param]
-                     , getBody      :: IO ByteString
+                     , getBody      :: IO LBS8.ByteString
                      , getBodyChunk :: IO BS.ByteString
                      , getFiles     :: [File]
                      }
 
 data RequestBodyState = BodyUntouched
-                      | BodyCached ByteString [BS.ByteString] -- whole body, chunks left to stream
+                      | BodyCached LBS8.ByteString [BS.ByteString] -- whole body, chunks left to stream
                       | BodyCorrupted
 
 data BodyPartiallyStreamed = BodyPartiallyStreamed deriving (Show, Typeable)
@@ -283,3 +295,6 @@ data RoutePattern = Capture   Text
 
 instance IsString RoutePattern where
     fromString = Capture . pack
+
+
+

Web/Scotty/Route.hs

@@ -1,40 +1,33 @@
-{-# LANGUAGE CPP, FlexibleContexts, FlexibleInstances,
+{-# LANGUAGE FlexibleContexts, FlexibleInstances,
              OverloadedStrings, RankNTypes, ScopedTypeVariables #-}
 module Web.Scotty.Route
     ( get, post, put, delete, patch, options, addroute, matchAny, notFound,
       capture, regex, function, literal
     ) where
-
-import           Blaze.ByteString.Builder (fromByteString)
+
 import           Control.Arrow ((***))
-import           Control.Concurrent.MVar
-import           Control.Exception (throw, catch)
 import           Control.Monad.IO.Class
 import qualified Control.Monad.State as MS
 
 import qualified Data.ByteString.Char8 as B
-import qualified Data.ByteString.Lazy.Char8 as BL
 
-import           Data.Maybe (fromMaybe, isJust)
+import           Data.Maybe (fromMaybe)
 import           Data.String (fromString)
 import qualified Data.Text.Lazy as T
 import qualified Data.Text as TS
 
 import           Network.HTTP.Types
-import           Network.Wai (Request(..), Response, responseBuilder)
-#if MIN_VERSION_wai(3,2,2)
-import           Network.Wai.Internal (getRequestBodyChunk)
-#endif
-import qualified Network.Wai.Parse as Parse hiding (parseRequestBody)
+import           Network.Wai (Request(..))
 
 import           Prelude ()
 import           Prelude.Compat
 
 import qualified Text.Regex as Regex
 
 import           Web.Scotty.Action
-import           Web.Scotty.Internal.Types
-import           Web.Scotty.Util
+import           Web.Scotty.Internal.Types (RoutePattern(..), RouteOptions, ActionEnv(..), ActionT, ScottyState(..), ScottyT(..), Middleware, BodyInfo, ScottyError(..), ErrorHandler, handler, addRoute)
+import           Web.Scotty.Util (strictByteStringToLazyText)
+import Web.Scotty.Body (cloneBodyInfo, getBodyAction, getBodyChunkAction, getFormParamsAndFilesAction)
 
 -- | get = 'addroute' 'GET'
 get :: (ScottyError e, MonadIO m) => RoutePattern -> ActionT e m () -> ScottyT e m ()
@@ -86,31 +79,34 @@ notFound action = matchAny (Function (\req -> Just [("path", path req)])) (statu
 addroute :: (ScottyError e, MonadIO m) => StdMethod -> RoutePattern -> ActionT e m () -> ScottyT e m ()
 addroute method pat action = ScottyT $ MS.modify $ \s -> addRoute (route (routeOptions s) (handler s) (Just method) pat action) s
 
-route :: (ScottyError e, MonadIO m) => RouteOptions -> ErrorHandler e m -> Maybe StdMethod -> RoutePattern -> ActionT e m () -> Middleware m
-route opts h method pat action app req =
-    let tryNext = app req
+route :: (ScottyError e, MonadIO m) =>
+         RouteOptions -> ErrorHandler e m -> Maybe StdMethod -> RoutePattern -> ActionT e m () -> BodyInfo -> Middleware m
+route opts h method pat action bodyInfo app req =
+  let tryNext = app req
         {- |
           We match all methods in the case where 'method' is 'Nothing'.
           See https://github.com/scotty-web/scotty/issues/196
         -}
-        methodMatches :: Bool
-        methodMatches =
-            case method of
-                Nothing -> True
-                Just m -> Right m == parseMethod (requestMethod req)
-    in if methodMatches
-       then case matchRoute pat req of
+      methodMatches :: Bool
+      methodMatches = maybe True (\x -> (Right x == parseMethod (requestMethod req))) method
+
+  in if methodMatches
+     then case matchRoute pat req of
             Just captures -> do
-                env <- liftIO $ catch (Right <$> mkEnv req captures opts) (\ex -> return . Left $ ex)
-                res <- evalAction h env action
-                maybe tryNext return res
+              -- The user-facing API for "body" and "bodyReader" involve an IO action that
+              -- reads the body/chunks thereof only once, so we shouldn't pass in our BodyInfo
+              -- directly; otherwise, the body might get consumed and then it would be unavailable
+              -- if `next` is called and we try to match further routes.
+              -- Instead, make a "cloned" copy of the BodyInfo that allows the IO actions to be called
+              -- without messing up the state of the original BodyInfo.
+              clonedBodyInfo <- cloneBodyInfo bodyInfo
+
+              env <- mkEnv clonedBodyInfo req captures opts
+              res <- runAction h env action
+              maybe tryNext return res
             Nothing -> tryNext
-       else tryNext
+     else tryNext
 
-evalAction :: (ScottyError e, Monad m) => ErrorHandler e m -> (Either ScottyException ActionEnv) -> ActionT e m () -> m (Maybe Response)
-evalAction _ (Left (RequestException msg s)) _ = return . Just $ responseBuilder s [("Content-Type","text/html")] $ fromByteString msg
-evalAction h (Right env) action = runAction h env action
-
 matchRoute :: RoutePattern -> Request -> Maybe [Param]
 matchRoute (Literal pat)  req | pat == path req = Just []
                               | otherwise       = Nothing
@@ -133,73 +129,15 @@ matchRoute (Capture pat)  req = go (T.split (=='/') pat) (compress $ T.split (==
 path :: Request -> T.Text
 path = T.fromStrict . TS.cons '/' . TS.intercalate "/" . pathInfo
 
--- Stolen from wai-extra's Network.Wai.Parse, modified to accept body as list of Bytestrings.
--- Reason: WAI's getRequestBodyChunk 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]
-                 -> Parse.BackEnd y
-                 -> Request
-                 -> m ([Parse.Param], [Parse.File y])
-parseRequestBody bl s r =
-    case Parse.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 $ Parse.sinkRequestBody s rbt provider
-
-mkEnv :: forall m. MonadIO m => Request -> [Param] -> RouteOptions ->m ActionEnv
-mkEnv req captures opts = do
-    bodyState <- liftIO $ newMVar BodyUntouched
-
-    let rbody = getRequestBodyChunk req
-
-        safeBodyReader :: IO B.ByteString
-        safeBodyReader =  do
-          state <- takeMVar bodyState
-          let direct = putMVar bodyState BodyCorrupted >> rbody
-          case state of
-            s@(BodyCached _ []) ->
-              do putMVar bodyState s
-                 return B.empty
-            BodyCached b (chunk:rest) ->
-              do putMVar bodyState $ BodyCached b rest
-                 return chunk
-            BodyUntouched -> direct
-            BodyCorrupted -> direct
 
-        bs :: IO BL.ByteString
-        bs = do
-          state <- takeMVar bodyState
-          case state of
-            s@(BodyCached b _) ->
-              do putMVar bodyState s
-                 return b
-            BodyCorrupted -> throw BodyPartiallyStreamed
-            BodyUntouched ->
-              do chunks <- readRequestBody rbody return (maxRequestBodySize opts)
-                 let b = BL.fromChunks chunks
-                 putMVar bodyState $ BodyCached b chunks
-                 return b
+mkEnv :: MonadIO m => BodyInfo -> Request -> [Param] -> RouteOptions -> m ActionEnv
+mkEnv bodyInfo req captureps opts = do
+  (formps, bodyFiles) <- liftIO $ getFormParamsAndFilesAction req bodyInfo opts
+  let
+    queryps = parseEncodedParams $ rawQueryString req
+    bodyFiles' = [ (strictByteStringToLazyText k, fi) | (k,fi) <- bodyFiles ]
+  return $ Env req captureps formps queryps (getBodyAction bodyInfo opts) (getBodyChunkAction bodyInfo) bodyFiles'
 
-        shouldParseBody = isJust $ Parse.getRequestBodyType req
-
-    (formparams, fs) <- if shouldParseBody
-      then liftIO $ do wholeBody <- BL.toChunks `fmap` bs
-                       parseRequestBody wholeBody Parse.lbsBackEnd req
-      else return ([], [])
-
-    let
-        convert (k, v) = (strictByteStringToLazyText k, strictByteStringToLazyText v)
-        formparams' = map convert formparams
-        queryparams = parseEncodedParams $ rawQueryString req
-
-    return $ Env req captures formparams' queryparams bs safeBodyReader [ (strictByteStringToLazyText k, fi) | (k,fi) <- fs ]
 
 parseEncodedParams :: B.ByteString -> [Param]
 parseEncodedParams bs = [ (T.fromStrict k, T.fromStrict $ fromMaybe "" v) | (k,v) <- parseQueryText bs ]
@@ -255,8 +193,3 @@ function = Function
 -- | Build a route that requires the requested path match exactly, without captures.
 literal :: String -> RoutePattern
 literal = Literal . T.pack
-
-#if !(MIN_VERSION_wai(3,2,2))
-getRequestBodyChunk :: Request -> IO B.ByteString
-getRequestBodyChunk = requestBody
-#endif

Web/Scotty/Trans.hs

@@ -62,6 +62,7 @@ import Web.Scotty.Route
 import Web.Scotty.Internal.Types hiding (Application, Middleware)
 import Web.Scotty.Util (socketDescription)
 import qualified Web.Scotty.Internal.Types as Scotty
+import Web.Scotty.Body (newBodyInfo)
 
 -- | Run a scotty application using the warp server.
 -- NB: scotty p === scottyT p id
@@ -108,7 +109,9 @@ scottyAppT :: (Monad m, Monad n)
            -> n Application
 scottyAppT runActionToIO defs = do
     let s = execState (runS defs) def
-    let rapp req callback = runActionToIO (foldl (flip ($)) notFoundApp (routes s) req) >>= callback
+    let rapp req callback = do
+          bodyInfo <- newBodyInfo req
+          runActionToIO (foldl (flip ($)) notFoundApp ([midd bodyInfo | midd <- routes s]) req) >>= callback
     return $ foldl (flip ($)) rapp (middlewares s)
 
 notFoundApp :: Monad m => Scotty.Application m