traefikoidc/pkce_test.go

package traefikoidc

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestPKCEGeneration(t *testing.T) {
	tests := []struct {
		test func(t *testing.T)
		name string
	}{
		{
			name: "generateCodeVerifier creates valid verifier",
			test: func(t *testing.T) {
				verifier, err := generateCodeVerifier()
				require.NoError(t, err)

				// RFC 7636: code_verifier must be 43-128 characters
				assert.GreaterOrEqual(t, len(verifier), 43)
				assert.LessOrEqual(t, len(verifier), 128)

				// Should be base64url encoded (no padding, no +/)
				assert.NotContains(t, verifier, "=")
				assert.NotContains(t, verifier, "+")
				assert.NotContains(t, verifier, "/")

				// Should be URL safe
				assert.Equal(t, url.QueryEscape(verifier), verifier)
			},
		},
		{
			name: "generateCodeVerifier creates unique values",
			test: func(t *testing.T) {
				verifiers := make(map[string]bool)
				for i := 0; i < 100; i++ {
					v, err := generateCodeVerifier()
					require.NoError(t, err)
					assert.False(t, verifiers[v], "Generated duplicate code verifier")
					verifiers[v] = true
				}
			},
		},
		{
			name: "deriveCodeChallenge creates valid S256 challenge",
			test: func(t *testing.T) {
				verifier := "dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk"
				challenge := deriveCodeChallenge(verifier)

				// Expected challenge for the test verifier (from RFC 7636 example)
				expected := "E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM"
				assert.Equal(t, expected, challenge)

				// Should be base64url encoded
				assert.NotContains(t, challenge, "=")
				assert.NotContains(t, challenge, "+")
				assert.NotContains(t, challenge, "/")
			},
		},
		{
			name: "deriveCodeChallenge handles empty verifier",
			test: func(t *testing.T) {
				challenge := deriveCodeChallenge("")

				// SHA256 of empty string, base64url encoded
				h := sha256.Sum256([]byte(""))
				expected := base64.RawURLEncoding.EncodeToString(h[:])
				assert.Equal(t, expected, challenge)
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tt.test(t)
		})
	}
}

func TestPKCEAuthorizationFlow(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping test in short mode")
	}

	tests := []struct {
		test       func(t *testing.T, authURL string)
		name       string
		enablePKCE bool
	}{
		{
			name:       "PKCE enabled adds code_challenge parameters",
			enablePKCE: true,
			test: func(t *testing.T, authURL string) {
				u, err := url.Parse(authURL)
				require.NoError(t, err)

				params := u.Query()

				// Should have code_challenge and code_challenge_method
				assert.NotEmpty(t, params.Get("code_challenge"))
				assert.Equal(t, "S256", params.Get("code_challenge_method"))

				// Code challenge should be properly formatted
				challenge := params.Get("code_challenge")
				assert.NotContains(t, challenge, "=")
				assert.NotContains(t, challenge, "+")
				assert.NotContains(t, challenge, "/")
				assert.Greater(t, len(challenge), 0)
			},
		},
		{
			name:       "PKCE disabled omits code_challenge parameters",
			enablePKCE: false,
			test: func(t *testing.T, authURL string) {
				u, err := url.Parse(authURL)
				require.NoError(t, err)

				params := u.Query()

				// Should not have PKCE parameters
				assert.Empty(t, params.Get("code_challenge"))
				assert.Empty(t, params.Get("code_challenge_method"))
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Setup test environment
			config := createTestConfig()
			config.EnablePKCE = tt.enablePKCE

			oidc, _ := setupTestOIDCMiddleware(t, config)
			oidc.enablePKCE = tt.enablePKCE

			// Create test request
			req := httptest.NewRequest("GET", "/protected", nil)
			rec := httptest.NewRecorder()

			// Trigger authentication
			oidc.ServeHTTP(rec, req)

			// Check redirect
			assert.Equal(t, http.StatusFound, rec.Code)
			location := rec.Header().Get("Location")
			assert.NotEmpty(t, location)

			// Run test specific checks
			tt.test(t, location)
		})
	}
}

func TestPKCESessionManagement(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping test in short mode")
	}

	tests := []struct {
		test func(t *testing.T)
		name string
	}{
		{
			name: "stores and retrieves code verifier in session",
			test: func(t *testing.T) {
				session := createTestSession()
				verifier, err := generateCodeVerifier()
				require.NoError(t, err)

				// Store verifier
				session.SetCodeVerifier(verifier)

				// Retrieve verifier
				retrieved := session.GetCodeVerifier()
				assert.Equal(t, verifier, retrieved)
			},
		},
		{
			name: "code verifier persists through session operations",
			test: func(t *testing.T) {
				session := createTestSession()
				verifier, err := generateCodeVerifier()
				require.NoError(t, err)

				// Store verifier and other data
				session.SetCodeVerifier(verifier)
				session.SetAccessToken("test-access-token")
				session.SetIDToken("test-id-token")

				// Verifier should still be there
				assert.Equal(t, verifier, session.GetCodeVerifier())
			},
		},
		{
			name: "code verifier cleared after token exchange",
			test: func(t *testing.T) {
				config := createTestConfig()
				config.EnablePKCE = true

				oidc, server := setupTestOIDCMiddleware(t, config)
				defer server.Close()

				oidc.enablePKCE = true

				// Create session with code verifier
				session := createTestSession()
				verifier, err := generateCodeVerifier()
				require.NoError(t, err)
				session.SetCodeVerifier(verifier)

				// Simulate callback with code
				req := httptest.NewRequest("GET", config.CallbackURL+"?code=test-code&state=test-state", nil)
				rec := httptest.NewRecorder()

				// Add session cookie
				// For testing, we would need to add the session to the request
				// This is a simplified approach - in real tests, use proper session injection

				// Handle callback
				oidc.ServeHTTP(rec, req)

				// Verify code verifier was used and cleared
				// Note: In real implementation, this would be cleared after successful exchange
				// This test verifies the session flow
				assert.NotNil(t, session)
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tt.test(t)
		})
	}
}

func TestPKCETokenExchange(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping test in short mode")
	}

	tests := []struct {
		name         string
		codeVerifier string
		enablePKCE   bool
		expectParam  bool
	}{
		{
			name:         "includes code_verifier when PKCE enabled",
			enablePKCE:   true,
			codeVerifier: "test-verifier-123",
			expectParam:  true,
		},
		{
			name:         "omits code_verifier when PKCE disabled",
			enablePKCE:   false,
			codeVerifier: "",
			expectParam:  false,
		},
		{
			name:         "omits code_verifier when empty even if PKCE enabled",
			enablePKCE:   true,
			codeVerifier: "",
			expectParam:  false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Create a test server to capture the token exchange request
			var capturedBody string
			tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				body, _ := io.ReadAll(r.Body)
				capturedBody = string(body)

				// Return mock tokens
				w.Header().Set("Content-Type", "application/json")
				w.Write([]byte(`{
					"access_token": "test-access-token",
					"id_token": "` + ValidIDToken + `",
					"token_type": "bearer",
					"expires_in": 3600
				}`))
			}))
			defer tokenServer.Close()

			// Setup OIDC with custom token endpoint
			config := createTestConfig()
			config.EnablePKCE = tt.enablePKCE

			oidc, _ := setupTestOIDCMiddleware(t, config)
			oidc.tokenURL = tokenServer.URL

			// Exchange tokens
			_, err := oidc.ExchangeCodeForToken(
				context.Background(),
				"authorization_code",
				"test-code",
				config.CallbackURL,
				tt.codeVerifier,
			)
			require.NoError(t, err)

			// Check if code_verifier was included
			if tt.expectParam {
				assert.Contains(t, capturedBody, "code_verifier="+tt.codeVerifier)
			} else {
				assert.NotContains(t, capturedBody, "code_verifier")
			}
		})
	}
}

func TestPKCEEndToEndFlow(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping test in short mode")
	}

	// Setup test environment
	config := createTestConfig()
	config.EnablePKCE = true

	oidc, server := setupTestOIDCMiddleware(t, config)
	defer server.Close()

	oidc.enablePKCE = true

	// Generate a code verifier for testing
	testCodeVerifier, err := generateCodeVerifier()
	require.NoError(t, err)
	testCodeChallenge := deriveCodeChallenge(testCodeVerifier)

	// Mock the token exchange to verify code_verifier is sent
	var receivedVerifier string
	tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		r.ParseForm()
		receivedVerifier = r.Form.Get("code_verifier")

		// Return mock tokens
		w.Header().Set("Content-Type", "application/json")
		w.Write([]byte(`{
			"access_token": "test-access-token",
			"id_token": "` + ValidIDToken + `",
			"token_type": "bearer",
			"expires_in": 3600
		}`))
	}))
	defer tokenServer.Close()

	oidc.tokenURL = tokenServer.URL

	// Mock the token verifier to avoid JWKS lookup
	oidc.tokenVerifier = &mockTokenVerifier{
		verifyFunc: func(token string) error {
			// Always return success for test tokens
			claims, err := extractClaims(token)
			if err != nil {
				return err
			}
			// Cache the claims for the token
			oidc.tokenCache.Set(token, claims, time.Hour)
			return nil
		},
	}

	// Step 1: Simulate the callback directly with a pre-configured session
	// This bypasses the session persistence issue in the test environment
	callbackReq := httptest.NewRequest("GET", config.CallbackURL+"?code=test-code&state=test-state", nil)
	callbackRec := httptest.NewRecorder()

	// Get a session and set it up as if the auth flow had started
	session, err := oidc.sessionManager.GetSession(callbackReq)
	require.NoError(t, err)

	// Set up the session as the auth initiation would have done
	session.SetCSRF("test-state")
	session.SetNonce("nonce123") // Must match the nonce in ValidIDToken
	session.SetCodeVerifier(testCodeVerifier)
	session.SetIncomingPath("/protected")

	// Save the session
	err = session.Save(callbackReq, callbackRec)
	require.NoError(t, err)

	// Create a new request with the session cookies
	callbackReq2 := httptest.NewRequest("GET", config.CallbackURL+"?code=test-code&state=test-state", nil)
	for _, cookie := range callbackRec.Result().Cookies() {
		callbackReq2.AddCookie(cookie)
	}
	callbackRec2 := httptest.NewRecorder()

	// Handle callback
	oidc.ServeHTTP(callbackRec2, callbackReq2)

	// Verify successful authentication
	assert.Equal(t, http.StatusFound, callbackRec2.Code)
	assert.Equal(t, testCodeVerifier, receivedVerifier, "Code verifier should be sent in token exchange")

	// Also test the authorization URL building with PKCE
	authURL := oidc.buildAuthURL("http://example.com/callback", "test-csrf", "test-nonce", testCodeChallenge)
	parsedURL, err := url.Parse(authURL)
	require.NoError(t, err)

	assert.Equal(t, testCodeChallenge, parsedURL.Query().Get("code_challenge"))
	assert.Equal(t, "S256", parsedURL.Query().Get("code_challenge_method"))
}

func TestPKCESecurityEdgeCases(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping test in short mode")
	}

	tests := []struct {
		test func(t *testing.T)
		name string
	}{
		{
			name: "rejects callback without matching state",
			test: func(t *testing.T) {
				config := createTestConfig()
				config.EnablePKCE = true

				oidc, _ := setupTestOIDCMiddleware(t, config)
				oidc.enablePKCE = true

				// Create callback request with wrong state
				req := httptest.NewRequest("GET", config.CallbackURL+"?code=test-code&state=wrong-state", nil)
				rec := httptest.NewRecorder()

				oidc.ServeHTTP(rec, req)

				// Should reject due to state mismatch
				assert.Equal(t, http.StatusBadRequest, rec.Code)
			},
		},
		{
			name: "handles missing code_verifier gracefully",
			test: func(t *testing.T) {
				config := createTestConfig()
				config.EnablePKCE = true

				oidc, server := setupTestOIDCMiddleware(t, config)
				defer server.Close()

				// Create session without code verifier
				session := createTestSession()
				session.mainSession.Values["state"] = "test-state"
				// Intentionally not setting code verifier

				req := httptest.NewRequest("GET", config.CallbackURL+"?code=test-code&state=test-state", nil)
				rec := httptest.NewRecorder()

				// Add session
				// For testing, we would need to add the session to the request
				// This is a simplified approach - in real tests, use proper session injection

				// Should handle gracefully even without verifier
				oidc.ServeHTTP(rec, req)

				// The actual behavior depends on provider - some may reject, others may accept
				// The important thing is no panic/crash
				assert.NotNil(t, rec)
			},
		},
		{
			name: "code verifier is single use",
			test: func(t *testing.T) {
				session := createTestSession()
				verifier, err := generateCodeVerifier()
				require.NoError(t, err)

				// Set verifier
				session.SetCodeVerifier(verifier)
				assert.Equal(t, verifier, session.GetCodeVerifier())

				// In real flow, it would be cleared after use
				// This test verifies the concept
				session.SetCodeVerifier("")
				assert.Empty(t, session.GetCodeVerifier())
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tt.test(t)
		})
	}
}

func TestPKCECompatibilityWithProviders(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping test in short mode")
	}

	providers := []struct {
		name         string
		providerType string
		supportsPKCE bool
	}{
		{"Google", "google", true},
		{"Azure", "azure", true},
		{"Generic", "generic", true},
	}

	for _, provider := range providers {
		t.Run(provider.name+" provider with PKCE", func(t *testing.T) {
			config := createTestConfig()
			config.EnablePKCE = true
			config.ProviderURL = "https://" + provider.providerType + ".example.com"

			oidc, _ := setupTestOIDCMiddleware(t, config)
			oidc.enablePKCE = true

			// Test auth URL generation
			req := httptest.NewRequest("GET", "/protected", nil)
			rec := httptest.NewRecorder()

			oidc.ServeHTTP(rec, req)

			if provider.supportsPKCE {
				location := rec.Header().Get("Location")
				assert.Contains(t, location, "code_challenge")
				assert.Contains(t, location, "code_challenge_method=S256")
			}
		})
	}
}