{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}

-- |
-- Sometimes incoming requests don't stick to the
-- "no duplicate headers" invariant, for a number
-- of possible reasons (e.g. proxy servers blindly
-- adding headers), or your application (or other
-- middleware) blindly adds headers.
--
-- In those cases, you can use this 'Middleware'
-- to make sure that headers that /can/ be combined
-- /are/ combined. (e.g. applications might only
-- check the first \"Accept\" header and fail, while
-- there might be another one that would match)
module Network.Wai.Middleware.CombineHeaders (
    combineHeaders,
    CombineSettings,
    defaultCombineSettings,
    HeaderMap,
    HandleType,
    defaultHeaderMap,

    -- * Adjusting the settings
    setHeader,
    removeHeader,
    setHeaderMap,
    regular,
    keepOnly,
    setRequestHeaders,
    setResponseHeaders,
) where

import qualified Data.ByteString as B
import qualified Data.List as L (foldl', reverse)
import qualified Data.Map.Strict as M
import Data.Word8 (_comma, _space, _tab)
import Network.HTTP.Types (Header, HeaderName, RequestHeaders)
import qualified Network.HTTP.Types.Header as H
import Network.Wai (Middleware, mapResponseHeaders, requestHeaders)
import Network.Wai.Util (dropWhileEnd)

-- | The mapping of 'HeaderName' to 'HandleType'
type HeaderMap = M.Map HeaderName HandleType

-- | These settings define which headers should be combined,
-- if the combining should happen on incoming (request) headers
-- and if it should happen on outgoing (response) headers.
--
-- Any header you put in the header map *will* be used to
-- combine those headers with commas. There's no check to see
-- if it is a header that allows comma-separated lists, so if
-- you want to combine custom headers, go ahead.
--
-- (You can check the documentation of 'defaultCombineSettings'
-- to see which standard headers are specified to be able to be
-- combined)
--
-- @since 3.1.13.0
data CombineSettings = CombineSettings
    { combineHeaderMap :: HeaderMap
    -- ^ Which headers should be combined? And how? (cf. 'HandleType')
    , combineRequestHeaders :: Bool
    -- ^ Should request headers be combined?
    , combineResponseHeaders :: Bool
    -- ^ Should response headers be combined?
    }
    deriving (Eq, Show)

-- | Settings that combine request headers,
-- but don't touch response headers.
--
-- All types of headers that /can/ be combined
-- (as defined in the spec) /will/ be combined.
--
-- To be exact, this is the list:
--
-- * Accept
-- * Accept-CH
-- * Accept-Charset
-- * Accept-Encoding
-- * Accept-Language
-- * Accept-Post
-- * Access-Control-Allow-Headers
-- * Access-Control-Allow-Methods
-- * Access-Control-Expose-Headers
-- * Access-Control-Request-Headers
-- * Allow
-- * Alt-Svc @(KeepOnly \"clear\"")@
-- * Cache-Control
-- * Clear-Site-Data @(KeepOnly \"*\")@
-- * Connection
-- * Content-Encoding
-- * Content-Language
-- * Digest
-- * If-Match
-- * If-None-Match @(KeepOnly \"*\")@
-- * Link
-- * Permissions-Policy
-- * TE
-- * Timing-Allow-Origin @(KeepOnly \"*\")@
-- * Trailer
-- * Transfer-Encoding
-- * Upgrade
-- * Via
-- * Vary @(KeepOnly \"*\")@
-- * Want-Digest
--
-- N.B. Any header name that has \"KeepOnly\" after it
-- will be combined like normal, unless one of the values
-- is the one mentioned (\"*\" most of the time), then
-- that value is used and all others are dropped.
--
-- @since 3.1.13.0
defaultCombineSettings :: CombineSettings
defaultCombineSettings =
    CombineSettings
        { combineHeaderMap = defaultHeaderMap
        , combineRequestHeaders = True
        , combineResponseHeaders = False
        }

-- | Override the 'HeaderMap' of the 'CombineSettings'
--  (default: 'defaultHeaderMap')
--
-- @since 3.1.13.0
setHeaderMap :: HeaderMap -> CombineSettings -> CombineSettings
setHeaderMap mp set = set{combineHeaderMap = mp}

-- | Set whether the combining of headers should be applied to
-- the incoming request headers. (default: True)
--
-- @since 3.1.13.0
setRequestHeaders :: Bool -> CombineSettings -> CombineSettings
setRequestHeaders b set = set{combineRequestHeaders = b}

-- | Set whether the combining of headers should be applied to
-- the outgoing response headers. (default: False)
--
-- @since 3.1.13.0
setResponseHeaders :: Bool -> CombineSettings -> CombineSettings
setResponseHeaders b set = set{combineResponseHeaders = b}

-- | Convenience function to add a header to the header map or,
-- if it is already in the map, to change the 'HandleType'.
--
-- @since 3.1.13.0
setHeader :: HeaderName -> HandleType -> CombineSettings -> CombineSettings
setHeader name typ settings =
    settings
        { combineHeaderMap = M.insert name typ $ combineHeaderMap settings
        }

-- | Convenience function to remove a header from the header map.
--
-- @since 3.1.13.0
removeHeader :: HeaderName -> CombineSettings -> CombineSettings
removeHeader name settings =
    settings
        { combineHeaderMap = M.delete name $ combineHeaderMap settings
        }

-- | This middleware will reorganize the incoming and/or outgoing
-- headers in such a way that it combines any duplicates of
-- headers that, on their own, can normally have more than one
-- value, and any other headers will stay untouched.
--
-- This middleware WILL change the global order of headers
-- (they will be put in alphabetical order), but keep the
-- order of the same type of header. I.e. if there are 3
-- \"Set-Cookie\" headers, the first one will still be first,
-- the second one will still be second, etc. But now they are
-- guaranteed to be next to each other.
--
-- N.B. This 'Middleware' assumes the headers it combines
-- are correctly formatted. If one of the to-be-combined
-- headers is malformed, the new combined header will also
-- (probably) be malformed.
--
-- @since 3.1.13.0
combineHeaders :: CombineSettings -> Middleware
combineHeaders CombineSettings{..} app req resFunc =
    app newReq $ resFunc . adjustRes
  where
    newReq
        | combineRequestHeaders = req{requestHeaders = mkNewHeaders oldHeaders}
        | otherwise = req
    oldHeaders = requestHeaders req
    adjustRes
        | combineResponseHeaders = mapResponseHeaders mkNewHeaders
        | otherwise = id
    mkNewHeaders =
        M.foldrWithKey' finishHeaders [] . L.foldl' go mempty
    go acc hdr@(name, _) =
        M.alter (checkHeader hdr) name acc
    checkHeader :: Header -> Maybe HeaderHandling -> Maybe HeaderHandling
    checkHeader (name, newVal) =
        Just . \case
            Nothing -> (name `M.lookup` combineHeaderMap, [newVal])
            -- Yes, this reverses the order of headers, but these
            -- will be reversed again in 'finishHeaders'
            Just (mHandleType, hdrs) -> (mHandleType, newVal : hdrs)

-- | Unpack 'HeaderHandling' back into 'Header's again
finishHeaders
    :: HeaderName -> HeaderHandling -> RequestHeaders -> RequestHeaders
finishHeaders name (shouldCombine, xs) hdrs =
    case shouldCombine of
        Just typ -> (name, combinedHeader typ) : hdrs
        Nothing ->
            -- Yes, this reverses the headers, but they
            -- were already reversed by 'checkHeader'
            L.foldl' (\acc el -> (name, el) : acc) hdrs xs
  where
    combinedHeader Regular = combineHdrs xs
    combinedHeader (KeepOnly val)
        | val `elem` xs = val
        | otherwise = combineHdrs xs
    -- headers were reversed, so do 'reverse' before combining
    combineHdrs = B.intercalate ", " . fmap clean . L.reverse
    clean = dropWhileEnd $ \w -> w == _comma || w == _space || w == _tab

type HeaderHandling = (Maybe HandleType, [B.ByteString])

-- | Both will concatenate with @,@ (commas), but 'KeepOnly' will drop all
-- values except the given one if present (e.g. in case of wildcards/special values)
--
-- For example: If there are multiple @"Clear-Site-Data"@ headers, but one of
-- them is the wildcard @\"*\"@ value, using @'KeepOnly' "*"@ will cause all
-- others to be dropped and only the wildcard value to remain.
-- (The @\"*\"@ wildcard in this case means /ALL site data/ should be cleared,
-- so no need to include more)
--
-- @since 3.1.13.0
data HandleType
    = Regular
    | KeepOnly B.ByteString
    deriving (Eq, Show)

-- | Use the regular strategy when combining headers.
-- (i.e. merge into one header and separate values with commas)
--
-- @since 3.1.13.0
regular :: HandleType
regular = Regular

-- | Use the regular strategy when combining headers,
-- but if the exact supplied 'ByteString' is encountered
-- then discard all other values and only keep that value.
--
-- e.g. @keepOnly "*"@ will drop all other encountered values
--
-- @since 3.1.13.0
keepOnly :: B.ByteString -> HandleType
keepOnly = KeepOnly

-- | The default collection of HTTP headers that can be combined
-- in case there are multiples in one request or response.
--
-- See the documentation of 'defaultCombineSettings' for the exact list.
--
-- @since 3.1.13.0
defaultHeaderMap :: HeaderMap
defaultHeaderMap =
    M.fromList
        [ (H.hAccept, Regular)
        , ("Accept-CH", Regular)
        , (H.hAcceptCharset, Regular)
        , (H.hAcceptEncoding, Regular)
        , (H.hAcceptLanguage, Regular)
        , ("Accept-Post", Regular)
        , ("Access-Control-Allow-Headers", Regular) -- wildcard? yes, but can just add to list
        , ("Access-Control-Allow-Methods", Regular) -- wildcard? yes, but can just add to list
        , ("Access-Control-Expose-Headers", Regular) -- wildcard? yes, but can just add to list
        , ("Access-Control-Request-Headers", Regular)
        , (H.hAllow, Regular)
        , ("Alt-Svc", KeepOnly "clear") -- special "clear" value (if any is "clear", only keep that one)
        , (H.hCacheControl, Regular)
        , ("Clear-Site-Data", KeepOnly "*") -- wildcard (if any is "*", only keep that one)
        , -- If "close" and anything else is used together, it's already F-ed,
          -- so just combine them.
          (H.hConnection, Regular)
        , (H.hContentEncoding, Regular)
        , (H.hContentLanguage, Regular)
        , ("Digest", Regular)
        , -- We could handle this, but it's experimental AND
          -- will be replaced by "Permissions-Policy"
          -- , "Feature-Policy" -- "semicolon ';' separated"

          (H.hIfMatch, Regular)
        , (H.hIfNoneMatch, KeepOnly "*") -- wildcard? (if any is "*", only keep that one)
        , ("Link", Regular)
        , ("Permissions-Policy", Regular)
        , (H.hTE, Regular)
        , ("Timing-Allow-Origin", KeepOnly "*") -- wildcard? (if any is "*", only keep that one)
        , (H.hTrailer, Regular)
        , (H.hTransferEncoding, Regular)
        , (H.hUpgrade, Regular)
        , (H.hVia, Regular)
        , (H.hVary, KeepOnly "*") -- wildcard? (if any is "*", only keep that one)
        , ("Want-Digest", Regular)
        ]