@@ -12,7 +12,6 @@ module Network.Wai.Parse
1212 , RequestBodyType (..)
1313 , getRequestBodyType
1414 , sinkRequestBody
15− , conduitRequestBody
1615 , BackEnd
1716 , lbsBackEnd
1817 , tempFileBackEnd
@@ -38,22 +37,14 @@ import qualified Data.ByteString.Char8 as S8
3837import Data.Word (Word8)
3938import Data.Maybe (fromMaybe)
4039import Data.List (sortBy)
41−import Data.Function (on)
40+import Data.Function (on, fix)
4241import System.Directory (removeFile, getTemporaryDirectory)
4342import System.IO (hClose, openBinaryTempFile)
4443import Network.Wai
45−import Data.Conduit
46−import Data.Conduit.Internal ()
47−import qualified Data.Conduit.List as CL
48−import qualified Data.Conduit.Binary as CB
49−import Control.Monad.IO.Class (liftIO)
5044import qualified Network.HTTP.Types as H
51−import Data.Either (partitionEithers)
5245import Control.Monad (when, unless)
53−import Control.Monad.Trans.Class (lift)
5446import Control.Monad.Trans.Resource (allocate, release, register, InternalState, runInternalState)
55−import Data.Conduit.Internal (Pipe (NeedInput, HaveOutput), (>+>), withUpstream, injectLeftovers, ConduitM (..))
56−import Data.Void (Void)
47+import Data.IORef
5748
5849breakDiscard :: Word8 -> S.ByteString -> (S.ByteString, S.ByteString)
5950breakDiscard w s =
@@ -85,15 +76,22 @@ parseHttpAccept = map fst
8576 _ -> 1.0
8677
8778-- | Store uploaded files in memory
88−lbsBackEnd :: Monad m => ignored1 -> ignored2 -> Sink S.ByteString m L.ByteString
89−lbsBackEnd _ _ = fmap L.fromChunks CL.consume
79+lbsBackEnd :: Monad m => ignored1 -> ignored2 -> m S.ByteString -> m L.ByteString
80+lbsBackEnd _ _ popper =
81+ loop id
82+ where
83+ loop front = do
84+ bs <- popper
85+ if S.null bs
86+ then return $ L.fromChunks $ front []
87+ else loop $ front . (bs:)
9088
9189-- | Save uploaded files on disk as temporary files
9290--
9391-- Note: starting with version 2.0, removal of temp files is registered with
9492-- the provided @InternalState@. It is the responsibility of the caller to
9593-- ensure that this @InternalState@ gets cleaned up.
96−tempFileBackEnd :: InternalState -> ignored1 -> ignored2 -> Sink S.ByteString IO FilePath
94+tempFileBackEnd :: InternalState -> ignored1 -> ignored2 -> IO S.ByteString -> IO FilePath
9795tempFileBackEnd = tempFileBackEndOpts getTemporaryDirectory "webenc.buf"
9896
9997-- | Same as 'tempFileSink', but use configurable temp folders and patterns.
@@ -102,14 +100,19 @@ tempFileBackEndOpts :: IO FilePath -- ^ get temporary directory
102100 -> InternalState
103101 -> ignored1
104102 -> ignored2
105− -> Sink S.ByteString IO FilePath
106−tempFileBackEndOpts getTmpDir pattern internalState _ _ = do
103+ -> IO S.ByteString
104+ -> IO FilePath
105+tempFileBackEndOpts getTmpDir pattern internalState _ _ popper = do
107106 (key, (fp, h)) <- flip runInternalState internalState $ allocate (do
108107 tempDir <- getTmpDir
109108 openBinaryTempFile tempDir pattern) (\(_, h) -> hClose h)
110109 _ <- runInternalState (register $ removeFile fp) internalState
111− CB.sinkHandle h
112− lift $ release key
110+ fix $ \loop -> do
111+ bs <- popper
112+ unless (S.null bs) $ do
113+ S.hPut h bs
114+ loop
115+ release key
113116 return fp
114117
115118-- | Information on an uploaded file.
@@ -126,11 +129,12 @@ type Param = (S.ByteString, S.ByteString)
126129-- | Post parameter name and associated file information.
127130type File y = (S.ByteString, FileInfo y)
128131
129−-- | A file uploading backend. Takes the parameter name, file name, and content
130−-- type, and returns a `Sink` for storing the contents.
132+-- | A file uploading backend. Takes the parameter name, file name, and a
133+-- stream of data.
131134type BackEnd a = S.ByteString -- ^ parameter name
132135 -> FileInfo ()
133− -> Sink S.ByteString IO a
136+ -> IO S.ByteString
137+ -> IO a
134138
135139data RequestBodyType = UrlEncoded | Multipart S.ByteString
136140
@@ -173,60 +177,99 @@ parseRequestBody :: BackEnd y
173177parseRequestBody s r =
174178 case getRequestBodyType r of
175179 Nothing -> return ([], [])
176− Just rbt -> fmap partitionEithers $ requestBody r $$ conduitRequestBody s rbt =$ CL.consume
180+ Just rbt -> sinkRequestBody s rbt (requestBody r)
177181
178182sinkRequestBody :: BackEnd y
179183 -> RequestBodyType
180− -> Sink S.ByteString IO ([Param], [File y])
181−sinkRequestBody s r = fmap partitionEithers $ conduitRequestBody s r =$ CL.consume
184+ -> IO S.ByteString
185+ -> IO ([Param], [File y])
186+sinkRequestBody s r body = do
187+ ref <- newIORef (id, id)
188+ let add x = atomicModifyIORef ref $ \(y, z) ->
189+ case x of
190+ Left y' -> ((y . (y':), z), ())
191+ Right z' -> ((y, z . (z':)), ())
192+ conduitRequestBody s r body add
193+ (x, y) <- readIORef ref
194+ return (x [], y [])
182195
183196conduitRequestBody :: BackEnd y
184197 -> RequestBodyType
185− -> Conduit S.ByteString IO (Either Param (File y))
186−conduitRequestBody _ UrlEncoded = do
198+ -> IO S.ByteString
199+ -> (Either Param (File y) -> IO ())
200+ -> IO ()
201+conduitRequestBody _ UrlEncoded rbody add = do
187202 -- NOTE: in general, url-encoded data will be in a single chunk.
188203 -- Therefore, I'm optimizing for the usual case by sticking with
189204 -- strict byte strings here.
190− bs <- CL.consume
191− mapM_ yield $ map Left $ H.parseSimpleQuery $ S.concat bs
192−conduitRequestBody backend (Multipart bound) =
193− parsePieces backend $ S8.pack "--" `S.append` bound
194−
195−takeLine :: Monad m => Consumer S.ByteString m (Maybe S.ByteString)
196−takeLine =
205+ let loop front = do
206+ bs <- rbody
207+ if S.null bs
208+ then return $ S.concat $ front []
209+ else loop $ front . (bs:)
210+ bs <- loop id
211+ mapM_ (add . Left) $ H.parseSimpleQuery bs
212+conduitRequestBody backend (Multipart bound) rbody add =
213+ parsePieces backend (S8.pack "--" `S.append` bound) rbody add
214+
215+takeLine :: Source -> IO (Maybe S.ByteString)
216+takeLine src =
197217 go id
198218 where
199− go front = await >>= maybe (close front) (push front)
219+ go front = do
220+ bs <- readSource src
221+ if S.null bs
222+ then close front
223+ else push front bs
200224
201− close front = leftover (front S.empty) >> return Nothing
225+ close front = leftover src (front S.empty) >> return Nothing
202226 push front bs = do
203227 let (x, y) = S.break (== 10) $ front bs -- LF
204228 in if S.null y
205229 then go $ S.append x
206230 else do
207− when (S.length y > 1) $ leftover $ S.drop 1 y
231+ when (S.length y > 1) $ leftover src $ S.drop 1 y
208232 return $ Just $ killCR x
209233
210−takeLines :: Consumer S.ByteString IO [S.ByteString]
211−takeLines = do
212− res <- takeLine
234+takeLines :: Source -> IO [S.ByteString]
235+takeLines src = do
236+ res <- takeLine src
213237 case res of
214238 Nothing -> return []
215239 Just l
216240 | S.null l -> return []
217241 | otherwise -> do
218− ls <- takeLines
242+ ls <- takeLines src
219243 return $ l : ls
220244
245+data Source = Source (IO S.ByteString) (IORef S.ByteString)
246+
247+mkSource :: IO S.ByteString -> IO Source
248+mkSource f = do
249+ ref <- newIORef S.empty
250+ return $ Source f ref
251+
252+readSource :: Source -> IO S.ByteString
253+readSource (Source f ref) = do
254+ bs <- atomicModifyIORef ref $ \bs -> (S.empty, bs)
255+ if S.null bs
256+ then f
257+ else return bs
258+
259+leftover :: Source -> S.ByteString -> IO ()
260+leftover (Source _ ref) bs = writeIORef ref bs
261+
221262parsePieces :: BackEnd y
222263 -> S.ByteString
223− -> ConduitM S.ByteString (Either Param (File y)) IO ()
224−parsePieces sink bound =
225− loop
264+ -> IO S.ByteString
265+ -> (Either Param (File y) -> IO ())
266+ -> IO ()
267+parsePieces sink bound rbody add =
268+ mkSource rbody >>= loop
226269 where
227− loop = do
228− _boundLine <- takeLine
229− res' <- takeLines
270+ loop src = do
271+ _boundLine <- takeLine src
272+ res' <- takeLines src
230273 unless (null res') $ do
231274 let ls' = map parsePair res'
232275 let x = do
@@ -239,23 +282,23 @@ parsePieces sink bound =
239282 Just (mct, name, Just filename) -> do
240283 let ct = fromMaybe "application/octet-stream" mct
241284 fi0 = FileInfo filename ct ()
242− (wasFound, y) <- sinkTillBound' bound name fi0 sink
243− yield $ Right (name, fi0 { fileContent = y })
244− when wasFound loop
285+ (wasFound, y) <- sinkTillBound' bound name fi0 sink src
286+ add $ Right (name, fi0 { fileContent = y })
287+ when wasFound (loop src)
245288 Just (_ct, name, Nothing) -> do
246289 let seed = id
247290 let iter front bs = return $ front . (:) bs
248− (wasFound, front) <- sinkTillBound bound iter seed
291+ (wasFound, front) <- sinkTillBound bound iter seed src
249292 let bs = S.concat $ front []
250293 let x' = (name, bs)
251− yield $ Left x'
252− when wasFound loop
294+ add $ Left x'
295+ when wasFound (loop src)
253296 _ -> do
254297 -- ignore this part
255298 let seed = ()
256299 iter () _ = return ()
257− (wasFound, ()) <- sinkTillBound bound iter seed
258− when wasFound loop
300+ (wasFound, ()) <- sinkTillBound bound iter seed src
301+ when wasFound (loop src)
259302 where
260303 contDisp = S8.pack "Content-Disposition"
261304 contType = S8.pack "Content-Type"
@@ -295,64 +338,78 @@ sinkTillBound' :: S.ByteString
295338 -> S.ByteString
296339 -> FileInfo ()
297340 -> BackEnd y
298− -> ConduitM S.ByteString o IO (Bool, y)
299−sinkTillBound' bound name fi sink =
300− ConduitM $ anyOutput $
301− conduitTillBound bound >+> withUpstream (fix $ sink name fi)
302− where
303− fix :: Sink S8.ByteString IO y -> Pipe Void S8.ByteString Void Bool IO y
304− fix p = ignoreTerm >+> injectLeftovers (unConduitM p)
305− ignoreTerm = await' >>= maybe (return ()) (\x -> yield' x >> ignoreTerm)
306− await' = NeedInput (return . Just) (const $ return Nothing)
307− yield' = HaveOutput (return ()) (return ())
308−
309− anyOutput p = p >+> dropInput
310− dropInput = NeedInput (const dropInput) return
311−
312−conduitTillBound :: Monad m
313− => S.ByteString -- bound
314− -> Pipe S.ByteString S.ByteString S.ByteString () m Bool
315−conduitTillBound bound =
316− unConduitM $
317− go id
341+ -> Source
342+ -> IO (Bool, y)
343+sinkTillBound' bound name fi sink src = do
344+ (next, final) <- wrapTillBound bound src
345+ y <- sink name fi next
346+ b <- final
347+ return (b, y)
348+
349+data WTB = WTBWorking (S.ByteString -> S.ByteString)
350+ | WTBDone Bool
351+wrapTillBound :: S.ByteString -- ^ bound
352+ -> Source
353+ -> IO (IO S.ByteString, IO Bool) -- ^ Bool indicates if the bound was found
354+wrapTillBound bound src = do
355+ ref <- newIORef $ WTBWorking id
356+ return (go ref, final ref)
318357 where
319− go front = await >>= maybe (close front) (push front)
320− close front = do
321− let bs = front S.empty
322− unless (S.null bs) $ yield bs
323− return False
324− push front bs' = do
325− let bs = front bs'
326− case findBound bound bs of
327− FoundBound before after -> do
328− let before' = killCRLF before
329− yield before'
330− leftover after
331− return True
332− NoBound -> do
333− -- don't emit newlines, in case it's part of a bound
334− let (toEmit, front') =
335− if not (S8.null bs) && S8.last bs `elem` "\r\n"
336− then let (x, y) = S.splitAt (S.length bs - 2) bs
337− in (x, S.append y)
338− else (bs, id)
339− yield toEmit
340− go front'
341− PartialBound -> go $ S.append bs
358+ final ref = do
359+ x <- readIORef ref
360+ case x of
361+ WTBWorking _ -> error "wrapTillBound did not finish"
362+ WTBDone y -> return y
363+
364+ go ref = do
365+ state <- readIORef ref
366+ case state of
367+ WTBDone _ -> return S.empty
368+ WTBWorking front -> do
369+ bs <- readSource src
370+ if S.null bs
371+ then do
372+ writeIORef ref $ WTBDone False
373+ return $ front bs
374+ else push $ front bs
375+ where
376+ push bs =
377+ case findBound bound bs of
378+ FoundBound before after -> do
379+ let before' = killCRLF before
380+ leftover src after
381+ writeIORef ref $ WTBDone True
382+ return before'
383+ NoBound -> do
384+ -- don't emit newlines, in case it's part of a bound
385+ let (toEmit, front') =
386+ if not (S8.null bs) && S8.last bs `elem` "\r\n"
387+ then let (x, y) = S.splitAt (S.length bs - 2) bs
388+ in (x, S.append y)
389+ else (bs, id)
390+ writeIORef ref $ WTBWorking front'
391+ if S.null toEmit
392+ then go ref
393+ else return toEmit
394+ PartialBound -> do
395+ writeIORef ref $ WTBWorking $ S.append bs
396+ go ref
342397
343398sinkTillBound :: S.ByteString
344399 -> (x -> S.ByteString -> IO x)
345400 -> x
346− -> Consumer S.ByteString IO (Bool, x)
347−sinkTillBound bound iter seed0 =
348− ConduitM $
349− (conduitTillBound bound >+> (withUpstream $ ij $ CL.foldM iter' seed0))
350− where
351− iter' a b = liftIO $ iter a b
352− ij (ConduitM p) = ignoreTerm >+> injectLeftovers p
353− ignoreTerm = await' >>= maybe (return ()) (\x -> yield' x >> ignoreTerm)
354− await' = NeedInput (return . Just) (const $ return Nothing)
355− yield' = HaveOutput (return ()) (return ())
401+ -> Source
402+ -> IO (Bool, x)
403+sinkTillBound bound iter seed0 src = do
404+ (next, final) <- wrapTillBound bound src
405+ let loop seed = do
406+ bs <- next
407+ if S.null bs
408+ then return seed
409+ else iter seed bs >>= loop
410+ seed <- loop seed0
411+ b <- final
412+ return (b, seed)
356413
357414parseAttrs :: S.ByteString -> [(S.ByteString, S.ByteString)]
358415parseAttrs = map go . S.split 59 -- semicolon