{-# LANGUAGE CPP #-}

-- | The functions in this module allow you to limit the total size of incoming request bodies.
--
-- Limiting incoming request body size helps protect your server against denial-of-service (DOS) attacks,
-- in which an attacker sends huge bodies to your server.
module Network.Wai.Middleware.RequestSizeLimit (
    -- * Middleware
    requestSizeLimitMiddleware,

    -- * Constructing 'RequestSizeLimitSettings'
    defaultRequestSizeLimitSettings,

    -- * 'RequestSizeLimitSettings' and accessors
    RequestSizeLimitSettings,
    setMaxLengthForRequest,
    setOnLengthExceeded,
) where

import Control.Exception (catch, try)
import qualified Data.ByteString.Lazy as BSL
import qualified Data.ByteString.Lazy.Char8 as LS8
#if __GLASGOW_HASKELL__ < 804
import Data.Monoid ((<>))
#endif
import Data.Word (Word64)
import Network.HTTP.Types.Status (requestEntityTooLarge413)
import Network.Wai

import Network.Wai.Middleware.RequestSizeLimit.Internal (
    RequestSizeLimitSettings (..),
    setMaxLengthForRequest,
    setOnLengthExceeded,
 )
import Network.Wai.Request

-- | Create a 'RequestSizeLimitSettings' with these settings:
--
-- * 2MB size limit for all requests
-- * When the limit is exceeded, return a plain text response describing the error, with a 413 status code.
--
-- @since 3.1.1
defaultRequestSizeLimitSettings :: RequestSizeLimitSettings
defaultRequestSizeLimitSettings :: RequestSizeLimitSettings
defaultRequestSizeLimitSettings =
    RequestSizeLimitSettings
        { maxLengthForRequest :: Request -> IO (Maybe Word64)
maxLengthForRequest = \Request
_req -> Maybe Word64 -> IO (Maybe Word64)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Word64 -> IO (Maybe Word64))
-> Maybe Word64 -> IO (Maybe Word64)
forall a b. (a -> b) -> a -> b
$ Word64 -> Maybe Word64
forall a. a -> Maybe a
Just (Word64 -> Maybe Word64) -> Word64 -> Maybe Word64
forall a b. (a -> b) -> a -> b
$ Word64
2 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
1024 Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
1024
        , onLengthExceeded :: Word64 -> Middleware
onLengthExceeded = \Word64
maxLen Application
_app Request
req Response -> IO ResponseReceived
sendResponse -> Response -> IO ResponseReceived
sendResponse (Word64 -> RequestBodyLength -> Response
tooLargeResponse Word64
maxLen (Request -> RequestBodyLength
requestBodyLength Request
req))
        }

-- | Middleware to limit request bodies to a certain size.
--
-- This uses 'requestSizeCheck' under the hood; see that function for details.
--
-- @since 3.1.1
requestSizeLimitMiddleware :: RequestSizeLimitSettings -> Middleware
requestSizeLimitMiddleware :: RequestSizeLimitSettings -> Middleware
requestSizeLimitMiddleware RequestSizeLimitSettings
settings Application
app Request
req Response -> IO ResponseReceived
sendResponse = do
    Maybe Word64
maybeMaxLen <- RequestSizeLimitSettings -> Request -> IO (Maybe Word64)
maxLengthForRequest RequestSizeLimitSettings
settings Request
req

    case Maybe Word64
maybeMaxLen of
        Maybe Word64
Nothing -> Application
app Request
req Response -> IO ResponseReceived
sendResponse
        Just Word64
maxLen -> do
            Either RequestSizeException Request
eitherSizeExceptionOrNewReq <- IO Request -> IO (Either RequestSizeException Request)
forall e a. Exception e => IO a -> IO (Either e a)
try (Word64 -> Request -> IO Request
requestSizeCheck Word64
maxLen Request
req)
            case Either RequestSizeException Request
eitherSizeExceptionOrNewReq of
                -- In the case of a known-length request, RequestSizeException will be thrown immediately
                Left (RequestSizeException Word64
_maxLen) -> Word64 -> IO ResponseReceived
handleLengthExceeded Word64
maxLen
                -- In the case of a chunked request (unknown length), RequestSizeException will be thrown during the processing of a body
                Right Request
newReq ->
                    Application
app Request
newReq Response -> IO ResponseReceived
sendResponse IO ResponseReceived
-> (RequestSizeException -> IO ResponseReceived)
-> IO ResponseReceived
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(RequestSizeException Word64
_maxLen) -> Word64 -> IO ResponseReceived
handleLengthExceeded Word64
maxLen
  where
    handleLengthExceeded :: Word64 -> IO ResponseReceived
handleLengthExceeded Word64
maxLen = RequestSizeLimitSettings -> Word64 -> Middleware
onLengthExceeded RequestSizeLimitSettings
settings Word64
maxLen Application
app Request
req Response -> IO ResponseReceived
sendResponse

tooLargeResponse :: Word64 -> RequestBodyLength -> Response
tooLargeResponse :: Word64 -> RequestBodyLength -> Response
tooLargeResponse Word64
maxLen RequestBodyLength
bodyLen =
    Status -> ResponseHeaders -> ByteString -> Response
responseLBS
        Status
requestEntityTooLarge413
        [(HeaderName
"Content-Type", ByteString
"text/plain")]
        ( [ByteString] -> ByteString
BSL.concat
            [ ByteString
"Request body too large to be processed. The maximum size is "
            , [Char] -> ByteString
LS8.pack (Word64 -> [Char]
forall a. Show a => a -> [Char]
show Word64
maxLen)
            , ByteString
" bytes; your request body was "
            , case RequestBodyLength
bodyLen of
                KnownLength Word64
knownLen -> [Char] -> ByteString
LS8.pack (Word64 -> [Char]
forall a. Show a => a -> [Char]
show Word64
knownLen) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
" bytes."
                RequestBodyLength
ChunkedBody -> ByteString
"split into chunks, whose total size is unknown, but exceeded the limit."
            , ByteString
" If you're the developer of this site, you can configure the maximum length with `requestSizeLimitMiddleware`."
            ]
        )