cmd/forwarder/run: add -H, --header and -R --response-header flags

saucelabs · Mar 27, 2023 · 15d0495 · 15d0495
1 parent 1739270
commit 15d0495
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 1 deletion.
diff --git a/bind/flag.go b/bind/flag.go
@@ -16,6 +16,7 @@ import (
 	"github.com/mmatczuk/anyflag"
 	"github.com/saucelabs/forwarder"
 	"github.com/saucelabs/forwarder/fileurl"
+	"github.com/saucelabs/forwarder/header"
 	"github.com/saucelabs/forwarder/httplog"
 	"github.com/saucelabs/forwarder/log"
 	"github.com/spf13/cobra"
@@ -62,6 +63,27 @@ func PAC(fs *pflag.FlagSet, pac **url.URL) {
 			"It can be a local file or a URL, you can also use '-' to read from stdin. ")
 }
 
+func RequestHeaders(fs *pflag.FlagSet, headers *[]header.Header) {
+	fs.VarP(anyflag.NewSliceValue[header.Header](*headers, headers, header.ParseHeader),
+		"header", "H", "<header>"+
+			"Add or remove HTTP request headers. "+
+			"Use the format \"name: value\" to add a header, "+
+			"\"name;\" to set the header to empty value, "+
+			"\"-name\" to remove the header, "+
+			"\"-name*\" to remove headers by prefix. "+
+			"The header name will be normalized to canonical form. "+
+			"The header value should not contain any newlines or carriage returns. "+
+			"The flag can be specified multiple times. "+
+			"Example: -H \"Host: example.com\" -H \"-User-Agent\" -H \"-X-*\". ")
+}
+
+func ResponseHeaders(fs *pflag.FlagSet, headers *[]header.Header) {
+	fs.VarP(anyflag.NewSliceValue[header.Header](*headers, headers, header.ParseHeader),
+		"response-header", "R", "<header>"+
+			"Add or remove HTTP headers on the received response before sending it to the client. "+
+			"See the documentation for the -H, --header flag for more details on the format. ")
+}
+
 func HTTPProxyConfig(fs *pflag.FlagSet, cfg *forwarder.HTTPProxyConfig, lcfg *log.Config) {
 	HTTPServerConfig(fs, &cfg.HTTPServerConfig, "", forwarder.HTTPScheme, forwarder.HTTPSScheme)
 	LogConfig(fs, lcfg)

diff --git a/cmd/forwarder/run/run.go b/cmd/forwarder/run/run.go
@@ -17,6 +17,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/saucelabs/forwarder"
 	"github.com/saucelabs/forwarder/bind"
+	"github.com/saucelabs/forwarder/header"
 	"github.com/saucelabs/forwarder/log"
 	"github.com/saucelabs/forwarder/log/stdlog"
 	"github.com/saucelabs/forwarder/pac"
@@ -31,6 +32,8 @@ type command struct {
 	httpTransportConfig *forwarder.HTTPTransportConfig
 	pac                 *url.URL
 	credentials         []*forwarder.HostPortUser
+	requestHeaders      []header.Header
+	responseHeaders     []header.Header
 	httpProxyConfig     *forwarder.HTTPProxyConfig
 	apiServerConfig     *forwarder.HTTPServerConfig
 	logConfig           *log.Config
@@ -95,6 +98,13 @@ func (c *command) RunE(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("credentials: %w", err)
 	}
 
+	if len(c.requestHeaders) > 0 {
+		c.httpProxyConfig.RequestModifiers = append(c.httpProxyConfig.RequestModifiers, header.Headers(c.requestHeaders))
+	}
+	if len(c.responseHeaders) > 0 {
+		c.httpProxyConfig.ResponseModifiers = append(c.httpProxyConfig.ResponseModifiers, header.Headers(c.responseHeaders))
+	}
+
 	var f runctx.Funcs
 	p, err := forwarder.NewHTTPProxy(c.httpProxyConfig, pr, cm, rt, logger.Named("proxy"))
 	if err != nil {
@@ -136,8 +146,10 @@ func Command() (cmd *cobra.Command) {
 
 	defer func() {
 		fs := cmd.Flags()
-		bind.HTTPProxyConfig(fs, c.httpProxyConfig, c.logConfig)
 		bind.Credentials(fs, &c.credentials)
+		bind.RequestHeaders(fs, &c.requestHeaders)
+		bind.ResponseHeaders(fs, &c.responseHeaders)
+		bind.HTTPProxyConfig(fs, c.httpProxyConfig, c.logConfig)
 		bind.PAC(fs, &c.pac)
 		bind.DNSConfig(fs, c.dnsConfig)
 		bind.HTTPServerConfig(fs, c.apiServerConfig, "api", forwarder.HTTPScheme)