From acc23b94916e286c0ef141ec6e330f89839c0fa2 Mon Sep 17 00:00:00 2001
From: Alfredo Di Napoli <alfredo@well-typed.com>
Date: Mon, 17 Jun 2024 11:53:41 +0200
Subject: [PATCH 1/2] Add wpsModifyResponseHeaders function

This commit adds a new function to the `WaiProxySettings` that allows
modifying the response headers before they are sent upstream.
---
 Network/HTTP/ReverseProxy.hs | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/Network/HTTP/ReverseProxy.hs b/Network/HTTP/ReverseProxy.hs
index 5d113d0..5f4a40a 100644
--- a/Network/HTTP/ReverseProxy.hs
+++ b/Network/HTTP/ReverseProxy.hs
@@ -27,6 +27,7 @@ module Network.HTTP.ReverseProxy
     , wpsUpgradeToRaw
     , wpsGetDest
     , wpsLogRequest
+    , wpsModifyResponseHeaders
     , SetIpHeader (..)
       -- *** Local settings
     , LocalWaiProxySettings
@@ -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.
@@ -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
@@ -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)
@@ -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

From 22c6027beb3ba20b6e039bd436cc6f6230f71c7a Mon Sep 17 00:00:00 2001
From: Alfredo Di Napoli <alfredo@well-typed.com>
Date: Mon, 24 Jun 2024 07:31:11 +0200
Subject: [PATCH 2/2] Update changelong, bump to 0.6.1.0

---
 ChangeLog.md                 | 4 ++++
 Network/HTTP/ReverseProxy.hs | 2 +-
 http-reverse-proxy.cabal     | 2 +-
 3 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/ChangeLog.md b/ChangeLog.md
index 65c474d..275a027 100644
--- a/ChangeLog.md
+++ b/ChangeLog.md
@@ -1,4 +1,8 @@
 
+## 0.6.1.0
+
+* Add the `wpsModifyResponseHeaders` option to `WaiProxySettings` to tweak response headers before they are returned upstream. [#48](https://github.com/fpco/http-reverse-proxy/pull/48)
+
 ## 0.6.0.3
 
 * Fix a regression introduced in 0.6.0.2: wrong 'Content-Length' header is preserved for responses with encoded content. [#47](https://github.com/fpco/http-reverse-proxy/pull/47)
diff --git a/Network/HTTP/ReverseProxy.hs b/Network/HTTP/ReverseProxy.hs
index 5f4a40a..db3ad68 100644
--- a/Network/HTTP/ReverseProxy.hs
+++ b/Network/HTTP/ReverseProxy.hs
@@ -279,7 +279,7 @@ data WaiProxySettings = WaiProxySettings
     --
     -- Default: no op
     --
-    -- @since 0.6.0.4
+    -- @since 0.6.1.0
     }
 
 -- | How to set the X-Real-IP request header.
diff --git a/http-reverse-proxy.cabal b/http-reverse-proxy.cabal
index 2614837..b952628 100644
--- a/http-reverse-proxy.cabal
+++ b/http-reverse-proxy.cabal
@@ -1,5 +1,5 @@
 name:                http-reverse-proxy
-version:             0.6.0.3
+version:             0.6.1.0
 synopsis:            Reverse proxy HTTP requests, either over raw sockets or with WAI
 description:         Provides a simple means of reverse-proxying HTTP requests. The raw approach uses the same technique as leveraged by keter, whereas the WAI approach performs full request/response parsing via WAI and http-conduit.
 homepage:            https://github.com/fpco/http-reverse-proxy