Skip to content

Commit

Permalink
Add wpsModifyResponseHeaders function
Browse files Browse the repository at this point in the history
This commit adds a new function to the `WaiProxySettings` that allows
modifying the response headers before they are sent upstream.
  • Loading branch information
adinapoli committed Jun 17, 2024
1 parent 77a936e commit acc23b9
Showing 1 changed file with 13 additions and 3 deletions.
16 changes: 13 additions & 3 deletions Network/HTTP/ReverseProxy.hs
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ module Network.HTTP.ReverseProxy
, wpsUpgradeToRaw
, wpsGetDest
, wpsLogRequest
, wpsModifyResponseHeaders
, SetIpHeader (..)
-- *** Local settings
, LocalWaiProxySettings
Expand Down Expand Up @@ -272,6 +273,13 @@ data WaiProxySettings = WaiProxySettings
-- Default: no op
--
-- @since 0.6.0.1
, wpsModifyResponseHeaders :: WAI.Request -> HC.Response () -> HT.ResponseHeaders -> HT.ResponseHeaders
-- ^ Allow to override the response headers before the response is returned upstream. Useful for example
-- to override overly-strict 'Content-Security-Policy' when the source is known to be trustworthy.
--
-- Default: no op
--
-- @since 0.6.0.4
}

-- | How to set the X-Real-IP request header.
Expand All @@ -294,6 +302,7 @@ defaultWaiProxySettings = WaiProxySettings
(CI.mk <$> lookup "upgrade" (WAI.requestHeaders req)) == Just "websocket"
, wpsGetDest = Nothing
, wpsLogRequest = const (pure ())
, wpsModifyResponseHeaders = \_ _ -> id
}

renderHeaders :: WAI.Request -> HT.RequestHeaders -> Builder
Expand Down Expand Up @@ -422,9 +431,10 @@ waiProxyToSettings getDest wps' manager req0 sendResponse = do
$ \case
Left e -> wpsOnExc wps e req sendResponse
Right res -> do
let conduit = fromMaybe
let res' = const () <$> res
conduit = fromMaybe
(awaitForever (\bs -> yield (Chunk $ fromByteString bs) >> yield Flush))
(wpsProcessBody wps req $ const () <$> res)
(wpsProcessBody wps req res')
src = bodyReaderSource $ HC.responseBody res
headers = HC.responseHeaders res
notEncoded = isNothing (lookup "content-encoding" headers)
Expand All @@ -433,7 +443,7 @@ waiProxyToSettings getDest wps' manager req0 sendResponse = do
(HC.responseStatus res)
(filter (\(key, v) -> not (key `Set.member` strippedHeaders) ||
key == "content-length" && (notEncoded && notChunked || v == "0"))
headers)
(wpsModifyResponseHeaders wps req res' headers))
(\sendChunk flush -> runConduit $ src .| conduit .| CL.mapM_ (\mb ->
case mb of
Flush -> flush
Expand Down

0 comments on commit acc23b9

Please sign in to comment.