@@ -19,6 +19,8 @@ module Network.Wai.Parse
1919 , File
2020 , FileInfo (..)
2121 , parseContentType
22+ , ParseRequestBodyOptions (..)
23+ , parseRequestBodyEx
2224#if TEST
2325 , Bound (..)
2426 , findBound
@@ -33,7 +35,9 @@ import qualified Data.ByteString.Search as Search
3335import qualified Data.ByteString as S
3436import qualified Data.ByteString.Lazy as L
3537import qualified Data.ByteString.Char8 as S8
38+import Data.Default (Default, def)
3639import Data.Word (Word8)
40+import Data.Int (Int64)
3741import Data.Maybe (fromMaybe)
3842import Data.List (sortBy)
3943import Data.Function (on, fix)
@@ -47,6 +51,8 @@ import Data.IORef
4751import Network.HTTP.Types (hContentType)
4852import Data.CaseInsensitive (mk)
4953
54+import Debug.Trace
55+
5056breakDiscard :: Word8 -> S.ByteString -> (S.ByteString, S.ByteString)
5157breakDiscard w s =
5258 let (x, y) = S.break (== w) s
@@ -103,10 +109,8 @@ tempFileBackEndOpts :: IO FilePath -- ^ get temporary directory
103109 -> ignored2
104110 -> IO S.ByteString
105111 -> IO FilePath
106−tempFileBackEndOpts getTmpDir pattern internalState _ _ popper = do
107− (key, (fp, h)) <- flip runInternalState internalState $ allocate (do
108− tempDir <- getTmpDir
109− openBinaryTempFile tempDir pattern) (\(_, h) -> hClose h)
112+tempFileBackEndOpts getTmpDir pattrn internalState _ _ popper = do
113+ (key, (fp, h)) <- flip runInternalState internalState $ allocate it (hClose . snd)
110114 _ <- runInternalState (register $ removeFile fp) internalState
111115 fix $ \loop -> do
112116 bs <- popper
@@ -115,6 +119,35 @@ tempFileBackEndOpts getTmpDir pattern internalState _ _ popper = do
115119 loop
116120 release key
117121 return fp
122+ where
123+ it = do
124+ tempDir <- getTmpDir
125+ openBinaryTempFile tempDir pattrn
126+
127+-- | A data structure that describes the behavior of
128+-- the parseRequestBodyEx function.
129+data ParseRequestBodyOptions = ParseRequestBodyOptions
130+ { -- | The maximum length of a filename
131+ prboKeyLength :: Int
132+ , -- | The maximum number of files.
133+ prboMaxNumFiles :: Int
134+ , -- | The maximum filesize per file.
135+ prboMaxFileSize :: Maybe Int64
136+ , -- | The maximum total filesize.
137+ prboMaxFilesSize :: Maybe Int64
138+ , -- | The maximum size of the sum of all parameters values
139+ prboMaxParmsValueSize :: Int
140+ , -- | The maximum header lines per entry
141+ prboMaxHeaderLines :: Int }
142+
143+instance Default ParseRequestBodyOptions where
144+ def = ParseRequestBodyOptions
145+ { prboKeyLength=32
146+ , prboMaxNumFiles=10
147+ , prboMaxFileSize=Nothing
148+ , prboMaxFilesSize=Nothing
149+ , prboMaxParmsValueSize=65336
150+ , prboMaxHeaderLines=32 }
118151
119152-- | Information on an uploaded file.
120153data FileInfo c = FileInfo
@@ -180,6 +213,15 @@ parseRequestBody s r =
180213 Nothing -> return ([], [])
181214 Just rbt -> sinkRequestBody s rbt (requestBody r)
182215
216+parseRequestBodyEx :: ParseRequestBodyOptions
217+ -> BackEnd y
218+ -> Request
219+ -> IO ([Param], [File y])
220+parseRequestBodyEx o s r =
221+ case getRequestBodyType r of
222+ Nothing -> return ([], [])
223+ Just rbt -> sinkRequestBodyEx o s rbt (requestBody r)
224+
183225sinkRequestBody :: BackEnd y
184226 -> RequestBodyType
185227 -> IO S.ByteString
@@ -194,6 +236,20 @@ sinkRequestBody s r body = do
194236 (x, y) <- readIORef ref
195237 return (x [], y [])
196238
239+sinkRequestBodyEx :: ParseRequestBodyOptions
240+ -> BackEnd y
241+ -> RequestBodyType
242+ -> IO S.ByteString
243+ -> IO ([Param], [File y])
244+sinkRequestBodyEx o s r body = do
245+ ref <- newIORef ([], [])
246+ let add x = atomicModifyIORef ref $ \(y, z) ->
247+ case x of
248+ Left y' -> ((y':y, z), ())
249+ Right z' -> ((y, z':z), ())
250+ conduitRequestBodyEx o s r body add
251+ readIORef ref
252+
197253conduitRequestBody :: BackEnd y
198254 -> RequestBodyType
199255 -> IO S.ByteString
@@ -213,6 +269,27 @@ conduitRequestBody _ UrlEncoded rbody add = do
213269conduitRequestBody backend (Multipart bound) rbody add =
214270 parsePieces backend (S8.pack "--" `S.append` bound) rbody add
215271
272+conduitRequestBodyEx :: ParseRequestBodyOptions
273+ -> BackEnd y
274+ -> RequestBodyType
275+ -> IO S.ByteString
276+ -> (Either Param (File y) -> IO ())
277+ -> IO ()
278+conduitRequestBodyEx o _ UrlEncoded rbody add = do
279+ -- NOTE: in general, url-encoded data will be in a single chunk.
280+ -- Therefore, I'm optimizing for the usual case by sticking with
281+ -- strict byte strings here.
282+ let loop front = do
283+ bs <- rbody
284+ if S.null bs
285+ then return $ S.concat $ front []
286+ else loop $ front . (bs:)
287+ bs <- loop id
288+ mapM_ (add . Left) $ H.parseSimpleQuery bs
289+conduitRequestBodyEx o backend (Multipart bound) rbody add =
290+ parsePiecesEx o backend (S8.pack "--" `S.append` bound) rbody add
291+
292+
216293takeLine :: Source -> IO (Maybe S.ByteString)
217294takeLine src =
218295 go id
@@ -243,6 +320,20 @@ takeLines src = do
243320 ls <- takeLines src
244321 return $ l : ls
245322
323+takeLines' :: Int -> Source -> IO [S.ByteString]
324+takeLines' a b = reverse <$> takeLines'' [] a b
325+
326+takeLines'' :: [S.ByteString] -> Int -> Source -> IO [S.ByteString]
327+takeLines'' lines maxLines src
328+ | length lines > maxLines = error "Too many lines in mime/multipart header"
329+ | otherwise = do
330+ res <- takeLine src
331+ case res of
332+ Nothing -> return lines
333+ Just l
334+ | S.null l -> return lines
335+ | otherwise -> takeLines'' (l:lines) maxLines src
336+
246337data Source = Source (IO S.ByteString) (IORef S.ByteString)
247338
248339mkSource :: IO S.ByteString -> IO Source
@@ -283,13 +374,13 @@ parsePieces sink bound rbody add =
283374 Just (mct, name, Just filename) -> do
284375 let ct = fromMaybe "application/octet-stream" mct
285376 fi0 = FileInfo filename ct ()
286− (wasFound, y) <- sinkTillBound' bound name fi0 sink src
377+ ((wasFound, _fileSize), y) <- sinkTillBound' bound name fi0 sink src
287378 add $ Right (name, fi0 { fileContent = y })
288379 when wasFound (loop src)
289380 Just (_ct, name, Nothing) -> do
290381 let seed = id
291382 let iter front bs = return $ front . (:) bs
292− (wasFound, front) <- sinkTillBound bound iter seed src
383+ ((wasFound, _fileSize), front) <- sinkTillBound bound iter seed src
293384 let bs = S.concat $ front []
294385 let x' = (name, bs)
295386 add $ Left x'
@@ -298,7 +389,7 @@ parsePieces sink bound rbody add =
298389 -- ignore this part
299390 let seed = ()
300391 iter () _ = return ()
301− (wasFound, ()) <- sinkTillBound bound iter seed src
392+ ((wasFound, _fileSize), ()) <- sinkTillBound bound iter seed src
302393 when wasFound (loop src)
303394 where
304395 contDisp = mk $ S8.pack "Content-Disposition"
@@ -307,6 +398,76 @@ parsePieces sink bound rbody add =
307398 let (x, y) = breakDiscard 58 s -- colon
308399 in (mk $ x, S.dropWhile (== 32) y) -- space
309400
401+parsePiecesEx :: ParseRequestBodyOptions
402+ -> BackEnd y
403+ -> S.ByteString
404+ -> IO S.ByteString
405+ -> (Either Param (File y) -> IO ())
406+ -> IO ()
407+parsePiecesEx o sink bound rbody add =
408+ mkSource rbody >>= loop 0 0 0 0
409+ where
410+ loop :: Int -> Int -> Int -> Int64 -> Source -> IO ()
411+ loop numParms numFiles parmSize filesSize src = do
412+ _boundLine <- takeLine src
413+ res' <- takeLines' (prboMaxHeaderLines o) src
414+ unless (null res') $ do
415+ let ls' = map parsePair res'
416+ let x = do
417+ cd <- lookup contDisp ls'
418+ let ct = lookup contType ls'
419+ let attrs = parseAttrs cd
420+ name <- lookup "name" attrs
421+ return (ct, name, lookup "filename" attrs)
422+ case x of
423+ Just (mct, name, Just filename) -> do
424+ when (S.length name > prboKeyLength o) $
425+ error "Filename is too long"
426+ when (numFiles >= prboMaxNumFiles o) $
427+ error "Maximum number of files exceeded"
428+ let ct = fromMaybe "application/octet-stream" mct
429+ fi0 = FileInfo filename ct ()
430+ ((wasFound, fileSize), y) <- sinkTillBound' bound name fi0 sink src
431+ case prboMaxFileSize o of
432+ Just maxFileSize -> when (fileSize > maxFileSize ) $
433+ error "Maximum file size exceeded"
434+ Nothing -> return ()
435+ let newFilesSize = filesSize + fileSize
436+ case prboMaxFilesSize o of
437+ Just maxFilesSize ->
438+ when (newFilesSize > maxFilesSize) $
439+ error "Maximum size of uploaded files exceeded"
440+ Nothing -> return ()
441+ add $ Right (name, fi0 { fileContent = y })
442+ when wasFound $ loop numParms (numFiles + 1) parmSize newFilesSize src
443+ Just (_ct, name, Nothing) -> do
444+ when (S.length name > prboKeyLength o) $
445+ error "Parameter name is too long"
446+ let seed = id
447+ let iter front bs = return $ front . (:) bs
448+ ((wasFound, _fileSize), front) <- sinkTillBound bound iter seed src
449+ let bs = S.concat $ front []
450+ let x' = (name, bs)
451+ let newParmSize = parmSize + S.length name + S.length bs
452+ when (newParmSize > prboMaxParmsValueSize o) $
453+ error "Maximum size of parameters exceeded"
454+ add $ Left x'
455+ when wasFound $ loop (numParms + 1) numFiles
456+ newParmSize filesSize src
457+ _ -> do
458+ -- ignore this part
459+ let seed = ()
460+ iter () _ = return ()
461+ ((wasFound, _fileSize), ()) <- sinkTillBound bound iter seed src
462+ when wasFound $ loop numParms numFiles parmSize filesSize src
463+ where
464+ contDisp = mk $ S8.pack "Content-Disposition"
465+ contType = mk $ S8.pack "Content-Type"
466+ parsePair s =
467+ let (x, y) = breakDiscard 58 s -- colon
468+ in (mk $ x, S.dropWhile (== 32) y) -- space
469+
470+
310471data Bound = FoundBound S.ByteString S.ByteString
311472 | NoBound
312473 | PartialBound
@@ -340,7 +501,7 @@ sinkTillBound' :: S.ByteString
340501 -> FileInfo ()
341502 -> BackEnd y
342503 -> Source
343− -> IO (Bool, y)
504+ -> IO ((Bool, Int64), y)
344505sinkTillBound' bound name fi sink src = do
345506 (next, final) <- wrapTillBound bound src
346507 y <- sink name fi next
@@ -351,23 +512,27 @@ data WTB = WTBWorking (S.ByteString -> S.ByteString)
351512 | WTBDone Bool
352513wrapTillBound :: S.ByteString -- ^ bound
353514 -> Source
354− -> IO (IO S.ByteString, IO Bool) -- ^ Bool indicates if the bound was found
515+ -> IO (IO S.ByteString, IO (Bool, Int64)) -- ^ Bool indicates if the bound was found
355516wrapTillBound bound src = do
356517 ref <- newIORef $ WTBWorking id
357− return (go ref, final ref)
518+ sref <- newIORef (0 :: Int64)
519+ return (go ref sref, final ref sref)
358520 where
359− final ref = do
521+ final ref sref = do
360522 x <- readIORef ref
361523 case x of
362524 WTBWorking _ -> error "wrapTillBound did not finish"
363− WTBDone y -> return y
525+ WTBDone y -> do
526+ siz <- readIORef sref
527+ return (y, siz)
364528
365− go ref = do
529+ go ref sref = do
366530 state <- readIORef ref
367531 case state of
368532 WTBDone _ -> return S.empty
369533 WTBWorking front -> do
370534 bs <- readSource src
535+ modifyIORef sref $ (+) (fromIntegral $S.length bs)
371536 if S.null bs
372537 then do
373538 writeIORef ref $ WTBDone False
@@ -390,17 +555,17 @@ wrapTillBound bound src = do
390555 else (bs, id)
391556 writeIORef ref $ WTBWorking front'
392557 if S.null toEmit
393− then go ref
558+ then go ref sref
394559 else return toEmit
395560 PartialBound -> do
396561 writeIORef ref $ WTBWorking $ S.append bs
397− go ref
562+ go ref sref
398563
399564sinkTillBound :: S.ByteString
400565 -> (x -> S.ByteString -> IO x)
401566 -> x
402567 -> Source
403− -> IO (Bool, x)
568+ -> IO ((Bool, Int64), x)
404569sinkTillBound bound iter seed0 src = do
405570 (next, final) <- wrapTillBound bound src
406571 let loop seed = do