From f9808a73baa585e7d70a954f82e734ea8cee8167 Mon Sep 17 00:00:00 2001
From: Kenneth Jenkins <51246568+kenjenkins@users.noreply.github.com>
Date: Tue, 6 Feb 2024 08:53:58 -0800
Subject: [PATCH] integration: unauthorized response Content-Type (#4956)

Modify the request 'Accept' header to behave more like a web browser,
and add an assertion to verify that Pomerium serves an HTML response for
the unauthorized error page.
---
 integration/authorization_test.go |  8 ++++++--
 integration/flows/flows.go        | 16 ++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/integration/authorization_test.go b/integration/authorization_test.go
index 7a10753ea..9379f6f79 100644
--- a/integration/authorization_test.go
+++ b/integration/authorization_test.go
@@ -15,6 +15,9 @@ func TestAuthorization(t *testing.T) {
 	ctx, clearTimeout := context.WithTimeout(context.Background(), time.Second*30)
 	defer clearTimeout()
 
+	withBrowserAcceptHeader := flows.WithRequestHeader("Accept",
+		"text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7")
+
 	accessType := []string{"direct", "api"}
 	for _, at := range accessType {
 		t.Run(at, func(t *testing.T) {
@@ -45,7 +48,7 @@ func TestAuthorization(t *testing.T) {
 				t.Run("allowed", func(t *testing.T) {
 					client := getClient(t)
 					res, err := flows.Authenticate(ctx, client, mustParseURL("https://httpdetails.localhost.pomerium.io/by-domain"),
-						withAPI, flows.WithEmail("user1@dogs.test"))
+						withAPI, flows.WithEmail("user1@dogs.test"), withBrowserAcceptHeader)
 					if assert.NoError(t, err) {
 						assert.Equal(t, http.StatusOK, res.StatusCode, "expected OK for dogs.test")
 					}
@@ -53,9 +56,10 @@ func TestAuthorization(t *testing.T) {
 				t.Run("not allowed", func(t *testing.T) {
 					client := getClient(t)
 					res, err := flows.Authenticate(ctx, client, mustParseURL("https://httpdetails.localhost.pomerium.io/by-domain"),
-						withAPI, flows.WithEmail("user1@cats.test"))
+						withAPI, flows.WithEmail("user1@cats.test"), withBrowserAcceptHeader)
 					if assert.NoError(t, err) {
 						assertDeniedAccess(t, res, "expected Forbidden for cats.test, but got: %d", res.StatusCode)
+						assert.Contains(t, res.Header.Get("Content-Type"), "text/html")
 					}
 				})
 			})
diff --git a/integration/flows/flows.go b/integration/flows/flows.go
index 1ae1f7477..1c6b9e552 100644
--- a/integration/flows/flows.go
+++ b/integration/flows/flows.go
@@ -28,6 +28,7 @@ type authenticateConfig struct {
 	groups          []string
 	tokenExpiration time.Duration
 	apiPath         string
+	requestHeaders  http.Header
 }
 
 // An AuthenticateOption is an option for authentication.
@@ -36,6 +37,7 @@ type AuthenticateOption func(cfg *authenticateConfig)
 func getAuthenticateConfig(options ...AuthenticateOption) *authenticateConfig {
 	cfg := &authenticateConfig{
 		tokenExpiration: time.Hour * 24,
+		requestHeaders:  http.Header{},
 	}
 	for _, option := range options {
 		if option != nil {
@@ -73,6 +75,12 @@ func WithAPI() AuthenticateOption {
 	}
 }
 
+func WithRequestHeader(name, value string) AuthenticateOption {
+	return func(cfg *authenticateConfig) {
+		cfg.requestHeaders.Set(name, value)
+	}
+}
+
 // Authenticate submits a request to a URL, expects a redirect to authenticate and then openid and logs in.
 // Finally it expects to redirect back to the original page.
 func Authenticate(ctx context.Context, client *http.Client, url *url.URL, options ...AuthenticateOption) (*http.Response, error) {
@@ -168,6 +176,7 @@ func Authenticate(ctx context.Context, client *http.Client, url *url.URL, option
 			if err != nil {
 				return nil, err
 			}
+			addRequestHeaders(req, cfg.requestHeaders)
 		} else {
 			req, err = requestFromRedirectResponse(ctx, res, req)
 			if err != nil {
@@ -238,5 +247,12 @@ func requestFromRedirectResponse(ctx context.Context, res *http.Response, req *h
 	if err != nil {
 		return nil, err
 	}
+	addRequestHeaders(newreq, req.Header)
 	return newreq, nil
 }
+
+func addRequestHeaders(req *http.Request, headers http.Header) {
+	for h := range headers {
+		req.Header[h] = headers[h]
+	}
+}