aca7bfb · typeable/wai

git.iko.soy/typeable/wai

on master help sign in

Using Matt's new takeHeaders code

aca7bfbMichael Snoyman16 years ago

Network/Wai/Handler/Warp.hs+68−4

compare with master

unifiedsplit

@@ -181,10 +181,11 @@ parseRequest port remoteHost' = do

181 parseRequest' port headers' remoteHost'

182

183-- FIXME come up with good values here

184−maxHeaders, maxHeaderLength, bytesPerRead :: Int

184+maxHeaders, maxHeaderLength, bytesPerRead, maxTotalHeaderLength :: Int

185maxHeaders = 30

186maxHeaderLength = 1024

187bytesPerRead = 4096

188+maxTotalHeaderLength = 50 * 1024

188

189

189sendFileCount :: Integer

190sendFileCount :: Integer

190sendFileCount = 65536

191sendFileCount = 65536

@@ -398,9 +399,6 @@ enumSocket th len socket =

398------ The functions below are not warp-specific and could be split out into a

399------ The functions below are not warp-specific and could be split out into a

399--separate package.

400--separate package.

400

401

401−takeHeaders :: E.Iteratee ByteString IO [ByteString]

402−takeHeaders = takeUntilBlank 0 id

403−

404takeUntilBlank :: Int

402takeUntilBlank :: Int

405 -> ([ByteString] -> [ByteString])

403 -> ([ByteString] -> [ByteString])

406 -> E.Iteratee S.ByteString IO [ByteString]

404 -> E.Iteratee S.ByteString IO [ByteString]

@@ -464,3 +462,69 @@ defaultSettings = Settings

464 where

462 where

465 go :: InvalidRequest -> IO ()

463 go :: InvalidRequest -> IO ()

466 go _ = return ()

464 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 #-}

runtests.hs+1−1

@@ -64,7 +64,7 @@ caseTakeLineMaxIncomplete = do

65caseTakeUntilBlankTooMany = do

66 x <- run $ (enumList 1 $ repeat "f\n") $$ takeHeaders

67− assertException TooManyHeaders x

67+ assertException OverLargeHeader x

69caseTakeUntilBlankTooLarge = do

70 x <- run $ (enumList 1 $ repeat "f") $$ takeHeaders