Headers response modifier is directly applied by headers middleware

Co-authored-by: Ludovic Fernandez <ldez@users.noreply.github.com>

pkg/middlewares/headers/responsewriter.go

@@ -0,0 +1,75 @@
+package headers
+
+import (
+	"net/http"
+
+	"github.com/containous/traefik/v2/pkg/log"
+)
+
+type responseModifier struct {
+	r *http.Request
+	w http.ResponseWriter
+
+	headersSent bool // whether headers have already been sent
+	code        int  // status code, must default to 200
+
+	modifier    func(*http.Response) error // can be nil
+	modified    bool                       // whether modifier has already been called for the current request
+	modifierErr error                      // returned by modifier call
+}
+
+// modifier can be nil.
+func newResponseModifier(w http.ResponseWriter, r *http.Request, modifier func(*http.Response) error) *responseModifier {
+	return &responseModifier{
+		r:        r,
+		w:        w,
+		modifier: modifier,
+		code:     http.StatusOK,
+	}
+}
+
+func (w *responseModifier) WriteHeader(code int) {
+	if w.headersSent {
+		return
+	}
+	defer func() {
+		w.code = code
+		w.headersSent = true
+	}()
+
+	if w.modifier == nil || w.modified {
+		w.w.WriteHeader(code)
+		return
+	}
+
+	resp := http.Response{
+		Header:  w.w.Header(),
+		Request: w.r,
+	}
+
+	if err := w.modifier(&resp); err != nil {
+		w.modifierErr = err
+		// we are propagating when we are called in Write, but we're logging anyway,
+		// because we could be called from another place which does not take care of
+		// checking w.modifierErr.
+		log.WithoutContext().Errorf("Error when applying response modifier: %v", err)
+		w.w.WriteHeader(http.StatusInternalServerError)
+		return
+	}
+
+	w.modified = true
+	w.w.WriteHeader(code)
+}
+
+func (w *responseModifier) Header() http.Header {
+	return w.w.Header()
+}
+
+func (w *responseModifier) Write(b []byte) (int, error) {
+	w.WriteHeader(w.code)
+	if w.modifierErr != nil {
+		return 0, w.modifierErr
+	}
+
+	return w.w.Write(b)
+}