-- |
-- Copyright        : (c) Raghu Kaippully, 2020
-- License          : MPL-2.0
-- Maintainer       : rkaippully@gmail.com
--
-- Middlewares related to HTTP body.
module WebGear.Middlewares.Body
  ( JSONRequestBody
  , jsonRequestBody
  , jsonResponseBody
  ) where

import Control.Arrow (Kleisli (..))
import Control.Monad ((>=>))
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.Aeson (FromJSON, ToJSON, eitherDecode', encode)
import Data.ByteString.Lazy (ByteString, fromChunks, fromStrict)
import Data.Kind (Type)
import Data.Text (Text, pack)
import Data.Text.Encoding (encodeUtf8)
import Network.HTTP.Types (hContentType)

import WebGear.Trait (Result (..), Trait (..), probe)
import WebGear.Types (MonadRouter (..), Request, RequestMiddleware', Response (..),
                      ResponseMiddleware', badRequest400, getRequestBodyChunk, setResponseHeader)
import WebGear.Util (takeWhileM)

-- | A 'Trait' for converting a JSON request body into a value.
data JSONRequestBody (t :: Type)

instance (FromJSON t, MonadIO m) => Trait (JSONRequestBody t) Request m where
  type Attribute (JSONRequestBody t) Request = t
  type Absence (JSONRequestBody t) Request = Text

  toAttribute :: Request -> m (Result (JSONRequestBody t) Request)
  toAttribute :: Request -> m (Result (JSONRequestBody t) Request)
toAttribute Request
r = do
    [ByteString]
chunks <- (ByteString -> Bool) -> [m ByteString] -> m [ByteString]
forall (m :: * -> *) a. Monad m => (a -> Bool) -> [m a] -> m [a]
takeWhileM (ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
/= ByteString
forall a. Monoid a => a
mempty) ([m ByteString] -> m [ByteString])
-> [m ByteString] -> m [ByteString]
forall a b. (a -> b) -> a -> b
$ m ByteString -> [m ByteString]
forall a. a -> [a]
repeat (m ByteString -> [m ByteString]) -> m ByteString -> [m ByteString]
forall a b. (a -> b) -> a -> b
$ IO ByteString -> m ByteString
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> m ByteString) -> IO ByteString -> m ByteString
forall a b. (a -> b) -> a -> b
$ Request -> IO ByteString
getRequestBodyChunk Request
r
    Result (JSONRequestBody t) Request
-> m (Result (JSONRequestBody t) Request)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Result (JSONRequestBody t) Request
 -> m (Result (JSONRequestBody t) Request))
-> Result (JSONRequestBody t) Request
-> m (Result (JSONRequestBody t) Request)
forall a b. (a -> b) -> a -> b
$ case ByteString -> Either String t
forall a. FromJSON a => ByteString -> Either String a
eitherDecode' ([ByteString] -> ByteString
fromChunks [ByteString]
chunks) of
             Left String
e  -> Absence (JSONRequestBody t) Request
-> Result (JSONRequestBody t) Request
forall k (t :: k) a. Absence t a -> Result t a
NotFound (String -> Text
pack String
e)
             Right t
t -> Attribute (JSONRequestBody t) Request
-> Result (JSONRequestBody t) Request
forall k (t :: k) a. Attribute t a -> Result t a
Found t
Attribute (JSONRequestBody t) Request
t

-- | A middleware to parse the request body as JSON and convert it to
-- a value via a 'FromJSON' instance.
--
-- Usage for a type @t@ which has a 'FromJSON' instance:
--
-- > jsonRequestBody @t handler
--
-- Returns a 400 Bad Request response on failure to parse body.
jsonRequestBody :: forall t m req a. (FromJSON t, MonadRouter m, MonadIO m)
                => RequestMiddleware' m req (JSONRequestBody t:req) a
jsonRequestBody :: RequestMiddleware' m req (JSONRequestBody t : req) a
jsonRequestBody Handler' m (JSONRequestBody t : req) a
handler = (Linked req Request -> m (Response a))
-> Kleisli m (Linked req Request) (Response a)
forall (m :: * -> *) a b. (a -> m b) -> Kleisli m a b
Kleisli ((Linked req Request -> m (Response a))
 -> Kleisli m (Linked req Request) (Response a))
-> (Linked req Request -> m (Response a))
-> Kleisli m (Linked req Request) (Response a)
forall a b. (a -> b) -> a -> b
$
  forall (ts :: [*]) a (m :: * -> *).
Trait (JSONRequestBody t) a m =>
Linked ts a
-> m (Either
        (Absence (JSONRequestBody t) a)
        (Linked (JSONRequestBody t : ts) a))
forall t (ts :: [*]) a (m :: * -> *).
Trait t a m =>
Linked ts a -> m (Either (Absence t a) (Linked (t : ts) a))
probe @(JSONRequestBody t) (Linked req Request
 -> m (Either Text (Linked (JSONRequestBody t : req) Request)))
-> (Either Text (Linked (JSONRequestBody t : req) Request)
    -> m (Response a))
-> Linked req Request
-> m (Response a)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> (Text -> m (Response a))
-> (Linked (JSONRequestBody t : req) Request -> m (Response a))
-> Either Text (Linked (JSONRequestBody t : req) Request)
-> m (Response a)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Response ByteString -> m (Response a)
forall (m :: * -> *) a. MonadRouter m => Response ByteString -> m a
errorResponse (Response ByteString -> m (Response a))
-> (Text -> Response ByteString) -> Text -> m (Response a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Response ByteString
mkError) (Handler' m (JSONRequestBody t : req) a
-> Linked (JSONRequestBody t : req) Request -> m (Response a)
forall (m :: * -> *) a b. Kleisli m a b -> a -> m b
runKleisli Handler' m (JSONRequestBody t : req) a
handler)
  where
    mkError :: Text -> Response ByteString
    mkError :: Text -> Response ByteString
mkError Text
e = ByteString -> Response ByteString
forall a. a -> Response a
badRequest400 (ByteString -> Response ByteString)
-> ByteString -> Response ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
fromStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Text -> ByteString
encodeUtf8 (Text -> ByteString) -> Text -> ByteString
forall a b. (a -> b) -> a -> b
$ Text
"Error parsing request body: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
e

-- | A middleware that converts the response that has a 'ToJSON'
-- instance to a 'ByteString' response.
--
-- This will also set the "Content-Type" header of the response to
-- "application/json".
--
-- Usage for a type @t@ which has a 'ToJSON' instance:
--
-- > jsonResponseBody @t handler
--
jsonResponseBody :: (ToJSON t, Monad m) => ResponseMiddleware' m req t ByteString
jsonResponseBody :: ResponseMiddleware' m req t ByteString
jsonResponseBody Handler' m req t
handler = (Linked req Request -> m (Response ByteString))
-> Kleisli m (Linked req Request) (Response ByteString)
forall (m :: * -> *) a b. (a -> m b) -> Kleisli m a b
Kleisli ((Linked req Request -> m (Response ByteString))
 -> Kleisli m (Linked req Request) (Response ByteString))
-> (Linked req Request -> m (Response ByteString))
-> Kleisli m (Linked req Request) (Response ByteString)
forall a b. (a -> b) -> a -> b
$ \Linked req Request
req -> do
  Response t
x <- Handler' m req t -> Linked req Request -> m (Response t)
forall (m :: * -> *) a b. Kleisli m a b -> a -> m b
runKleisli Handler' m req t
handler Linked req Request
req
  Response ByteString -> m (Response ByteString)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Response ByteString -> m (Response ByteString))
-> Response ByteString -> m (Response ByteString)
forall a b. (a -> b) -> a -> b
$ HeaderName
-> ByteString -> Response ByteString -> Response ByteString
forall a. HeaderName -> ByteString -> Response a -> Response a
setResponseHeader HeaderName
hContentType ByteString
"application/json" (Response ByteString -> Response ByteString)
-> Response ByteString -> Response ByteString
forall a b. (a -> b) -> a -> b
$ t -> ByteString
forall a. ToJSON a => a -> ByteString
encode (t -> ByteString) -> Response t -> Response ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Response t
x