module Network.Wai.Handler.Warp.Conduit where import Control.Applicative import Control.Exception import Control.Monad (when, unless) import Data.ByteString (ByteString) import Data.ByteString.Lazy.Char8 (pack) import qualified Data.ByteString as S import qualified Data.ByteString.Char8 as S8 import qualified Data.ByteString.Lazy as L import qualified Data.IORef as I import Data.Word (Word, Word8) import Network.Wai.Handler.Warp.Types ---------------------------------------------------------------- -- | Contains a @Source@ and a byte count that is still to be read in. data ISource = ISource !Source !(I.IORef Int) mkISource :: Source -> Int -> IO ISource mkISource src cnt = do ref <- I.newIORef cnt return $! ISource src ref -- | Given an @IsolatedBSSource@ provide a @Source@ that only allows up to the -- specified number of bytes to be passed downstream. All leftovers should be -- retained within the @Source@. If there are not enough bytes available, -- throws a @ConnectionClosedByPeer@ exception. readISource :: ISource -> IO ByteString readISource (ISource src ref) = do count <- I.readIORef ref if count == 0 then return S.empty else do bs <- readSource src -- If no chunk available, then there aren't enough bytes in the -- stream. Throw a ConnectionClosedByPeer when (S.null bs) $ throwIO ConnectionClosedByPeer let -- How many of the bytes in this chunk to send downstream toSend = min count (S.length bs) -- How many bytes will still remain to be sent downstream count' = count - toSend case () of () -- The expected count is greater than the size of the -- chunk we just read. Send the entire chunk -- downstream, and then loop on this function for the -- next chunk. | count' > 0 -> do I.writeIORef ref count' return bs -- Some of the bytes in this chunk should not be sent -- downstream. Split up the chunk into the sent and -- not-sent parts, add the not-sent parts onto the new -- source, and send the rest of the chunk downstream. | otherwise -> do let (x, y) = S.splitAt toSend bs leftoverSource src y assert (count' == 0) $ I.writeIORef ref count' return x -- | Extract the underlying @Source@ from an @IsolatedBSSource@, which will not -- perform any more isolation. isourceDone :: ISource -> Source isourceDone (ISource src _) = src ---------------------------------------------------------------- data CSource = CSource !Source !(I.IORef ChunkState) data ChunkState = NeedLen | NeedLenNewline | HaveLen Word | DoneChunking deriving Show bsCRLF :: L.ByteString bsCRLF = pack "\r\n" mkCSource :: Source -> IO CSource mkCSource src = do ref <- I.newIORef NeedLen return $! CSource src ref readCSource :: CSource -> IO ByteString readCSource (CSource src ref) = do mlen <- I.readIORef ref go mlen where withLen 0 bs = do leftoverSource src bs dropCRLF yield' S.empty DoneChunking withLen len bs | S.null bs = do -- FIXME should this throw an exception if len > 0? I.writeIORef ref DoneChunking return S.empty | otherwise = case S.length bs `compare` fromIntegral len of EQ -> yield' bs NeedLenNewline LT -> yield' bs $ HaveLen $ len - fromIntegral (S.length bs) GT -> do let (x, y) = S.splitAt (fromIntegral len) bs leftoverSource src y yield' x NeedLenNewline yield' bs mlen = do I.writeIORef ref mlen return bs dropCRLF = do bs <- readSource src case S.uncons bs of Nothing -> return () Just (13, bs') -> dropLF bs' Just (10, bs') -> leftoverSource src bs' Just _ -> leftoverSource src bs dropLF bs = case S.uncons bs of Nothing -> do bs2 <- readSource' src unless (S.null bs2) $ dropLF bs2 Just (10, bs') -> leftoverSource src bs' Just _ -> leftoverSource src bs go NeedLen = getLen go NeedLenNewline = dropCRLF >> getLen go (HaveLen 0) = do -- Drop the final CRLF dropCRLF I.writeIORef ref DoneChunking return S.empty go (HaveLen len) = do bs <- readSource src withLen len bs go DoneChunking = return S.empty -- Get the length from the source, and then pass off control to withLen getLen = do bs <- readSource src if S.null bs then do I.writeIORef ref $ assert False $ HaveLen 0 return S.empty else do (x, y) <- case S.breakByte 10 bs of (x, y) | S.null y -> do bs2 <- readSource' src if S.null bs2 then return (x, y) else return $ S.breakByte 10 $ bs `S.append` bs2 | otherwise -> return (x, y) let w = S.foldl' (\i c -> i * 16 + fromIntegral (hexToWord c)) 0 $ S.takeWhile isHexDigit x let y' = S.drop 1 y y'' <- if S.null y' then readSource src else return y' withLen w y'' hexToWord w | w < 58 = w - 48 | w < 71 = w - 55 | otherwise = w - 87 isHexDigit :: Word8 -> Bool isHexDigit w = w >= 48 && w <= 57 || w >= 65 && w <= 70 || w >= 97 && w <= 102