traefikoidc/templated_header_integration_test.go

package traefikoidc

import (
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"
	"text/template"
	"time"

	"golang.org/x/time/rate"
)

// TestTemplatedHeadersIntegration tests that templated headers are correctly added to requests
// in the actual middleware flow
func TestTemplatedHeadersIntegration(t *testing.T) {
	// Create a TestSuite to use its helper methods and fields
	ts := &TestSuite{t: t}
	ts.Setup()

	tests := []struct {
		name               string
		headers            []TemplatedHeader
		sessionSetup       func(*SessionData)
		claims             map[string]interface{}
		expectedHeaders    map[string]string
		interceptedHeaders map[string]string
	}{
		{
			name: "Basic Email Header",
			headers: []TemplatedHeader{
				{Name: "X-User-Email", Value: "{{.Claims.email}}"},
			},
			claims: map[string]interface{}{
				"email": "user@example.com",
			},
			expectedHeaders: map[string]string{
				"X-User-Email": "user@example.com",
			},
		},
		{
			name: "Multiple Headers",
			headers: []TemplatedHeader{
				{Name: "X-User-Email", Value: "{{.Claims.email}}"},
				{Name: "X-User-ID", Value: "{{.Claims.sub}}"},
				{Name: "X-User-Name", Value: "{{.Claims.given_name}} {{.Claims.family_name}}"},
			},
			claims: map[string]interface{}{
				"email":       "user@example.com",
				"sub":         "user123",
				"given_name":  "John",
				"family_name": "Doe",
			},
			expectedHeaders: map[string]string{
				"X-User-Email": "user@example.com",
				"X-User-ID":    "user123",
				"X-User-Name":  "John Doe",
			},
		},
		{
			name: "Authorization Header with Bearer Token",
			headers: []TemplatedHeader{
				{Name: "Authorization", Value: "Bearer {{.AccessToken}}"},
			},
			expectedHeaders: map[string]string{
				// We'll update this dynamically after generating the token
				"Authorization": "",
			},
		},
		{
			name: "ID Token Header",
			headers: []TemplatedHeader{
				{Name: "X-ID-Token", Value: "{{.IdToken}}"},
			},
			expectedHeaders: map[string]string{
				// We'll update this dynamically after generating the token
				"X-ID-Token": "",
			},
		},
		{
			name: "Both Token Types",
			headers: []TemplatedHeader{
				{Name: "X-Access-Token", Value: "{{.AccessToken}}"},
				{Name: "X-ID-Token", Value: "{{.IdToken}}"},
			},
			expectedHeaders: map[string]string{
				// We'll update these dynamically after generating the tokens
				"X-Access-Token": "",
				"X-ID-Token":     "",
			},
		},
		{
			name: "Missing Claim",
			headers: []TemplatedHeader{
				{Name: "X-User-Role", Value: "{{.Claims.role}}"},
			},
			claims: map[string]interface{}{
				"email": "user@example.com",
				// role claim is missing
			},
			expectedHeaders: map[string]string{
				"X-User-Role": "<no value>", // Go templates provide <no value> for missing fields
			},
		},
		{
			name: "Conditional Header",
			headers: []TemplatedHeader{
				{Name: "X-User-Admin", Value: "{{if .Claims.is_admin}}true{{else}}false{{end}}"},
			},
			claims: map[string]interface{}{
				"email":    "admin@example.com",
				"is_admin": true,
			},
			expectedHeaders: map[string]string{
				"X-User-Admin": "true",
			},
		},
		{
			name: "Combined Token and Claim",
			headers: []TemplatedHeader{
				{Name: "X-Auth-Info", Value: "User={{.Claims.email}}, Token={{.AccessToken}}"},
			},
			claims: map[string]interface{}{
				"email": "user@example.com",
			},
			expectedHeaders: map[string]string{
				// We'll update this dynamically after generating the token
				"X-Auth-Info": "",
			},
		},
		{
			name: "Opaque Access Token with AccessTokenField",
			headers: []TemplatedHeader{
				{Name: "X-User-AccessToken", Value: "{{.AccessToken}}"},
			},
			claims: map[string]interface{}{ // For ID Token
				"email": "opaque_user@example.com",
				"sub":   "opaque_sub_for_id_token",
			},
			expectedHeaders: map[string]string{
				"X-User-AccessToken": "this_is_an_opaque_access_token",
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			// Create token with the test claims
			token := ts.token
			if len(tc.claims) > 0 {
				var err error
				baseClaims := map[string]interface{}{
					"iss":   "https://test-issuer.com",
					"aud":   "test-client-id",
					"exp":   float64(3000000000), // Far future timestamp
					"iat":   float64(1000000000),
					"nbf":   float64(1000000000),
					"sub":   "test-subject",
					"nonce": "test-nonce",
					"jti":   generateRandomString(16),
				}

				// Add the test-specific claims
				for k, v := range tc.claims {
					baseClaims[k] = v
				}

				token, err = createTestJWT(ts.rsaPrivateKey, "RS256", "test-key-id", baseClaims)
				if err != nil {
					t.Fatalf("Failed to create test JWT: %v", err)
				}
			}

			// Update expectedHeaders for the token-based tests after token generation
			if tc.name == "Authorization Header with Bearer Token" {
				tc.expectedHeaders["Authorization"] = "Bearer " + token
			}

			if tc.name == "Combined Token and Claim" {
				// If this test case uses specific ID/Access tokens, 'token' here might be just the ID token.
				// This part might need adjustment if AccessToken is different and opaque.
				// For now, assuming 'token' is the one to be used if not overridden later.
				// The specific test "Opaque Access Token with AccessTokenField" will handle its AccessToken.
				// This generic 'token' is used as a fallback if specific logic isn't hit.
				// Let's ensure this test case uses the JWT access token if not otherwise specified.
				accessTokenForHeader := token                        // Default to the generated JWT 'token'
				if sessionVal, ok := tc.claims["_accessToken"]; ok { // Check if a specific access token is provided for this test
					accessTokenForHeader = sessionVal.(string)
				}
				tc.expectedHeaders["X-Auth-Info"] = "User=" + tc.claims["email"].(string) + ", Token=" + accessTokenForHeader
			}

			// Store intercepted headers for verification
			interceptedHeaders := make(map[string]string)

			// Create a test next handler that captures the headers
			nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				// Capture headers for verification
				for name := range tc.expectedHeaders {
					if value := r.Header.Get(name); value != "" {
						interceptedHeaders[name] = value
					}
				}
				w.WriteHeader(http.StatusOK)
			})

			tOidc := &TraefikOidc{
				next:               nextHandler,
				name:               "test",
				redirURLPath:       "/callback",
				logoutURLPath:      "/callback/logout",
				issuerURL:          "https://test-issuer.com",
				clientID:           "test-client-id",
				clientSecret:       "test-client-secret",
				jwkCache:           ts.mockJWKCache,
				jwksURL:            "https://test-jwks-url.com",
				tokenBlacklist:     NewCache(),
				tokenCache:         NewTokenCache(),
				limiter:            rate.NewLimiter(rate.Every(time.Second), 10),
				logger:             NewLogger("debug"),
				allowedUserDomains: map[string]struct{}{"example.com": {}, "opaque_user@example.com": {}}, // Ensure domain for opaque test is allowed
				excludedURLs:       map[string]struct{}{"/favicon": {}},
				httpClient:         &http.Client{},
				initComplete:       make(chan struct{}),
				sessionManager:     ts.sessionManager,
				extractClaimsFunc:  extractClaims,
				headerTemplates:    make(map[string]*template.Template),
				// Default to true, which means PopulateSessionWithIdTokenClaims is true
				// UseIdTokenForSession: true, // Explicitly can be set if needed
			}

			// Initialize and parse header templates
			for _, header := range tc.headers {
				tmpl, err := template.New(header.Name).Parse(header.Value)
				if err != nil {
					t.Fatalf("Failed to parse header template for %s: %v", header.Name, err)
				}
				tOidc.headerTemplates[header.Name] = tmpl
			}

			close(tOidc.initComplete)

			req := httptest.NewRequest("GET", "/protected", nil)
			req.Header.Set("X-Forwarded-Proto", "https")
			req.Header.Set("X-Forwarded-Host", "example.com")
			rr := httptest.NewRecorder()

			session, err := tOidc.sessionManager.GetSession(req)
			if err != nil {
				t.Fatalf("Failed to get session: %v", err)
			}

			session.SetAuthenticated(true)
			// Set a default email; specific tests might override or rely on ID token population
			defaultEmail := "user@example.com"
			if emailClaim, ok := tc.claims["email"].(string); ok {
				defaultEmail = emailClaim // Use email from claims if available for initial setup
			}
			session.SetEmail(defaultEmail)

			// Default token setup (can be overridden by specific test cases below)
			session.SetIDToken(token)
			session.SetAccessToken(token)
			session.SetRefreshToken("test-refresh-token")

			if tc.name == "ID Token Header" || tc.name == "Both Token Types" {
				idTokenClaims := map[string]interface{}{
					"iss": "https://test-issuer.com", "aud": "test-client-id", "exp": float64(3000000000),
					"iat": float64(1000000000), "nbf": float64(1000000000), "sub": "test-subject",
					"nonce": "test-nonce", "jti": generateRandomString(16), "type": "id_token",
					"email": tc.claims["email"], // Ensure email from test case claims is in ID token
				}
				// Add other claims from tc.claims to idTokenClaims
				for k, v := range tc.claims {
					if _, exists := idTokenClaims[k]; !exists {
						idTokenClaims[k] = v
					}
				}

				idTokenForSession, idErr := createTestJWT(ts.rsaPrivateKey, "RS256", "test-key-id", idTokenClaims)
				if idErr != nil {
					t.Fatalf("Failed to create test ID JWT: %v", idErr)
				}

				accessTokenClaims := map[string]interface{}{
					"iss": "https://test-issuer.com", "aud": "test-client-id", "exp": float64(3000000000),
					"iat": float64(1000000000), "nbf": float64(1000000000), "sub": "test-subject",
					"jti": generateRandomString(16), "type": "access_token", "scope": "openid email profile",
					"email": tc.claims["email"], // Include email in access token too for these tests
				}
				accessTokenForSession, accessErr := createTestJWT(ts.rsaPrivateKey, "RS256", "test-key-id", accessTokenClaims)
				if accessErr != nil {
					t.Fatalf("Failed to create test access JWT: %v", accessErr)
				}

				session.SetIDToken(idTokenForSession)
				session.SetAccessToken(accessTokenForSession)

				tOidc.tokenExchanger = &MockTokenExchanger{
					RefreshTokenFunc: func(refreshToken string) (*TokenResponse, error) {
						return &TokenResponse{
							IDToken: idTokenForSession, AccessToken: accessTokenForSession,
							RefreshToken: refreshToken, ExpiresIn: 3600,
						}, nil
					},
				}
				tOidc.tokenVerifier = &MockTokenVerifier{VerifyFunc: func(token string) error { return nil }}

				if tc.name == "ID Token Header" {
					tc.expectedHeaders["X-ID-Token"] = idTokenForSession
				} else if tc.name == "Both Token Types" {
					tc.expectedHeaders["X-ID-Token"] = idTokenForSession
					tc.expectedHeaders["X-Access-Token"] = accessTokenForSession
				}
			} else if tc.name == "Opaque Access Token with AccessTokenField" {
				idTokenClaims := map[string]interface{}{
					"iss": "https://test-issuer.com", "aud": "test-client-id", "exp": float64(3000000000),
					"iat": float64(1000000000), "nbf": float64(1000000000), "sub": "test-subject", // Default sub
					"nonce": "test-nonce", "jti": generateRandomString(16), "type": "id_token",
				}
				// Populate ID token claims from tc.claims
				for k, v := range tc.claims {
					idTokenClaims[k] = v
				}
				// Ensure email from tc.claims is used for the ID token
				session.SetEmail(tc.claims["email"].(string)) // Also set it directly for initial session state

				idTokenForSession, err := createTestJWT(ts.rsaPrivateKey, "RS256", "test-key-id", idTokenClaims)
				if err != nil {
					t.Fatalf("Failed to create test ID JWT for opaque test: %v", err)
				}

				opaqueAccessToken := "this_is_an_opaque_access_token"

				session.SetIDToken(idTokenForSession)
				session.SetAccessToken(opaqueAccessToken)

				tOidc.tokenExchanger = &MockTokenExchanger{
					RefreshTokenFunc: func(refreshToken string) (*TokenResponse, error) {
						return &TokenResponse{
							IDToken:      idTokenForSession,
							AccessToken:  opaqueAccessToken,
							RefreshToken: refreshToken,
							ExpiresIn:    3600,
						}, nil
					},
				}
				tOidc.tokenVerifier = &MockTokenVerifier{
					VerifyFunc: func(tokenToVerify string) error {
						if tokenToVerify == idTokenForSession {
							return nil // ID token is expected to be verified
						}
						if tokenToVerify == opaqueAccessToken {
							t.Errorf("TokenVerifier was incorrectly called with the opaque access token.")
							return errors.New("opaque access token should not be verified by this path")
						}
						t.Logf("TokenVerifier called with unexpected token: %s", tokenToVerify)
						return errors.New("unexpected token passed to verifier for this test case")
					},
				}
				// Expected header X-User-AccessToken is already set in tc.expectedHeaders
			}

			if err := session.Save(req, rr); err != nil {
				t.Fatalf("Failed to save session: %v", err)
			}

			for _, cookie := range rr.Result().Cookies() {
				req.AddCookie(cookie)
			}

			rr = httptest.NewRecorder()
			tOidc.ServeHTTP(rr, req)

			if rr.Code != http.StatusOK {
				t.Errorf("Expected status code %d, got %d. Body: %s", http.StatusOK, rr.Code, rr.Body.String())
			}

			for name, expectedValue := range tc.expectedHeaders {
				if value, exists := interceptedHeaders[name]; !exists {
					// For <no value> case, it might not be set if template resolves to empty and header is omitted.
					// However, Go templates usually insert "<no value>" string.
					if expectedValue == "<no value>" && tc.name == "Missing Claim" { // Special handling for <no value>
						// If the template {{.Claims.role}} results in an empty string because role is missing,
						// and the header is not set, this is also acceptable for "<no value>".
						// The current test expects the literal string "<no value>".
						// Let's assume for now that if it's missing, it's an error unless specifically handled.
						// The test as written expects "<no value>" to be present.
					}
					t.Errorf("Expected header %s was not set", name)

				} else if value != expectedValue {
					t.Errorf("Header %s expected value %q, got %q", name, expectedValue, value)
				}
			}

			if tc.name == "Opaque Access Token with AccessTokenField" {
				postReq := httptest.NewRequest("GET", "/protected", nil)
				for _, cookie := range rr.Result().Cookies() {
					postReq.AddCookie(cookie)
				}
				updatedSession, err := tOidc.sessionManager.GetSession(postReq)
				if err != nil {
					t.Fatalf("Failed to get updated session for opaque test: %v", err)
				}

				expectedEmail := tc.claims["email"].(string)
				if updatedSession.GetEmail() != expectedEmail {
					t.Errorf("Expected session email to be %q (from ID token), got %q", expectedEmail, updatedSession.GetEmail())
				}
				if !updatedSession.GetAuthenticated() {
					t.Errorf("Session should be authenticated after successful flow for opaque test")
				}
			}
		})
	}
}

// TestEdgeCaseTemplatedHeaders tests edge cases for templated headers
func TestEdgeCaseTemplatedHeaders(t *testing.T) {
	// Create a TestSuite to use its helper methods and fields
	ts := &TestSuite{t: t}
	ts.Setup()

	tests := []struct {
		name               string
		headers            []TemplatedHeader
		claims             map[string]interface{}
		shouldExecuteCheck bool
	}{
		{
			name: "Very Large Template",
			headers: []TemplatedHeader{
				{
					Name:  "X-Large-Header",
					Value: createLargeTemplate(500), // Template with 500 variable references
				},
			},
			claims:             createLargeClaims(500), // Map with 500 claims
			shouldExecuteCheck: true,
		},
		{
			name: "Array Claim Access",
			headers: []TemplatedHeader{
				{Name: "X-Roles", Value: "{{range $i, $e := .Claims.roles}}{{if $i}},{{end}}{{$e}}{{end}}"},
			},
			claims: map[string]interface{}{
				"roles": []interface{}{"admin", "user", "manager"},
			},
			shouldExecuteCheck: true,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			// Create token with the test claims
			claims := map[string]interface{}{
				"iss":   "https://test-issuer.com",
				"aud":   "test-client-id",
				"exp":   float64(3000000000), // Far future timestamp
				"iat":   float64(1000000000),
				"nbf":   float64(1000000000),
				"sub":   "test-subject",
				"nonce": "test-nonce",
				"jti":   generateRandomString(16),
			}

			// Add the test-specific claims
			for k, v := range tc.claims {
				claims[k] = v
			}

			token, err := createTestJWT(ts.rsaPrivateKey, "RS256", "test-key-id", claims)
			if err != nil {
				t.Fatalf("Failed to create test JWT: %v", err)
			}

			// Create a test next handler
			nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.WriteHeader(http.StatusOK)
			})

			tOidc := &TraefikOidc{
				next:               nextHandler,
				name:               "test",
				redirURLPath:       "/callback",
				logoutURLPath:      "/callback/logout",
				issuerURL:          "https://test-issuer.com",
				clientID:           "test-client-id",
				clientSecret:       "test-client-secret",
				jwkCache:           ts.mockJWKCache,
				jwksURL:            "https://test-jwks-url.com",
				tokenBlacklist:     NewCache(),
				tokenCache:         NewTokenCache(),
				limiter:            rate.NewLimiter(rate.Every(time.Second), 10),
				logger:             NewLogger("debug"),
				allowedUserDomains: map[string]struct{}{"example.com": {}},
				excludedURLs:       map[string]struct{}{"/favicon": {}},
				httpClient:         &http.Client{},
				initComplete:       make(chan struct{}),
				sessionManager:     ts.sessionManager,
				extractClaimsFunc:  extractClaims,
				headerTemplates:    make(map[string]*template.Template),
			}

			// Initialize and parse header templates
			for _, header := range tc.headers {
				tmpl, err := template.New(header.Name).Parse(header.Value)
				if err != nil {
					t.Fatalf("Failed to parse header template for %s: %v", header.Name, err)
				}
				tOidc.headerTemplates[header.Name] = tmpl
			}

			close(tOidc.initComplete)

			req := httptest.NewRequest("GET", "/protected", nil)
			req.Header.Set("X-Forwarded-Proto", "https")
			req.Header.Set("X-Forwarded-Host", "example.com")
			rr := httptest.NewRecorder()

			session, err := tOidc.sessionManager.GetSession(req)
			if err != nil {
				t.Fatalf("Failed to get session: %v", err)
			}

			session.SetAuthenticated(true)
			session.SetEmail("user@example.com")
			session.SetIDToken(token)     // Use the new method
			session.SetAccessToken(token) // Also set access token to match
			session.SetRefreshToken("test-refresh-token")

			tOidc.extractClaimsFunc = extractClaims
			tOidc.tokenExchanger = &MockTokenExchanger{
				RefreshTokenFunc: func(refreshToken string) (*TokenResponse, error) {
					return &TokenResponse{
						IDToken:      token,
						AccessToken:  token,
						RefreshToken: refreshToken,
						ExpiresIn:    3600,
					}, nil
				},
			}

			if err := session.Save(req, rr); err != nil {
				t.Fatalf("Failed to save session: %v", err)
			}

			for _, cookie := range rr.Result().Cookies() {
				req.AddCookie(cookie)
			}

			rr = httptest.NewRecorder()
			tOidc.ServeHTTP(rr, req)

			if rr.Code != http.StatusOK {
				t.Errorf("Expected status code %d, got %d", http.StatusOK, rr.Code)
			}

			// The "Array Claim Access" check previously here was problematic as it didn't correctly
			// intercept headers in TestEdgeCaseTemplatedHeaders. The primary goal of this
			// function is to test edge cases for panics/errors, and robust header value
			// checking is already covered in TestTemplatedHeadersIntegration.
			// Removing the ineffective check to resolve the "declared and not used" error.
		})
	}
}

// Helper functions for edge case tests

// createLargeTemplate creates a template with many variable references
func createLargeTemplate(size int) string {
	template := "{{with .Claims}}"
	for i := 0; i < size; i++ {
		if i > 0 {
			template += ","
		}
		template += "{{.field" + string(rune('a'+i%26)) + string(rune('0'+i%10)) + "}}"
	}
	template += "{{end}}"
	return template
}

// createLargeClaims creates a map with many claims for testing large templates
func createLargeClaims(size int) map[string]interface{} {
	claims := make(map[string]interface{})
	for i := 0; i < size; i++ {
		key := "field" + string(rune('a'+i%26)) + string(rune('0'+i%10))
		claims[key] = "value" + string(rune('a'+i%26)) + string(rune('0'+i%10))
	}
	return claims
}