Add templated headers sent to the downstream service. (#40)

2026-06-05 22:44:17 +00:00 · 2025-04-14 00:45:26 +01:00
parent 983585e96e
commit a895333964
7 changed files with 1091 additions and 3 deletions
@@ -0,0 +1,424 @@
+package traefikoidc
+
+import (
+	"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: "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": "",
+			},
+		},
+	}
+
+	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
+				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)
+				}
+			}
+
+			// 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" {
+				tc.expectedHeaders["X-Auth-Info"] = "User=user@example.com, Token=" + token
+			}
+
+			// 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)
+			})
+
+			// Instead of using New(), we'll directly create a TraefikOidc instance
+			// similar to how it's done in TestSuite.Setup()
+			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 the initComplete channel to bypass the waiting
+			close(tOidc.initComplete)
+
+			// Create a test request
+			req := httptest.NewRequest("GET", "/protected", nil)
+			req.Header.Set("X-Forwarded-Proto", "https")
+			req.Header.Set("X-Forwarded-Host", "example.com")
+			rr := httptest.NewRecorder()
+
+			// Create a session
+			session, err := tOidc.sessionManager.GetSession(req)
+			if err != nil {
+				t.Fatalf("Failed to get session: %v", err)
+			}
+
+			// Setup the session with authentication data
+			session.SetAuthenticated(true)
+			session.SetEmail("user@example.com")
+			session.SetAccessToken(token)
+			session.SetRefreshToken("test-refresh-token")
+
+			if err := session.Save(req, rr); err != nil {
+				t.Fatalf("Failed to save session: %v", err)
+			}
+
+			// Add session cookies to the request
+			for _, cookie := range rr.Result().Cookies() {
+				req.AddCookie(cookie)
+			}
+
+			// Reset the response recorder for the main test
+			rr = httptest.NewRecorder()
+
+			// Process the request
+			tOidc.ServeHTTP(rr, req)
+
+			// Check status code
+			if rr.Code != http.StatusOK {
+				t.Errorf("Expected status code %d, got %d", http.StatusOK, rr.Code)
+			}
+
+			// Verify headers were set correctly
+			for name, expectedValue := range tc.expectedHeaders {
+				if value, exists := interceptedHeaders[name]; !exists {
+					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)
+				}
+			}
+		})
+	}
+}
+
+// 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)
+			})
+
+			// Instead of using New(), we'll directly create a TraefikOidc instance
+			// similar to how it's done in TestSuite.Setup()
+			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 the initComplete channel to bypass the waiting
+			close(tOidc.initComplete)
+
+			// Create a test request
+			req := httptest.NewRequest("GET", "/protected", nil)
+			req.Header.Set("X-Forwarded-Proto", "https")
+			req.Header.Set("X-Forwarded-Host", "example.com")
+			rr := httptest.NewRecorder()
+
+			// Create a session
+			session, err := tOidc.sessionManager.GetSession(req)
+			if err != nil {
+				t.Fatalf("Failed to get session: %v", err)
+			}
+
+			// Setup the session with authentication data
+			session.SetAuthenticated(true)
+			session.SetEmail("user@example.com")
+			session.SetAccessToken(token)
+			session.SetRefreshToken("test-refresh-token")
+
+			if err := session.Save(req, rr); err != nil {
+				t.Fatalf("Failed to save session: %v", err)
+			}
+
+			// Add session cookies to the request
+			for _, cookie := range rr.Result().Cookies() {
+				req.AddCookie(cookie)
+			}
+
+			// Reset the response recorder for the main test
+			rr = httptest.NewRecorder()
+
+			// Process the request
+			tOidc.ServeHTTP(rr, req)
+
+			// Check status code
+			if rr.Code != http.StatusOK {
+				t.Errorf("Expected status code %d, got %d", http.StatusOK, rr.Code)
+			}
+
+			// We are primarily checking that these edge cases don't cause panics or errors
+			// For the array test, we can verify the content
+			if tc.name == "Array Claim Access" {
+				// Check if the header was set
+				headerValue := req.Header.Get("X-Roles")
+				expectedValue := "admin,user,manager"
+				if headerValue != expectedValue {
+					t.Errorf("Expected X-Roles header to be %q, got %q", expectedValue, headerValue)
+				}
+			}
+		})
+	}
+}
+
+// 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
+}