{-# LANGUAGE CPP #-}

module Network.Wai.Handler.Warp.Conduit where

import Control.Exception (assert, throwIO)
import qualified Data.ByteString as S
import qualified Data.IORef as I
import Data.Word8 (_0, _9, _A, _F, _a, _cr, _f, _lf, _semicolon)

import Network.Wai.Handler.Warp.Imports
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 :: Source -> Int -> IO ISource
mkISource Source
src Int
cnt = do
    IORef Int
ref <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
I.newIORef Int
cnt
    ISource -> IO ISource
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ISource -> IO ISource) -> ISource -> IO ISource
forall a b. (a -> b) -> a -> b
$! Source -> IORef Int -> ISource
ISource Source
src IORef Int
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 -> IO ByteString
readISource (ISource Source
src IORef Int
ref) = do
    Int
count <- IORef Int -> IO Int
forall a. IORef a -> IO a
I.readIORef IORef Int
ref
    if Int
count Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
S.empty
        else do
            ByteString
bs <- Source -> IO ByteString
readSource Source
src

            -- If no chunk available, then there aren't enough bytes in the
            -- stream. Throw a ConnectionClosedByPeer
            Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ByteString -> Bool
S.null ByteString
bs) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ InvalidRequest -> IO ()
forall e a. Exception e => e -> IO a
throwIO InvalidRequest
ConnectionClosedByPeer

            let -- How many of the bytes in this chunk to send downstream
                toSend :: Int
toSend = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
count (ByteString -> Int
S.length ByteString
bs)
                -- How many bytes will still remain to be sent downstream
                count' :: Int
count' = Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
toSend

            -- [WRITE_IOREF_NOTE]
            -- This doesn't need to be "atomic", since it is only used in
            -- 'recvRequest', which creates the 'Source' and doesn't fork it,
            -- so the 'IORef' is not shared outside of 'recvRequest'.
            IORef Int -> Int -> IO ()
forall a. IORef a -> a -> IO ()
I.writeIORef IORef Int
ref Int
count'

            if Int
count' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
                then -- 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.
                    ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
bs
                else do
                    -- 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.
                    let (ByteString
x, ByteString
y) = Int -> ByteString -> (ByteString, ByteString)
S.splitAt Int
toSend ByteString
bs
                    Source -> ByteString -> IO ()
leftoverSource Source
src ByteString
y
                    Bool -> IO ByteString -> IO ByteString
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
count' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (IO ByteString -> IO ByteString) -> IO ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
x

----------------------------------------------------------------

data CSource = CSource !Source !(I.IORef ChunkState)

data ChunkState
    = NeedLen
    | NeedLenNewline
    | HaveLen Word
    | DoneChunking
    deriving (Int -> ChunkState -> ShowS
[ChunkState] -> ShowS
ChunkState -> String
(Int -> ChunkState -> ShowS)
-> (ChunkState -> String)
-> ([ChunkState] -> ShowS)
-> Show ChunkState
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ChunkState -> ShowS
showsPrec :: Int -> ChunkState -> ShowS
$cshow :: ChunkState -> String
show :: ChunkState -> String
$cshowList :: [ChunkState] -> ShowS
showList :: [ChunkState] -> ShowS
Show)

mkCSource :: Source -> IO CSource
mkCSource :: Source -> IO CSource
mkCSource Source
src = do
    IORef ChunkState
ref <- ChunkState -> IO (IORef ChunkState)
forall a. a -> IO (IORef a)
I.newIORef ChunkState
NeedLen
    CSource -> IO CSource
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (CSource -> IO CSource) -> CSource -> IO CSource
forall a b. (a -> b) -> a -> b
$! Source -> IORef ChunkState -> CSource
CSource Source
src IORef ChunkState
ref

-- | A chunk-size line is a few bytes of hex, and chunk extensions are
--   all but unused; a trailer section is a handful of fields. Neither is
--   a reason to hold on to everything a client cares to send.
maxChunkLineLength :: Int
maxChunkLineLength :: Int
maxChunkLineLength = Int
1024

maxTrailerLength :: Int
maxTrailerLength :: Int
maxTrailerLength = Int
8192

readCSource :: CSource -> IO ByteString
readCSource :: CSource -> IO ByteString
readCSource (CSource Source
src IORef ChunkState
ref) = do
    ChunkState
mlen <- IORef ChunkState -> IO ChunkState
forall a. IORef a -> IO a
I.readIORef IORef ChunkState
ref
    ChunkState -> IO ByteString
go ChunkState
mlen
  where
    go :: ChunkState -> IO ByteString
go ChunkState
NeedLen = IO ByteString
getLen
    go ChunkState
NeedLenNewline = IO ()
dropCRLF IO () -> IO ByteString -> IO ByteString
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ByteString
getLen
    go (HaveLen Word
len) = Source -> IO ByteString
readSource Source
src IO ByteString -> (ByteString -> IO ByteString) -> IO ByteString
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Word -> ByteString -> IO ByteString
withLen Word
len
    go ChunkState
DoneChunking = ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
S.empty

    getLen :: IO ByteString
getLen = do
        Word
len <- ByteString -> IO Word
forall {a}. Num a => ByteString -> IO a
chunkSize (ByteString -> IO Word) -> IO ByteString -> IO Word
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO ByteString
requireLine
        if Word
len Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word
0
            then do
                -- Everything up to the empty line which ends the trailer
                -- section belongs to this request. Warp used to drop one
                -- CRLF and leave the trailers in the source, where they
                -- were read as the next request on the connection.
                Int -> IO ()
dropTrailers Int
0
                ByteString -> ChunkState -> IO ByteString
forall {b}. b -> ChunkState -> IO b
yield' ByteString
S.empty ChunkState
DoneChunking
            else Source -> IO ByteString
readSource Source
src IO ByteString -> (ByteString -> IO ByteString) -> IO ByteString
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Word -> ByteString -> IO ByteString
withLen Word
len

    withLen :: Word -> ByteString -> IO ByteString
withLen Word
len ByteString
bs
        -- The chunk promised more than the client went on to send.
        | ByteString -> Bool
S.null ByteString
bs = InvalidRequest -> IO ByteString
forall e a. Exception e => e -> IO a
throwIO InvalidRequest
ConnectionClosedByPeer
        | Bool
otherwise =
            case ByteString -> Int
S.length ByteString
bs Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
len of
                Ordering
EQ -> ByteString -> ChunkState -> IO ByteString
forall {b}. b -> ChunkState -> IO b
yield' ByteString
bs ChunkState
NeedLenNewline
                Ordering
LT -> ByteString -> ChunkState -> IO ByteString
forall {b}. b -> ChunkState -> IO b
yield' ByteString
bs (ChunkState -> IO ByteString) -> ChunkState -> IO ByteString
forall a b. (a -> b) -> a -> b
$ Word -> ChunkState
HaveLen (Word -> ChunkState) -> Word -> ChunkState
forall a b. (a -> b) -> a -> b
$ Word
len Word -> Word -> Word
forall a. Num a => a -> a -> a
- Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
S.length ByteString
bs)
                Ordering
GT -> do
                    let (ByteString
x, ByteString
y) = Int -> ByteString -> (ByteString, ByteString)
S.splitAt (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
len) ByteString
bs
                    Source -> ByteString -> IO ()
leftoverSource Source
src ByteString
y
                    ByteString -> ChunkState -> IO ByteString
forall {b}. b -> ChunkState -> IO b
yield' ByteString
x ChunkState
NeedLenNewline

    yield' :: b -> ChunkState -> IO b
yield' b
bs ChunkState
mlen = do
        IORef ChunkState -> ChunkState -> IO ()
forall a. IORef a -> a -> IO ()
I.writeIORef IORef ChunkState
ref ChunkState
mlen -- [WRITE_IOREF_NOTE]
        b -> IO b
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return b
bs

    -- A chunk size is 1*HEXDIG, with nothing after it but chunk
    -- extensions. Warp used to fold whatever hex digits it could find
    -- and take 0 for the rest, so a line which is no chunk size at all
    -- -- a request line, say -- ended the body and left the rest of it
    -- to be read as the next request.
    chunkSize :: ByteString -> IO a
chunkSize ByteString
line
        | ByteString -> Bool
S.null ByteString
digits = InvalidRequest -> IO a
forall e a. Exception e => e -> IO a
throwIO InvalidRequest
BadChunkedBody
        | Bool -> Bool
not (ByteString -> Bool
S.null ByteString
ext), (?callStack::CallStack) => ByteString -> Word8
ByteString -> Word8
S.head ByteString
ext Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
_semicolon = InvalidRequest -> IO a
forall e a. Exception e => e -> IO a
throwIO InvalidRequest
BadChunkedBody
        | Bool
otherwise =
            IO a -> (Int -> IO a) -> Maybe Int -> IO a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (InvalidRequest -> IO a
forall e a. Exception e => e -> IO a
throwIO InvalidRequest
BadChunkedBody) (a -> IO a
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> IO a) -> (Int -> a) -> Int -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral) (Maybe Int -> IO a) -> Maybe Int -> IO a
forall a b. (a -> b) -> a -> b
$
                (Maybe Int -> Word8 -> Maybe Int)
-> Maybe Int -> ByteString -> Maybe Int
forall a. (a -> Word8 -> a) -> a -> ByteString -> a
S.foldl' Maybe Int -> Word8 -> Maybe Int
forall {a}. (Integral a, Bounded a) => Maybe a -> Word8 -> Maybe a
accumulate (Int -> Maybe Int
forall a. a -> Maybe a
Just (Int
0 :: Int)) ByteString
digits
      where
        (ByteString
digits, ByteString
ext) = (Word8 -> Bool) -> ByteString -> (ByteString, ByteString)
S.span Word8 -> Bool
isHexDigit ByteString
line
        accumulate :: Maybe a -> Word8 -> Maybe a
accumulate Maybe a
Nothing Word8
_ = Maybe a
forall a. Maybe a
Nothing
        -- A chunk warp cannot measure is a chunk warp cannot read.
        accumulate (Just a
n) Word8
w
            | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> (a
forall a. Bounded a => a
maxBound a -> a -> a
forall a. Num a => a -> a -> a
- a
d) a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
16 = Maybe a
forall a. Maybe a
Nothing
            | Bool
otherwise = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$! a
n a -> a -> a
forall a. Num a => a -> a -> a
* a
16 a -> a -> a
forall a. Num a => a -> a -> a
+ a
d
          where
            d :: a
d = Word8 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Word8
hexToWord Word8
w)

    -- The CRLF which ends chunk-data. Warp used to put anything else
    -- back and read it as the next chunk size, so a chunk that was not
    -- closed off ended the body early.
    dropCRLF :: IO ()
dropCRLF = do
        ByteString
line <- IO ByteString
requireLine
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ByteString -> Bool
S.null ByteString
line) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ InvalidRequest -> IO ()
forall e a. Exception e => e -> IO a
throwIO InvalidRequest
BadChunkedBody

    dropTrailers :: Int -> IO ()
dropTrailers Int
total = do
        Maybe ByteString
mline <- IO (Maybe ByteString)
readLine
        case Maybe ByteString
mline of
            -- A client which leaves out the last CRLF and closes has
            -- nothing left to hide behind it.
            Maybe ByteString
Nothing -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            Just ByteString
line -> Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ByteString -> Bool
S.null ByteString
line) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                let total' :: Int
total' = Int
total Int -> Int -> Int
forall a. Num a => a -> a -> a
+ ByteString -> Int
S.length ByteString
line
                Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
total' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxTrailerLength) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ InvalidRequest -> IO ()
forall e a. Exception e => e -> IO a
throwIO InvalidRequest
OverLargeHeader
                Int -> IO ()
dropTrailers Int
total'

    requireLine :: IO ByteString
requireLine = IO (Maybe ByteString)
readLine IO (Maybe ByteString)
-> (Maybe ByteString -> IO ByteString) -> IO ByteString
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO ByteString
-> (ByteString -> IO ByteString)
-> Maybe ByteString
-> IO ByteString
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (InvalidRequest -> IO ByteString
forall e a. Exception e => e -> IO a
throwIO InvalidRequest
ConnectionClosedByPeer) ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return

    -- One line of the chunked framing: up to the next LF, without it or
    -- the CR before it, leaving the rest of what was read in the source.
    -- 'Nothing' when the connection ends where a line was to begin.
    -- Warp used to give up one read short of the LF, which mis-read
    -- every size line that arrived in more than two pieces.
    readLine :: IO (Maybe ByteString)
readLine = Int -> ([ByteString] -> [ByteString]) -> IO (Maybe ByteString)
loop Int
0 [ByteString] -> [ByteString]
forall a. a -> a
id
      where
        loop :: Int -> ([ByteString] -> [ByteString]) -> IO (Maybe ByteString)
loop Int
total [ByteString] -> [ByteString]
front = do
            ByteString
bs <- Source -> IO ByteString
readSource Source
src
            if ByteString -> Bool
S.null ByteString
bs
                then
                    if Int
total Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
                        then Maybe ByteString -> IO (Maybe ByteString)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ByteString
forall a. Maybe a
Nothing
                        else InvalidRequest -> IO (Maybe ByteString)
forall e a. Exception e => e -> IO a
throwIO InvalidRequest
ConnectionClosedByPeer
                else case Word8 -> ByteString -> Maybe Int
S.elemIndex Word8
_lf ByteString
bs of
                    Maybe Int
Nothing -> do
                        let total' :: Int
total' = Int
total Int -> Int -> Int
forall a. Num a => a -> a -> a
+ ByteString -> Int
S.length ByteString
bs
                        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
total' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxChunkLineLength) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ InvalidRequest -> IO ()
forall e a. Exception e => e -> IO a
throwIO InvalidRequest
OverLargeHeader
                        Int -> ([ByteString] -> [ByteString]) -> IO (Maybe ByteString)
loop Int
total' ([ByteString] -> [ByteString]
front ([ByteString] -> [ByteString])
-> ([ByteString] -> [ByteString]) -> [ByteString] -> [ByteString]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString
bs ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:))
                    Just Int
ix -> do
                        Source -> ByteString -> IO ()
leftoverSource Source
src (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> ByteString
S.drop (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ByteString
bs
                        let line :: ByteString
line = [ByteString] -> ByteString
S.concat ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString] -> [ByteString]
front [Int -> ByteString -> ByteString
S.take Int
ix ByteString
bs]
                        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ByteString -> Int
S.length ByteString
line Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxChunkLineLength) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
                            InvalidRequest -> IO ()
forall e a. Exception e => e -> IO a
throwIO InvalidRequest
OverLargeHeader
                        Maybe ByteString -> IO (Maybe ByteString)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ByteString -> IO (Maybe ByteString))
-> Maybe ByteString -> IO (Maybe ByteString)
forall a b. (a -> b) -> a -> b
$ ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Maybe ByteString) -> ByteString -> Maybe ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
dropCR ByteString
line
        dropCR :: ByteString -> ByteString
dropCR ByteString
bs = case ByteString -> Maybe (ByteString, Word8)
S.unsnoc ByteString
bs of
            Just (ByteString
bs', Word8
w) | Word8
w Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
_cr -> ByteString
bs'
            Maybe (ByteString, Word8)
_ -> ByteString
bs

    hexToWord :: Word8 -> Word8
hexToWord Word8
w
        | Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word8
_9 = Word8
w Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
- Word8
_0
        | Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word8
_F = Word8
w Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
- Word8
55
        | Bool
otherwise = Word8
w Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
- Word8
87

isHexDigit :: Word8 -> Bool
isHexDigit :: Word8 -> Bool
isHexDigit Word8
w =
    Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word8
_0 Bool -> Bool -> Bool
&& Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word8
_9
        Bool -> Bool -> Bool
|| Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word8
_A Bool -> Bool -> Bool
&& Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word8
_F
        Bool -> Bool -> Bool
|| Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word8
_a Bool -> Bool -> Bool
&& Word8
w Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word8
_f