Limit request body sizes for chunked bodies

yesod-core/Yesod/Handler.hs

@@ -172,7 +172,6 @@ import Control.Monad.Logger
 
 import qualified Yesod.Internal.Cache as Cache
 import Yesod.Internal.Cache (mkCacheKey, CacheKey)
-import Data.Typeable (Typeable)
 import qualified Data.IORef as I
 import Control.Exception.Lifted (catch)
 import Control.Monad.Trans.Control
@@ -303,19 +302,6 @@ data YesodAppResult
     = YARWai W.Response
     | YARPlain H.Status [Header] ContentType Content SessionMap
 
-data HandlerContents =
-      HCContent H.Status ChooseRep
-    | HCError ErrorResponse
-    | HCSendFile ContentType FilePath (Maybe W.FilePart) -- FIXME replace FilePath with opaque type from system-filepath?
-    | HCRedirect H.Status Text
-    | HCCreated Text
-    | HCWai W.Response
-    deriving Typeable
-
-instance Show HandlerContents where
-    show _ = "Cannot show a HandlerContents"
-instance Exception HandlerContents
-
 getRequest :: GHandler s m Request
 getRequest = handlerRequest `liftM` ask
 

yesod-core/Yesod/Internal.hs

@@ -9,6 +9,7 @@
 module Yesod.Internal
     ( -- * Error responses
       ErrorResponse (..)
+    , HandlerContents (..)
       -- * Header
     , Header (..)
       -- * Cookie names
@@ -46,6 +47,8 @@ import qualified Data.Map as Map
 import Data.Text.Lazy.Builder (Builder)
 import Web.Cookie (SetCookie (..))
 import Data.ByteString (ByteString)
+import qualified Network.Wai as W
+import Yesod.Content (ChooseRep, ContentType)
 
 -- | Responses to indicate some form of an error occurred. These are different
 -- from 'SpecialResponse' in that they allow for custom error pages.
@@ -120,3 +123,16 @@ instance Monoid (GWData a) where
         (Map.unionWith mappend a5 b5)
         (a6 `mappend` b6)
         (a7 `mappend` b7)
+
+data HandlerContents =
+      HCContent H.Status ChooseRep
+    | HCError ErrorResponse
+    | HCSendFile ContentType FilePath (Maybe W.FilePart) -- FIXME replace FilePath with opaque type from system-filepath?
+    | HCRedirect H.Status Text
+    | HCCreated Text
+    | HCWai W.Response
+    deriving Typeable
+
+instance Show HandlerContents where
+    show _ = "Cannot show a HandlerContents"
+instance Exception HandlerContents

yesod-core/Yesod/Internal/Core.hs

@@ -359,7 +359,7 @@ $doctype 5
 
     -- | How to store uploaded files.
     --
-    -- Default: Whe nthe request body is greater than 50kb, store in a temp
+    -- Default: When the request body is greater than 50kb, store in a temp
     -- file. Otherwise, store in memory.
     fileUpload :: a
                -> Word64 -- ^ request body size
@@ -433,17 +433,13 @@ defaultYesodRunner :: Yesod master
                    -> Maybe (SessionBackend master)
                    -> W.Application
 defaultYesodRunner logger handler' master sub murl toMasterRoute msb req
-  | maximumContentLength master (fmap toMasterRoute murl) < len =
-        return $ W.responseLBS
-            (H.Status 413 "Too Large")
-            [("Content-Type", "text/plain")]
-            "Request body too large to be processed."
+  | 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
+    rr <- liftIO $ parseWaiRequest req session (isJust msb) len maxLen
     let h = {-# SCC "h" #-} do
           case murl of
             Nothing -> handler
@@ -477,6 +473,7 @@ defaultYesodRunner logger handler' master sub murl toMasterRoute msb req
         _ -> return []
     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

yesod-core/Yesod/Internal/Request.hs

@@ -13,6 +13,7 @@ module Yesod.Internal.Request
     , mkFileInfoFile
     , mkFileInfoSource
     , FileUpload (..)
+    , tooLargeResponse
     -- The below are exported for testing.
     , randomString
     , parseWaiRequest'
@@ -28,7 +29,7 @@ import Web.Cookie (parseCookiesText)
 import Data.ByteString (ByteString)
 import qualified Data.ByteString.Char8 as S8
 import Data.Text (Text, pack)
-import Network.HTTP.Types (queryToQueryText)
+import Network.HTTP.Types (queryToQueryText, Status (Status))
 import Control.Monad (join)
 import Data.Maybe (fromMaybe, catMaybes)
 import qualified Data.ByteString.Lazy as L
@@ -40,6 +41,8 @@ import Data.Conduit
 import Data.Conduit.List (sourceList)
 import Data.Conduit.Binary (sourceFile, sinkFile)
 import Data.Word (Word64)
+import Control.Monad.IO.Class (liftIO)
+import Control.Exception (throwIO)
 
 -- | The parsed request information.
 data Request = Request
@@ -51,26 +54,56 @@ data Request = Request
       -- | A random, session-specific token used to prevent CSRF attacks.
     , reqToken :: Maybe Text
       -- | Size of the request body.
-    , reqBodySize :: Word64
+      --
+      -- 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
+                -> Word64 -- ^ actual length... might be meaningless, see 'reqBodySize'
+                -> Word64 -- ^ maximum allowed body size
                 -> IO Request
-parseWaiRequest env session' useToken bodySize =
-    parseWaiRequest' env session' useToken bodySize <$> newStdGen
+parseWaiRequest env session' useToken bodySize maxBodySize =
+    parseWaiRequest' env session' useToken bodySize maxBodySize <$> newStdGen
+
+-- | Impose a limit on the size of the request body.
+limitRequestBody :: Word64 -> W.Request -> W.Request
+limitRequestBody maxLen req =
+    req { W.requestBody = W.requestBody req $= limit maxLen }
+  where
+    tooLarge = liftIO $ throwIO $ HCWai tooLargeResponse
+
+    limit 0 = tooLarge
+    limit remaining =
+        await >>= maybe (return ()) go
+      where
+        go bs = do
+            let len = fromIntegral $ S8.length bs
+            if len > remaining
+                then tooLarge
+                else do
+                    yield bs
+                    limit $ remaining - len
+
+tooLargeResponse :: W.Response
+tooLargeResponse = W.responseLBS
+    (Status 413 "Too Large")
+    [("Content-Type", "text/plain")]
+    "Request body too large to be processed."
 
 parseWaiRequest' :: RandomGen g
                  => W.Request
                  -> [(Text, ByteString)] -- ^ session
                  -> Bool
                  -> Word64
+                 -> Word64 -- ^ max body size
                  -> g
                  -> Request
-parseWaiRequest' env session' useToken bodySize gen =
-    Request gets'' cookies' env langs'' token bodySize
+parseWaiRequest' env session' useToken bodySize maxBodySize gen =
+    Request gets'' cookies' (limitRequestBody maxBodySize env) langs'' token bodySize
   where
     gets' = queryToQueryText $ W.queryString env
     gets'' = map (second $ fromMaybe "") gets'

yesod-core/test/YesodCoreTest.hs

@@ -12,6 +12,7 @@ import YesodCoreTest.Cache
 import qualified YesodCoreTest.WaiSubsite as WaiSubsite
 import qualified YesodCoreTest.Redirect as Redirect
 import qualified YesodCoreTest.JsLoader as JsLoader
+import qualified YesodCoreTest.RequestBodySize as RequestBodySize
 
 import Test.Hspec
 
@@ -29,3 +30,4 @@ specs = do
       WaiSubsite.specs
       Redirect.specs
       JsLoader.specs
+      RequestBodySize.specs

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 g
+  r = parseWaiRequest' defaultRequest [] False 0 1000 g
 
 ignoreDisabledToken :: Bool
 ignoreDisabledToken = reqToken r == Nothing where
-  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] False 0 g
+  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] False 0 1000 g
 
 useOldToken :: Bool
 useOldToken = reqToken r == Just "old" where
-  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] True 0 g
+  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] True 0 1000 g
 
 generateToken :: Bool
 generateToken = reqToken r /= Nothing where
-  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] True 0 g
+  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] True 0 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 g
+        { requestHeaders = [("Accept-Language", "en-US, es")] } [] False 0 1000 g
 
 respectSessionLang :: Bool
 respectSessionLang = reqLangs r == ["en"] where
-  r = parseWaiRequest' defaultRequest [("_LANG", "en")] False 0 g
+  r = parseWaiRequest' defaultRequest [("_LANG", "en")] False 0 1000 g
 
 respectCookieLang :: Bool
 respectCookieLang = reqLangs r == ["en"] where
   r = parseWaiRequest' defaultRequest
         { requestHeaders = [("Cookie", "_LANG=en")]
-        } [] False 0 g
+        } [] False 0 1000 g
 
 respectQueryLang :: Bool
 respectQueryLang = reqLangs r == ["en-US", "en"] where
-  r = parseWaiRequest' defaultRequest { queryString = [("_LANG", Just "en-US")] } [] False 0 g
+  r = parseWaiRequest' defaultRequest { queryString = [("_LANG", Just "en-US")] } [] False 0 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 g
+        } [("_LANG", "en-SESSION")] False 0 10000 g
 
 
 internalRequestTest :: Spec

yesod-core/test/YesodCoreTest/RequestBodySize.hs

@@ -0,0 +1,89 @@
+{-# LANGUAGE QuasiQuotes, TypeFamilies, TemplateHaskell, MultiParamTypeClasses #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE FlexibleInstances #-}
+module YesodCoreTest.RequestBodySize (specs, Widget) where
+
+import Test.Hspec
+
+import Yesod.Core hiding (Request)
+
+import Network.Wai
+import Network.Wai.Test
+import Data.ByteString (ByteString)
+import qualified Data.ByteString as S
+import qualified Data.ByteString.Lazy as L
+import qualified Data.ByteString.Char8 as S8
+import Data.Text (Text)
+import qualified Data.Text as T
+import Data.Conduit
+import Data.Conduit.List (consume)
+import Data.Conduit.Binary (isolate)
+
+data Y = Y
+
+mkYesod "Y" [parseRoutes|
+/post PostR POST
+/consume ConsumeR POST
+/partial-consume PartialConsumeR POST
+/unused UnusedR POST
+|]
+
+instance Yesod Y where
+    maximumContentLength _ _ = 10
+
+postPostR, postConsumeR, postPartialConsumeR, postUnusedR :: Handler RepPlain
+
+postPostR = do
+    val <- lookupPostParams "foobarbaz"
+    return $ RepPlain $ toContent $ T.concat val
+
+postConsumeR = do
+    req <- waiRequest
+    body <- lift $ requestBody req $$ consume
+    return $ RepPlain $ toContent $ S.concat body
+
+postPartialConsumeR = do
+    req <- waiRequest
+    body <- lift $ requestBody req $$ isolate 5 =$ consume
+    return $ RepPlain $ toContent $ S.concat body
+
+postUnusedR = return $ RepPlain ""
+
+runner :: Session () -> IO ()
+runner f = toWaiApp Y >>= runSession f
+
+caseHelper :: String -- ^ name
+           -> Text -- ^ pathinfo
+           -> ByteString -- ^ request body
+           -> Int -- ^ expected status code, chunked
+           -> Int -- ^ expected status code, non-chunked
+           -> Spec
+caseHelper name path body statusChunked statusNonChunked = describe name $ do
+    it "chunked" $ runner $ do
+        res <- mkRequest False
+        assertStatus statusChunked res
+    it "non-chunked" $ runner $ do
+        res <- mkRequest True
+        assertStatus statusNonChunked res
+  where
+    mkRequest includeLength = srequest $ SRequest defaultRequest
+        { pathInfo = [path]
+        , requestHeaders =
+            ("content-type", "application/x-www-form-urlencoded") :
+
+            if includeLength
+                then [("content-length", S8.pack $ show $ S.length body)]
+                else []
+        , requestMethod = "POST"
+        } $ L.fromChunks $ map S.singleton $ S.unpack body
+
+specs :: Spec
+specs = describe "Test.RequestBodySize" $ do
+    caseHelper "lookupPostParam- large" "post" "foobarbaz=bin" 413 413
+    caseHelper "lookupPostParam- small" "post" "foo=bin" 200 200
+    caseHelper "consume- large" "consume" "this is longer than 10" 413 413
+    caseHelper "consume- small" "consume" "smaller" 200 200
+    caseHelper "partial consume- large" "partial-consume" "this is longer than 10" 200 413
+    caseHelper "partial consume- small" "partial-consume" "smaller" 200 200
+    caseHelper "unused- large" "unused" "this is longer than 10" 200 413
+    caseHelper "unused- small" "unused" "smaller" 200 200

yesod-core/yesod-core.cabal

@@ -30,6 +30,7 @@ extra-source-files:
   test/YesodCoreTest/MediaData.hs
   test/YesodCoreTest/NoOverloadedStrings.hs
   test/YesodCoreTest/Redirect.hs
+  test/YesodCoreTest/RequestBodySize.hs
   test/YesodCoreTest/WaiSubsite.hs
   test/YesodCoreTest/Widget.hs
   test/YesodCoreTest/YesodTest.hs
@@ -123,6 +124,7 @@ test-suite tests
                   ,HUnit
                   ,QuickCheck >= 2 && < 3
                   ,transformers
+                  , conduit
     ghc-options:     -Wall
 
 source-repository head