CODEWITHSUNDEEP
haskellscotty
tauli
tauli

Reputation: 1420

How do i limit the size of a file upload in scotty?

I'm currently looking into Scotty for web development, and so far it looks pretty good. I'm worried though, that there seems to be no way to discard a file upload (or better yet an arbitrary POST body) where the file size is above a certain limit without receiving the whole file first. The example at https://github.com/scotty-web/scotty/blob/master/examples/upload.hs doesn't mention file size limits and i can't find anything in the documentation.

I could of course do a length on the ByteString, but i can't see how that would work until the whole file is already loaded into memory.

Upvotes: 1

Views: 615

Answers (2)

user933161
user933161

Reputation:

From https://github.com/scotty-web/scotty/issues/203

As a workaround I prevent Scotty from parsing the body by putting away the Content-Type header:

{-# LANGUAGE OverloadedStrings #-}

module Main
  ( main
  ) where

import Control.Exception (bracket)
import Control.Exception.Base (catch, throwIO)
import Control.Monad.Trans (liftIO)
import qualified Data.ByteString as BS
import Data.CaseInsensitive (CI)
import Network.HTTP.Types.Header (hContentType)
import Network.Wai (Middleware, Request, requestHeaders)
import Network.Wai.Parse
       (BackEnd, FileInfo(..), getRequestBodyType, parseRequestBody)
import System.FilePath ((</>))
import System.IO (hClose)
import System.IO.Error (isDoesNotExistError)
import System.Posix.Files (removeLink)
import System.Posix.Temp (mkstemp)
import Web.Scotty

data UploadState = UploadState
  { size :: !Int
  }

removeIfExists :: FilePath -> IO ()
removeIfExists path = removeLink path `catch` handleExists
  where
    handleExists e
      | isDoesNotExistError e = return ()
      | otherwise = throwIO e

fileBackend :: BackEnd UploadState
fileBackend _ (FileInfo _fname _cntType ()) reader = bracket start stop work
  where
    st0 = UploadState {size = 0}
    start = mkstemp ("uploads" </> "tmp-")
    stop (p, h) = do
      hClose h
      removeIfExists p
    work (_p, h) = do
      st <- loop h st0
      return st
    loop h st = do
      bs <- reader
      if BS.null bs
        then return st
        else do
          BS.hPut h bs
          loop h st {size = size st + BS.length bs}

scottyHack :: Middleware
scottyHack app req resp =
  case getRequestBodyType req of
    Nothing -> app req resp
    Just _ -> app (fixRequest req) resp

xContentType :: CI BS.ByteString
xContentType = "X-Content-Type"

fixRequest :: Request -> Request
fixRequest req = req {requestHeaders = map putaway $ requestHeaders req}
  where
    putaway (h, v) =
      if h == hContentType
        then (xContentType, v)
        else (h, v)

unFixRequest :: Request -> Request
unFixRequest req = req {requestHeaders = map putback $ requestHeaders req}
  where
    putback (h, v) =
      if h == xContentType
        then (hContentType, v)
        else (h, v)

main :: IO ()
main =
  scotty 3000 $ do
    middleware scottyHack
    post "/upload" $ do
      req <- request
      (_, docs) <- liftIO $ parseRequestBody fileBackend (unFixRequest req)
      json $ map (size . fileContent . snd) docs

Upvotes: 0

Benjamin Kovach
Benjamin Kovach

Reputation: 3260

You should be able to set some maxBytes parameter, take maxBytes from each file contents lazily, partition your file uploads into failures and successes, then handle each of them. Here's some untested code to illustrate what I mean in the context of your application:

post "/upload" $ do
 fs <- files
 let maxBytes = 9000 -- etc
     fs' = [ (fieldName, BS.unpack (fileName fi), B.take (maxBytes + 1) (fileContent fi)) | (fieldName,fi) <- fs ]
     (oks, fails) = partition ((<= maxBytes) . B.length) fs' -- separate out failures
 liftIO $ sequence_ [ B.writeFile ("uploads" </> fn) fc | (_,fn,fc) <- oks ]
 -- do something with 'fails'
 -- and continue...

It's also entirely possible to just filter out failures "on the fly" but that solution is more specific to what you want to do with the failures -- this should illustrate the idea though. This solution should take care of your concerns; since you're using lazy ByteStrings, B.take shouldn't have to read in the full contents of any of the files to be tagged as a failed upload.

Upvotes: 7

Related Questions