module Network.Wai.Handler.Warp.IO where

import Control.Exception (mask_)
import qualified Data.ByteString as B (length)
import Data.ByteString.Builder (Builder)
import Data.ByteString.Builder.Extra (Next (Chunk, Done, More), runBuilder)
import Data.IORef (IORef, readIORef, writeIORef)
import Foreign.Ptr (plusPtr)
import Network.Wai.Handler.Warp.Buffer
import Network.Wai.Handler.Warp.Imports
import Network.Wai.Handler.Warp.Types

toBufIOWith
    :: Int -> IORef WriteBuffer -> (ByteString -> IO ()) -> Builder -> IO Integer
toBufIOWith :: Int
-> IORef WriteBuffer
-> (ByteString -> IO ())
-> Builder
-> IO Integer
toBufIOWith = Int
-> Int
-> IORef WriteBuffer
-> (ByteString -> IO ())
-> Builder
-> IO Integer
unsafeToBufIOWithOffset Int
0

-- | Like 'toBufIOWith' but the first @offset@ bytes of the write buffer
-- are assumed to be already filled (e.g. with a response header composed
-- directly into the buffer). They are flushed together with the first
-- batch of builder output and included in the returned total.
--
-- === WARNING: @offset@ MUST NOT exceed the current buffer size!!!
--
-- This function performs NO bounds checking on @offset@. The builder is
-- handed the pointer @buffer + offset@ with @bufSize - offset@ bytes of
-- claimed free space, so an oversized @offset@ points past the end of the
-- allocation and advertises negative capacity, i.e. out-of-bounds writes
-- and memory corruption. Every caller MUST verify
-- @offset < bufSize@ of the current write buffer first (see the
-- @hdrLen@ check in 'Network.Wai.Handler.Warp.Response.sendRsp').
unsafeToBufIOWithOffset
    :: Int
    -> Int
    -> IORef WriteBuffer
    -> (ByteString -> IO ())
    -> Builder
    -> IO Integer
unsafeToBufIOWithOffset :: Int
-> Int
-> IORef WriteBuffer
-> (ByteString -> IO ())
-> Builder
-> IO Integer
unsafeToBufIOWithOffset Int
offset0 Int
maxRspBufSize IORef WriteBuffer
writeBufferRef ByteString -> IO ()
io Builder
builder = do
    WriteBuffer
writeBuffer <- IORef WriteBuffer -> IO WriteBuffer
forall a. IORef a -> IO a
readIORef IORef WriteBuffer
writeBufferRef
    WriteBuffer -> Int -> BufferWriter -> Integer -> IO Integer
loop WriteBuffer
writeBuffer Int
offset0 BufferWriter
firstWriter Integer
0
  where
    firstWriter :: BufferWriter
firstWriter = Builder -> BufferWriter
runBuilder Builder
builder
    loop :: WriteBuffer -> Int -> BufferWriter -> Integer -> IO Integer
loop WriteBuffer
writeBuffer Int
offset BufferWriter
writer Integer
bytesSent = do
        let buf :: Buffer
buf = WriteBuffer -> Buffer
bufBuffer WriteBuffer
writeBuffer
            size :: Int
size = WriteBuffer -> Int
bufSize WriteBuffer
writeBuffer
        (Int
len, Next
signal) <- BufferWriter
writer (Buffer
buf Buffer -> Int -> Buffer
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
offset) (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
offset)
        Buffer -> Int -> (ByteString -> IO ()) -> IO ()
bufferIO Buffer
buf (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len) ByteString -> IO ()
io
        let totalBytesSent :: Integer
totalBytesSent = Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
bytesSent
        case Next
signal of
            Next
Done -> Integer -> IO Integer
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Integer
totalBytesSent
            More Int
minSize BufferWriter
next
                | Int
size Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
minSize -> do
                    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
minSize Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxRspBufSize) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
                        [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$
                            [Char]
"Sending a Builder response required a buffer of size "
                                [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
minSize
                                [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" which is bigger than the specified maximum of "
                                [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
maxRspBufSize
                                [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"!"
                    -- The current WriteBuffer is too small to fit the next
                    -- batch of bytes from the Builder so we free it and
                    -- create a new bigger one. Freeing the current buffer,
                    -- creating a new one and writing it to the IORef need
                    -- to be performed atomically to prevent both double
                    -- frees and missed frees. So we mask async exceptions:
                    WriteBuffer
biggerWriteBuffer <- IO WriteBuffer -> IO WriteBuffer
forall a. IO a -> IO a
mask_ (IO WriteBuffer -> IO WriteBuffer)
-> IO WriteBuffer -> IO WriteBuffer
forall a b. (a -> b) -> a -> b
$ do
                        WriteBuffer -> IO ()
bufFree WriteBuffer
writeBuffer
                        WriteBuffer
biggerWriteBuffer <- Int -> IO WriteBuffer
createWriteBuffer Int
minSize
                        -- This doesn't need to be "atomic", since these two
                        -- functions are only used in 'sendResponse', which is
                        -- ultimately only used in 'serveConnection', which does
                        -- not share nor fork the created 'Connection'.
                        IORef WriteBuffer -> WriteBuffer -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef WriteBuffer
writeBufferRef WriteBuffer
biggerWriteBuffer
                        WriteBuffer -> IO WriteBuffer
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return WriteBuffer
biggerWriteBuffer
                    WriteBuffer -> Int -> BufferWriter -> Integer -> IO Integer
loop WriteBuffer
biggerWriteBuffer Int
0 BufferWriter
next Integer
totalBytesSent
                | Bool
otherwise -> WriteBuffer -> Int -> BufferWriter -> Integer -> IO Integer
loop WriteBuffer
writeBuffer Int
0 BufferWriter
next Integer
totalBytesSent
            Chunk ByteString
bs BufferWriter
next -> do
                ByteString -> IO ()
io ByteString
bs
                WriteBuffer -> Int -> BufferWriter -> Integer -> IO Integer
loop WriteBuffer
writeBuffer Int
0 BufferWriter
next (Integer -> IO Integer) -> Integer -> IO Integer
forall a b. (a -> b) -> a -> b
$
                    Integer
totalBytesSent Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
B.length ByteString
bs)