ab68123 · typeable/wai

git.iko.soy/typeable/wai

Chunking

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

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

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

7878 )

7979import Control.Concurrent (forkIO)

8080import Data.Maybe (fromMaybe, isJust)

81+import Data.Char (toLower, isHexDigit)

82+import Data.Word (Word)

8183

8284import Data.Typeable (Typeable)

8385

8486import Data.Conduit (ResourceT, runResourceT)

8587import qualified Data.Conduit as C

8688import qualified Data.Conduit.List as CL

89+import qualified Data.Conduit.Binary as CB

8790import Data.Conduit.Blaze (builderToByteString)

8891import Control.Exception.Lifted (throwIO)

8992import Blaze.ByteString.Builder.HTTP

@@ -97,6 +100,7 @@ import Network.Sendfile

97100import qualified System.PosixCompat.Files as P

98101

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

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

100104import qualified Timeout as T

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

102106import Data.Word (Word8)

@@ -108,7 +112,6 @@ import System.IO (hPrint, stderr)

108112import ReadInt (readInt64)

109113import qualified Data.IORef as I

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

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

112115

113116#if WINDOWS

114117import Control.Concurrent (threadDelay)

@@ -272,17 +275,6 @@ ibsIsolate ibs@(IsolatedBSSource ref) =

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

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

274277

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

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

277− where

278− pull = do

279− src <- liftIO $ I.readIORef ref

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

281− liftIO $ I.writeIORef ref src'

282− case ma of

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

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

285−

286278serveConnection :: Settings

287279 -> T.Handle

288280 -> Port -> Application -> Connection -> SockAddr-> IO ()

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

373365 Nothing -> 0

374366 Just bs -> readInt bs

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

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

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

377369 $ lookup "transfer-encoding" heads

378370 (rbody, getSource) <- liftIO $

379371 if chunked

380372 then do

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

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

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

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

383375 else do

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

385377 return (ibsIsolate ibs, ibsDone ibs)

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

400392 , vault = mempty

401393 }, getSource)

402394

395+data ChunkState = NeedLen

396+ | NeedLenNewline

397+ | HaveLen Word

398+

399+chunkedSource :: MonadIO m

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

401+ -> C.Source m ByteString

402+chunkedSource ipair = do

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

404+ go src mlen

405+ where

406+ go src NeedLen = do

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

408+ go src' $ HaveLen len

409+ go src NeedLenNewline = do

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

411+ go src' $ HaveLen len

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

413+ go src (HaveLen len) = do

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

415+ case mbs of

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

417+ Just bs ->

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

419+ EQ -> yield' src' NeedLenNewline bs

420+ LT -> do

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

422+ yield' src' mlen bs

423+ GT -> do

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

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

426+ yield' src'' NeedLenNewline x

427+

428+ yield' src mlen bs = do

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

430+ C.yield bs

431+ go src mlen

432+

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

434+ getLen = do

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

436+ _ <- CB.take 1

437+ return $

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

439+ $ B.takeWhile isHexDigit

440+ $ B.concat bss

441+

442+ hexToWord w

443+ | w < 58 = w - 48

444+ | w < 71 = w - 55

445+ | otherwise = w - 87

403446

404447takeUntil :: Word8 -> ByteString -> ByteString

405448takeUntil 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