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 {
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 {

View File

@ -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) {