serve: support custom http response headers

Co-authored-by: Tim Schumacher <tim@tschumacher.net>

fs/config/configflags/configflags.go

@@ -61,23 +61,6 @@ func AddFlags(ci *fs.ConfigInfo, flagSet *pflag.FlagSet) {
 	flags.StringVarP(flagSet, &dscp, "dscp", "", "", "Set DSCP value to connections, value or name, e.g. CS1, LE, DF, AF21", "Networking")
 }
 
-// ParseHeaders converts the strings passed in via the header flags into HTTPOptions
-func ParseHeaders(headers []string) []*fs.HTTPOption {
-	opts := []*fs.HTTPOption{}
-	for _, header := range headers {
-		parts := strings.SplitN(header, ":", 2)
-		if len(parts) == 1 {
-			fs.Fatalf(nil, "Failed to parse '%s' as an HTTP header. Expecting a string like: 'Content-Encoding: gzip'", header)
-		}
-		option := &fs.HTTPOption{
-			Key:   strings.TrimSpace(parts[0]),
-			Value: strings.TrimSpace(parts[1]),
-		}
-		opts = append(opts, option)
-	}
-	return opts
-}
-
 // SetFlags sets flags which aren't part of the config system
 func SetFlags(ci *fs.ConfigInfo) {
 	// Process obsolete --dump-headers and --dump-bodies flags
@@ -153,13 +136,13 @@ func SetFlags(ci *fs.ConfigInfo) {
 
 	// Process --headers-upload, --headers-download, --headers
 	if len(uploadHeaders) != 0 {
-		ci.UploadHeaders = ParseHeaders(uploadHeaders)
+		ci.UploadHeaders = fs.MustParseHeaders(uploadHeaders)
 	}
 	if len(downloadHeaders) != 0 {
-		ci.DownloadHeaders = ParseHeaders(downloadHeaders)
+		ci.DownloadHeaders = fs.MustParseHeaders(downloadHeaders)
 	}
 	if len(headers) != 0 {
-		ci.Headers = ParseHeaders(headers)
+		ci.Headers = fs.MustParseHeaders(headers)
 	}
 
 	// Process --metadata-set

fs/open_options.go

@@ -198,6 +198,32 @@ func (o *SeekOption) Mandatory() bool {
 	return true
 }
 
+// ParseHeaders converts the strings passed in via the header flags into HTTPOptions
+func ParseHeaders(headers []string) ([]*HTTPOption, error) {
+	opts := []*HTTPOption{}
+	for _, header := range headers {
+		parts := strings.SplitN(header, ":", 2)
+		if len(parts) != 2 {
+			return nil, fmt.Errorf("failed to parse %q as an HTTP header. Expecting a string like: 'Cache-Control: no-store'", header)
+		}
+		option := &HTTPOption{
+			Key:   strings.TrimSpace(parts[0]),
+			Value: strings.TrimSpace(parts[1]),
+		}
+		opts = append(opts, option)
+	}
+	return opts, nil
+}
+
+// MustParseHeaders converts the strings passed in via the header flags into HTTPOptions or exits.
+func MustParseHeaders(headers []string) []*HTTPOption {
+	opts, err := ParseHeaders(headers)
+	if err != nil {
+		Fatalf(nil, "%v", err)
+	}
+	return opts
+}
+
 // HTTPOption defines a general purpose HTTP option
 type HTTPOption struct {
 	Key   string

lib/http/middleware.go

@@ -199,6 +199,18 @@ func MiddlewareCORS(allowOrigin string) Middleware {
 	}
 }
 
+// MiddlewareResponseHeaders instantiates middleware that apply custom headers to every response
+func MiddlewareResponseHeaders(headers []*fs.HTTPOption) Middleware {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			for _, header := range headers {
+				w.Header().Set(header.Key, header.Value)
+			}
+			next.ServeHTTP(w, r)
+		})
+	}
+}
+
 // MiddlewareStripPrefix instantiates middleware that removes the BaseURL from the path
 func MiddlewareStripPrefix(prefix string) Middleware {
 	return func(next http.Handler) http.Handler {

lib/http/middleware_test.go

@@ -583,3 +583,66 @@ func TestMiddlewareCORSWithAuth(t *testing.T) {
 		})
 	}
 }
+
+func TestMiddlewareResponseHeaders(t *testing.T) {
+	servers := []struct {
+		name   string
+		http   Config
+		header http.Header
+	}{
+		{
+			name: "SingleHeader",
+			http: Config{
+				ListenAddr:      []string{"127.0.0.1:0"},
+				ResponseHeaders: []string{"X-Test-Header: test-value"},
+			},
+			header: http.Header{
+				"X-Test-Header": []string{"test-value"},
+			},
+		},
+		{
+			name: "MultipleHeaders",
+			http: Config{
+				ListenAddr:      []string{"127.0.0.1:0"},
+				ResponseHeaders: []string{"X-Header-One: one", "X-Header-Two: two"},
+			},
+			header: http.Header{
+				"X-Header-One": []string{"one"},
+				"X-Header-Two": []string{"two"},
+			},
+		},
+	}
+
+	for _, ss := range servers {
+		t.Run(ss.name, func(t *testing.T) {
+			s, err := NewServer(context.Background(), WithConfig(ss.http))
+			require.NoError(t, err)
+			defer func() {
+				require.NoError(t, s.Shutdown())
+			}()
+
+			expected := []byte("header-test")
+			s.Router().Mount("/", testEchoHandler(expected))
+			s.Serve()
+
+			url := testGetServerURL(t, s)
+			client := &http.Client{}
+			req, err := http.NewRequest("GET", url, nil)
+			require.NoError(t, err)
+
+			resp, err := client.Do(req)
+			require.NoError(t, err)
+			defer func() {
+				_ = resp.Body.Close()
+			}()
+
+			require.Equal(t, http.StatusOK, resp.StatusCode, "should return ok")
+			testExpectRespBody(t, resp, expected)
+
+			for key, vals := range ss.header {
+				require.Contains(t, resp.Header, key, "response should contain custom header")
+				require.Equal(t, vals, resp.Header.Values(key), "header value should match")
+			}
+		})
+	}
+}

lib/http/server.go

@@ -53,6 +53,10 @@ for a transfer.
 ` + "`--{{ .Prefix }}max-header-bytes`" + ` controls the maximum number of bytes the server will
 accept in the HTTP header.
 
+` + "`--{{ .Prefix }}response-header`" + ` can be used to set an HTTP header for all responses,
+will overriding existing values. The flag may be repeated to add multiple
+headers. Use the format ` + "`Header-Name: value`" + `.
+
 ` + "`--{{ .Prefix }}baseurl`" + ` controls the URL prefix that rclone serves from.  By default
 rclone will serve from the root.  If you used ` + "`--{{ .Prefix }}baseurl \"/rclone\"`" + ` then
 rclone would serve from a URL starting with "/rclone/".  This is
@@ -167,6 +171,10 @@ var ConfigInfo = fs.Options{{
 	Name:    "allow_origin",
 	Default: "",
 	Help:    "Origin which cross-domain request (CORS) can be executed from",
+}, {
+	Name:    "response_header",
+	Default: []string{},
+	Help:    "Set HTTP header for all responses, overriding existing values",
 }}
 
 // Config contains options for the http Server
@@ -181,8 +189,9 @@ type Config struct {
 	TLSCertBody        []byte      `config:"-"`                    // TLS PEM public key certificate body (can also include intermediate/CA certificates), ignores TLSCert
 	TLSKeyBody         []byte      `config:"-"`                    // TLS PEM private key body, ignores TLSKey
 	ClientCA           string      `config:"client_ca"`            // Path to TLS PEM CA file with certificate authorities to verify clients with
-	MinTLSVersion      string      `config:"min_tls_version"`      // MinTLSVersion contains the minimum TLS version that is acceptable.
+	MinTLSVersion      string      `config:"min_tls_version"`      // MinTLSVersion contains the minimum TLS version that is acceptable
 	AllowOrigin        string      `config:"allow_origin"`         // AllowOrigin sets the Access-Control-Allow-Origin header
+	ResponseHeaders    []string    `config:"response_header"`      // Set HTTP header for all responses, overriding existing values
 }
 
 // AddFlagsPrefix adds flags for the httplib
@@ -197,6 +206,7 @@ func (cfg *Config) AddFlagsPrefix(flagSet *pflag.FlagSet, prefix string) {
 	flags.StringVarP(flagSet, &cfg.BaseURL, prefix+"baseurl", "", cfg.BaseURL, "Prefix for URLs - leave blank for root", prefix)
 	flags.StringVarP(flagSet, &cfg.MinTLSVersion, prefix+"min-tls-version", "", cfg.MinTLSVersion, "Minimum TLS version that is acceptable", prefix)
 	flags.StringVarP(flagSet, &cfg.AllowOrigin, prefix+"allow-origin", "", cfg.AllowOrigin, "Origin which cross-domain request (CORS) can be executed from", prefix)
+	flags.StringArrayVarP(flagSet, &cfg.ResponseHeaders, prefix+"response-header", "", cfg.ResponseHeaders, "Set HTTP header for all responses, overriding existing values", prefix)
 }
 
 // AddHTTPFlagsPrefix adds flags for the httplib
@@ -348,7 +358,13 @@ func NewServer(ctx context.Context, options ...Option) (*Server, error) {
 		return nil, err
 	}
 
+	responseHeaders, err := fs.ParseHeaders(s.cfg.ResponseHeaders)
+	if err != nil {
+		return nil, err
+	}
+
 	s.mux.Use(MiddlewareCORS(s.cfg.AllowOrigin))
+	s.mux.Use(MiddlewareResponseHeaders(responseHeaders))
 
 	s.initAuth()