on master help sign in
Simplified takeHeaders
commitb408c0cb6ab0b6b532e91c8c0c4c50b1d0d0c5ea authorMichael Snoyman <michael@snoyman.com> dateSun 25 Mar 2012, 09:42 (15 years ago) parent6715687
1 file changed+46−42
warp/Network/Wai/Handler/Warp.hs+46−42
warp/Network/Wai/Handler/Warp.hs+46−42
@@ -626,51 +626,55 @@ data THStatus = THStatus
626626
627627takeHeaders :: C.Sink ByteString (ResourceT IO) [ByteString]
628628takeHeaders =
629− C.sinkState (THStatus 0 id id) takeHeadersPush close
629+ C.Processing (push (THStatus 0 id id)) close
630630 where
631− close _ = throwIO IncompleteHeaders
631+ close = throwIO IncompleteHeaders
632+
633+ push (THStatus len lines prepend) bs
634+ -- Too many bytes
635+ | len > maxTotalHeaderLength = throwIO OverLargeHeader
636+ | otherwise =
637+ case mnl of
638+ -- No newline find in this chunk. Add it to the prepend,
639+ -- update the length, and continue processing.
640+ Nothing ->
641+ let len' = len + bsLen
642+ prepend' = prepend . S.append bs
643+ status = THStatus len' lines prepend'
644+ in C.Processing (push status) close
645+ -- Found a newline at position end.
646+ Just end ->
647+ let start = end + 1 -- start of next chunk
648+ line
649+ -- There were some bytes before the newline, get them
650+ | end > 0 = prepend $ SU.unsafeTake (checkCR bs end) bs
651+ -- No bytes before the newline
652+ | otherwise = prepend S.empty
653+ in if S.null line
654+ -- no more headers
655+ then
656+ let lines' = lines []
657+ -- leftover
658+ rest = if start < bsLen
659+ then Just (SU.unsafeDrop start bs)
660+ else Nothing
661+ in C.Done rest lines'
662+ -- more headers
663+ else
664+ let len' = len + start
665+ lines' = lines . (line:)
666+ status = THStatus len' lines' id
667+ in if start < bsLen
668+ -- more bytes in this chunk, push again
669+ then let bs' = SU.unsafeDrop start bs
670+ in push status bs'
671+ -- no more bytes in this chunk, ask for more
672+ else C.Processing (push status) close
673+ where
674+ bsLen = S.length bs
675+ mnl = S.elemIndex 10 bs
632676{-# INLINE takeHeaders #-}
633677
634−takeHeadersPush :: THStatus
635− -> ByteString
636− -> ResourceT IO (C.SinkStateResult THStatus ByteString [ByteString])
637−takeHeadersPush (THStatus len _ _ ) _
638− | len > maxTotalHeaderLength = throwIO OverLargeHeader
639−takeHeadersPush (THStatus len lines prepend) bs =
640− case mnl of
641− -- no newline. prepend entire bs to next line
642− Nothing -> do
643− let len' = len + bsLen
644− return $ C.StateProcessing $ THStatus len' lines (prepend . S.append bs)
645− Just nl -> do
646− let end = nl
647− start = nl + 1
648− line = prepend (if end > 0
649− then SU.unsafeTake (checkCR bs end) bs
650− else S.empty)
651− if S.null line
652− -- no more headers
653− then do
654− let lines' = lines []
655− if start < bsLen
656− then do
657− let rest = SU.unsafeDrop start bs
658− return $ C.StateDone (Just rest) lines'
659− else return $ C.StateDone Nothing lines'
660− -- more headers
661− else do
662− let len' = len + start
663− lines' = lines . (:) line
664− if start < bsLen
665− then do
666− let more = SU.unsafeDrop start bs
667− takeHeadersPush (THStatus len' lines' id) more
668− else return $ C.StateProcessing $ THStatus len' lines' id
669− where
670− bsLen = S.length bs
671− mnl = S.elemIndex 10 bs
672−{-# INLINE takeHeadersPush #-}
673−
674678checkCR :: ByteString -> Int -> Int
675679checkCR bs pos =
676680 let !p = pos - 1