a5d5037 · typeable/wai

git.iko.soy/typeable/wai

Chunking

commita5d5037c98c968fd0de05db2305a4e6dfb0c63de authorMichael Snoyman <michael@snoyman.com> dateSun 13 May 2012, 16:53 (14 years ago) parent69f48d3

3 files changed+109−16

compare with master

unifiedsplit

warp/Network/Wai/Handler/Warp.hs+58−15 warp/test/main.hs+51 warp/warp.cabal−1

warp/Network/Wai/Handler/Warp.hs+58−15

@@ -77,12 +77,15 @@ import Control.Exception

7777 )

7878import Control.Concurrent (forkIO)

7979import Data.Maybe (fromMaybe, isJust)

80+import Data.Char (toLower, isHexDigit)

81+import Data.Word (Word)

8082

8183import Data.Typeable (Typeable)

8284

8385import Data.Conduit (ResourceT, runResourceT)

8486import qualified Data.Conduit as C

8587import qualified Data.Conduit.List as CL

88+import qualified Data.Conduit.Binary as CB

8689import Data.Conduit.Blaze (builderToByteString)

8790import Control.Exception.Lifted (throwIO)

8891import Blaze.ByteString.Builder.HTTP

@@ -96,6 +99,7 @@ import Network.Sendfile

9699import qualified System.PosixCompat.Files as P

97100

98101import Control.Monad.IO.Class (MonadIO, liftIO)

102+import Control.Monad.Trans.Class (lift)

99103import qualified Timeout as T

100104import Timeout (Manager, registerKillThread, pause, resume)

101105import Data.Word (Word8)

@@ -107,7 +111,6 @@ import System.IO (hPrint, stderr)

107111import ReadInt (readInt64)

108112import qualified Data.IORef as I

109113import Data.Conduit.Network (bindPort, HostPreference (HostIPv4))

110−import Network.HTTP.Conduit.Chunk (chunkedConduit)

111114

112115#if WINDOWS

113116import Control.Concurrent (threadDelay)

@@ -270,17 +273,6 @@ ibsIsolate ibs@(IsolatedBSSource ref) =

270273ibsDone :: IsolatedBSSource -> IO (C.Source (ResourceT IO) ByteString)

271274ibsDone (IsolatedBSSource ref) = fmap snd $ I.readIORef ref

272275

273−refSource :: MonadIO m => I.IORef (C.Source m a) -> C.Source m a

274−refSource ref = C.PipeM pull (return ())

275− where

276− pull = do

277− src <- liftIO $ I.readIORef ref

278− (src', ma) <- src C.$$+ CL.head

279− liftIO $ I.writeIORef ref src'

280− case ma of

281− Nothing -> return $ C.Done Nothing ()

282− Just a -> return $ C.HaveOutput (refSource ref) (return ()) a

283−

284276serveConnection :: Settings

285277 -> T.Handle

286278 -> (SomeException -> IO ())

@@ -376,13 +368,13 @@ parseRequest' conn port (firstLine:otherLines) remoteHost' src = do

376368 Nothing -> 0

377369 Just bs -> readInt bs

378370 let serverName' = takeUntil 58 host -- ':'

379− let chunked = maybe False (=="chunked")

371+ let chunked = maybe False ((== "chunked") . B.map toLower)

380372 $ lookup "transfer-encoding" heads

381373 (rbody, getSource) <- liftIO $

382374 if chunked

383375 then do

384− ref <- I.newIORef (src C.$= chunkedConduit False)

385− return (refSource ref, I.readIORef ref)

376+ ref <- I.newIORef (src, NeedLen)

377+ return (chunkedSource ref, fmap fst $ I.readIORef ref)

386378 else do

387379 ibs <- fmap IsolatedBSSource $ I.newIORef (len0, src)

388380 return (ibsIsolate ibs, ibsDone ibs)

@@ -403,6 +395,57 @@ parseRequest' conn port (firstLine:otherLines) remoteHost' src = do

403395 , vault = mempty

404396 }, getSource)

405397

398+data ChunkState = NeedLen

399+ | NeedLenNewline

400+ | HaveLen Word

401+

402+chunkedSource :: MonadIO m

403+ => I.IORef (C.Source m ByteString, ChunkState)

404+ -> C.Source m ByteString

405+chunkedSource ipair = do

406+ (src, mlen) <- liftIO $ I.readIORef ipair

407+ go src mlen

408+ where

409+ go src NeedLen = do

410+ (src', len) <- lift $ src C.$$+ getLen

411+ go src' $ HaveLen len

412+ go src NeedLenNewline = do

413+ (src', len) <- lift $ src C.$$+ (CB.take 2 >> getLen)

414+ go src' $ HaveLen len

415+ go src (HaveLen 0) = liftIO $ I.writeIORef ipair (src, HaveLen 0)

416+ go src (HaveLen len) = do

417+ (src', mbs) <- lift $ src C.$$+ CL.head

418+ case mbs of

419+ Nothing -> liftIO $ I.writeIORef ipair (src', HaveLen 0)

420+ Just bs ->

421+ case S.length bs `compare` fromIntegral len of

422+ EQ -> yield' src' NeedLenNewline bs

423+ LT -> do

424+ let mlen = HaveLen $ len - fromIntegral (S.length bs)

425+ yield' src' mlen bs

426+ GT -> do

427+ let (x, y) = S.splitAt (fromIntegral len) bs

428+ let src'' = C.yield y >> src'

429+ yield' src'' NeedLenNewline x

430+

431+ yield' src mlen bs = do

432+ liftIO $ I.writeIORef ipair (src, mlen)

433+ C.yield bs

434+ go src mlen

435+

436+ getLen :: MonadIO m => C.Sink ByteString m Word

437+ getLen = do

438+ bss <- CB.takeWhile (/= 10) C.=$ CL.consume -- FIXME avoid attack

439+ _ <- CB.take 1

440+ return $

441+ S.foldl' (\i c -> i * 16 + fromIntegral (hexToWord c)) 0

442+ $ B.takeWhile isHexDigit

443+ $ B.concat bss

444+

445+ hexToWord w

446+ | w < 58 = w - 48

447+ | w < 71 = w - 55

448+ | otherwise = w - 87

406449

407450takeUntil :: Word8 -> ByteString -> ByteString

408451takeUntil c bs =

warp/test/main.hs+51

@@ -146,3 +146,54 @@ main = hspecX $ do

146146 describe "connection termination" $ do

147147 it "ConnectionClosedByPeer" $ runTerminateTest ConnectionClosedByPeer "GET / HTTP/1.1\r\ncontent-length: 10\r\n\r\nhello"

148148 it "IncompleteHeaders" $ runTerminateTest IncompleteHeaders "GET / HTTP/1.1\r\ncontent-length: 10\r\n"

149+

150+ describe "chunked bodies" $ do

151+ it "works" $ do

152+ ifront <- I.newIORef id

153+ port <- getPort

154+ tid <- forkIO $ run port $ \req -> do

155+ bss <- requestBody req $$ Data.Conduit.List.consume

156+ liftIO $ I.atomicModifyIORef ifront $ \front -> (front . (S.concat bss:), ())

157+ return $ responseLBS status200 [] ""

158+ threadDelay 1000

159+ handle <- connectTo "127.0.0.1" $ PortNumber $ fromIntegral port

160+ let input = S.concat

161+ [ "POST / HTTP/1.1\r\nTransfer-Encoding: Chunked\r\n\r\n"

162+ , "c\r\nHello World\n\r\n3\r\nBye\r\n0\r\n"

163+ , "POST / HTTP/1.1\r\nTransfer-Encoding: Chunked\r\n\r\n"

164+ , "b\r\nHello World\r\n0\r\n"

165+ ]

166+ hPutStr handle input

167+ hFlush handle

168+ hClose handle

169+ threadDelay 1000

170+ killThread tid

171+ front <- I.readIORef ifront

172+ front [] @?=

173+ [ "Hello World\nBye"

174+ , "Hello World"

175+ ]

176+ it "in chunks" $ do

177+ ifront <- I.newIORef id

178+ port <- getPort

179+ tid <- forkIO $ run port $ \req -> do

180+ bss <- requestBody req $$ Data.Conduit.List.consume

181+ liftIO $ I.atomicModifyIORef ifront $ \front -> (front . (S.concat bss:), ())

182+ return $ responseLBS status200 [] ""

183+ threadDelay 1000

184+ handle <- connectTo "127.0.0.1" $ PortNumber $ fromIntegral port

185+ let input = S.concat

186+ [ "POST / HTTP/1.1\r\nTransfer-Encoding: Chunked\r\n\r\n"

187+ , "c\r\nHello World\n\r\n3\r\nBye\r\n0\r\n"

188+ , "POST / HTTP/1.1\r\nTransfer-Encoding: Chunked\r\n\r\n"

189+ , "b\r\nHello World\r\n0\r\n"

190+ ]

191+ mapM_ (\bs -> hPutStr handle bs >> hFlush handle) $ map S.singleton $ S.unpack input

192+ hClose handle

193+ threadDelay 1000

194+ killThread tid

195+ front <- I.readIORef ifront

196+ front [] @?=

197+ [ "Hello World\nBye"

198+ , "Hello World"

199+ ]

warp/warp.cabal−1

@@ -24,7 +24,6 @@ Library

2424 , conduit >= 0.4 && < 0.5

2525 , network-conduit >= 0.4 && < 0.5

2626 , blaze-builder-conduit >= 0.4 && < 0.5

27− , http-conduit >= 1.4 && < 1.5

2827 , lifted-base >= 0.1 && < 0.2

2928 , blaze-builder >= 0.2.1.4 && < 0.4

3029 , simple-sendfile >= 0.2.2 && < 0.3