{-# 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)
import Network.Wai.Handler.Warp.Imports
import Network.Wai.Handler.Warp.Types
data ISource = ISource !Source !(I.IORef Int)
mkISource :: Source -> Int -> IO ISource
mkISource :: Source -> Int -> IO ISource
mkISource Source
src Int
cnt = do
ref <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
I.newIORef Int
cnt
return $! ISource src ref
readISource :: ISource -> IO ByteString
readISource :: ISource -> IO ByteString
readISource (ISource Source
src IORef Int
ref) = do
count <- IORef Int -> IO Int
forall a. IORef a -> IO a
I.readIORef IORef Int
ref
if count == 0
then return S.empty
else do
bs <- readSource src
when (S.null bs) $ throwIO ConnectionClosedByPeer
let
toSend = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
count (ByteString -> Int
S.length ByteString
bs)
count' = Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
toSend
I.writeIORef ref count'
if count' > 0
then
return bs
else do
let (x, y) = S.splitAt toSend bs
leftoverSource src y
assert (count' == 0) $ return 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
ref <- ChunkState -> IO (IORef ChunkState)
forall a. a -> IO (IORef a)
I.newIORef ChunkState
NeedLen
return $! CSource src ref
readCSource :: CSource -> IO ByteString
readCSource :: CSource -> IO ByteString
readCSource (CSource Source
src IORef ChunkState
ref) = do
mlen <- IORef ChunkState -> IO ChunkState
forall a. IORef a -> IO a
I.readIORef IORef ChunkState
ref
go mlen
where
withLen :: Word -> ByteString -> IO ByteString
withLen Word
0 ByteString
bs = do
Source -> ByteString -> IO ()
leftoverSource Source
src ByteString
bs
IO ()
dropCRLF
ByteString -> ChunkState -> IO ByteString
forall {b}. b -> ChunkState -> IO b
yield' ByteString
S.empty ChunkState
DoneChunking
withLen Word
len ByteString
bs
| ByteString -> Bool
S.null ByteString
bs = do
IORef ChunkState -> ChunkState -> IO ()
forall a. IORef a -> a -> IO ()
I.writeIORef IORef ChunkState
ref ChunkState
DoneChunking
ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
S.empty
| 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
b -> IO b
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return b
bs
dropCRLF :: IO ()
dropCRLF = do
bs <- Source -> IO ByteString
readSource Source
src
case S.uncons bs of
Maybe (Word8, ByteString)
Nothing -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
Just (Word8
w8, ByteString
bs')
| Word8
w8 Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
_cr -> ByteString -> IO ()
dropLF ByteString
bs'
| Word8
w8 Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
_lf -> Source -> ByteString -> IO ()
leftoverSource Source
src ByteString
bs'
| Bool
otherwise -> Source -> ByteString -> IO ()
leftoverSource Source
src ByteString
bs
dropLF :: ByteString -> IO ()
dropLF ByteString
bs =
case ByteString -> Maybe (Word8, ByteString)
S.uncons ByteString
bs of
Maybe (Word8, ByteString)
Nothing -> do
bs2 <- Source -> IO ByteString
readSource' Source
src
unless (S.null bs2) $ dropLF bs2
Just (Word8
w8, ByteString
bs') ->
Source -> ByteString -> IO ()
leftoverSource Source
src (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$
if Word8
w8 Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
_lf then ByteString
bs' else ByteString
bs
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
0) = do
IO ()
dropCRLF
IORef ChunkState -> ChunkState -> IO ()
forall a. IORef a -> a -> IO ()
I.writeIORef IORef ChunkState
ref ChunkState
DoneChunking
ByteString -> IO ByteString
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
S.empty
go (HaveLen Word
len) = do
bs <- Source -> IO ByteString
readSource Source
src
withLen len bs
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
bs <- Source -> IO ByteString
readSource Source
src
if S.null bs
then do
I.writeIORef ref $ assert False $ HaveLen 0
return S.empty
else do
(x, y) <-
case S.break (== _lf) bs of
(ByteString
x, ByteString
y)
| ByteString -> Bool
S.null ByteString
y -> do
bs2 <- Source -> IO ByteString
readSource' Source
src
return $
if S.null bs2
then (x, y)
else S.break (== _lf) $ bs `S.append` bs2
| Bool
otherwise -> (ByteString, ByteString) -> IO (ByteString, ByteString)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
x, ByteString
y)
let w =
(Word -> Word8 -> Word) -> Word -> ByteString -> Word
forall a. (a -> Word8 -> a) -> a -> ByteString -> a
S.foldl' (\Word
i Word8
c -> Word
i Word -> Word -> Word
forall a. Num a => a -> a -> a
* Word
16 Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word8 -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Word8
hexToWord Word8
c)) Word
0 (ByteString -> Word) -> ByteString -> Word
forall a b. (a -> b) -> a -> b
$
(Word8 -> Bool) -> ByteString -> ByteString
S.takeWhile Word8 -> Bool
isHexDigit ByteString
x
let y' = Int -> ByteString -> ByteString
S.drop Int
1 ByteString
y
y'' <-
if S.null y'
then readSource src
else return y'
withLen w y''
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