-- | 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 Network.Wai
import Network.Wai.Request
import qualified Data.ByteString.Lazy as BSL
import Data.Word (Word64)
import Network.HTTP.Types.Status (requestEntityTooLarge413)
import qualified Data.ByteString.Lazy.Char8 as LS8
import Control.Exception (try, catch)
import Data.Monoid ((<>))
import Network.Wai.Middleware.RequestSizeLimit.Internal (RequestSizeLimitSettings(..), setMaxLengthForRequest, setOnLengthExceeded)

-- | 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 :: (Request -> IO (Maybe Word64))
-> (Word64 -> Middleware) -> RequestSizeLimitSettings
RequestSizeLimitSettings
    { maxLengthForRequest :: Request -> IO (Maybe Word64)
maxLengthForRequest = \Request
_req -> Maybe Word64 -> IO (Maybe Word64)
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`."
        ])