{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}

module Network.Wai.Handler.Warp.IO where

import Control.Exception (mask_)
import Data.ByteString.Builder (Builder)
import Data.ByteString.Builder.Extra (Next (Chunk, Done, More), runBuilder)
import Data.IORef (IORef, readIORef, writeIORef)
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 ()
toBufIOWith :: Int
-> IORef WriteBuffer -> (ByteString -> IO ()) -> Builder -> IO ()
toBufIOWith 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 -> BufferWriter -> IO ()
loop WriteBuffer
writeBuffer BufferWriter
firstWriter
  where
    firstWriter :: BufferWriter
firstWriter = Builder -> BufferWriter
runBuilder Builder
builder
    loop :: WriteBuffer -> BufferWriter -> IO ()
loop WriteBuffer
writeBuffer BufferWriter
writer = 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 Int
size
      Buffer -> Int -> (ByteString -> IO ()) -> IO ()
bufferIO Buffer
buf Int
len ByteString -> IO ()
io
      case Next
signal of
        Next
Done -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        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
                IORef WriteBuffer -> WriteBuffer -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef WriteBuffer
writeBufferRef WriteBuffer
biggerWriteBuffer
                WriteBuffer -> IO WriteBuffer
forall (m :: * -> *) a. Monad m => a -> m a
return WriteBuffer
biggerWriteBuffer
              WriteBuffer -> BufferWriter -> IO ()
loop WriteBuffer
biggerWriteBuffer BufferWriter
next
          | Bool
otherwise -> WriteBuffer -> BufferWriter -> IO ()
loop WriteBuffer
writeBuffer BufferWriter
next
        Chunk ByteString
bs BufferWriter
next -> do
          ByteString -> IO ()
io ByteString
bs
          WriteBuffer -> BufferWriter -> IO ()
loop WriteBuffer
writeBuffer BufferWriter
next