From 7b4a74c99123b9357b779d89776943bf0c75e340 Mon Sep 17 00:00:00 2001
From: Andrew Dunstall <andydunstall@hotmail.co.uk>
Date: Fri, 15 Nov 2024 06:06:13 +0000
Subject: [PATCH] server: support x-piko-authorization header

---
 pkg/middleware/auth.go      | 20 +++++++++++++++++---
 pkg/middleware/auth_test.go | 33 ++++++++++++++++++++++++++++++++-
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/pkg/middleware/auth.go b/pkg/middleware/auth.go
index 2f25eb9a..44b14405 100644
--- a/pkg/middleware/auth.go
+++ b/pkg/middleware/auth.go
@@ -76,9 +76,14 @@ func (m *Auth) Verify(c *gin.Context) {
 }
 
 func (m *Auth) parseToken(c *gin.Context) (string, bool) {
-	authorization := c.Request.Header.Get("Authorization")
-	authType, tokenString, ok := strings.Cut(authorization, " ")
-	if !ok {
+	// Support both x-piko-authorization and authorization, where
+	// x-piko-authorization takes precedence. x-piko-authorization can be used
+	// to avoid conflicts with the upstream authorization header.
+	authorization := c.Request.Header.Get("x-piko-authorization")
+	if authorization == "" {
+		authorization = c.Request.Header.Get("Authorization")
+	}
+	if authorization == "" {
 		m.logger.Warn("missing authorization header")
 		c.AbortWithStatusJSON(
 			http.StatusUnauthorized,
@@ -86,6 +91,15 @@ func (m *Auth) parseToken(c *gin.Context) (string, bool) {
 		)
 		return "", false
 	}
+	authType, tokenString, ok := strings.Cut(authorization, " ")
+	if !ok {
+		m.logger.Warn("invalid authorization header")
+		c.AbortWithStatusJSON(
+			http.StatusUnauthorized,
+			gin.H{"error": "invalid authorization"},
+		)
+		return "", false
+	}
 	if authType != "Bearer" {
 		m.logger.Warn(
 			"unsupported auth type",
diff --git a/pkg/middleware/auth_test.go b/pkg/middleware/auth_test.go
index caf11532..41f78ebf 100644
--- a/pkg/middleware/auth_test.go
+++ b/pkg/middleware/auth_test.go
@@ -30,7 +30,7 @@ type errorMessage struct {
 }
 
 func TestAuth(t *testing.T) {
-	t.Run("ok", func(t *testing.T) {
+	t.Run("authorization ok", func(t *testing.T) {
 		verifier := &fakeVerifier{
 			handler: func(token string) (*auth.Token, error) {
 				assert.Equal(t, "123", token)
@@ -58,6 +58,37 @@ func TestAuth(t *testing.T) {
 		assert.Equal(t, []string{"e1", "e2", "e3"}, token.(*auth.Token).Endpoints)
 	})
 
+	t.Run("x-piko-authorization ok", func(t *testing.T) {
+		verifier := &fakeVerifier{
+			handler: func(token string) (*auth.Token, error) {
+				assert.Equal(t, "123", token)
+				return &auth.Token{
+					Expiry:    time.Now().Add(time.Hour),
+					Endpoints: []string{"e1", "e2", "e3"},
+				}, nil
+			},
+		}
+		m := NewAuth(verifier, log.NewNopLogger())
+
+		w := httptest.NewRecorder()
+		c, _ := gin.CreateTestContext(w)
+		c.Request = httptest.NewRequest("GET", "http://example.com/foo", nil)
+		// Add both x-piko-authorization and Authorization, where
+		// x-piko-authorization should take precedence.
+		c.Request.Header.Add("x-piko-authorization", "Bearer 123")
+		c.Request.Header.Add("Authorization", "Bearer xyz")
+
+		m.Verify(c)
+
+		resp := w.Result()
+		assert.Equal(t, http.StatusOK, resp.StatusCode)
+
+		// Verify the token was added to context.
+		token, ok := c.Get(TokenContextKey)
+		assert.True(t, ok)
+		assert.Equal(t, []string{"e1", "e2", "e3"}, token.(*auth.Token).Endpoints)
+	})
+
 	t.Run("invalid token", func(t *testing.T) {
 		verifier := &fakeVerifier{
 			handler: func(token string) (*auth.Token, error) {