add requestSizeCheck

wai-extra/Network/Wai/Request.hs

@@ -3,15 +3,22 @@
 module Network.Wai.Request
     ( appearsSecure
     , guessApproot
+    , RequestSizeException(..)
+    , requestSizeCheck
     ) where
 
 import Data.ByteString (ByteString)
 import Data.Maybe (fromMaybe)
 import Network.HTTP.Types (HeaderName)
-import Network.Wai (Request, isSecure, requestHeaders, requestHeaderHost)
+import Network.Wai
 
 import qualified Data.ByteString as S
 import qualified Data.ByteString.Char8 as C
+import Control.Exception (Exception, throwIO)
+import Data.Typeable (Typeable)
+import Data.Word (Word64)
+import Data.IORef (atomicModifyIORef', newIORef)
+
 
 -- | Does this request appear to have been made over an SSL connection?
 --
@@ -49,3 +56,41 @@ guessApproot :: Request -> ByteString
 guessApproot req =
     (if appearsSecure req then "https://" else "http://") `S.append`
     (fromMaybe "localhost" $ requestHeaderHost req)
+
+-- | see 'requestSizeCheck'
+data RequestSizeException
+    = RequestSizeException Word64
+    deriving (Eq, Ord, Typeable)
+
+instance Exception RequestSizeException
+
+instance Show RequestSizeException where
+    showsPrec p (RequestSizeException limit) =
+        showString ("Request Body is larger than ") . showsPrec p limit . showString " bytes."
+
+-- | Check request body size to avoid server crash when request is too large.
+--
+-- This function first checks @'requestBodyLength'@, if content-length is known
+-- but larger than limit, or it's unknown but we have received too many chunks,
+-- a 'RequestSizeException' are thrown when user use @'requestBody'@ to extract
+-- request body inside IO.
+requestSizeCheck :: Word64 -> Request -> IO Request
+requestSizeCheck maxSize req =
+    case requestBodyLength req of
+        KnownLength len   ->
+            if len > maxSize
+                then return $ req { requestBody = throwIO (RequestSizeException maxSize) }
+                else return req
+        ChunkedBody	      -> do
+            currentSize <- newIORef 0
+            return $ req
+                { requestBody = do
+                    bs <- requestBody req
+                    total <-
+                        atomicModifyIORef' currentSize $ \sz ->
+                            let nextSize = sz + fromIntegral (S.length bs)
+                            in (nextSize, nextSize)
+                    if total > maxSize
+                    then throwIO (RequestSizeException maxSize)
+                    else return bs
+                }

wai-extra/test/Network/Wai/RequestSpec.hs

@@ -8,52 +8,82 @@ import Test.Hspec
 
 import Data.ByteString (ByteString)
 import Network.HTTP.Types (HeaderName)
-import Network.Wai (Request(..), defaultRequest)
+import Network.Wai (Request(..), defaultRequest, RequestBodyLength(..))
 
 import Network.Wai.Request
+import Control.Exception (try)
+import Control.Monad (forever)
 
 main :: IO ()
 main = hspec spec
 
 spec :: Spec
-spec = describe "appearsSecure" $ do
-    let insecureRequest = defaultRequest
-            { isSecure = False
-            , requestHeaders =
-                [ ("HTTPS", "off")
-                , ("HTTP_X_FORWARDED_SSL", "off")
-                , ("HTTP_X_FORWARDED_SCHEME", "http")
-                , ("HTTP_X_FORWARDED_PROTO", "http,xyz")
-                ]
-            }
+spec = do
+    describe "requestSizeCheck" $ do
+        it "too large content length should throw RequestSizeException" $ do
+            let limit = 1024
+                largeRequest = defaultRequest
+                    { isSecure = False
+                    , requestBodyLength = KnownLength (limit + 1)
+                    , requestBody = return "repeat this chunk"
+                    }
+            checkedRequest <- requestSizeCheck limit largeRequest
+            body <- try (requestBody checkedRequest)
+            case body of
+                Left (RequestSizeException l) -> l `shouldBe` limit
+                Right _   -> expectationFailure "request size check failed"
 
-    it "returns False for an insecure request" $
-        insecureRequest `shouldSatisfy` not . appearsSecure
+        it "too many chunks should throw RequestSizeException" $ do
+            let limit = 1024
+                largeRequest = defaultRequest
+                    { isSecure = False
+                    , requestBodyLength = ChunkedBody
+                    , requestBody = return "repeat this chunk"
+                    }
+            checkedRequest <- requestSizeCheck limit largeRequest
+            body <- try (forever $ requestBody checkedRequest)
+            case body of
+                Left (RequestSizeException l) -> l `shouldBe` limit
+                Right _   -> expectationFailure "request size check failed"
 
-    it "checks if the Request is actually secure" $ do
-        let req = insecureRequest { isSecure = True }
+    describe "appearsSecure" $ do
+        let insecureRequest = defaultRequest
+                { isSecure = False
+                , requestHeaders =
+                    [ ("HTTPS", "off")
+                    , ("HTTP_X_FORWARDED_SSL", "off")
+                    , ("HTTP_X_FORWARDED_SCHEME", "http")
+                    , ("HTTP_X_FORWARDED_PROTO", "http,xyz")
+                    ]
+                }
 
-        req `shouldSatisfy` appearsSecure
+        it "returns False for an insecure request" $
+            insecureRequest `shouldSatisfy` not . appearsSecure
 
-    it "checks for HTTP: on" $ do
-        let req = addHeader "HTTPS" "on" insecureRequest
+        it "checks if the Request is actually secure" $ do
+            let req = insecureRequest { isSecure = True }
 
-        req `shouldSatisfy` appearsSecure
+            req `shouldSatisfy` appearsSecure
 
-    it "checks for HTTP_X_FORWARDED_SSL: on" $ do
-        let req = addHeader "HTTP_X_FORWARDED_SSL" "on" insecureRequest
+        it "checks for HTTP: on" $ do
+            let req = addHeader "HTTPS" "on" insecureRequest
 
-        req `shouldSatisfy` appearsSecure
+            req `shouldSatisfy` appearsSecure
 
-    it "checks for HTTP_X_FORWARDED_SCHEME: https" $ do
-        let req = addHeader "HTTP_X_FORWARDED_SCHEME" "https" insecureRequest
+        it "checks for HTTP_X_FORWARDED_SSL: on" $ do
+            let req = addHeader "HTTP_X_FORWARDED_SSL" "on" insecureRequest
 
-        req `shouldSatisfy` appearsSecure
+            req `shouldSatisfy` appearsSecure
 
-    it "checks for HTTP_X_FORWARDED_PROTO: https,..." $ do
-        let req = addHeader "HTTP_X_FORWARDED_PROTO" "https,xyz" insecureRequest
+        it "checks for HTTP_X_FORWARDED_SCHEME: https" $ do
+            let req = addHeader "HTTP_X_FORWARDED_SCHEME" "https" insecureRequest
 
-        req `shouldSatisfy` appearsSecure
+            req `shouldSatisfy` appearsSecure
+
+        it "checks for HTTP_X_FORWARDED_PROTO: https,..." $ do
+            let req = addHeader "HTTP_X_FORWARDED_PROTO" "https,xyz" insecureRequest
+
+            req `shouldSatisfy` appearsSecure
 
 addHeader :: HeaderName -> ByteString -> Request -> Request
 addHeader name value req = req

wai-extra/wai-extra.cabal

@@ -151,6 +151,7 @@ Library
                      Network.Wai.EventSource
                      Network.Wai.EventSource.EventStream
   other-modules:     Network.Wai.Middleware.RequestLogger.Internal
+  default-language:          Haskell2010
   ghc-options:       -Wall
 
 test-suite spec