on master help sign in
All wai-extra tests pass
commitb3cb523f78a70af8155db0c5d5ec6865de3861b6 authorMichael Snoyman <michael@snoyman.com> dateThu 22 May 2014, 17:03 (12 years ago) parent02ae797
1 file changed+125−95
wai-extra/Network/Wai/Parse.hs+125−95
wai-extra/Network/Wai/Parse.hs+125−95
@@ -11,10 +11,7 @@ module Network.Wai.Parse
1111 , parseRequestBody
1212 , RequestBodyType (..)
1313 , getRequestBodyType
14− , sinkRequestBody -- FIXME bad name
15− {- FIXME
16− , conduitRequestBody
17− -}
14+ , sinkRequestBody
1815 , BackEnd
1916 , lbsBackEnd
2017 , tempFileBackEnd
@@ -44,13 +41,9 @@ import Data.Function (on, fix)
4441import System.Directory (removeFile, getTemporaryDirectory)
4542import System.IO (hClose, openBinaryTempFile)
4643import Network.Wai
47−import Control.Monad.IO.Class (liftIO)
4844import qualified Network.HTTP.Types as H
49−import Data.Either (partitionEithers)
5045import Control.Monad (when, unless)
51−import Control.Monad.Trans.Class (lift)
5246import Control.Monad.Trans.Resource (allocate, release, register, InternalState, runInternalState)
53−import Data.Void (Void)
5447import Data.IORef
5548
5649breakDiscard :: Word8 -> S.ByteString -> (S.ByteString, S.ByteString)
@@ -205,52 +198,78 @@ conduitRequestBody :: BackEnd y
205198 -> IO S.ByteString
206199 -> (Either Param (File y) -> IO ())
207200 -> IO ()
208−conduitRequestBody = error "conduitRequestBody"
209−{-
210−conduitRequestBody _ UrlEncoded = do
201+conduitRequestBody _ UrlEncoded rbody add = do
211202 -- NOTE: in general, url-encoded data will be in a single chunk.
212203 -- Therefore, I'm optimizing for the usual case by sticking with
213204 -- strict byte strings here.
214− bs <- CL.consume
215− mapM_ yield $ map Left $ H.parseSimpleQuery $ S.concat bs
216−conduitRequestBody backend (Multipart bound) =
217− parsePieces backend $ S8.pack "--" `S.append` bound
218−
219−takeLine :: Monad m => Consumer S.ByteString m (Maybe S.ByteString)
220−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 =
221217 go id
222218 where
223− 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
224224
225− close front = leftover (front S.empty) >> return Nothing
225+ close front = leftover src (front S.empty) >> return Nothing
226226 push front bs = do
227227 let (x, y) = S.break (== 10) $ front bs -- LF
228228 in if S.null y
229229 then go $ S.append x
230230 else do
231− when (S.length y > 1) $ leftover $ S.drop 1 y
231+ when (S.length y > 1) $ leftover src $ S.drop 1 y
232232 return $ Just $ killCR x
233233
234−takeLines :: Consumer S.ByteString IO [S.ByteString]
235−takeLines = do
236− res <- takeLine
234+takeLines :: Source -> IO [S.ByteString]
235+takeLines src = do
236+ res <- takeLine src
237237 case res of
238238 Nothing -> return []
239239 Just l
240240 | S.null l -> return []
241241 | otherwise -> do
242− ls <- takeLines
242+ ls <- takeLines src
243243 return $ l : ls
244244
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+
245262parsePieces :: BackEnd y
246263 -> S.ByteString
247− -> ConduitM S.ByteString (Either Param (File y)) IO ()
248−parsePieces sink bound =
249− loop
264+ -> IO S.ByteString
265+ -> (Either Param (File y) -> IO ())
266+ -> IO ()
267+parsePieces sink bound rbody add =
268+ mkSource rbody >>= loop
250269 where
251− loop = do
252− _boundLine <- takeLine
253− res' <- takeLines
270+ loop src = do
271+ _boundLine <- takeLine src
272+ res' <- takeLines src
254273 unless (null res') $ do
255274 let ls' = map parsePair res'
256275 let x = do
@@ -263,30 +282,29 @@ parsePieces sink bound =
263282 Just (mct, name, Just filename) -> do
264283 let ct = fromMaybe "application/octet-stream" mct
265284 fi0 = FileInfo filename ct ()
266− (wasFound, y) <- sinkTillBound' bound name fi0 sink
267− yield $ Right (name, fi0 { fileContent = y })
268− when wasFound loop
285+ (wasFound, y) <- sinkTillBound' bound name fi0 sink src
286+ add $ Right (name, fi0 { fileContent = y })
287+ when wasFound (loop src)
269288 Just (_ct, name, Nothing) -> do
270289 let seed = id
271290 let iter front bs = return $ front . (:) bs
272− (wasFound, front) <- sinkTillBound bound iter seed
291+ (wasFound, front) <- sinkTillBound bound iter seed src
273292 let bs = S.concat $ front []
274293 let x' = (name, bs)
275− yield $ Left x'
276− when wasFound loop
294+ add $ Left x'
295+ when wasFound (loop src)
277296 _ -> do
278297 -- ignore this part
279298 let seed = ()
280299 iter () _ = return ()
281− (wasFound, ()) <- sinkTillBound bound iter seed
282− when wasFound loop
300+ (wasFound, ()) <- sinkTillBound bound iter seed src
301+ when wasFound (loop src)
283302 where
284303 contDisp = S8.pack "Content-Disposition"
285304 contType = S8.pack "Content-Type"
286305 parsePair s =
287306 let (x, y) = breakDiscard 58 s -- colon
288307 in (x, S.dropWhile (== 32) y) -- space
289−-}
290308
291309data Bound = FoundBound S.ByteString S.ByteString
292310 | NoBound
@@ -316,70 +334,82 @@ findBound b bs = handleBreak $ Search.breakOn b bs
316334 | S.index b x == S.index bs y = mismatch xs ys
317335 | otherwise = True
318336
319−{- FIXME
320337sinkTillBound' :: S.ByteString
321338 -> S.ByteString
322339 -> FileInfo ()
323340 -> BackEnd y
324− -> ConduitM S.ByteString o IO (Bool, y)
325−sinkTillBound' bound name fi sink =
326− ConduitM $ anyOutput $
327− conduitTillBound bound >+> withUpstream (fix $ sink name fi)
328− where
329− fix :: Sink S8.ByteString IO y -> Pipe Void S8.ByteString Void Bool IO y
330− fix p = ignoreTerm >+> injectLeftovers (unConduitM p)
331− ignoreTerm = await' >>= maybe (return ()) (\x -> yield' x >> ignoreTerm)
332− await' = NeedInput (return . Just) (const $ return Nothing)
333− yield' = HaveOutput (return ()) (return ())
334−
335− anyOutput p = p >+> dropInput
336− dropInput = NeedInput (const dropInput) return
337−
338−conduitTillBound :: Monad m
339− => S.ByteString -- bound
340− -> Pipe S.ByteString S.ByteString S.ByteString () m Bool
341−conduitTillBound bound =
342− unConduitM $
343− 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)
344357 where
345− go front = await >>= maybe (close front) (push front)
346− close front = do
347− let bs = front S.empty
348− unless (S.null bs) $ yield bs
349− return False
350− push front bs' = do
351− let bs = front bs'
352− case findBound bound bs of
353− FoundBound before after -> do
354− let before' = killCRLF before
355− yield before'
356− leftover after
357− return True
358− NoBound -> do
359− -- don't emit newlines, in case it's part of a bound
360− let (toEmit, front') =
361− if not (S8.null bs) && S8.last bs `elem` "\r\n"
362− then let (x, y) = S.splitAt (S.length bs - 2) bs
363− in (x, S.append y)
364− else (bs, id)
365− yield toEmit
366− go front'
367− 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
368397
369398sinkTillBound :: S.ByteString
370399 -> (x -> S.ByteString -> IO x)
371400 -> x
372− -> Consumer S.ByteString IO (Bool, x)
373−sinkTillBound bound iter seed0 =
374− ConduitM $
375− (conduitTillBound bound >+> (withUpstream $ ij $ CL.foldM iter' seed0))
376− where
377− iter' a b = liftIO $ iter a b
378− ij (ConduitM p) = ignoreTerm >+> injectLeftovers p
379− ignoreTerm = await' >>= maybe (return ()) (\x -> yield' x >> ignoreTerm)
380− await' = NeedInput (return . Just) (const $ return Nothing)
381− yield' = HaveOutput (return ()) (return ())
382−-}
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)
383413
384414parseAttrs :: S.ByteString -> [(S.ByteString, S.ByteString)]
385415parseAttrs = map go . S.split 59 -- semicolon