Merge pull request #48 from adinapoli/adinapoli/wpsModifyResponseHeaders

Add `wpsModifyResponseHeaders` function
fpco · Jun 25, 2024 · 08d3b8e · 08d3b8e
2 parents 77a936e + 22c6027
commit 08d3b8e
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 4 deletions.
diff --git 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
@@ -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.1.0
     }
 
 -- | 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

diff --git 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