From 4aa349609276a5c2d599030d37010d7751bdf477 Mon Sep 17 00:00:00 2001
From: Romain <rtribotte@users.noreply.github.com>
Date: Fri, 17 Mar 2023 16:46:06 +0100
Subject: [PATCH] Add HTTP 103 early hints unit test

Co-authored-by: Mathieu Lonjaret <mathieu.lonjaret@gmail.com>
---
 pkg/server/service/service_test.go | 93 ++++++++++++++++++++++++++++++
 1 file changed, 93 insertions(+)

diff --git a/pkg/server/service/service_test.go b/pkg/server/service/service_test.go
index 961f47091..617b3569a 100644
--- a/pkg/server/service/service_test.go
+++ b/pkg/server/service/service_test.go
@@ -2,8 +2,11 @@ package service
 
 import (
 	"context"
+	"io"
 	"net/http"
 	"net/http/httptest"
+	"net/http/httptrace"
+	"net/textproto"
 	"strings"
 	"testing"
 
@@ -385,6 +388,96 @@ func TestGetLoadBalancerServiceHandler(t *testing.T) {
 	}
 }
 
+// This test is an adapted version of net/http/httputil.Test1xxResponses test.
+func Test1xxResponses(t *testing.T) {
+	sm := NewManager(nil, nil, nil, &RoundTripperManager{
+		roundTrippers: map[string]http.RoundTripper{
+			"default@internal": http.DefaultTransport,
+		},
+	})
+
+	backend := httptest.NewServer(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"))
+	}))
+	t.Cleanup(backend.Close)
+
+	config := &dynamic.ServersLoadBalancer{
+		Servers: []dynamic.Server{
+			{
+				URL: backend.URL,
+			},
+		},
+	}
+	handler, err := sm.getLoadBalancerServiceHandler(context.Background(), "foobar", config)
+	assert.Nil(t, err)
+
+	frontend := httptest.NewServer(handler)
+	t.Cleanup(frontend.Close)
+	frontendClient := frontend.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, frontend.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)
+	}
+}
+
 func TestManager_Build(t *testing.T) {
 	testCases := []struct {
 		desc         string