@@ -278,7 +278,7 @@ parseRequest :: Port -> SockAddr
278278 -> C.BufferedSource IO S.ByteString
279279 -> ResourceT IO Request
280280parseRequest port remoteHost' src = do
281− headers' <- takeHeaders src
281+ headers' <- src C.$$ takeHeaders
282282 parseRequest' port headers' remoteHost' src
283283
284284-- FIXME come up with good values here
@@ -291,7 +291,6 @@ data InvalidRequest =
291291 | BadFirstLine String
292292 | NonHttp
293293 | IncompleteHeaders
294− | ConnectionClosedByPeer
295294 | OverLargeHeader
296295 deriving (Show, Typeable, Eq)
297296instance Exception InvalidRequest
@@ -587,73 +586,63 @@ defaultSettings = Settings
587586 go' (Just ThreadKilled) = False
588587 go' _ = True
589588
590−takeHeaders :: C.BufferedSource IO ByteString -> ResourceT IO [ByteString]
591−takeHeaders src = do
592− !x <- forceHead ConnectionClosedByPeer src
593− takeHeaders' 0 id id x src
589+type BSEndo = ByteString -> ByteString
590+type BSEndoList = [ByteString] -> [ByteString]
594591
592+data THStatus = THStatus
593+ {-# UNPACK #-} !Int -- running total byte count
594+ BSEndoList -- previously parsed lines
595+ BSEndo -- bytestrings to be prepended
596+
597+takeHeaders :: C.Sink ByteString IO [ByteString]
598+takeHeaders =
599+ C.sinkState (THStatus 0 id id) takeHeadersPush close
600+ where
601+ close _ = throwIO IncompleteHeaders
595602{-# INLINE takeHeaders #-}
596603
597−takeHeaders' :: Int
598− -> ([ByteString] -> [ByteString])
599− -> ([ByteString] -> [ByteString])
600− -> ByteString
601− -> C.BufferedSource IO ByteString
602− -> ResourceT IO [ByteString]
603−takeHeaders' !len _ _ _ _ | len > maxTotalHeaderLength = throwIO OverLargeHeader
604−takeHeaders' !len !lines !prepend !bs src = do
605− let !bsLen = {-# SCC "takeHeaders'.bsLen" #-} S.length bs
606− !mnl = {-# SCC "takeHeaders'.mnl" #-} S.elemIndex 10 bs
607− case mnl of
608− -- no newline. prepend entire bs to next line
609− !Nothing -> {-# SCC "takeHeaders'.noNewline" #-} do
610− let !len' = len + bsLen
611− !more <- forceHead IncompleteHeaders src
612− takeHeaders' len' lines (prepend . (:) bs) more src
613− Just !nl -> {-# SCC "takeHeaders'.newline" #-} do
614− let !end = nl
615− !start = nl + 1
616− !line = {-# SCC "takeHeaders'.line" #-}
617− if end > 0
618− -- line data included in this chunk
619− then S.concat $! prepend [SU.unsafeTake (checkCR bs end) bs]
620− --then S.concat $! prepend [SU.unsafeTake (end-1) bs]
621− -- no line data in this chunk (all in prepend, or empty line)
622− else S.concat $! prepend []
623− if S.null line
624− -- no more headers
625− then {-# SCC "takeHeaders'.noMoreHeaders" #-} do
626− let !lines' = {-# SCC "takeHeaders'.noMoreHeaders.lines'" #-} lines []
627− if start < bsLen
628− then {-# SCC "takeHeaders'.noMoreHeaders.yield" #-} do
629− let !rest = {-# SCC "takeHeaders'.noMoreHeaders.yield.rest" #-} SU.unsafeDrop start bs
630− C.bsourceUnpull src rest
631− return lines'
632− else return lines'
633−
634− -- more headers
635− else {-# SCC "takeHeaders'.moreHeaders" #-} do
636− let !len' = len + start
637− !lines' = {-# SCC "takeHeaders.lines'" #-} lines . (:) line
638− !more <- {-# SCC "takeHeaders'.more" #-}
639− if start < bsLen
640− then return $! SU.unsafeDrop start bs
641− else forceHead IncompleteHeaders src
642− {-# SCC "takeHeaders'.takeMore" #-} takeHeaders' len' lines' id more src
643−{-# INLINE takeHeaders' #-}
644−
645−forceHead :: InvalidRequest
646− -> C.BufferedSource IO ByteString
647− -> ResourceT IO ByteString
648−forceHead err src = do
649− !mx <- src C.$$ CL.head
650− case mx of
651− Just !x -> return x
652− Nothing -> throwIO err
653−{-# INLINE forceHead #-}
604+takeHeadersPush :: THStatus
605+ -> ByteString
606+ -> ResourceT IO (THStatus, C.SinkResult ByteString [ByteString])
607+takeHeadersPush (THStatus len _ _ ) _
608+ | len > maxTotalHeaderLength = throwIO OverLargeHeader
609+takeHeadersPush (THStatus len lines prepend) bs =
610+ case mnl of
611+ -- no newline. prepend entire bs to next line
612+ Nothing -> do
613+ let len' = len + bsLen
614+ return (THStatus len' lines (prepend . S.append bs), C.Processing)
615+ Just nl -> do
616+ let end = nl
617+ start = nl + 1
618+ line = prepend (if end > 0
619+ then SU.unsafeTake (checkCR bs end) bs
620+ else S.empty)
621+ if S.null line
622+ -- no more headers
623+ then do
624+ let lines' = lines []
625+ if start < bsLen
626+ then do
627+ let rest = SU.unsafeDrop start bs
628+ return (undefined, C.Done (Just rest) lines')
629+ else return (undefined, C.Done Nothing lines')
630+ -- more headers
631+ else do
632+ let len' = len + start
633+ lines' = lines . (:) line
634+ if start < bsLen
635+ then do
636+ let more = SU.unsafeDrop start bs
637+ takeHeadersPush (THStatus len' lines' id) more
638+ else return (THStatus len' lines' id, C.Processing)
639+ where
640+ bsLen = S.length bs
641+ mnl = S.elemIndex 10 bs
642+{-# INLINE takeHeadersPush #-}
654643
655644checkCR :: ByteString -> Int -> Int
656−checkCR bs pos =
645+checkCR bs pos =
657646 let !p = pos - 1
658647 in if '\r' == B.index bs p
659648 then p