@@ -181,10 +181,11 @@ parseRequest port remoteHost' = do
181181 parseRequest' port headers' remoteHost'
182182
183183-- FIXME come up with good values here
184−maxHeaders, maxHeaderLength, bytesPerRead :: Int
184+maxHeaders, maxHeaderLength, bytesPerRead, maxTotalHeaderLength :: Int
185185maxHeaders = 30
186186maxHeaderLength = 1024
187187bytesPerRead = 4096
188+maxTotalHeaderLength = 50 * 1024
188189
189190sendFileCount :: Integer
190191sendFileCount = 65536
@@ -398,9 +399,6 @@ enumSocket th len socket =
398399------ The functions below are not warp-specific and could be split out into a
399400--separate package.
400401
401−takeHeaders :: E.Iteratee ByteString IO [ByteString]
402−takeHeaders = takeUntilBlank 0 id
403−
404402takeUntilBlank :: Int
405403 -> ([ByteString] -> [ByteString])
406404 -> E.Iteratee S.ByteString IO [ByteString]
@@ -464,3 +462,69 @@ defaultSettings = Settings
464462 where
465463 go :: InvalidRequest -> IO ()
466464 go _ = return ()
465+
466+takeHeaders = do
467+ !x <- forceHead
468+ takeHeaders' 0 id id x
469+
470+{-# INLINE takeHeaders #-}
471+
472+takeHeaders' :: Int
473+ -> ([ByteString] -> [ByteString])
474+ -> ([ByteString] -> [ByteString])
475+ -> ByteString
476+ -> E.Iteratee S.ByteString IO [ByteString]
477+takeHeaders' !len _ _ _ | len > maxTotalHeaderLength = E.throwError OverLargeHeader
478+takeHeaders' !len !lines !prepend !bs = do
479+ let !bsLen = {-# SCC "takeHeaders'.bsLen" #-} S.length bs
480+ !mnl = {-# SCC "takeHeaders'.mnl" #-} S.elemIndex 10 bs
481+ case mnl of
482+ -- no newline. prepend entire bs to next line
483+ !Nothing -> {-# SCC "takeHeaders'.noNewline" #-} do
484+ let !len' = len + bsLen
485+ !more <- forceHead
486+ takeHeaders' len' lines (prepend . (:) bs) more
487+ Just !nl -> {-# SCC "takeHeaders'.newline" #-} do
488+ let !end = nl
489+ !start = nl + 1
490+ !line = {-# SCC "takeHeaders'.line" #-}
491+ if end > 0
492+ -- line data included in this chunk
493+ then S.concat $! prepend [SU.unsafeTake (checkCR bs end) bs]
494+ --then S.concat $! prepend [SU.unsafeTake (end-1) bs]
495+ -- no line data in this chunk (all in prepend, or empty line)
496+ else S.concat $! prepend []
497+ if S.null line
498+ -- no more headers
499+ then {-# SCC "takeHeaders'.noMoreHeaders" #-} do
500+ let !lines' = {-# SCC "takeHeaders'.noMoreHeaders.lines'" #-} lines []
501+ if start < bsLen
502+ then {-# SCC "takeHeaders'.noMoreHeaders.yield" #-} do
503+ let !rest = {-# SCC "takeHeaders'.noMoreHeaders.yield.rest" #-} SU.unsafeDrop start bs
504+ E.yield lines' $! E.Chunks [rest]
505+ else return lines'
506+
507+ -- more headers
508+ else {-# SCC "takeHeaders'.moreHeaders" #-} do
509+ let !len' = len + start
510+ !lines' = {-# SCC "takeHeaders.lines'" #-} lines . (:) line
511+ !more <- {-# SCC "takeHeaders'.more" #-}
512+ if start < bsLen
513+ then return $! SU.unsafeDrop start bs
514+ else forceHead
515+ {-# SCC "takeHeaders'.takeMore" #-} takeHeaders' len' lines' id more
516+{-# INLINE takeHeaders' #-}
517+
518+forceHead = do
519+ !mx <- E.head
520+ case mx of
521+ !Nothing -> E.throwError IncompleteHeaders
522+ Just !x -> return x
523+{-# INLINE forceHead #-}
524+
525+checkCR bs pos =
526+ let !p = pos - 1
527+ in if '\r' == B.index bs p
528+ then p
529+ else pos
530+{-# INLINE checkCR #-}