Use RequestBodyLength

yesod-core/Yesod/Handler.hs

@@ -179,7 +179,6 @@ import Control.Monad.Trans.Control
 import Control.Monad.Trans.Resource
 import Control.Monad.Base
 import Yesod.Routes.Class
-import Data.Word (Word64)
 import Language.Haskell.TH.Syntax (Loc)
 
 class YesodSubRoute s y where
@@ -193,7 +192,7 @@ data HandlerData sub master = HandlerData
     , handlerRender   :: Route master -> [(Text, Text)] -> Text
     , handlerToMaster :: Route sub -> Route master
     , handlerState    :: I.IORef GHState
-    , handlerUpload   :: Word64 -> FileUpload
+    , handlerUpload   :: W.RequestBodyLength -> FileUpload
     , handlerLog      :: Loc -> LogSource -> LogLevel -> LogStr -> IO ()
     }
 
@@ -313,7 +312,9 @@ runRequestBody :: GHandler s m RequestBodyContents
 runRequestBody = do
     hd <- ask
     let getUpload = handlerUpload hd
-        len = reqBodySize $ handlerRequest hd
+        len = W.requestBodyLength
+            $ reqWaiRequest
+            $ handlerRequest hd
         upload = getUpload len
     x <- get
     case ghsRBC x of
@@ -422,9 +423,10 @@ handlerToIO =
     -- Let go of the request body, cache and response headers.
     let oldReq    = handlerRequest oldHandlerData
         oldWaiReq = reqWaiRequest oldReq
-        newWaiReq = oldWaiReq { W.requestBody = mempty }
+        newWaiReq = oldWaiReq { W.requestBody = mempty
-        newReq    = oldReq { reqWaiRequest = newWaiReq
+                              , W.requestBodyLength = W.KnownLength 0
-                           , reqBodySize   = 0 }
+                              }
+        newReq    = oldReq { reqWaiRequest = newWaiReq }
         clearedOldHandlerData =
           oldHandlerData { handlerRequest = err "handlerRequest never here"
                          , handlerState   = err "handlerState never here" }
@@ -457,7 +459,7 @@ runHandler :: HasReps c
            -> (Route sub -> Route master)
            -> master
            -> sub
-           -> (Word64 -> FileUpload)
+           -> (W.RequestBodyLength -> FileUpload)
            -> (Loc -> LogSource -> LogLevel -> LogStr -> IO ())
            -> YesodApp
 runHandler handler mrender sroute tomr master sub upload log' =
@@ -872,7 +874,7 @@ getSession = liftM ghsSession get
 handlerToYAR :: (HasReps a, HasReps b)
              => master -- ^ master site foundation
              -> sub    -- ^ sub site foundation
-             -> (Word64 -> FileUpload)
+             -> (W.RequestBodyLength -> FileUpload)
              -> (Loc -> LogSource -> LogLevel -> LogStr -> IO ())
              -> (Route sub -> Route master)
              -> (Route master -> [(Text, Text)] -> Text) -- route renderer

yesod-core/Yesod/Internal/Core.hs

@@ -4,6 +4,7 @@
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE PatternGuards #-}
 -- | The basic typeclass for a Yesod application.
 module Yesod.Internal.Core
     ( -- * Type classes
@@ -360,13 +361,12 @@ $doctype 5
     -- | How to store uploaded files.
     --
     -- Default: When the request body is greater than 50kb, store in a temp
-    -- file. Otherwise, store in memory.
+    -- file. For chunked request bodies, store in a temp file. Otherwise, store
-    fileUpload :: a
+    -- in memory.
-               -> Word64 -- ^ request body size
+    fileUpload :: a -> W.RequestBodyLength -> FileUpload
-               -> FileUpload
+    fileUpload _ (W.KnownLength size)
-    fileUpload _ size
+        | size <= 50000 = FileUploadMemory lbsBackEnd
-        | size > 50000 = FileUploadDisk tempFileBackEnd
+    fileUpload _ _ = FileUploadDisk tempFileBackEnd
-        | otherwise = FileUploadMemory lbsBackEnd
 
     -- | Should we log the given log source/level combination.
     --
@@ -433,13 +433,13 @@ defaultYesodRunner :: Yesod master
                    -> Maybe (SessionBackend master)
                    -> W.Application
 defaultYesodRunner logger handler' master sub murl toMasterRoute msb req
-  | maxLen < len = return tooLargeResponse
+  | W.KnownLength len <- W.requestBodyLength req, maxLen < len = return tooLargeResponse
   | otherwise = do
     let dontSaveSession _ _ = return []
     now <- liftIO getCurrentTime -- FIXME remove in next major version bump
     (session, saveSession) <- liftIO $ do
         maybe (return ([], dontSaveSession)) (\sb -> sbLoadSession sb master req now) msb
-    rr <- liftIO $ parseWaiRequest req session (isJust msb) len maxLen
+    rr <- liftIO $ parseWaiRequest req session (isJust msb) maxLen
     let h = {-# SCC "h" #-} do
           case murl of
             Nothing -> handler
@@ -474,11 +474,6 @@ defaultYesodRunner logger handler' master sub murl toMasterRoute msb req
     return $ yarToResponse yar extraHeaders
   where
     maxLen = maximumContentLength master $ fmap toMasterRoute murl
-    len = fromMaybe 0 $ lookup "content-length" (W.requestHeaders req) >>= readMay
-    readMay s =
-        case reads $ S8.unpack s of
-            [] -> Nothing
-            (x, _):_ -> Just x
     handler = yesodMiddleware handler'
 
 data AuthResult = Authorized | AuthenticationRequired | Unauthorized Text
@@ -920,7 +915,6 @@ runFakeHandler fakeSessionMap logger master handler = liftIO $ do
           , reqWaiRequest = fakeWaiRequest
           , reqLangs      = []
           , reqToken      = Just "NaN" -- not a nonce =)
-          , reqBodySize   = 0
           }
       fakeContentType = []
   _ <- runResourceT $ yapp errHandler fakeRequest fakeContentType fakeSessionMap

yesod-core/Yesod/Internal/Request.hs

@@ -53,21 +53,15 @@ data Request = Request
     , reqLangs :: [Text]
       -- | A random, session-specific token used to prevent CSRF attacks.
     , reqToken :: Maybe Text
-      -- | Size of the request body.
-      --
-      -- Note: in the presence of chunked request bodies, this value will be 0,
-      -- even though data is available.
-    , reqBodySize :: Word64 -- FIXME Consider in the future using a Maybe to represent chunked bodies
     }
 
 parseWaiRequest :: W.Request
                 -> [(Text, ByteString)] -- ^ session
                 -> Bool
-                -> Word64 -- ^ actual length... might be meaningless, see 'reqBodySize'
                 -> Word64 -- ^ maximum allowed body size
                 -> IO Request
-parseWaiRequest env session' useToken bodySize maxBodySize =
+parseWaiRequest env session' useToken maxBodySize =
-    parseWaiRequest' env session' useToken bodySize maxBodySize <$> newStdGen
+    parseWaiRequest' env session' useToken maxBodySize <$> newStdGen
 
 -- | Impose a limit on the size of the request body.
 limitRequestBody :: Word64 -> W.Request -> W.Request
@@ -98,12 +92,11 @@ parseWaiRequest' :: RandomGen g
                  => W.Request
                  -> [(Text, ByteString)] -- ^ session
                  -> Bool
-                 -> Word64
                  -> Word64 -- ^ max body size
                  -> g
                  -> Request
-parseWaiRequest' env session' useToken bodySize maxBodySize gen =
+parseWaiRequest' env session' useToken maxBodySize gen =
-    Request gets'' cookies' (limitRequestBody maxBodySize env) langs'' token bodySize
+    Request gets'' cookies' (limitRequestBody maxBodySize env) langs'' token
   where
     gets' = queryToQueryText $ W.queryString env
     gets'' = map (second $ fromMaybe "") gets'

yesod-core/test/YesodCoreTest/InternalRequest.hs

@@ -38,19 +38,19 @@ tokenSpecs = describe "Yesod.Internal.Request.parseWaiRequest (reqToken)" $ do
 
 noDisabledToken :: Bool
 noDisabledToken = reqToken r == Nothing where
-  r = parseWaiRequest' defaultRequest [] False 0 1000 g
+  r = parseWaiRequest' defaultRequest [] False 1000 g
 
 ignoreDisabledToken :: Bool
 ignoreDisabledToken = reqToken r == Nothing where
-  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] False 0 1000 g
+  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] False 1000 g
 
 useOldToken :: Bool
 useOldToken = reqToken r == Just "old" where
-  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] True 0 1000 g
+  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] True 1000 g
 
 generateToken :: Bool
 generateToken = reqToken r /= Nothing where
-  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] True 0 1000 g
+  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] True 1000 g
 
 
 langSpecs :: Spec
@@ -64,21 +64,21 @@ langSpecs = describe "Yesod.Internal.Request.parseWaiRequest (reqLangs)" $ do
 respectAcceptLangs :: Bool
 respectAcceptLangs = reqLangs r == ["en-US", "es", "en"] where
   r = parseWaiRequest' defaultRequest
-        { requestHeaders = [("Accept-Language", "en-US, es")] } [] False 0 1000 g
+        { requestHeaders = [("Accept-Language", "en-US, es")] } [] False 1000 g
 
 respectSessionLang :: Bool
 respectSessionLang = reqLangs r == ["en"] where
-  r = parseWaiRequest' defaultRequest [("_LANG", "en")] False 0 1000 g
+  r = parseWaiRequest' defaultRequest [("_LANG", "en")] False 1000 g
 
 respectCookieLang :: Bool
 respectCookieLang = reqLangs r == ["en"] where
   r = parseWaiRequest' defaultRequest
         { requestHeaders = [("Cookie", "_LANG=en")]
-        } [] False 0 1000 g
+        } [] False 1000 g
 
 respectQueryLang :: Bool
 respectQueryLang = reqLangs r == ["en-US", "en"] where
-  r = parseWaiRequest' defaultRequest { queryString = [("_LANG", Just "en-US")] } [] False 0 1000 g
+  r = parseWaiRequest' defaultRequest { queryString = [("_LANG", Just "en-US")] } [] False 1000 g
 
 prioritizeLangs :: Bool
 prioritizeLangs = reqLangs r == ["en-QUERY", "en-COOKIE", "en-SESSION", "en", "es"] where
@@ -87,7 +87,7 @@ prioritizeLangs = reqLangs r == ["en-QUERY", "en-COOKIE", "en-SESSION", "en", "e
                            , ("Cookie", "_LANG=en-COOKIE")
                            ]
         , queryString = [("_LANG", Just "en-QUERY")]
-        } [("_LANG", "en-SESSION")] False 0 10000 g
+        } [("_LANG", "en-SESSION")] False 10000 g
 
 
 internalRequestTest :: Spec

yesod-core/test/YesodCoreTest/RequestBodySize.hs

@@ -75,6 +75,10 @@ caseHelper name path body statusChunked statusNonChunked = describe name $ do
                 then [("content-length", S8.pack $ show $ S.length body)]
                 else []
         , requestMethod = "POST"
+        , requestBodyLength =
+            if includeLength
+                then KnownLength $ fromIntegral $ S.length body
+                else ChunkedBody
         } $ L.fromChunks $ map S.singleton $ S.unpack body
 
 specs :: Spec

yesod-core/yesod-core.cabal

@@ -50,7 +50,7 @@ library
     build-depends:   base                  >= 4.3      && < 5
                    , time                  >= 1.1.4
                    , yesod-routes          >= 1.1      && < 1.2
-                   , wai                   >= 1.3      && < 1.5
+                   , wai                   >= 1.4      && < 1.5
                    , wai-extra             >= 1.3      && < 1.4
                    , bytestring            >= 0.9.1.4
                    , text                  >= 0.7      && < 0.12