Support informational headers in middlewares redefining the response writer.

Co-authored-by: LandryBe <lbenguigui@gmail.com>

pkg/middlewares/compress/compress_test.go

@@ -5,6 +5,8 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"net/http/httptrace"
+	"net/textproto"
 	"testing"
 
 	"github.com/klauspost/compress/gzhttp"
@@ -356,6 +358,86 @@ func TestMinResponseBodyBytes(t *testing.T) {
 	}
 }
 
+// This test is an adapted version of net/http/httputil.Test1xxResponses test.
+func Test1xxResponses(t *testing.T) {
+	fakeBody := generateBytes(100000)
+
+	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		h := w.Header()
+		h.Add("Link", "</style.css>; rel=preload; as=style")
+		h.Add("Link", "</script.js>; rel=preload; as=script")
+		w.WriteHeader(http.StatusEarlyHints)
+
+		h.Add("Link", "</foo.js>; rel=preload; as=script")
+		w.WriteHeader(http.StatusProcessing)
+
+		if _, err := w.Write(fakeBody); err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+		}
+	})
+
+	compress, err := New(context.Background(), next, dynamic.Compress{MinResponseBodyBytes: 1024}, "testing")
+	require.NoError(t, err)
+
+	server := httptest.NewServer(compress)
+	t.Cleanup(server.Close)
+	frontendClient := server.Client()
+
+	checkLinkHeaders := func(t *testing.T, expected, got []string) {
+		t.Helper()
+
+		if len(expected) != len(got) {
+			t.Errorf("Expected %d link headers; got %d", len(expected), len(got))
+		}
+
+		for i := range expected {
+			if i >= len(got) {
+				t.Errorf("Expected %q link header; got nothing", expected[i])
+
+				continue
+			}
+
+			if expected[i] != got[i] {
+				t.Errorf("Expected %q link header; got %q", expected[i], got[i])
+			}
+		}
+	}
+
+	var respCounter uint8
+	trace := &httptrace.ClientTrace{
+		Got1xxResponse: func(code int, header textproto.MIMEHeader) error {
+			switch code {
+			case http.StatusEarlyHints:
+				checkLinkHeaders(t, []string{"</style.css>; rel=preload; as=style", "</script.js>; rel=preload; as=script"}, header["Link"])
+			case http.StatusProcessing:
+				checkLinkHeaders(t, []string{"</style.css>; rel=preload; as=style", "</script.js>; rel=preload; as=script", "</foo.js>; rel=preload; as=script"}, header["Link"])
+			default:
+				t.Error("Unexpected 1xx response")
+			}
+
+			respCounter++
+
+			return nil
+		},
+	}
+	req, _ := http.NewRequestWithContext(httptrace.WithClientTrace(context.Background(), trace), http.MethodGet, server.URL, nil)
+	req.Header.Add(acceptEncodingHeader, gzipValue)
+
+	res, err := frontendClient.Do(req)
+	assert.Nil(t, err)
+
+	defer res.Body.Close()
+
+	if respCounter != 2 {
+		t.Errorf("Expected 2 1xx responses; got %d", respCounter)
+	}
+	checkLinkHeaders(t, []string{"</style.css>; rel=preload; as=style", "</script.js>; rel=preload; as=script", "</foo.js>; rel=preload; as=script"}, res.Header["Link"])
+
+	assert.Equal(t, gzipValue, res.Header.Get(contentEncodingHeader))
+	body, _ := io.ReadAll(res.Body)
+	assert.NotEqualValues(t, body, fakeBody)
+}
+
 func BenchmarkCompress(b *testing.B) {
 	testCases := []struct {
 		name     string

pkg/middlewares/customerrors/custom_errors.go

@@ -193,11 +193,25 @@ func (cc *codeCatcher) Write(buf []byte) (int, error) {
 	return cc.responseWriter.Write(buf)
 }
 
+// WriteHeader is, in the specific case of 1xx status codes, a direct call to the wrapped ResponseWriter, without marking headers as sent,
+// allowing so further calls.
 func (cc *codeCatcher) WriteHeader(code int) {
 	if cc.headersSent || cc.caughtFilteredCode {
 		return
 	}
 
+	// Handling informational headers.
+	if code >= 100 && code <= 199 {
+		// Multiple informational status codes can be used,
+		// so here the copy is not appending the values to not repeat them.
+		for k, v := range cc.Header() {
+			cc.responseWriter.Header()[k] = v
+		}
+
+		cc.responseWriter.WriteHeader(code)
+		return
+	}
+
 	cc.code = code
 	for _, block := range cc.httpCodeRanges {
 		if cc.code >= block[0] && cc.code <= block[1] {
@@ -208,7 +222,11 @@ func (cc *codeCatcher) WriteHeader(code int) {
 		}
 	}
 
-	utils.CopyHeaders(cc.responseWriter.Header(), cc.Header())
+	// The copy is not appending the values,
+	// to not repeat them in case any informational status code has been written.
+	for k, v := range cc.Header() {
+		cc.responseWriter.Header()[k] = v
+	}
 	cc.responseWriter.WriteHeader(cc.code)
 	cc.headersSent = true
 }
@@ -303,12 +321,28 @@ func (r *codeModifierWithoutCloseNotify) Write(buf []byte) (int, error) {
 
 // WriteHeader sends the headers, with the enforced code (the code in argument is always ignored),
 // if it hasn't already been done.
-func (r *codeModifierWithoutCloseNotify) WriteHeader(_ int) {
+// WriteHeader is, in the specific case of 1xx status codes, a direct call to the wrapped ResponseWriter, without marking headers as sent,
+// allowing so further calls.
+func (r *codeModifierWithoutCloseNotify) WriteHeader(code int) {
 	if r.headerSent {
 		return
 	}
 
-	utils.CopyHeaders(r.responseWriter.Header(), r.Header())
+	// Handling informational headers.
+	if code >= 100 && code <= 199 {
+		// Multiple informational status codes can be used,
+		// so here the copy is not appending the values to not repeat them.
+		for k, v := range r.headerMap {
+			r.responseWriter.Header()[k] = v
+		}
+
+		r.responseWriter.WriteHeader(code)
+		return
+	}
+
+	for k, v := range r.headerMap {
+		r.responseWriter.Header()[k] = v
+	}
 	r.responseWriter.WriteHeader(r.code)
 	r.headerSent = true
 }

pkg/middlewares/customerrors/custom_errors_test.go

@@ -3,8 +3,11 @@ package customerrors
 import (
 	"context"
 	"fmt"
+	"io"
 	"net/http"
 	"net/http/httptest"
+	"net/http/httptrace"
+	"net/textproto"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -181,6 +184,88 @@ func TestHandler(t *testing.T) {
 	}
 }
 
+// This test is an adapted version of net/http/httputil.Test1xxResponses test.
+func Test1xxResponses(t *testing.T) {
+	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		h := w.Header()
+		h.Add("Link", "</style.css>; rel=preload; as=style")
+		h.Add("Link", "</script.js>; rel=preload; as=script")
+		w.WriteHeader(http.StatusEarlyHints)
+
+		h.Add("Link", "</foo.js>; rel=preload; as=script")
+		w.WriteHeader(http.StatusProcessing)
+
+		h.Add("User-Agent", "foobar")
+		_, _ = w.Write([]byte("Hello"))
+		w.WriteHeader(http.StatusBadGateway)
+	})
+
+	serviceBuilderMock := &mockServiceBuilder{handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_, _ = fmt.Fprintln(w, "My error page.")
+	})}
+
+	config := dynamic.ErrorPage{Service: "error", Query: "/", Status: []string{"200"}}
+
+	errorPageHandler, err := New(context.Background(), next, config, serviceBuilderMock, "test")
+	require.NoError(t, err)
+
+	server := httptest.NewServer(errorPageHandler)
+	t.Cleanup(server.Close)
+	frontendClient := server.Client()
+
+	checkLinkHeaders := func(t *testing.T, expected, got []string) {
+		t.Helper()
+
+		if len(expected) != len(got) {
+			t.Errorf("Expected %d link headers; got %d", len(expected), len(got))
+		}
+
+		for i := range expected {
+			if i >= len(got) {
+				t.Errorf("Expected %q link header; got nothing", expected[i])
+
+				continue
+			}
+
+			if expected[i] != got[i] {
+				t.Errorf("Expected %q link header; got %q", expected[i], got[i])
+			}
+		}
+	}
+
+	var respCounter uint8
+	trace := &httptrace.ClientTrace{
+		Got1xxResponse: func(code int, header textproto.MIMEHeader) error {
+			switch code {
+			case http.StatusEarlyHints:
+				checkLinkHeaders(t, []string{"</style.css>; rel=preload; as=style", "</script.js>; rel=preload; as=script"}, header["Link"])
+			case http.StatusProcessing:
+				checkLinkHeaders(t, []string{"</style.css>; rel=preload; as=style", "</script.js>; rel=preload; as=script", "</foo.js>; rel=preload; as=script"}, header["Link"])
+			default:
+				t.Error("Unexpected 1xx response")
+			}
+
+			respCounter++
+
+			return nil
+		},
+	}
+	req, _ := http.NewRequestWithContext(httptrace.WithClientTrace(context.Background(), trace), http.MethodGet, server.URL, nil)
+
+	res, err := frontendClient.Do(req)
+	assert.Nil(t, err)
+
+	defer res.Body.Close()
+
+	if respCounter != 2 {
+		t.Errorf("Expected 2 1xx responses; got %d", respCounter)
+	}
+	checkLinkHeaders(t, []string{"</style.css>; rel=preload; as=style", "</script.js>; rel=preload; as=script", "</foo.js>; rel=preload; as=script"}, res.Header["Link"])
+
+	body, _ := io.ReadAll(res.Body)
+	assert.Equal(t, "My error page.\n", string(body))
+}
+
 type mockServiceBuilder struct {
 	handler http.Handler
 }

pkg/middlewares/headers/headers_test.go

@@ -4,8 +4,11 @@ package headers
 
 import (
 	"context"
+	"io"
 	"net/http"
 	"net/http/httptest"
+	"net/http/httptrace"
+	"net/textproto"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -111,3 +114,87 @@ func Test_headers_getTracingInformation(t *testing.T) {
 	assert.Equal(t, "testing", name)
 	assert.Equal(t, tracing.SpanKindNoneEnum, trace)
 }
+
+// This test is an adapted version of net/http/httputil.Test1xxResponses test.
+func Test1xxResponses(t *testing.T) {
+	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		h := w.Header()
+		h.Add("Link", "</style.css>; rel=preload; as=style")
+		h.Add("Link", "</script.js>; rel=preload; as=script")
+		w.WriteHeader(http.StatusEarlyHints)
+
+		h.Add("Link", "</foo.js>; rel=preload; as=script")
+		w.WriteHeader(http.StatusProcessing)
+
+		_, _ = w.Write([]byte("Hello"))
+	})
+
+	cfg := dynamic.Headers{
+		CustomResponseHeaders: map[string]string{
+			"X-Custom-Response-Header": "test_response",
+		},
+	}
+
+	mid, err := New(context.Background(), next, cfg, "testing")
+	require.NoError(t, err)
+
+	server := httptest.NewServer(mid)
+	t.Cleanup(server.Close)
+	frontendClient := server.Client()
+
+	checkLinkHeaders := func(t *testing.T, expected, got []string) {
+		t.Helper()
+
+		if len(expected) != len(got) {
+			t.Errorf("Expected %d link headers; got %d", len(expected), len(got))
+		}
+
+		for i := range expected {
+			if i >= len(got) {
+				t.Errorf("Expected %q link header; got nothing", expected[i])
+
+				continue
+			}
+
+			if expected[i] != got[i] {
+				t.Errorf("Expected %q link header; got %q", expected[i], got[i])
+			}
+		}
+	}
+
+	var respCounter uint8
+	trace := &httptrace.ClientTrace{
+		Got1xxResponse: func(code int, header textproto.MIMEHeader) error {
+			switch code {
+			case http.StatusEarlyHints:
+				checkLinkHeaders(t, []string{"</style.css>; rel=preload; as=style", "</script.js>; rel=preload; as=script"}, header["Link"])
+			case http.StatusProcessing:
+				checkLinkHeaders(t, []string{"</style.css>; rel=preload; as=style", "</script.js>; rel=preload; as=script", "</foo.js>; rel=preload; as=script"}, header["Link"])
+			default:
+				t.Error("Unexpected 1xx response")
+			}
+
+			respCounter++
+
+			return nil
+		},
+	}
+	req, _ := http.NewRequestWithContext(httptrace.WithClientTrace(context.Background(), trace), http.MethodGet, server.URL, nil)
+
+	res, err := frontendClient.Do(req)
+	assert.Nil(t, err)
+
+	defer res.Body.Close()
+
+	if respCounter != 2 {
+		t.Errorf("Expected 2 1xx responses; got %d", respCounter)
+	}
+	checkLinkHeaders(t, []string{"</style.css>; rel=preload; as=style", "</script.js>; rel=preload; as=script", "</foo.js>; rel=preload; as=script"}, res.Header["Link"])
+
+	body, _ := io.ReadAll(res.Body)
+	if string(body) != "Hello" {
+		t.Errorf("Read body %q; want Hello", body)
+	}
+
+	assert.Equal(t, "test_response", res.Header.Get("X-Custom-Response-Header"))
+}

pkg/middlewares/headers/responsewriter.go

@@ -36,10 +36,19 @@ func newResponseModifier(w http.ResponseWriter, r *http.Request, modifier func(*
 	return rm
 }
 
+// WriteHeader is, in the specific case of 1xx status codes, a direct call to the wrapped ResponseWriter, without marking headers as sent,
+// allowing so further calls.
 func (r *responseModifier) WriteHeader(code int) {
 	if r.headersSent {
 		return
 	}
+
+	// Handling informational headers.
+	if code >= 100 && code <= 199 {
+		r.rw.WriteHeader(code)
+		return
+	}
+
 	defer func() {
 		r.code = code
 		r.headersSent = true

pkg/middlewares/pipelining/pipelining_test.go

@@ -2,8 +2,11 @@ package pipelining
 
 import (
 	"context"
+	"io"
 	"net/http"
 	"net/http/httptest"
+	"net/http/httptrace"
+	"net/textproto"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -68,3 +71,80 @@ func TestNew(t *testing.T) {
 		})
 	}
 }
+
+// This test is an adapted version of net/http/httputil.Test1xxResponses test.
+// This test is only here to guarantee that there would not be any regression in the future,
+// because the pipelining middleware is already supporting informational headers.
+func Test1xxResponses(t *testing.T) {
+	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		h := w.Header()
+		h.Add("Link", "</style.css>; rel=preload; as=style")
+		h.Add("Link", "</script.js>; rel=preload; as=script")
+		w.WriteHeader(http.StatusEarlyHints)
+
+		h.Add("Link", "</foo.js>; rel=preload; as=script")
+		w.WriteHeader(http.StatusProcessing)
+
+		_, _ = w.Write([]byte("Hello"))
+	})
+
+	pipe := New(context.Background(), next, "pipe")
+
+	server := httptest.NewServer(pipe)
+	t.Cleanup(server.Close)
+	frontendClient := server.Client()
+
+	checkLinkHeaders := func(t *testing.T, expected, got []string) {
+		t.Helper()
+
+		if len(expected) != len(got) {
+			t.Errorf("Expected %d link headers; got %d", len(expected), len(got))
+		}
+
+		for i := range expected {
+			if i >= len(got) {
+				t.Errorf("Expected %q link header; got nothing", expected[i])
+
+				continue
+			}
+
+			if expected[i] != got[i] {
+				t.Errorf("Expected %q link header; got %q", expected[i], got[i])
+			}
+		}
+	}
+
+	var respCounter uint8
+	trace := &httptrace.ClientTrace{
+		Got1xxResponse: func(code int, header textproto.MIMEHeader) error {
+			switch code {
+			case http.StatusEarlyHints:
+				checkLinkHeaders(t, []string{"</style.css>; rel=preload; as=style", "</script.js>; rel=preload; as=script"}, header["Link"])
+			case http.StatusProcessing:
+				checkLinkHeaders(t, []string{"</style.css>; rel=preload; as=style", "</script.js>; rel=preload; as=script", "</foo.js>; rel=preload; as=script"}, header["Link"])
+			default:
+				t.Error("Unexpected 1xx response")
+			}
+
+			respCounter++
+
+			return nil
+		},
+	}
+	req, _ := http.NewRequestWithContext(httptrace.WithClientTrace(context.Background(), trace), http.MethodGet, server.URL, nil)
+
+	res, err := frontendClient.Do(req)
+	assert.Nil(t, err)
+
+	defer res.Body.Close()
+
+	if respCounter != 2 {
+		t.Errorf("Expected 2 1xx responses; got %d", respCounter)
+	}
+	checkLinkHeaders(t, []string{"</style.css>; rel=preload; as=style", "</script.js>; rel=preload; as=script", "</foo.js>; rel=preload; as=script"}, res.Header["Link"])
+
+	body, _ := io.ReadAll(res.Body)
+	if string(body) != "Hello" {
+		t.Errorf("Read body %q; want Hello", body)
+	}
+}

pkg/middlewares/retry/retry.go

@@ -209,7 +209,7 @@ func (r *responseWriterWithoutCloseNotify) WriteHeader(code int) {
 		r.DisableRetries()
 	}
 
-	if r.ShouldRetry() {
+	if r.ShouldRetry() || r.written {
 		return
 	}
 
@@ -223,6 +223,13 @@ func (r *responseWriterWithoutCloseNotify) WriteHeader(code int) {
 	}
 
 	r.responseWriter.WriteHeader(code)
+
+	// Handling informational headers.
+	// This allows to keep writing to r.headers map until a final status code is written.
+	if code >= 100 && code <= 199 {
+		return
+	}
+
 	r.written = true
 }
 

pkg/middlewares/retry/retry_test.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/http/httptrace"
+	"net/textproto"
 	"strings"
 	"testing"
 	"time"
@@ -309,3 +310,82 @@ func TestRetryWebsocket(t *testing.T) {
 		})
 	}
 }
+
+// This test is an adapted version of net/http/httputil.Test1xxResponses test.
+func Test1xxResponses(t *testing.T) {
+	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		h := w.Header()
+		h.Add("Link", "</style.css>; rel=preload; as=style")
+		h.Add("Link", "</script.js>; rel=preload; as=script")
+		w.WriteHeader(http.StatusEarlyHints)
+
+		h.Add("Link", "</foo.js>; rel=preload; as=script")
+		w.WriteHeader(http.StatusProcessing)
+
+		_, _ = w.Write([]byte("Hello"))
+	})
+
+	retryListener := &countingRetryListener{}
+	retry, err := New(context.Background(), next, dynamic.Retry{Attempts: 1}, retryListener, "traefikTest")
+	require.NoError(t, err)
+
+	server := httptest.NewServer(retry)
+	t.Cleanup(server.Close)
+	frontendClient := server.Client()
+
+	checkLinkHeaders := func(t *testing.T, expected, got []string) {
+		t.Helper()
+
+		if len(expected) != len(got) {
+			t.Errorf("Expected %d link headers; got %d", len(expected), len(got))
+		}
+
+		for i := range expected {
+			if i >= len(got) {
+				t.Errorf("Expected %q link header; got nothing", expected[i])
+
+				continue
+			}
+
+			if expected[i] != got[i] {
+				t.Errorf("Expected %q link header; got %q", expected[i], got[i])
+			}
+		}
+	}
+
+	var respCounter uint8
+	trace := &httptrace.ClientTrace{
+		Got1xxResponse: func(code int, header textproto.MIMEHeader) error {
+			switch code {
+			case http.StatusEarlyHints:
+				checkLinkHeaders(t, []string{"</style.css>; rel=preload; as=style", "</script.js>; rel=preload; as=script"}, header["Link"])
+			case http.StatusProcessing:
+				checkLinkHeaders(t, []string{"</style.css>; rel=preload; as=style", "</script.js>; rel=preload; as=script", "</foo.js>; rel=preload; as=script"}, header["Link"])
+			default:
+				t.Error("Unexpected 1xx response")
+			}
+
+			respCounter++
+
+			return nil
+		},
+	}
+	req, _ := http.NewRequestWithContext(httptrace.WithClientTrace(context.Background(), trace), http.MethodGet, server.URL, nil)
+
+	res, err := frontendClient.Do(req)
+	assert.Nil(t, err)
+
+	defer res.Body.Close()
+
+	if respCounter != 2 {
+		t.Errorf("Expected 2 1xx responses; got %d", respCounter)
+	}
+	checkLinkHeaders(t, []string{"</style.css>; rel=preload; as=style", "</script.js>; rel=preload; as=script", "</foo.js>; rel=preload; as=script"}, res.Header["Link"])
+
+	body, _ := io.ReadAll(res.Body)
+	if string(body) != "Hello" {
+		t.Errorf("Read body %q; want Hello", body)
+	}
+
+	assert.Equal(t, 0, retryListener.timesCalled)
+}