{-# LANGUAGE LambdaCase #-}

-- |
-- Copyright: © 2018-2020 IOHK
-- License: Apache-2.0
--
-- Middleware between Wai <-> Servant to accommodate raw error responses
-- returned by servant. See also 'handleRawError'.

module Network.Wai.Middleware.ServerError
    ( handleRawError
    ) where

import Prelude

import Control.Monad
    ( guard )
import Data.ByteString.Lazy
    ( ByteString )
import Network.HTTP.Types.Status
    ( statusCode, statusMessage )
import Network.Wai
    ( Middleware, responseHeaders, responseStatus )
import Network.Wai.Internal
    ( Request, Response (..) )
import Servant.Server.Internal.ServerError
    ( ServerError (..), responseServerError )

import qualified Data.Binary.Builder as Binary
import qualified Data.ByteString.Char8 as B8

-- | Make sure every error is converted to a suitable application-level error.
--
-- There are many cases where Servant will handle errors itself and reply to a
-- client without even disturbing the application. This is both handy and clunky
-- since our application return errors in a specific format (e.g. JSON, XML
-- ...).
--
-- This is the case for instance if the client hits a non-exiting endpoint of
-- the API, or if the client requests an invalid content-type, etc ...
--
-- Ideally, we would like clients to be able to expect one and only one format,
-- so this middleware allows for manipulating the response returned by a Wai
-- application (what servant boils down to) and adjust the response when
-- necessary. So, any response with or without payload but no content-type will
-- trigger the 'convert' function and offer the caller to adjust the response as
-- needed.
handleRawError
    :: (Request -> ServerError -> ServerError)
    -- ^ Convert a raw response into something that better fits the application
    -- error
    -> Middleware
handleRawError :: (Request -> ServerError -> ServerError) -> Middleware
handleRawError Request -> ServerError -> ServerError
adjust Application
app Request
req Response -> IO ResponseReceived
send =
    Application
app Request
req (Response -> IO ResponseReceived
send (Response -> IO ResponseReceived)
-> (Response -> Response) -> Response -> IO ResponseReceived
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ServerError -> Response)
-> (Response -> Response)
-> Either ServerError Response
-> Response
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (ServerError -> Response
responseServerError (ServerError -> Response)
-> (ServerError -> ServerError) -> ServerError -> Response
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Request -> ServerError -> ServerError
adjust Request
req) Response -> Response
forall a. a -> a
id (Either ServerError Response -> Response)
-> (Response -> Either ServerError Response)
-> Response
-> Response
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Response -> Either ServerError Response
eitherRawError)

-- | Analyze whether a given error is a raw error thrown by Servant before
-- reaching our application layer, or one from our application layer.
eitherRawError
    :: Response
    -> Either ServerError Response
eitherRawError :: Response -> Either ServerError Response
eitherRawError Response
res =
    let
        status :: Status
status = Response -> Status
responseStatus Response
res
        code :: Int
code = Status -> Int
statusCode Status
status
        reason :: [Char]
reason = ByteString -> [Char]
B8.unpack (Status -> ByteString
statusMessage Status
status)
        headers :: ResponseHeaders
headers = Response -> ResponseHeaders
responseHeaders Response
res
        body :: Maybe ByteString
body = Response -> Maybe ByteString
responseBody Response
res
        maybeToEither :: Maybe ByteString -> Either ServerError Response
maybeToEither = Either ServerError Response
-> (ByteString -> Either ServerError Response)
-> Maybe ByteString
-> Either ServerError Response
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
            (Response -> Either ServerError Response
forall a b. b -> Either a b
Right Response
res)
            (ServerError -> Either ServerError Response
forall a b. a -> Either a b
Left (ServerError -> Either ServerError Response)
-> (ByteString -> ServerError)
-> ByteString
-> Either ServerError Response
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString -> ResponseHeaders -> ServerError)
-> ResponseHeaders -> ByteString -> ServerError
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Int -> [Char] -> ByteString -> ResponseHeaders -> ServerError
ServerError Int
code [Char]
reason) ResponseHeaders
headers)
    in
        Maybe ByteString -> Either ServerError Response
maybeToEither (Maybe ByteString -> Either ServerError Response)
-> Maybe ByteString -> Either ServerError Response
forall a b. (a -> b) -> a -> b
$ Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Int
code Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
400) Maybe () -> Maybe ByteString -> Maybe ByteString
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Maybe ByteString
body

-- | Extract raw body of a response, only if it suitables for transformation.
-- Servant doesn't return files or streams by default, so if one of the two is
-- met, it means it comes from our application layer anyway.
responseBody
  :: Response
  -> Maybe ByteString
responseBody :: Response -> Maybe ByteString
responseBody = \case
    ResponseBuilder Status
_ ResponseHeaders
_ Builder
b ->
        ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (Builder -> ByteString
Binary.toLazyByteString Builder
b)
    ResponseRaw IO ByteString -> (ByteString -> IO ()) -> IO ()
_ Response
r ->
        Response -> Maybe ByteString
responseBody Response
r
    ResponseFile{} ->
        Maybe ByteString
forall a. Maybe a
Nothing
    ResponseStream{} ->
        Maybe ByteString
forall a. Maybe a
Nothing