From 43bfe338ab667f62a3bcf943fc63c1ef7fe19442 Mon Sep 17 00:00:00 2001
From: Michael Quigley <michael@quigley.com>
Date: Wed, 7 Aug 2024 13:49:31 -0400
Subject: [PATCH] new '--response-header' flag in 'zrok access private' (#522)

---
 cmd/zrok/accessPrivate.go   | 11 +++++++----
 endpoints/proxy/frontend.go | 20 +++++++++++++++-----
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/cmd/zrok/accessPrivate.go b/cmd/zrok/accessPrivate.go
index 3d1d5c39..1bb41be1 100644
--- a/cmd/zrok/accessPrivate.go
+++ b/cmd/zrok/accessPrivate.go
@@ -28,9 +28,10 @@ func init() {
 }
 
 type accessPrivateCommand struct {
-	bindAddress string
-	headless    bool
-	cmd         *cobra.Command
+	bindAddress     string
+	headless        bool
+	responseHeaders []string
+	cmd             *cobra.Command
 }
 
 func newAccessPrivateCommand() *accessPrivateCommand {
@@ -41,8 +42,9 @@ func newAccessPrivateCommand() *accessPrivateCommand {
 	}
 	command := &accessPrivateCommand{cmd: cmd}
 	cmd.Flags().BoolVar(&command.headless, "headless", false, "Disable TUI and run headless")
-	cmd.Run = command.run
 	cmd.Flags().StringVarP(&command.bindAddress, "bind", "b", "127.0.0.1:9191", "The address to bind the private frontend")
+	cmd.Flags().StringArrayVar(&command.responseHeaders, "response-header", []string{}, "Add a response header ('key:value')")
+	cmd.Run = command.run
 	return command
 }
 
@@ -194,6 +196,7 @@ func (cmd *accessPrivateCommand) run(_ *cobra.Command, args []string) {
 		cfg := proxy.DefaultFrontendConfig(env.EnvironmentIdentityName())
 		cfg.ShrToken = shrToken
 		cfg.Address = cmd.bindAddress
+		cfg.ResponseHeaders = cmd.responseHeaders
 		cfg.RequestsChan = requests
 		fe, err := proxy.NewFrontend(cfg)
 		if err != nil {
diff --git a/endpoints/proxy/frontend.go b/endpoints/proxy/frontend.go
index 9660270e..f868f0aa 100644
--- a/endpoints/proxy/frontend.go
+++ b/endpoints/proxy/frontend.go
@@ -15,15 +15,17 @@ import (
 	"net/http"
 	"net/http/httputil"
 	"net/url"
+	"strings"
 	"time"
 )
 
 type FrontendConfig struct {
-	IdentityName string
-	ShrToken     string
-	Address      string
-	Tls          *endpoints.TlsConfig
-	RequestsChan chan *endpoints.Request
+	IdentityName    string
+	ShrToken        string
+	Address         string
+	ResponseHeaders []string
+	Tls             *endpoints.TlsConfig
+	RequestsChan    chan *endpoints.Request
 }
 
 func DefaultFrontendConfig(identityName string) *FrontendConfig {
@@ -112,6 +114,14 @@ func newServiceProxy(cfg *FrontendConfig, ctx ziti.Context) (*httputil.ReversePr
 		req.Header.Set("X-Proxy", "zrok")
 	}
 	proxy.ModifyResponse = func(resp *http.Response) error {
+		for _, responseHeader := range cfg.ResponseHeaders {
+			tokens := strings.Split(responseHeader, ":")
+			if len(tokens) == 2 {
+				resp.Header.Set(strings.TrimSpace(tokens[0]), strings.TrimSpace(tokens[1]))
+			} else {
+				logrus.Errorf("invalid response header '%v' (expecting header:value", responseHeader)
+			}
+		}
 		return nil
 	}
 	proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {