new '--response-header' flag in 'zrok access private' (#522)

This commit is contained in:
Michael Quigley 2024-08-07 13:49:31 -04:00
parent 6561a3b6bb
commit 43bfe338ab
No known key found for this signature in database
GPG Key ID: 9B60314A9DD20A62
2 changed files with 22 additions and 9 deletions

View File

@ -28,9 +28,10 @@ func init() {
} }
type accessPrivateCommand struct { type accessPrivateCommand struct {
bindAddress string bindAddress string
headless bool headless bool
cmd *cobra.Command responseHeaders []string
cmd *cobra.Command
} }
func newAccessPrivateCommand() *accessPrivateCommand { func newAccessPrivateCommand() *accessPrivateCommand {
@ -41,8 +42,9 @@ func newAccessPrivateCommand() *accessPrivateCommand {
} }
command := &accessPrivateCommand{cmd: cmd} command := &accessPrivateCommand{cmd: cmd}
cmd.Flags().BoolVar(&command.headless, "headless", false, "Disable TUI and run headless") 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().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 return command
} }
@ -194,6 +196,7 @@ func (cmd *accessPrivateCommand) run(_ *cobra.Command, args []string) {
cfg := proxy.DefaultFrontendConfig(env.EnvironmentIdentityName()) cfg := proxy.DefaultFrontendConfig(env.EnvironmentIdentityName())
cfg.ShrToken = shrToken cfg.ShrToken = shrToken
cfg.Address = cmd.bindAddress cfg.Address = cmd.bindAddress
cfg.ResponseHeaders = cmd.responseHeaders
cfg.RequestsChan = requests cfg.RequestsChan = requests
fe, err := proxy.NewFrontend(cfg) fe, err := proxy.NewFrontend(cfg)
if err != nil { if err != nil {

View File

@ -15,15 +15,17 @@ import (
"net/http" "net/http"
"net/http/httputil" "net/http/httputil"
"net/url" "net/url"
"strings"
"time" "time"
) )
type FrontendConfig struct { type FrontendConfig struct {
IdentityName string IdentityName string
ShrToken string ShrToken string
Address string Address string
Tls *endpoints.TlsConfig ResponseHeaders []string
RequestsChan chan *endpoints.Request Tls *endpoints.TlsConfig
RequestsChan chan *endpoints.Request
} }
func DefaultFrontendConfig(identityName string) *FrontendConfig { func DefaultFrontendConfig(identityName string) *FrontendConfig {
@ -112,6 +114,14 @@ func newServiceProxy(cfg *FrontendConfig, ctx ziti.Context) (*httputil.ReversePr
req.Header.Set("X-Proxy", "zrok") req.Header.Set("X-Proxy", "zrok")
} }
proxy.ModifyResponse = func(resp *http.Response) error { 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 return nil
} }
proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) { proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {