Implement copy_response_headers, with include/exclude list support

caddyconfig/httpcaddyfile/directives.go

@@ -41,6 +41,7 @@ var directiveOrder = []string{
 	"root",
 
 	"header",
+	"copy_response_headers",
 	"request_body",
 
 	"redir",

caddytest/integration/caddyfile_adapt/reverse_proxy_handle_response.txt

@@ -19,8 +19,7 @@ reverse_proxy 127.0.0.1:65535 {
 	}
 
 	handle_response {
-		header Foo "This should be last in the JSON!"
-		copy_response 404
+		respond "Any! This should be last in the JSON!"
 	}
 
 	@403 {
@@ -42,6 +41,20 @@ reverse_proxy 127.0.0.1:65535 {
 
 	@changeStatus status 500
 	handle_response @changeStatus 400
+
+	@200 status 200
+	handle_response @200 {
+		copy_response_headers {
+			include Foo Bar
+		}
+		respond "Copied headers from the response"
+	}
+
+	@201 status 201
+	handle_response @201 {
+		header Foo "Copying the response"
+		copy_response 404
+	}
 }
 ----------
 {
@@ -165,6 +178,35 @@ reverse_proxy 127.0.0.1:65535 {
 											"status_code": 400
 										},
 										{
+											"match": {
+												"status_code": [
+													200
+												]
+											},
+											"routes": [
+												{
+													"handle": [
+														{
+															"handler": "copy_response_headers",
+															"include": [
+																"Foo",
+																"Bar"
+															]
+														},
+														{
+															"body": "Copied headers from the response",
+															"handler": "static_response"
+														}
+													]
+												}
+											]
+										},
+										{
+											"match": {
+												"status_code": [
+													201
+												]
+											},
 											"routes": [
 												{
 													"handle": [
@@ -173,7 +215,7 @@ reverse_proxy 127.0.0.1:65535 {
 															"response": {
 																"set": {
 																	"Foo": [
-																		"This should be last in the JSON!"
+																		"Copying the response"
 																	]
 																}
 															}
@@ -185,6 +227,18 @@ reverse_proxy 127.0.0.1:65535 {
 													]
 												}
 											]
+										},
+										{
+											"routes": [
+												{
+													"handle": [
+														{
+															"body": "Any! This should be last in the JSON!",
+															"handler": "static_response"
+														}
+													]
+												}
+											]
 										}
 									],
 									"handler": "reverse_proxy",

cmd/caddy/Caddyfile.test2

@@ -0,0 +1,58 @@
+:8884
+
+reverse_proxy 127.0.0.1:65535 {
+	@accel header X-Accel-Redirect *
+	handle_response @accel {
+		respond "Header X-Accel-Redirect!"
+	}
+
+	@another {
+		header X-Another *
+	}
+	handle_response @another {
+		respond "Header X-Another!"
+	}
+
+	@401 status 401
+	handle_response @401 {
+		respond "Status 401!"
+	}
+
+	handle_response {
+		respond "Any! This should be last in the JSON!"
+	}
+
+	@403 {
+		status 403
+	}
+	handle_response @403 {
+		respond "Status 403!"
+	}
+
+	@multi {
+		status 401 403
+		status 404
+		header Foo *
+		header Bar *
+	}
+	handle_response @multi {
+		respond "Headers Foo, Bar AND statuses 401, 403 and 404!"
+	}
+
+	@changeStatus status 500
+	handle_response @changeStatus 400
+
+	@200 status 200
+	handle_response @200 {
+		copy_response_headers {
+			include Foo Bar
+		}
+		respond "Copied headers from the response"
+	}
+
+	@201 status 201
+	handle_response @201 {
+		header Foo "Copying the response"
+		copy_response 404
+	}
+}

modules/caddyhttp/reverseproxy/caddyfile.go

@@ -35,6 +35,7 @@ import (
 func init() {
 	httpcaddyfile.RegisterHandlerDirective("reverse_proxy", parseCaddyfile)
 	httpcaddyfile.RegisterHandlerDirective("copy_response", parseCopyResponseCaddyfile)
+	httpcaddyfile.RegisterHandlerDirective("copy_response_headers", parseCopyResponseHeadersCaddyfile)
 }
 
 func parseCaddyfile(h httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, error) {
@@ -1065,6 +1066,44 @@ func (h *CopyResponseHandler) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
 	return nil
 }
 
+func parseCopyResponseHeadersCaddyfile(h httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, error) {
+	crh := new(CopyResponseHeadersHandler)
+	err := crh.UnmarshalCaddyfile(h.Dispenser)
+	if err != nil {
+		return nil, err
+	}
+	return crh, nil
+}
+
+// UnmarshalCaddyfile sets up the handler from Caddyfile tokens. Syntax:
+//
+//    copy_response_headers [<matcher>] {
+//        exclude <fields...>
+//    }
+//
+func (h *CopyResponseHeadersHandler) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
+	for d.Next() {
+		args := d.RemainingArgs()
+		if len(args) > 0 {
+			return d.ArgErr()
+		}
+
+		for d.NextBlock(0) {
+			switch d.Val() {
+			case "include":
+				h.Include = append(h.Include, d.RemainingArgs()...)
+
+			case "exclude":
+				h.Exclude = append(h.Exclude, d.RemainingArgs()...)
+
+			default:
+				return d.Errf("unrecognized subdirective '%s'", d.Val())
+			}
+		}
+	}
+	return nil
+}
+
 const matcherPrefix = "@"
 
 // Interface guards

modules/caddyhttp/reverseproxy/copyresponse.go

@@ -26,11 +26,12 @@ import (
 
 func init() {
 	caddy.RegisterModule(CopyResponseHandler{})
+	caddy.RegisterModule(CopyResponseHeadersHandler{})
 }
 
-// CopyResponseHandler is a special HTTP handler which may only be used
-// within reverse_proxy's handle_response routes, to copy the response
-// from the
+// CopyResponseHandler is a special HTTP handler which may
+// only be used within reverse_proxy's handle_response routes,
+// to copy the proxy response.
 type CopyResponseHandler struct {
 	// To write the upstream response's body but with a different
 	// status code, set this field to the desired status code.
@@ -53,6 +54,7 @@ func (h *CopyResponseHandler) Provision(ctx caddy.Context) error {
 	return nil
 }
 
+// ServeHTTP implements the Handler interface.
 func (h CopyResponseHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request, _ caddyhttp.Handler) error {
 	repl := req.Context().Value(caddy.ReplacerCtxKey).(*caddy.Replacer)
 	hrc, ok := req.Context().Value(proxyHandleResponseContextCtxKey).(*handleResponseContext)
@@ -81,8 +83,107 @@ func (h CopyResponseHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request
 	return hrc.handler.finalizeResponse(rw, req, hrc.response, repl, hrc.start, hrc.logger, false)
 }
 
+// CopyResponseHeadersHandler is a special HTTP handler which may
+// only be used within reverse_proxy's handle_response routes,
+// to copy headers from the proxy response.
+type CopyResponseHeadersHandler struct {
+	// A list of header fields to copy from the response.
+	// Cannot be defined at the same time as Exclude.
+	Include []string `json:"include,omitempty"`
+
+	// A list of header fields to skip copying from the response.
+	// Cannot be defined at the same time as Include.
+	Exclude []string `json:"exclude,omitempty"`
+
+	includeMap map[string]struct{}
+	excludeMap map[string]struct{}
+	ctx        caddy.Context
+}
+
+// CaddyModule returns the Caddy module information.
+func (CopyResponseHeadersHandler) CaddyModule() caddy.ModuleInfo {
+	return caddy.ModuleInfo{
+		ID:  "http.handlers.copy_response_headers",
+		New: func() caddy.Module { return new(CopyResponseHeadersHandler) },
+	}
+}
+
+// Validate ensures the h's configuration is valid.
+func (h *CopyResponseHeadersHandler) Validate() error {
+	if len(h.Exclude) > 0 && len(h.Include) > 0 {
+		return fmt.Errorf("cannot define both 'exclude' and 'include' lists at the same time")
+	}
+
+	return nil
+}
+
+// Provision ensures that h is set up properly before use.
+func (h *CopyResponseHeadersHandler) Provision(ctx caddy.Context) error {
+	h.ctx = ctx
+
+	// Optimize the include list by converting it to a map
+	for _, field := range h.Include {
+		if h.includeMap == nil {
+			h.includeMap = map[string]struct{}{}
+		}
+		h.includeMap[http.CanonicalHeaderKey(field)] = struct{}{}
+	}
+
+	// Optimize the exclude list by converting it to a map
+	for _, field := range h.Exclude {
+		if h.excludeMap == nil {
+			h.excludeMap = map[string]struct{}{}
+		}
+		h.excludeMap[http.CanonicalHeaderKey(field)] = struct{}{}
+	}
+
+	return nil
+}
+
+// ServeHTTP implements the Handler interface.
+func (h CopyResponseHeadersHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request, next caddyhttp.Handler) error {
+	hrc, ok := req.Context().Value(proxyHandleResponseContextCtxKey).(*handleResponseContext)
+
+	// don't allow this to be used outside of handle_response routes
+	if !ok {
+		return caddyhttp.Error(http.StatusInternalServerError,
+			fmt.Errorf("cannot use 'copy_response_headers' outside of reverse_proxy's handle_response routes"))
+	}
+
+	for field, values := range hrc.response.Header {
+		// Check the include list first, skip
+		// the header if it's _not_ in this list.
+		if len(h.includeMap) > 0 {
+			if _, ok := h.includeMap[field]; !ok {
+				continue
+			}
+		}
+
+		// Then, check the exclude list, skip
+		// the header if it _is_ in this list.
+		if len(h.excludeMap) > 0 {
+			if _, ok := h.excludeMap[field]; ok {
+				continue
+			}
+		}
+
+		// Copy all the values for the header.
+		for _, value := range values {
+			rw.Header().Add(field, value)
+		}
+	}
+
+	return next.ServeHTTP(rw, req)
+}
+
 // Interface guards
 var (
 	_ caddyhttp.MiddlewareHandler = (*CopyResponseHandler)(nil)
 	_ caddyfile.Unmarshaler       = (*CopyResponseHandler)(nil)
+	_ caddy.Provisioner           = (*CopyResponseHandler)(nil)
+
+	_ caddyhttp.MiddlewareHandler = (*CopyResponseHeadersHandler)(nil)
+	_ caddyfile.Unmarshaler       = (*CopyResponseHeadersHandler)(nil)
+	_ caddy.Provisioner           = (*CopyResponseHeadersHandler)(nil)
+	_ caddy.Validator             = (*CopyResponseHeadersHandler)(nil)
 )