-- | A light-weight wrapper around @Network.Wai@ to provide easy conduit support.
module Network.Wai.Conduit
    ( -- * Request body
      sourceRequestBody
      -- * Response body
    , responseSource
    , responseRawSource
      -- * Re-export
    , module Network.Wai
    ) where

import Network.Wai
import Data.Conduit
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.ByteString (ByteString)
import qualified Data.ByteString as S
import Control.Monad (unless)
import Network.HTTP.Types
import Data.ByteString.Builder (Builder)
import Data.IORef
import qualified Data.Conduit.List as CL

-- | Stream the request body.
--
-- Since 3.0.0
sourceRequestBody :: MonadIO m => Request -> Source m ByteString
sourceRequestBody :: forall (m :: * -> *). MonadIO m => Request -> Source m ByteString
sourceRequestBody Request
req =
    ConduitT () ByteString m ()
loop
  where
    go :: ConduitT () ByteString m ByteString
go = IO ByteString -> ConduitT () ByteString m ByteString
forall a. IO a -> ConduitT () ByteString m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Request -> IO ByteString
requestBody Request
req)

    loop :: ConduitT () ByteString m ()
loop = do
        ByteString
bs <- ConduitT () ByteString m ByteString
go
        Bool -> ConduitT () ByteString m () -> ConduitT () ByteString m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ByteString -> Bool
S.null ByteString
bs) (ConduitT () ByteString m () -> ConduitT () ByteString m ())
-> ConduitT () ByteString m () -> ConduitT () ByteString m ()
forall a b. (a -> b) -> a -> b
$ do
            ByteString -> ConduitT () ByteString m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield ByteString
bs
            ConduitT () ByteString m ()
loop

-- | Create an HTTP response out of a @Source@.
--
-- Since 3.0.0
responseSource :: Status -> ResponseHeaders -> Source IO (Flush Builder) -> Response
responseSource :: Status -> ResponseHeaders -> Source IO (Flush Builder) -> Response
responseSource Status
s ResponseHeaders
hs Source IO (Flush Builder)
src = Status -> ResponseHeaders -> StreamingBody -> Response
responseStream Status
s ResponseHeaders
hs (StreamingBody -> Response) -> StreamingBody -> Response
forall a b. (a -> b) -> a -> b
$ \Builder -> IO ()
send IO ()
flush ->
    Source IO (Flush Builder)
src Source IO (Flush Builder) -> Sink (Flush Builder) IO () -> IO ()
forall (m :: * -> *) a b.
Monad m =>
Source m a -> Sink a m b -> m b
$$ (Flush Builder -> IO ()) -> Sink (Flush Builder) IO ()
forall (m :: * -> *) a o.
Monad m =>
(a -> m ()) -> ConduitT a o m ()
CL.mapM_ (\Flush Builder
mbuilder ->
        case Flush Builder
mbuilder of
            Chunk Builder
b -> Builder -> IO ()
send Builder
b
            Flush Builder
Flush -> IO ()
flush)

-- | Create a raw response using a @Source@ and @Sink@ to represent the input
-- and output, respectively.
--
-- Since 3.0.0
responseRawSource :: (MonadIO m, MonadIO n)
                  => (Source m ByteString -> Sink ByteString n () -> IO ())
                  -> Response
                  -> Response
responseRawSource :: forall (m :: * -> *) (n :: * -> *).
(MonadIO m, MonadIO n) =>
(Source m ByteString -> Sink ByteString n () -> IO ())
-> Response -> Response
responseRawSource Source m ByteString -> Sink ByteString n () -> IO ()
app =
    (IO ByteString -> (ByteString -> IO ()) -> IO ())
-> Response -> Response
responseRaw IO ByteString -> (ByteString -> IO ()) -> IO ()
app'
  where
    app' :: IO ByteString -> (ByteString -> IO ()) -> IO ()
app' IO ByteString
recv ByteString -> IO ()
send =
        Source m ByteString -> Sink ByteString n () -> IO ()
app Source m ByteString
forall {i}. ConduitT i ByteString m ()
src Sink ByteString n ()
forall {o}. ConduitT ByteString o n ()
sink
      where
        src :: ConduitT i ByteString m ()
src = do
            ByteString
bs <- IO ByteString -> ConduitT i ByteString m ByteString
forall a. IO a -> ConduitT i ByteString m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ByteString
recv
            Bool -> ConduitT i ByteString m () -> ConduitT i ByteString m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ByteString -> Bool
S.null ByteString
bs) (ConduitT i ByteString m () -> ConduitT i ByteString m ())
-> ConduitT i ByteString m () -> ConduitT i ByteString m ()
forall a b. (a -> b) -> a -> b
$ do
                ByteString -> ConduitT i ByteString m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield ByteString
bs
                ConduitT i ByteString m ()
src
        sink :: ConduitT ByteString o n ()
sink = (ByteString -> n ()) -> ConduitT ByteString o n ()
forall (m :: * -> *) a o.
Monad m =>
(a -> m ()) -> ConduitT a o m ()
CL.mapM_ ((ByteString -> n ()) -> ConduitT ByteString o n ())
-> (ByteString -> n ()) -> ConduitT ByteString o n ()
forall a b. (a -> b) -> a -> b
$ IO () -> n ()
forall a. IO a -> n a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> n ()) -> (ByteString -> IO ()) -> ByteString -> n ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> IO ()
send