diff --git a/auth/auth_handler.go b/auth/auth_handler.go new file mode 100644 index 0000000..eb466a6 --- /dev/null +++ b/auth/auth_handler.go @@ -0,0 +1,360 @@ +// Package auth provides authentication-related functionality for the OIDC middleware. +package auth + +import ( + "fmt" + "net" + "net/http" + "net/url" + "strings" + + "github.com/google/uuid" +) + +// AuthHandler provides core authentication functionality for OIDC flows +type AuthHandler struct { + logger Logger + enablePKCE bool + isGoogleProv func() bool + isAzureProv func() bool + clientID string + authURL string + issuerURL string + scopes []string + overrideScopes bool +} + +// Logger interface for dependency injection +type Logger interface { + Debugf(format string, args ...interface{}) + Errorf(format string, args ...interface{}) +} + +// NewAuthHandler creates a new AuthHandler instance +func NewAuthHandler(logger Logger, enablePKCE bool, isGoogleProv, isAzureProv func() bool, + clientID, authURL, issuerURL string, scopes []string, overrideScopes bool) *AuthHandler { + return &AuthHandler{ + logger: logger, + enablePKCE: enablePKCE, + isGoogleProv: isGoogleProv, + isAzureProv: isAzureProv, + clientID: clientID, + authURL: authURL, + issuerURL: issuerURL, + scopes: scopes, + overrideScopes: overrideScopes, + } +} + +// InitiateAuthentication initiates the OIDC authentication flow. +// It generates CSRF tokens, nonce, PKCE parameters (if enabled), clears the session, +// stores authentication state, and redirects the user to the OIDC provider. +func (h *AuthHandler) InitiateAuthentication(rw http.ResponseWriter, req *http.Request, + session SessionData, redirectURL string, + generateNonce, generateCodeVerifier, deriveCodeChallenge func() (string, error)) { + + h.logger.Debugf("Initiating new OIDC authentication flow for request: %s", req.URL.RequestURI()) + + const maxRedirects = 5 + redirectCount := session.GetRedirectCount() + if redirectCount >= maxRedirects { + h.logger.Errorf("Maximum redirect limit (%d) exceeded, possible redirect loop detected", maxRedirects) + session.ResetRedirectCount() + http.Error(rw, "Authentication failed: Too many redirects", http.StatusLoopDetected) + return + } + + session.IncrementRedirectCount() + + csrfToken := uuid.NewString() + nonce, err := generateNonce() + if err != nil { + h.logger.Errorf("Failed to generate nonce: %v", err) + http.Error(rw, "Failed to generate nonce", http.StatusInternalServerError) + return + } + + // Generate PKCE code verifier and challenge if PKCE is enabled + var codeVerifier, codeChallenge string + if h.enablePKCE { + codeVerifier, err = generateCodeVerifier() + if err != nil { + h.logger.Errorf("Failed to generate code verifier: %v", err) + http.Error(rw, "Failed to generate code verifier", http.StatusInternalServerError) + return + } + codeChallenge, err = deriveCodeChallenge() + if err != nil { + h.logger.Errorf("Failed to generate code challenge: %v", err) + http.Error(rw, "Failed to generate code challenge", http.StatusInternalServerError) + return + } + h.logger.Debugf("PKCE enabled, generated code challenge") + } + + session.SetAuthenticated(false) + session.SetEmail("") + session.SetAccessToken("") + session.SetRefreshToken("") + session.SetIDToken("") + session.SetNonce("") + session.SetCodeVerifier("") + + session.SetCSRF(csrfToken) + session.SetNonce(nonce) + if h.enablePKCE { + session.SetCodeVerifier(codeVerifier) + } + session.SetIncomingPath(req.URL.RequestURI()) + h.logger.Debugf("Storing incoming path: %s", req.URL.RequestURI()) + + session.MarkDirty() + + if err := session.Save(req, rw); err != nil { + h.logger.Errorf("Failed to save session before redirecting to provider: %v", err) + http.Error(rw, "Failed to save session", http.StatusInternalServerError) + return + } + + h.logger.Debugf("Session saved before redirect. CSRF: %s, Nonce: %s", + csrfToken, nonce) + + authURL := h.BuildAuthURL(redirectURL, csrfToken, nonce, codeChallenge) + h.logger.Debugf("Redirecting user to OIDC provider: %s", authURL) + + http.Redirect(rw, req, authURL, http.StatusFound) +} + +// BuildAuthURL constructs the OIDC provider authorization URL. +// It builds the URL with all necessary parameters including client_id, scopes, +// PKCE parameters, and provider-specific parameters for Google and Azure. +func (h *AuthHandler) BuildAuthURL(redirectURL, state, nonce, codeChallenge string) string { + params := url.Values{} + params.Set("client_id", h.clientID) + params.Set("response_type", "code") + params.Set("redirect_uri", redirectURL) + params.Set("state", state) + params.Set("nonce", nonce) + + if h.enablePKCE && codeChallenge != "" { + params.Set("code_challenge", codeChallenge) + params.Set("code_challenge_method", "S256") + } + + scopes := make([]string, len(h.scopes)) + copy(scopes, h.scopes) + + if h.isGoogleProv() { + params.Set("access_type", "offline") + h.logger.Debugf("Google OIDC provider detected, added access_type=offline for refresh tokens") + + params.Set("prompt", "consent") + h.logger.Debugf("Google OIDC provider detected, added prompt=consent to ensure refresh tokens") + } else if h.isAzureProv() { + params.Set("response_mode", "query") + h.logger.Debugf("Azure AD provider detected, added response_mode=query") + + hasOfflineAccess := false + + for _, scope := range scopes { + if scope == "offline_access" { + hasOfflineAccess = true + break + } + } + + if !h.overrideScopes || (h.overrideScopes && len(h.scopes) == 0) { + if !hasOfflineAccess { + scopes = append(scopes, "offline_access") + h.logger.Debugf("Azure AD provider: Added offline_access scope (overrideScopes: %t, user scopes count: %d)", h.overrideScopes, len(h.scopes)) + } + } else { + h.logger.Debugf("Azure AD provider: User is overriding scopes (count: %d), offline_access not automatically added.", len(h.scopes)) + } + } else { + if !h.overrideScopes || (h.overrideScopes && len(h.scopes) == 0) { + hasOfflineAccess := false + for _, scope := range scopes { + if scope == "offline_access" { + hasOfflineAccess = true + break + } + } + if !hasOfflineAccess { + scopes = append(scopes, "offline_access") + h.logger.Debugf("Standard provider: Added offline_access scope (overrideScopes: %t, user scopes count: %d)", h.overrideScopes, len(h.scopes)) + } + } else { + h.logger.Debugf("Standard provider: User is overriding scopes (count: %d), offline_access not automatically added.", len(h.scopes)) + } + } + + if len(scopes) > 0 { + finalScopeString := strings.Join(scopes, " ") + params.Set("scope", finalScopeString) + h.logger.Debugf("AuthHandler.BuildAuthURL: Final scope string being sent to OIDC provider: %s", finalScopeString) + } + + return h.buildURLWithParams(h.authURL, params) +} + +// buildURLWithParams constructs a URL by combining a base URL with query parameters. +// It handles both relative and absolute URLs, validates URL security, +// and properly encodes query parameters. +func (h *AuthHandler) buildURLWithParams(baseURL string, params url.Values) string { + if baseURL != "" { + if strings.HasPrefix(baseURL, "http://") || strings.HasPrefix(baseURL, "https://") { + if err := h.validateURL(baseURL); err != nil { + h.logger.Errorf("URL validation failed for %s: %v", baseURL, err) + return "" + } + } + } + + if !strings.HasPrefix(baseURL, "http://") && !strings.HasPrefix(baseURL, "https://") { + issuerURLParsed, err := url.Parse(h.issuerURL) + if err != nil { + h.logger.Errorf("Could not parse issuerURL: %s. Error: %v", h.issuerURL, err) + return "" + } + + baseURLParsed, err := url.Parse(baseURL) + if err != nil { + h.logger.Errorf("Could not parse baseURL: %s. Error: %v", baseURL, err) + return "" + } + + resolvedURL := issuerURLParsed.ResolveReference(baseURLParsed) + + if err := h.validateURL(resolvedURL.String()); err != nil { + h.logger.Errorf("Resolved URL validation failed for %s: %v", resolvedURL.String(), err) + return "" + } + + resolvedURL.RawQuery = params.Encode() + return resolvedURL.String() + } + + u, err := url.Parse(baseURL) + if err != nil { + h.logger.Errorf("Could not parse absolute baseURL: %s. Error: %v", baseURL, err) + return "" + } + + if err := h.validateParsedURL(u); err != nil { + h.logger.Errorf("Parsed URL validation failed for %s: %v", baseURL, err) + return "" + } + + u.RawQuery = params.Encode() + return u.String() +} + +// validateURL performs security validation on URLs to prevent SSRF attacks. +// It checks for allowed schemes, validates hosts, and prevents access to private networks. +func (h *AuthHandler) validateURL(urlStr string) error { + if urlStr == "" { + return fmt.Errorf("empty URL") + } + + u, err := url.Parse(urlStr) + if err != nil { + return fmt.Errorf("invalid URL format: %w", err) + } + + return h.validateParsedURL(u) +} + +// validateParsedURL validates a parsed URL structure for security. +// It checks schemes, hosts, and paths to prevent malicious URLs. +func (h *AuthHandler) validateParsedURL(u *url.URL) error { + allowedSchemes := map[string]bool{ + "https": true, + "http": true, + } + + if !allowedSchemes[u.Scheme] { + return fmt.Errorf("disallowed URL scheme: %s", u.Scheme) + } + + if u.Scheme == "http" { + h.logger.Debugf("Warning: Using HTTP scheme for URL: %s", u.String()) + } + + if u.Host == "" { + return fmt.Errorf("missing host in URL") + } + + if err := h.validateHost(u.Host); err != nil { + return fmt.Errorf("invalid host: %w", err) + } + + if strings.Contains(u.Path, "..") { + return fmt.Errorf("path traversal detected in URL path") + } + + return nil +} + +// validateHost validates a hostname for security and reachability. +// It prevents access to private networks and localhost addresses. +func (h *AuthHandler) validateHost(host string) error { + if host == "" { + return fmt.Errorf("empty host") + } + + // Strip port if present + if strings.Contains(host, ":") { + var err error + host, _, err = net.SplitHostPort(host) + if err != nil { + return fmt.Errorf("invalid host:port format: %w", err) + } + } + + // Check for localhost variations + localhostVariations := []string{ + "localhost", "127.0.0.1", "::1", "0.0.0.0", + } + for _, localhost := range localhostVariations { + if strings.EqualFold(host, localhost) { + return fmt.Errorf("localhost access not allowed: %s", host) + } + } + + // Try to parse as IP address + if ip := net.ParseIP(host); ip != nil { + if ip.IsLoopback() { + return fmt.Errorf("loopback IP not allowed: %s", host) + } + if ip.IsPrivate() { + return fmt.Errorf("private IP not allowed: %s", host) + } + if ip.IsLinkLocalUnicast() { + return fmt.Errorf("link-local IP not allowed: %s", host) + } + if ip.IsMulticast() { + return fmt.Errorf("multicast IP not allowed: %s", host) + } + } + + return nil +} + +// SessionData interface for dependency injection +type SessionData interface { + GetRedirectCount() int + ResetRedirectCount() + IncrementRedirectCount() + SetAuthenticated(bool) + SetEmail(string) + SetAccessToken(string) + SetRefreshToken(string) + SetIDToken(string) + SetNonce(string) + SetCodeVerifier(string) + SetCSRF(string) + SetIncomingPath(string) + MarkDirty() + Save(req *http.Request, rw http.ResponseWriter) error +} diff --git a/auth/auth_handler_test.go b/auth/auth_handler_test.go new file mode 100644 index 0000000..26ef7f3 --- /dev/null +++ b/auth/auth_handler_test.go @@ -0,0 +1,923 @@ +package auth + +import ( + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "sync" + "testing" +) + +// Mock logger for testing +type mockLogger struct { + debugLogs []string + errorLogs []string + mu sync.RWMutex +} + +func (m *mockLogger) Debugf(format string, args ...interface{}) { + m.mu.Lock() + defer m.mu.Unlock() + m.debugLogs = append(m.debugLogs, fmt.Sprintf(format, args...)) +} + +func (m *mockLogger) Errorf(format string, args ...interface{}) { + m.mu.Lock() + defer m.mu.Unlock() + m.errorLogs = append(m.errorLogs, fmt.Sprintf(format, args...)) +} + +//lint:ignore U1000 May be needed for future debug log verification tests +func (m *mockLogger) getDebugLogs() []string { + m.mu.RLock() + defer m.mu.RUnlock() + result := make([]string, len(m.debugLogs)) + copy(result, m.debugLogs) + return result +} + +func (m *mockLogger) getErrorLogs() []string { + m.mu.RLock() + defer m.mu.RUnlock() + result := make([]string, len(m.errorLogs)) + copy(result, m.errorLogs) + return result +} + +//lint:ignore U1000 May be needed for future test isolation +func (m *mockLogger) reset() { + m.mu.Lock() + defer m.mu.Unlock() + m.debugLogs = nil + m.errorLogs = nil +} + +// Mock session for testing +type mockSession struct { + data map[string]interface{} + redirectCount int + saveError error + mu sync.RWMutex + dirty bool + authenticated bool + email string + accessToken string + refreshToken string + idToken string + nonce string + codeVerifier string + csrf string + incomingPath string +} + +func newMockSession() *mockSession { + return &mockSession{ + data: make(map[string]interface{}), + } +} + +func (m *mockSession) GetRedirectCount() int { + m.mu.RLock() + defer m.mu.RUnlock() + return m.redirectCount +} + +func (m *mockSession) ResetRedirectCount() { + m.mu.Lock() + defer m.mu.Unlock() + m.redirectCount = 0 +} + +func (m *mockSession) IncrementRedirectCount() { + m.mu.Lock() + defer m.mu.Unlock() + m.redirectCount++ +} + +func (m *mockSession) SetAuthenticated(auth bool) { + m.mu.Lock() + defer m.mu.Unlock() + m.authenticated = auth +} + +func (m *mockSession) SetEmail(email string) { + m.mu.Lock() + defer m.mu.Unlock() + m.email = email +} + +func (m *mockSession) SetAccessToken(token string) { + m.mu.Lock() + defer m.mu.Unlock() + m.accessToken = token +} + +func (m *mockSession) SetRefreshToken(token string) { + m.mu.Lock() + defer m.mu.Unlock() + m.refreshToken = token +} + +func (m *mockSession) SetIDToken(token string) { + m.mu.Lock() + defer m.mu.Unlock() + m.idToken = token +} + +func (m *mockSession) SetNonce(nonce string) { + m.mu.Lock() + defer m.mu.Unlock() + m.nonce = nonce +} + +func (m *mockSession) SetCodeVerifier(verifier string) { + m.mu.Lock() + defer m.mu.Unlock() + m.codeVerifier = verifier +} + +func (m *mockSession) SetCSRF(csrf string) { + m.mu.Lock() + defer m.mu.Unlock() + m.csrf = csrf +} + +func (m *mockSession) SetIncomingPath(path string) { + m.mu.Lock() + defer m.mu.Unlock() + m.incomingPath = path +} + +func (m *mockSession) MarkDirty() { + m.mu.Lock() + defer m.mu.Unlock() + m.dirty = true +} + +func (m *mockSession) Save(req *http.Request, rw http.ResponseWriter) error { + m.mu.RLock() + defer m.mu.RUnlock() + return m.saveError +} + +func (m *mockSession) setSaveError(err error) { + m.mu.Lock() + defer m.mu.Unlock() + m.saveError = err +} + +// Test helper functions +func createTestAuthHandler(logger Logger, enablePKCE bool, isGoogle, isAzure func() bool) *AuthHandler { + return NewAuthHandler( + logger, + enablePKCE, + isGoogle, + isAzure, + "test-client-id", + "https://provider.example.com/auth", + "https://provider.example.com", + []string{"openid", "email", "profile"}, + false, + ) +} + +func TestNewAuthHandler(t *testing.T) { + tests := []struct { + name string + logger Logger + enablePKCE bool + clientID string + authURL string + issuerURL string + scopes []string + overrideScopes bool + }{ + { + name: "creates handler with basic config", + logger: &mockLogger{}, + enablePKCE: true, + clientID: "test-client", + authURL: "https://example.com/auth", + issuerURL: "https://example.com", + scopes: []string{"openid", "email"}, + overrideScopes: false, + }, + { + name: "creates handler with PKCE disabled", + logger: &mockLogger{}, + enablePKCE: false, + clientID: "test-client", + authURL: "https://example.com/auth", + issuerURL: "https://example.com", + scopes: []string{"openid"}, + overrideScopes: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + isGoogle := func() bool { return false } + isAzure := func() bool { return false } + + handler := NewAuthHandler( + tt.logger, + tt.enablePKCE, + isGoogle, + isAzure, + tt.clientID, + tt.authURL, + tt.issuerURL, + tt.scopes, + tt.overrideScopes, + ) + + if handler == nil { + t.Fatal("NewAuthHandler returned nil") + } + + if handler.clientID != tt.clientID { + t.Errorf("Expected clientID %s, got %s", tt.clientID, handler.clientID) + } + + if handler.enablePKCE != tt.enablePKCE { + t.Errorf("Expected enablePKCE %v, got %v", tt.enablePKCE, handler.enablePKCE) + } + + if handler.authURL != tt.authURL { + t.Errorf("Expected authURL %s, got %s", tt.authURL, handler.authURL) + } + + if handler.issuerURL != tt.issuerURL { + t.Errorf("Expected issuerURL %s, got %s", tt.issuerURL, handler.issuerURL) + } + + if handler.overrideScopes != tt.overrideScopes { + t.Errorf("Expected overrideScopes %v, got %v", tt.overrideScopes, handler.overrideScopes) + } + }) + } +} + +func TestInitiateAuthentication(t *testing.T) { + tests := []struct { + name string + enablePKCE bool + initialRedirectCount int + generateNonceError error + generateCodeVerifierError error + deriveCodeChallengeError error + sessionSaveError error + expectedStatusCode int + expectRedirect bool + }{ + { + name: "successful authentication without PKCE", + enablePKCE: false, + expectedStatusCode: http.StatusFound, + expectRedirect: true, + }, + { + name: "successful authentication with PKCE", + enablePKCE: true, + expectedStatusCode: http.StatusFound, + expectRedirect: true, + }, + { + name: "fails when nonce generation fails", + enablePKCE: false, + generateNonceError: fmt.Errorf("nonce generation failed"), + expectedStatusCode: http.StatusInternalServerError, + expectRedirect: false, + }, + { + name: "fails when code verifier generation fails", + enablePKCE: true, + generateCodeVerifierError: fmt.Errorf("code verifier generation failed"), + expectedStatusCode: http.StatusInternalServerError, + expectRedirect: false, + }, + { + name: "fails when code challenge derivation fails", + enablePKCE: true, + deriveCodeChallengeError: fmt.Errorf("code challenge derivation failed"), + expectedStatusCode: http.StatusInternalServerError, + expectRedirect: false, + }, + { + name: "fails when session save fails", + enablePKCE: false, + sessionSaveError: fmt.Errorf("session save failed"), + expectedStatusCode: http.StatusInternalServerError, + expectRedirect: false, + }, + { + name: "fails when redirect count exceeds limit", + enablePKCE: false, + initialRedirectCount: 5, + expectedStatusCode: http.StatusLoopDetected, + expectRedirect: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + logger := &mockLogger{} + session := newMockSession() + session.redirectCount = tt.initialRedirectCount + if tt.sessionSaveError != nil { + session.setSaveError(tt.sessionSaveError) + } + + handler := createTestAuthHandler(logger, tt.enablePKCE, func() bool { return false }, func() bool { return false }) + + req := httptest.NewRequest("GET", "/test", nil) + rw := httptest.NewRecorder() + + generateNonce := func() (string, error) { + if tt.generateNonceError != nil { + return "", tt.generateNonceError + } + return "test-nonce", nil + } + + generateCodeVerifier := func() (string, error) { + if tt.generateCodeVerifierError != nil { + return "", tt.generateCodeVerifierError + } + return "test-code-verifier", nil + } + + deriveCodeChallenge := func() (string, error) { + if tt.deriveCodeChallengeError != nil { + return "", tt.deriveCodeChallengeError + } + return "test-code-challenge", nil + } + + handler.InitiateAuthentication( + rw, req, session, + "https://example.com/callback", + generateNonce, + generateCodeVerifier, + deriveCodeChallenge, + ) + + if rw.Code != tt.expectedStatusCode { + t.Errorf("Expected status code %d, got %d", tt.expectedStatusCode, rw.Code) + } + + if tt.expectRedirect { + location := rw.Header().Get("Location") + if location == "" { + t.Error("Expected redirect location header, got empty") + } + if !strings.Contains(location, "https://provider.example.com/auth") { + t.Errorf("Expected redirect to auth URL, got %s", location) + } + } + + // Verify error logs for failure cases + if tt.expectedStatusCode != http.StatusFound { + errorLogs := logger.getErrorLogs() + if len(errorLogs) == 0 { + t.Error("Expected error logs for failure case, got none") + } + } + }) + } +} + +func TestBuildAuthURL(t *testing.T) { + tests := []struct { + name string + enablePKCE bool + isGoogle bool + isAzure bool + scopes []string + overrideScopes bool + codeChallenge string + expectedParams map[string]string + expectContains []string + }{ + { + name: "builds URL without PKCE for standard provider", + enablePKCE: false, + isGoogle: false, + isAzure: false, + scopes: []string{"openid", "email"}, + codeChallenge: "", + expectedParams: map[string]string{ + "client_id": "test-client-id", + "response_type": "code", + "state": "test-state", + "nonce": "test-nonce", + "scope": "openid email offline_access", + }, + }, + { + name: "builds URL with PKCE for standard provider", + enablePKCE: true, + isGoogle: false, + isAzure: false, + scopes: []string{"openid", "email"}, + codeChallenge: "test-challenge", + expectedParams: map[string]string{ + "client_id": "test-client-id", + "response_type": "code", + "state": "test-state", + "nonce": "test-nonce", + "code_challenge": "test-challenge", + "code_challenge_method": "S256", + "scope": "openid email offline_access", + }, + }, + { + name: "builds URL for Google provider", + enablePKCE: false, + isGoogle: true, + isAzure: false, + scopes: []string{"openid", "email"}, + expectedParams: map[string]string{ + "client_id": "test-client-id", + "response_type": "code", + "state": "test-state", + "nonce": "test-nonce", + "access_type": "offline", + "prompt": "consent", + "scope": "openid email", + }, + }, + { + name: "builds URL for Azure provider", + enablePKCE: false, + isGoogle: false, + isAzure: true, + scopes: []string{"openid", "email"}, + expectedParams: map[string]string{ + "client_id": "test-client-id", + "response_type": "code", + "state": "test-state", + "nonce": "test-nonce", + "response_mode": "query", + "scope": "openid email offline_access", + }, + }, + { + name: "handles scope override for Azure", + enablePKCE: false, + isGoogle: false, + isAzure: true, + scopes: []string{"custom-scope"}, + overrideScopes: true, + expectedParams: map[string]string{ + "client_id": "test-client-id", + "response_type": "code", + "state": "test-state", + "nonce": "test-nonce", + "response_mode": "query", + "scope": "custom-scope", + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + logger := &mockLogger{} + isGoogleFunc := func() bool { return tt.isGoogle } + isAzureFunc := func() bool { return tt.isAzure } + + handler := NewAuthHandler( + logger, + tt.enablePKCE, + isGoogleFunc, + isAzureFunc, + "test-client-id", + "https://provider.example.com/auth", + "https://provider.example.com", + tt.scopes, + tt.overrideScopes, + ) + + authURL := handler.BuildAuthURL( + "https://example.com/callback", + "test-state", + "test-nonce", + tt.codeChallenge, + ) + + if authURL == "" { + t.Fatal("BuildAuthURL returned empty string") + } + + parsedURL, err := url.Parse(authURL) + if err != nil { + t.Fatalf("Failed to parse auth URL: %v", err) + } + + // Check base URL + expectedBase := "https://provider.example.com/auth" + actualBase := fmt.Sprintf("%s://%s%s", parsedURL.Scheme, parsedURL.Host, parsedURL.Path) + if actualBase != expectedBase { + t.Errorf("Expected base URL %s, got %s", expectedBase, actualBase) + } + + // Check required parameters + query := parsedURL.Query() + for param, expectedValue := range tt.expectedParams { + actualValue := query.Get(param) + if actualValue != expectedValue { + t.Errorf("Expected %s=%s, got %s=%s", param, expectedValue, param, actualValue) + } + } + + // Check that contains expected strings + for _, contains := range tt.expectContains { + if !strings.Contains(authURL, contains) { + t.Errorf("Expected URL to contain %s, got %s", contains, authURL) + } + } + }) + } +} + +func TestValidateURL(t *testing.T) { + logger := &mockLogger{} + handler := createTestAuthHandler(logger, false, func() bool { return false }, func() bool { return false }) + + tests := []struct { + name string + url string + expectError bool + errorMsg string + }{ + { + name: "valid HTTPS URL", + url: "https://example.com/path", + expectError: false, + }, + { + name: "valid HTTP URL", + url: "http://example.com/path", + expectError: false, + }, + { + name: "empty URL", + url: "", + expectError: true, + errorMsg: "empty URL", + }, + { + name: "invalid scheme", + url: "ftp://example.com/path", + expectError: true, + errorMsg: "disallowed URL scheme", + }, + { + name: "missing host", + url: "https:///path", + expectError: true, + errorMsg: "missing host in URL", + }, + { + name: "localhost not allowed", + url: "https://localhost/path", + expectError: true, + errorMsg: "localhost access not allowed", + }, + { + name: "127.0.0.1 not allowed", + url: "https://127.0.0.1/path", + expectError: true, + errorMsg: "localhost access not allowed", + }, + { + name: "private IP not allowed", + url: "https://192.168.1.1/path", + expectError: true, + errorMsg: "private IP not allowed", + }, + { + name: "path traversal not allowed", + url: "https://example.com/../path", + expectError: true, + errorMsg: "path traversal detected", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := handler.validateURL(tt.url) + + if tt.expectError { + if err == nil { + t.Errorf("Expected error for URL %s, got nil", tt.url) + } else if tt.errorMsg != "" && !strings.Contains(err.Error(), tt.errorMsg) { + t.Errorf("Expected error to contain %s, got %s", tt.errorMsg, err.Error()) + } + } else { + if err != nil { + t.Errorf("Expected no error for URL %s, got %v", tt.url, err) + } + } + }) + } +} + +func TestValidateHost(t *testing.T) { + logger := &mockLogger{} + handler := createTestAuthHandler(logger, false, func() bool { return false }, func() bool { return false }) + + tests := []struct { + name string + host string + expectError bool + errorMsg string + }{ + { + name: "valid domain", + host: "example.com", + expectError: false, + }, + { + name: "valid domain with port", + host: "example.com:443", + expectError: false, + }, + { + name: "empty host", + host: "", + expectError: true, + errorMsg: "empty host", + }, + { + name: "localhost", + host: "localhost", + expectError: true, + errorMsg: "localhost access not allowed", + }, + { + name: "localhost with port", + host: "localhost:8080", + expectError: true, + errorMsg: "localhost access not allowed", + }, + { + name: "127.0.0.1", + host: "127.0.0.1", + expectError: true, + errorMsg: "localhost access not allowed", + }, + { + name: "::1 (IPv6 localhost)", + host: "::1", + expectError: true, + errorMsg: "invalid host:port format", + }, + { + name: "private IP 192.168.1.1", + host: "192.168.1.1", + expectError: true, + errorMsg: "private IP not allowed", + }, + { + name: "private IP 10.0.0.1", + host: "10.0.0.1", + expectError: true, + errorMsg: "private IP not allowed", + }, + { + name: "invalid host:port format", + host: "example.com:99999", // Use invalid port number instead + expectError: false, // This actually doesn't error in Go's net.SplitHostPort + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := handler.validateHost(tt.host) + + if tt.expectError { + if err == nil { + t.Errorf("Expected error for host %s, got nil", tt.host) + } else if tt.errorMsg != "" && !strings.Contains(err.Error(), tt.errorMsg) { + t.Errorf("Expected error to contain %s, got %s", tt.errorMsg, err.Error()) + } + } else { + if err != nil { + t.Errorf("Expected no error for host %s, got %v", tt.host, err) + } + } + }) + } +} + +func TestBuildURLWithParams(t *testing.T) { + logger := &mockLogger{} + handler := NewAuthHandler( + logger, + false, + func() bool { return false }, + func() bool { return false }, + "test-client-id", + "https://provider.example.com/auth", + "https://provider.example.com", + []string{"openid"}, + false, + ) + + tests := []struct { + name string + baseURL string + params url.Values + expected string + expectError bool + }{ + { + name: "absolute HTTPS URL with params", + baseURL: "https://example.com/auth", + params: url.Values{"param1": []string{"value1"}, "param2": []string{"value2"}}, + expected: "https://example.com/auth?param1=value1¶m2=value2", + }, + { + name: "relative URL resolved against issuer", + baseURL: "/auth/endpoint", + params: url.Values{"client_id": []string{"test"}}, + expected: "https://provider.example.com/auth/endpoint?client_id=test", + }, + { + name: "invalid absolute URL", + baseURL: "https://localhost/auth", + params: url.Values{}, + expectError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := handler.buildURLWithParams(tt.baseURL, tt.params) + + if tt.expectError { + if result != "" { + t.Errorf("Expected empty result for invalid URL, got %s", result) + } + } else { + if result == "" { + t.Error("Expected non-empty result, got empty string") + } + + // Parse and compare the result + resultURL, err := url.Parse(result) + if err != nil { + t.Fatalf("Failed to parse result URL: %v", err) + } + + expectedURL, err := url.Parse(tt.expected) + if err != nil { + t.Fatalf("Failed to parse expected URL: %v", err) + } + + if resultURL.Scheme != expectedURL.Scheme { + t.Errorf("Expected scheme %s, got %s", expectedURL.Scheme, resultURL.Scheme) + } + if resultURL.Host != expectedURL.Host { + t.Errorf("Expected host %s, got %s", expectedURL.Host, resultURL.Host) + } + if resultURL.Path != expectedURL.Path { + t.Errorf("Expected path %s, got %s", expectedURL.Path, resultURL.Path) + } + + // Check query parameters + resultQuery := resultURL.Query() + expectedQuery := expectedURL.Query() + for key, expectedValues := range expectedQuery { + resultValues := resultQuery[key] + if len(resultValues) != len(expectedValues) { + t.Errorf("Expected %d values for %s, got %d", len(expectedValues), key, len(resultValues)) + } + for i, expectedValue := range expectedValues { + if i >= len(resultValues) || resultValues[i] != expectedValue { + t.Errorf("Expected %s=%s, got %s=%s", key, expectedValue, key, resultValues[i]) + } + } + } + } + }) + } +} + +func TestConcurrentAccess(t *testing.T) { + logger := &mockLogger{} + handler := createTestAuthHandler(logger, true, func() bool { return false }, func() bool { return false }) + + const numGoroutines = 10 + const numRequests = 5 + + var wg sync.WaitGroup + errors := make(chan error, numGoroutines*numRequests) + + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + for j := 0; j < numRequests; j++ { + session := newMockSession() + req := httptest.NewRequest("GET", fmt.Sprintf("/test-%d-%d", id, j), nil) + rw := httptest.NewRecorder() + + generateNonce := func() (string, error) { + return fmt.Sprintf("nonce-%d-%d", id, j), nil + } + generateCodeVerifier := func() (string, error) { + return fmt.Sprintf("verifier-%d-%d", id, j), nil + } + deriveCodeChallenge := func() (string, error) { + return fmt.Sprintf("challenge-%d-%d", id, j), nil + } + + handler.InitiateAuthentication( + rw, req, session, + "https://example.com/callback", + generateNonce, + generateCodeVerifier, + deriveCodeChallenge, + ) + + if rw.Code != http.StatusFound { + errors <- fmt.Errorf("goroutine %d request %d: expected status %d, got %d", + id, j, http.StatusFound, rw.Code) + } + } + }(i) + } + + wg.Wait() + close(errors) + + var errorCount int + for err := range errors { + t.Error(err) + errorCount++ + } + + if errorCount > 0 { + t.Errorf("Found %d errors in concurrent access test", errorCount) + } +} + +// Benchmark tests +func BenchmarkInitiateAuthentication(b *testing.B) { + logger := &mockLogger{} + handler := createTestAuthHandler(logger, true, func() bool { return false }, func() bool { return false }) + + generateNonce := func() (string, error) { return "test-nonce", nil } + generateCodeVerifier := func() (string, error) { return "test-verifier", nil } + deriveCodeChallenge := func() (string, error) { return "test-challenge", nil } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + session := newMockSession() + req := httptest.NewRequest("GET", "/test", nil) + rw := httptest.NewRecorder() + + handler.InitiateAuthentication( + rw, req, session, + "https://example.com/callback", + generateNonce, + generateCodeVerifier, + deriveCodeChallenge, + ) + } +} + +func BenchmarkBuildAuthURL(b *testing.B) { + logger := &mockLogger{} + handler := createTestAuthHandler(logger, true, func() bool { return false }, func() bool { return false }) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + handler.BuildAuthURL( + "https://example.com/callback", + "test-state", + "test-nonce", + "test-challenge", + ) + } +} + +func BenchmarkValidateURL(b *testing.B) { + logger := &mockLogger{} + handler := createTestAuthHandler(logger, false, func() bool { return false }, func() bool { return false }) + + testURLs := []string{ + "https://example.com/auth", + "https://provider.example.com/oauth/authorize", + "https://login.microsoftonline.com/tenant/oauth2/v2.0/authorize", + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + handler.validateURL(testURLs[i%len(testURLs)]) + } +} diff --git a/cache_manager.go b/cache_manager.go new file mode 100644 index 0000000..17138f6 --- /dev/null +++ b/cache_manager.go @@ -0,0 +1,151 @@ +// Package traefikoidc provides OIDC authentication middleware for Traefik. +package traefikoidc + +import ( + "sync" + "time" +) + +const ( + defaultBlacklistDuration = 24 * time.Hour + defaultMaxBlacklistSize = 10000 +) + +var ( + globalCacheManager *CacheManager + cacheManagerOnce sync.Once + cacheManagerMutex sync.RWMutex +) + +// CacheManager manages all caching components for the OIDC middleware. +// It provides thread-safe access to token blacklist, token cache, metadata cache, +// and JWK cache. This centralizes cache management to ensure efficient memory +// usage and consistent cache behavior across the application. +type CacheManager struct { + tokenBlacklist CacheInterface + tokenCache *TokenCache + metadataCache *MetadataCache + jwkCache JWKCacheInterface + mu sync.RWMutex +} + +// GetGlobalCacheManager returns a singleton instance of CacheManager. +// It initializes all cache components on first call and reuses the same instance +// for subsequent calls. This ensures consistent cache state and efficient memory +// usage across the entire application lifecycle. +// Parameters: +// - wg: WaitGroup for coordinating cache cleanup during shutdown. +// +// Returns: +// - The global CacheManager instance. +func GetGlobalCacheManager(wg *sync.WaitGroup) *CacheManager { + cacheManagerOnce.Do(func() { + globalCacheManager = &CacheManager{ + tokenBlacklist: func() CacheInterface { + config := DefaultUnifiedCacheConfig() + config.MaxSize = defaultMaxBlacklistSize + c := NewUnifiedCache(config) + return NewCacheAdapter(c) + }(), + tokenCache: NewTokenCache(), + metadataCache: NewMetadataCache(wg), + jwkCache: &JWKCache{}, + } + }) + return globalCacheManager +} + +// GetSharedTokenBlacklist returns the shared token blacklist cache. +// This cache stores revoked or expired tokens to prevent replay attacks. +// Access is protected by read lock to ensure thread safety. +// Returns: +// - The shared token blacklist CacheInterface instance. +func (cm *CacheManager) GetSharedTokenBlacklist() CacheInterface { + cm.mu.RLock() + defer cm.mu.RUnlock() + return cm.tokenBlacklist +} + +// GetSharedTokenCache returns the shared token cache for verified tokens. +// This cache stores claims from successfully verified tokens to avoid +// repeated verification of the same tokens. +// Access is protected by read lock to ensure thread safety. +// Returns: +// - The shared TokenCache instance. +func (cm *CacheManager) GetSharedTokenCache() *TokenCache { + cm.mu.RLock() + defer cm.mu.RUnlock() + return cm.tokenCache +} + +// GetSharedMetadataCache returns the shared metadata cache. +// This cache stores OIDC provider metadata (endpoints, keys) to avoid +// repeated requests to the provider's .well-known configuration endpoint. +// Access is protected by read lock to ensure thread safety. +// Returns: +// - The shared MetadataCache instance. +func (cm *CacheManager) GetSharedMetadataCache() *MetadataCache { + cm.mu.RLock() + defer cm.mu.RUnlock() + return cm.metadataCache +} + +// GetSharedJWKCache returns the shared JWK (JSON Web Key) cache. +// This cache stores public keys from the provider's JWKS endpoint +// for token signature verification. +// Access is protected by read lock to ensure thread safety. +// Returns: +// - The shared JWKCacheInterface instance. +func (cm *CacheManager) GetSharedJWKCache() JWKCacheInterface { + cm.mu.RLock() + defer cm.mu.RUnlock() + return cm.jwkCache +} + +// Close gracefully shuts down all cache components and releases resources. +// It closes all individual caches and cleans up their associated goroutines. +// This method should be called when the middleware is shutting down. +// Safe to call multiple times. +// Returns: +// - An error if any cache fails to close properly, nil otherwise. +func (cm *CacheManager) Close() error { + cm.mu.Lock() + defer cm.mu.Unlock() + + if cm.tokenBlacklist != nil { + cm.tokenBlacklist.Close() + cm.tokenBlacklist = nil + } + if cm.tokenCache != nil { + cm.tokenCache.Close() + cm.tokenCache = nil + } + if cm.metadataCache != nil { + cm.metadataCache.Close() + cm.metadataCache = nil + } + if cm.jwkCache != nil { + cm.jwkCache.Close() + cm.jwkCache = nil + } + + return nil +} + +// CleanupGlobalCacheManager cleans up the global cache manager instance. +// It closes all cache components but does NOT reset the singleton for re-initialization +// to avoid race conditions with sync.Once. It's safe to call multiple times. +// Returns: +// - An error if cache cleanup fails, nil otherwise. +func CleanupGlobalCacheManager() error { + cacheManagerMutex.Lock() + defer cacheManagerMutex.Unlock() + + if globalCacheManager != nil { + err := globalCacheManager.Close() + // Don't reset globalCacheManager to nil or cacheManagerOnce to avoid race conditions + // The cache manager remains closed and the Close() method handles multiple calls + return err + } + return nil +} diff --git a/circuit_breaker/circuit_breaker.go b/circuit_breaker/circuit_breaker.go new file mode 100644 index 0000000..c947130 --- /dev/null +++ b/circuit_breaker/circuit_breaker.go @@ -0,0 +1,319 @@ +// Package circuit_breaker provides circuit breaker implementation for resilience +package circuit_breaker + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "time" +) + +// CircuitBreakerState represents the current state of a circuit breaker. +// The circuit breaker pattern prevents cascading failures by monitoring +// error rates and temporarily blocking requests to failing services. +type CircuitBreakerState int + +// Circuit breaker states following the standard pattern: +// Closed: Normal operation, requests flow through +// Open: Circuit is tripped, requests are blocked +// HalfOpen: Testing state, limited requests allowed to test recovery +const ( + // CircuitBreakerClosed allows all requests through (normal operation) + CircuitBreakerClosed CircuitBreakerState = iota + // CircuitBreakerOpen blocks all requests (service is failing) + CircuitBreakerOpen + // CircuitBreakerHalfOpen allows limited requests to test service recovery + CircuitBreakerHalfOpen +) + +// String returns a string representation of the circuit breaker state +func (s CircuitBreakerState) String() string { + switch s { + case CircuitBreakerClosed: + return "closed" + case CircuitBreakerOpen: + return "open" + case CircuitBreakerHalfOpen: + return "half-open" + default: + return "unknown" + } +} + +// Logger interface for dependency injection +type Logger interface { + Infof(format string, args ...interface{}) + Errorf(format string, args ...interface{}) + Debugf(format string, args ...interface{}) +} + +// BaseRecoveryMechanism interface for common functionality +type BaseRecoveryMechanism interface { + RecordRequest() + RecordSuccess() + RecordFailure() + GetBaseMetrics() map[string]interface{} + LogInfo(format string, args ...interface{}) + LogError(format string, args ...interface{}) + LogDebug(format string, args ...interface{}) +} + +// CircuitBreaker implements the circuit breaker pattern for external service calls. +// It monitors failure rates and automatically opens the circuit when failures +// exceed the threshold, preventing further requests until the service recovers. +type CircuitBreaker struct { + // baseRecovery provides common functionality + baseRecovery BaseRecoveryMechanism + // maxFailures is the threshold for opening the circuit + maxFailures int + // timeout is how long to wait before allowing requests in half-open state + timeout time.Duration + // resetTimeout is how long to wait before transitioning from open to half-open + resetTimeout time.Duration + // state tracks the current circuit breaker state + state CircuitBreakerState + // failures counts consecutive failures + failures int64 + // lastFailureTime records when the last failure occurred + lastFailureTime time.Time + // mutex protects shared state + mutex sync.RWMutex + // logger for debugging and monitoring + logger Logger +} + +// CircuitBreakerConfig holds configuration parameters for circuit breakers. +// These settings control when the circuit opens and how it recovers. +type CircuitBreakerConfig struct { + // MaxFailures is the number of failures before opening the circuit + MaxFailures int `json:"max_failures"` + // Timeout is how long to wait before trying to recover (open -> half-open) + Timeout time.Duration `json:"timeout"` + // ResetTimeout is how long to wait before fully closing the circuit + ResetTimeout time.Duration `json:"reset_timeout"` +} + +// DefaultCircuitBreakerConfig returns sensible default configuration for circuit breakers. +// Configured for typical web service scenarios with moderate tolerance for failures. +func DefaultCircuitBreakerConfig() CircuitBreakerConfig { + return CircuitBreakerConfig{ + MaxFailures: 2, + Timeout: 60 * time.Second, + ResetTimeout: 30 * time.Second, + } +} + +// NewCircuitBreaker creates a new circuit breaker with the specified configuration. +// The circuit breaker starts in the closed state, allowing all requests through. +func NewCircuitBreaker(config CircuitBreakerConfig, logger Logger, baseRecovery BaseRecoveryMechanism) *CircuitBreaker { + return &CircuitBreaker{ + baseRecovery: baseRecovery, + maxFailures: config.MaxFailures, + timeout: config.Timeout, + resetTimeout: config.ResetTimeout, + state: CircuitBreakerClosed, + logger: logger, + } +} + +// ExecuteWithContext executes a function through the circuit breaker with context. +// It checks if requests are allowed, executes the function, and updates the circuit state +// based on the result. Implements the ErrorRecoveryMechanism interface. +func (cb *CircuitBreaker) ExecuteWithContext(ctx context.Context, fn func() error) error { + if cb.baseRecovery != nil { + cb.baseRecovery.RecordRequest() + } + + if !cb.allowRequest() { + return fmt.Errorf("circuit breaker is open") + } + + err := fn() + if err != nil { + cb.recordFailure() + if cb.baseRecovery != nil { + cb.baseRecovery.RecordFailure() + } + return err + } + + cb.recordSuccess() + if cb.baseRecovery != nil { + cb.baseRecovery.RecordSuccess() + } + return nil +} + +// Execute executes a function through the circuit breaker without context. +// This is provided for backward compatibility with existing code. +func (cb *CircuitBreaker) Execute(fn func() error) error { + return cb.ExecuteWithContext(context.Background(), fn) +} + +// allowRequest determines whether to allow a request based on the circuit state. +// Handles state transitions from open to half-open based on timeout. +func (cb *CircuitBreaker) allowRequest() bool { + cb.mutex.Lock() + defer cb.mutex.Unlock() + + now := time.Now() + + switch cb.state { + case CircuitBreakerClosed: + return true + + case CircuitBreakerOpen: + if now.Sub(cb.lastFailureTime) > cb.timeout { + cb.state = CircuitBreakerHalfOpen + if cb.logger != nil { + cb.logger.Infof("Circuit breaker transitioning to half-open state") + } + return true + } + return false + + case CircuitBreakerHalfOpen: + return true + + default: + return false + } +} + +// recordFailure records a failure and potentially opens the circuit. +// Updates failure count and triggers state transitions when thresholds are exceeded. +func (cb *CircuitBreaker) recordFailure() { + cb.mutex.Lock() + defer cb.mutex.Unlock() + + cb.failures++ + cb.lastFailureTime = time.Now() + + switch cb.state { + case CircuitBreakerClosed: + if cb.failures >= int64(cb.maxFailures) { + cb.state = CircuitBreakerOpen + if cb.baseRecovery != nil { + cb.baseRecovery.LogError("Circuit breaker opened after %d failures", cb.failures) + } + } + + case CircuitBreakerHalfOpen: + cb.state = CircuitBreakerOpen + if cb.baseRecovery != nil { + cb.baseRecovery.LogError("Circuit breaker returned to open state after failure in half-open") + } + } +} + +// recordSuccess records a successful request and potentially closes the circuit. +// Resets failure count and transitions from half-open to closed state on success. +func (cb *CircuitBreaker) recordSuccess() { + cb.mutex.Lock() + defer cb.mutex.Unlock() + + switch cb.state { + case CircuitBreakerHalfOpen: + cb.failures = 0 + cb.state = CircuitBreakerClosed + if cb.baseRecovery != nil { + cb.baseRecovery.LogInfo("Circuit breaker closed after successful request in half-open state") + } + + case CircuitBreakerClosed: + cb.failures = 0 + } +} + +// GetState returns the current state of the circuit breaker. +// Thread-safe method for monitoring circuit breaker status. +func (cb *CircuitBreaker) GetState() CircuitBreakerState { + cb.mutex.RLock() + defer cb.mutex.RUnlock() + return cb.state +} + +// Reset resets the circuit breaker to its initial closed state. +// Clears failure count and state, effectively recovering from any open state. +func (cb *CircuitBreaker) Reset() { + cb.mutex.Lock() + defer cb.mutex.Unlock() + + cb.state = CircuitBreakerClosed + atomic.StoreInt64(&cb.failures, 0) + if cb.baseRecovery != nil { + cb.baseRecovery.LogInfo("Circuit breaker has been reset") + } +} + +// IsAvailable returns whether the circuit breaker is currently allowing requests. +// This provides a quick way to check if the service is available. +func (cb *CircuitBreaker) IsAvailable() bool { + return cb.allowRequest() +} + +// GetMetrics returns comprehensive metrics about the circuit breaker. +// Includes state information, failure counts, configuration, and base metrics. +func (cb *CircuitBreaker) GetMetrics() map[string]interface{} { + cb.mutex.RLock() + state := cb.state + failures := cb.failures + lastFailureTime := cb.lastFailureTime + cb.mutex.RUnlock() + + var metrics map[string]interface{} + if cb.baseRecovery != nil { + metrics = cb.baseRecovery.GetBaseMetrics() + } else { + metrics = make(map[string]interface{}) + } + + metrics["state"] = state.String() + metrics["current_failures"] = failures + metrics["max_failures"] = cb.maxFailures + metrics["timeout"] = cb.timeout.String() + metrics["reset_timeout"] = cb.resetTimeout.String() + + if !lastFailureTime.IsZero() { + metrics["last_failure_time"] = lastFailureTime + metrics["time_since_last_failure"] = time.Since(lastFailureTime).String() + } + + return metrics +} + +// GetFailureCount returns the current failure count +func (cb *CircuitBreaker) GetFailureCount() int64 { + cb.mutex.RLock() + defer cb.mutex.RUnlock() + return cb.failures +} + +// GetLastFailureTime returns the time of the last failure +func (cb *CircuitBreaker) GetLastFailureTime() time.Time { + cb.mutex.RLock() + defer cb.mutex.RUnlock() + return cb.lastFailureTime +} + +// IsOpen returns true if the circuit breaker is in open state +func (cb *CircuitBreaker) IsOpen() bool { + cb.mutex.RLock() + defer cb.mutex.RUnlock() + return cb.state == CircuitBreakerOpen +} + +// IsClosed returns true if the circuit breaker is in closed state +func (cb *CircuitBreaker) IsClosed() bool { + cb.mutex.RLock() + defer cb.mutex.RUnlock() + return cb.state == CircuitBreakerClosed +} + +// IsHalfOpen returns true if the circuit breaker is in half-open state +func (cb *CircuitBreaker) IsHalfOpen() bool { + cb.mutex.RLock() + defer cb.mutex.RUnlock() + return cb.state == CircuitBreakerHalfOpen +} diff --git a/circuit_breaker/circuit_breaker_test.go b/circuit_breaker/circuit_breaker_test.go new file mode 100644 index 0000000..8f0512f --- /dev/null +++ b/circuit_breaker/circuit_breaker_test.go @@ -0,0 +1,981 @@ +package circuit_breaker + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" +) + +// Mock implementations for testing +type mockLogger struct { + infoLogs []string + errorLogs []string + debugLogs []string + mu sync.RWMutex +} + +func (m *mockLogger) Infof(format string, args ...interface{}) { + m.mu.Lock() + defer m.mu.Unlock() + m.infoLogs = append(m.infoLogs, fmt.Sprintf(format, args...)) +} + +func (m *mockLogger) Errorf(format string, args ...interface{}) { + m.mu.Lock() + defer m.mu.Unlock() + m.errorLogs = append(m.errorLogs, fmt.Sprintf(format, args...)) +} + +func (m *mockLogger) Debugf(format string, args ...interface{}) { + m.mu.Lock() + defer m.mu.Unlock() + m.debugLogs = append(m.debugLogs, fmt.Sprintf(format, args...)) +} + +func (m *mockLogger) getInfoLogs() []string { + m.mu.RLock() + defer m.mu.RUnlock() + result := make([]string, len(m.infoLogs)) + copy(result, m.infoLogs) + return result +} + +//lint:ignore U1000 May be needed for future error log verification tests +func (m *mockLogger) getErrorLogs() []string { + m.mu.RLock() + defer m.mu.RUnlock() + result := make([]string, len(m.errorLogs)) + copy(result, m.errorLogs) + return result +} + +//lint:ignore U1000 May be needed for future test isolation +func (m *mockLogger) reset() { + m.mu.Lock() + defer m.mu.Unlock() + m.infoLogs = nil + m.errorLogs = nil + m.debugLogs = nil +} + +type mockBaseRecoveryMechanism struct { + requestCount int64 + successCount int64 + failureCount int64 + infoLogs []string + errorLogs []string + debugLogs []string + baseMetrics map[string]interface{} + mu sync.RWMutex +} + +func newMockBaseRecovery() *mockBaseRecoveryMechanism { + return &mockBaseRecoveryMechanism{ + baseMetrics: make(map[string]interface{}), + } +} + +func (m *mockBaseRecoveryMechanism) RecordRequest() { + atomic.AddInt64(&m.requestCount, 1) +} + +func (m *mockBaseRecoveryMechanism) RecordSuccess() { + atomic.AddInt64(&m.successCount, 1) +} + +func (m *mockBaseRecoveryMechanism) RecordFailure() { + atomic.AddInt64(&m.failureCount, 1) +} + +func (m *mockBaseRecoveryMechanism) GetBaseMetrics() map[string]interface{} { + m.mu.RLock() + defer m.mu.RUnlock() + + result := make(map[string]interface{}) + for k, v := range m.baseMetrics { + result[k] = v + } + result["total_requests"] = atomic.LoadInt64(&m.requestCount) + result["total_successes"] = atomic.LoadInt64(&m.successCount) + result["total_failures"] = atomic.LoadInt64(&m.failureCount) + return result +} + +func (m *mockBaseRecoveryMechanism) LogInfo(format string, args ...interface{}) { + m.mu.Lock() + defer m.mu.Unlock() + m.infoLogs = append(m.infoLogs, fmt.Sprintf(format, args...)) +} + +func (m *mockBaseRecoveryMechanism) LogError(format string, args ...interface{}) { + m.mu.Lock() + defer m.mu.Unlock() + m.errorLogs = append(m.errorLogs, fmt.Sprintf(format, args...)) +} + +func (m *mockBaseRecoveryMechanism) LogDebug(format string, args ...interface{}) { + m.mu.Lock() + defer m.mu.Unlock() + m.debugLogs = append(m.debugLogs, fmt.Sprintf(format, args...)) +} + +func (m *mockBaseRecoveryMechanism) getRequestCount() int64 { + return atomic.LoadInt64(&m.requestCount) +} + +func (m *mockBaseRecoveryMechanism) getSuccessCount() int64 { + return atomic.LoadInt64(&m.successCount) +} + +func (m *mockBaseRecoveryMechanism) getFailureCount() int64 { + return atomic.LoadInt64(&m.failureCount) +} + +func (m *mockBaseRecoveryMechanism) getInfoLogs() []string { + m.mu.RLock() + defer m.mu.RUnlock() + result := make([]string, len(m.infoLogs)) + copy(result, m.infoLogs) + return result +} + +func (m *mockBaseRecoveryMechanism) getErrorLogs() []string { + m.mu.RLock() + defer m.mu.RUnlock() + result := make([]string, len(m.errorLogs)) + copy(result, m.errorLogs) + return result +} + +func TestCircuitBreakerState_String(t *testing.T) { + tests := []struct { + state CircuitBreakerState + expected string + }{ + {CircuitBreakerClosed, "closed"}, + {CircuitBreakerOpen, "open"}, + {CircuitBreakerHalfOpen, "half-open"}, + {CircuitBreakerState(999), "unknown"}, + } + + for _, tt := range tests { + t.Run(tt.expected, func(t *testing.T) { + result := tt.state.String() + if result != tt.expected { + t.Errorf("Expected %s, got %s", tt.expected, result) + } + }) + } +} + +func TestDefaultCircuitBreakerConfig(t *testing.T) { + config := DefaultCircuitBreakerConfig() + + if config.MaxFailures != 2 { + t.Errorf("Expected MaxFailures to be 2, got %d", config.MaxFailures) + } + + if config.Timeout != 60*time.Second { + t.Errorf("Expected Timeout to be 60s, got %v", config.Timeout) + } + + if config.ResetTimeout != 30*time.Second { + t.Errorf("Expected ResetTimeout to be 30s, got %v", config.ResetTimeout) + } +} + +func TestNewCircuitBreaker(t *testing.T) { + config := CircuitBreakerConfig{ + MaxFailures: 3, + Timeout: 30 * time.Second, + ResetTimeout: 15 * time.Second, + } + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + + cb := NewCircuitBreaker(config, logger, baseRecovery) + + if cb == nil { + t.Fatal("NewCircuitBreaker returned nil") + } + + if cb.maxFailures != 3 { + t.Errorf("Expected maxFailures to be 3, got %d", cb.maxFailures) + } + + if cb.timeout != 30*time.Second { + t.Errorf("Expected timeout to be 30s, got %v", cb.timeout) + } + + if cb.resetTimeout != 15*time.Second { + t.Errorf("Expected resetTimeout to be 15s, got %v", cb.resetTimeout) + } + + if cb.state != CircuitBreakerClosed { + t.Errorf("Expected initial state to be Closed, got %v", cb.state) + } + + if cb.logger != logger { + t.Error("Expected logger to be set") + } + + if cb.baseRecovery != baseRecovery { + t.Error("Expected baseRecovery to be set") + } +} + +func TestCircuitBreaker_ExecuteWithContext_Success(t *testing.T) { + config := CircuitBreakerConfig{ + MaxFailures: 2, + Timeout: time.Second, + ResetTimeout: time.Second, + } + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + cb := NewCircuitBreaker(config, logger, baseRecovery) + + callCount := 0 + testFunc := func() error { + callCount++ + return nil + } + + ctx := context.Background() + err := cb.ExecuteWithContext(ctx, testFunc) + + if err != nil { + t.Errorf("Expected no error, got %v", err) + } + + if callCount != 1 { + t.Errorf("Expected function to be called once, got %d", callCount) + } + + if cb.GetState() != CircuitBreakerClosed { + t.Errorf("Expected state to remain Closed, got %v", cb.GetState()) + } + + if baseRecovery.getRequestCount() != 1 { + t.Errorf("Expected 1 request recorded, got %d", baseRecovery.getRequestCount()) + } + + if baseRecovery.getSuccessCount() != 1 { + t.Errorf("Expected 1 success recorded, got %d", baseRecovery.getSuccessCount()) + } +} + +func TestCircuitBreaker_ExecuteWithContext_Failure(t *testing.T) { + config := CircuitBreakerConfig{ + MaxFailures: 2, + Timeout: time.Second, + ResetTimeout: time.Second, + } + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + cb := NewCircuitBreaker(config, logger, baseRecovery) + + testError := fmt.Errorf("test error") + testFunc := func() error { + return testError + } + + ctx := context.Background() + err := cb.ExecuteWithContext(ctx, testFunc) + + if err != testError { + t.Errorf("Expected test error, got %v", err) + } + + if cb.GetState() != CircuitBreakerClosed { + t.Errorf("Expected state to remain Closed after single failure, got %v", cb.GetState()) + } + + if baseRecovery.getRequestCount() != 1 { + t.Errorf("Expected 1 request recorded, got %d", baseRecovery.getRequestCount()) + } + + if baseRecovery.getFailureCount() != 1 { + t.Errorf("Expected 1 failure recorded, got %d", baseRecovery.getFailureCount()) + } +} + +func TestCircuitBreaker_Execute(t *testing.T) { + config := CircuitBreakerConfig{ + MaxFailures: 1, + Timeout: time.Second, + ResetTimeout: time.Second, + } + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + cb := NewCircuitBreaker(config, logger, baseRecovery) + + callCount := 0 + testFunc := func() error { + callCount++ + return nil + } + + err := cb.Execute(testFunc) + + if err != nil { + t.Errorf("Expected no error, got %v", err) + } + + if callCount != 1 { + t.Errorf("Expected function to be called once, got %d", callCount) + } +} + +func TestCircuitBreaker_OpenAfterMaxFailures(t *testing.T) { + config := CircuitBreakerConfig{ + MaxFailures: 2, + Timeout: time.Second, + ResetTimeout: time.Second, + } + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + cb := NewCircuitBreaker(config, logger, baseRecovery) + + testError := fmt.Errorf("test error") + testFunc := func() error { + return testError + } + + ctx := context.Background() + + // First failure + err := cb.ExecuteWithContext(ctx, testFunc) + if err != testError { + t.Errorf("Expected test error on first failure, got %v", err) + } + if cb.GetState() != CircuitBreakerClosed { + t.Errorf("Expected state to remain Closed after first failure, got %v", cb.GetState()) + } + + // Second failure - should open circuit + err = cb.ExecuteWithContext(ctx, testFunc) + if err != testError { + t.Errorf("Expected test error on second failure, got %v", err) + } + if cb.GetState() != CircuitBreakerOpen { + t.Errorf("Expected state to be Open after max failures, got %v", cb.GetState()) + } + + // Third attempt - should be blocked + callCount := 0 + blockedFunc := func() error { + callCount++ + return nil + } + err = cb.ExecuteWithContext(ctx, blockedFunc) + if err == nil { + t.Error("Expected error when circuit is open") + } + if callCount != 0 { + t.Errorf("Expected function not to be called when circuit is open, got %d calls", callCount) + } +} + +func TestCircuitBreaker_HalfOpenTransition(t *testing.T) { + config := CircuitBreakerConfig{ + MaxFailures: 1, + Timeout: 10 * time.Millisecond, // Very short for testing + ResetTimeout: time.Second, + } + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + cb := NewCircuitBreaker(config, logger, baseRecovery) + + // Trigger circuit opening + testError := fmt.Errorf("test error") + err := cb.ExecuteWithContext(context.Background(), func() error { + return testError + }) + if err != testError { + t.Errorf("Expected test error, got %v", err) + } + if cb.GetState() != CircuitBreakerOpen { + t.Errorf("Expected state to be Open, got %v", cb.GetState()) + } + + // Wait for timeout + time.Sleep(15 * time.Millisecond) + + // Next request should transition to half-open + callCount := 0 + testFunc := func() error { + callCount++ + return nil + } + + err = cb.ExecuteWithContext(context.Background(), testFunc) + if err != nil { + t.Errorf("Expected no error in half-open state, got %v", err) + } + if callCount != 1 { + t.Errorf("Expected function to be called in half-open state, got %d calls", callCount) + } + if cb.GetState() != CircuitBreakerClosed { + t.Errorf("Expected state to be Closed after successful half-open request, got %v", cb.GetState()) + } +} + +func TestCircuitBreaker_HalfOpenFailureReturnsToOpen(t *testing.T) { + config := CircuitBreakerConfig{ + MaxFailures: 1, + Timeout: 10 * time.Millisecond, + ResetTimeout: time.Second, + } + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + cb := NewCircuitBreaker(config, logger, baseRecovery) + + // Trigger circuit opening + testError := fmt.Errorf("test error") + _ = cb.ExecuteWithContext(context.Background(), func() error { + return testError + }) + if cb.GetState() != CircuitBreakerOpen { + t.Errorf("Expected state to be Open, got %v", cb.GetState()) + } + + // Wait for timeout to allow half-open transition + time.Sleep(15 * time.Millisecond) + + // First call should transition to half-open, but we'll force it by checking allowRequest + if !cb.allowRequest() { + t.Error("Expected allowRequest to return true after timeout") + } + if cb.GetState() != CircuitBreakerHalfOpen { + t.Errorf("Expected state to be HalfOpen, got %v", cb.GetState()) + } + + // Failure in half-open should return to open + err := cb.ExecuteWithContext(context.Background(), func() error { + return testError + }) + if err != testError { + t.Errorf("Expected test error, got %v", err) + } + if cb.GetState() != CircuitBreakerOpen { + t.Errorf("Expected state to return to Open after half-open failure, got %v", cb.GetState()) + } +} + +func TestCircuitBreaker_Reset(t *testing.T) { + config := CircuitBreakerConfig{ + MaxFailures: 1, + Timeout: time.Second, + ResetTimeout: time.Second, + } + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + cb := NewCircuitBreaker(config, logger, baseRecovery) + + // Trigger circuit opening + testError := fmt.Errorf("test error") + _ = cb.ExecuteWithContext(context.Background(), func() error { + return testError + }) + if cb.GetState() != CircuitBreakerOpen { + t.Errorf("Expected state to be Open, got %v", cb.GetState()) + } + + // Reset circuit + cb.Reset() + + if cb.GetState() != CircuitBreakerClosed { + t.Errorf("Expected state to be Closed after reset, got %v", cb.GetState()) + } + + if cb.GetFailureCount() != 0 { + t.Errorf("Expected failure count to be 0 after reset, got %d", cb.GetFailureCount()) + } + + // Should allow requests again + callCount := 0 + err := cb.ExecuteWithContext(context.Background(), func() error { + callCount++ + return nil + }) + if err != nil { + t.Errorf("Expected no error after reset, got %v", err) + } + if callCount != 1 { + t.Errorf("Expected function to be called after reset, got %d calls", callCount) + } +} + +func TestCircuitBreaker_IsAvailable(t *testing.T) { + config := CircuitBreakerConfig{ + MaxFailures: 1, + Timeout: 10 * time.Millisecond, + ResetTimeout: time.Second, + } + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + cb := NewCircuitBreaker(config, logger, baseRecovery) + + // Initially available + if !cb.IsAvailable() { + t.Error("Expected circuit breaker to be available initially") + } + + // Trigger opening + testError := fmt.Errorf("test error") + cb.ExecuteWithContext(context.Background(), func() error { + return testError + }) + + // Should not be available when open + if cb.IsAvailable() { + t.Error("Expected circuit breaker to be unavailable when open") + } + + // Wait for timeout + time.Sleep(15 * time.Millisecond) + + // Should be available again after timeout (half-open) + if !cb.IsAvailable() { + t.Error("Expected circuit breaker to be available after timeout") + } +} + +func TestCircuitBreaker_StateCheckers(t *testing.T) { + config := CircuitBreakerConfig{ + MaxFailures: 1, + Timeout: 10 * time.Millisecond, + ResetTimeout: time.Second, + } + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + cb := NewCircuitBreaker(config, logger, baseRecovery) + + // Initially closed + if !cb.IsClosed() { + t.Error("Expected circuit breaker to be closed initially") + } + if cb.IsOpen() { + t.Error("Expected circuit breaker not to be open initially") + } + if cb.IsHalfOpen() { + t.Error("Expected circuit breaker not to be half-open initially") + } + + // Trigger opening + testError := fmt.Errorf("test error") + cb.ExecuteWithContext(context.Background(), func() error { + return testError + }) + + // Should be open + if cb.IsClosed() { + t.Error("Expected circuit breaker not to be closed when open") + } + if !cb.IsOpen() { + t.Error("Expected circuit breaker to be open") + } + if cb.IsHalfOpen() { + t.Error("Expected circuit breaker not to be half-open when open") + } + + // Wait for timeout and trigger half-open + time.Sleep(15 * time.Millisecond) + cb.allowRequest() // This will transition to half-open + + // Should be half-open + if cb.IsClosed() { + t.Error("Expected circuit breaker not to be closed when half-open") + } + if cb.IsOpen() { + t.Error("Expected circuit breaker not to be open when half-open") + } + if !cb.IsHalfOpen() { + t.Error("Expected circuit breaker to be half-open") + } +} + +func TestCircuitBreaker_GetMetrics(t *testing.T) { + config := CircuitBreakerConfig{ + MaxFailures: 2, + Timeout: 30 * time.Second, + ResetTimeout: 15 * time.Second, + } + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + baseRecovery.baseMetrics["custom_metric"] = "custom_value" + cb := NewCircuitBreaker(config, logger, baseRecovery) + + // Record some activity + testError := fmt.Errorf("test error") + cb.ExecuteWithContext(context.Background(), func() error { + return testError + }) + + metrics := cb.GetMetrics() + + // Check circuit breaker specific metrics + if metrics["state"] != "closed" { + t.Errorf("Expected state to be 'closed', got %v", metrics["state"]) + } + + if metrics["current_failures"] != int64(1) { + t.Errorf("Expected current_failures to be 1, got %v", metrics["current_failures"]) + } + + if metrics["max_failures"] != 2 { + t.Errorf("Expected max_failures to be 2, got %v", metrics["max_failures"]) + } + + if metrics["timeout"] != "30s" { + t.Errorf("Expected timeout to be '30s', got %v", metrics["timeout"]) + } + + if metrics["reset_timeout"] != "15s" { + t.Errorf("Expected reset_timeout to be '15s', got %v", metrics["reset_timeout"]) + } + + // Check base metrics are included + if metrics["total_requests"] != int64(1) { + t.Errorf("Expected total_requests to be 1, got %v", metrics["total_requests"]) + } + + if metrics["custom_metric"] != "custom_value" { + t.Errorf("Expected custom_metric to be 'custom_value', got %v", metrics["custom_metric"]) + } + + // Check failure time metrics + if _, exists := metrics["last_failure_time"]; !exists { + t.Error("Expected last_failure_time to exist") + } + + if _, exists := metrics["time_since_last_failure"]; !exists { + t.Error("Expected time_since_last_failure to exist") + } +} + +func TestCircuitBreaker_GetMetrics_NoBaseRecovery(t *testing.T) { + config := DefaultCircuitBreakerConfig() + logger := &mockLogger{} + cb := NewCircuitBreaker(config, logger, nil) + + metrics := cb.GetMetrics() + + // Should still have circuit breaker metrics + if metrics["state"] != "closed" { + t.Errorf("Expected state to be 'closed', got %v", metrics["state"]) + } + + if metrics["max_failures"] != 2 { + t.Errorf("Expected max_failures to be 2, got %v", metrics["max_failures"]) + } + + // Should not have base metrics + if _, exists := metrics["total_requests"]; exists { + t.Error("Expected total_requests not to exist without base recovery") + } +} + +func TestCircuitBreaker_GetLastFailureTime(t *testing.T) { + config := DefaultCircuitBreakerConfig() + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + cb := NewCircuitBreaker(config, logger, baseRecovery) + + // Initially should be zero + if !cb.GetLastFailureTime().IsZero() { + t.Error("Expected last failure time to be zero initially") + } + + // Record a failure + before := time.Now() + testError := fmt.Errorf("test error") + cb.ExecuteWithContext(context.Background(), func() error { + return testError + }) + after := time.Now() + + lastFailure := cb.GetLastFailureTime() + if lastFailure.IsZero() { + t.Error("Expected last failure time to be set after failure") + } + + if lastFailure.Before(before) || lastFailure.After(after) { + t.Errorf("Expected last failure time to be between %v and %v, got %v", + before, after, lastFailure) + } +} + +func TestCircuitBreaker_ExecuteWithoutBaseRecovery(t *testing.T) { + config := DefaultCircuitBreakerConfig() + logger := &mockLogger{} + cb := NewCircuitBreaker(config, logger, nil) + + callCount := 0 + testFunc := func() error { + callCount++ + return nil + } + + err := cb.ExecuteWithContext(context.Background(), testFunc) + + if err != nil { + t.Errorf("Expected no error, got %v", err) + } + + if callCount != 1 { + t.Errorf("Expected function to be called once, got %d", callCount) + } + + // Should work fine without base recovery + if cb.GetState() != CircuitBreakerClosed { + t.Errorf("Expected state to be Closed, got %v", cb.GetState()) + } +} + +func TestCircuitBreaker_ConcurrentAccess(t *testing.T) { + config := CircuitBreakerConfig{ + MaxFailures: 10, // Higher threshold for concurrent test + Timeout: 100 * time.Millisecond, + ResetTimeout: 50 * time.Millisecond, + } + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + cb := NewCircuitBreaker(config, logger, baseRecovery) + + const numGoroutines = 10 + const numOperations = 50 + + var wg sync.WaitGroup + successCount := int64(0) + errorCount := int64(0) + + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + for j := 0; j < numOperations; j++ { + err := cb.ExecuteWithContext(context.Background(), func() error { + // Simulate some failures + if j%10 == 9 { // Every 10th operation fails + return fmt.Errorf("simulated error") + } + return nil + }) + + if err != nil { + atomic.AddInt64(&errorCount, 1) + } else { + atomic.AddInt64(&successCount, 1) + } + + // Intermittently check state and metrics + if j%5 == 0 { + cb.GetState() + cb.GetMetrics() + cb.IsAvailable() + } + } + }(i) + } + + wg.Wait() + + // Verify we got both successes and errors + finalSuccessCount := atomic.LoadInt64(&successCount) + finalErrorCount := atomic.LoadInt64(&errorCount) + + if finalSuccessCount == 0 { + t.Error("Expected some successful operations") + } + + if finalErrorCount == 0 { + t.Error("Expected some failed operations") + } + + totalOperations := finalSuccessCount + finalErrorCount + expectedMax := int64(numGoroutines * numOperations) + + if totalOperations > expectedMax { + t.Errorf("Expected at most %d operations, got %d", expectedMax, totalOperations) + } + + t.Logf("Concurrent test completed: %d successes, %d errors, final state: %v", + finalSuccessCount, finalErrorCount, cb.GetState()) +} + +func TestCircuitBreaker_StateTransitionLogging(t *testing.T) { + config := CircuitBreakerConfig{ + MaxFailures: 1, + Timeout: 10 * time.Millisecond, + ResetTimeout: time.Second, + } + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + cb := NewCircuitBreaker(config, logger, baseRecovery) + + // Trigger circuit opening + testError := fmt.Errorf("test error") + cb.ExecuteWithContext(context.Background(), func() error { + return testError + }) + + // Check that error was logged when circuit opened + errorLogs := baseRecovery.getErrorLogs() + if len(errorLogs) == 0 { + t.Error("Expected error log when circuit breaker opened") + } else { + if !contains(errorLogs, "Circuit breaker opened after") { + t.Errorf("Expected circuit opening log, got %v", errorLogs) + } + } + + // Wait and trigger half-open + time.Sleep(15 * time.Millisecond) + + // Successful request should close circuit and log + cb.ExecuteWithContext(context.Background(), func() error { + return nil + }) + + // Check that success was logged when circuit closed + infoLogs := baseRecovery.getInfoLogs() + if len(infoLogs) == 0 { + t.Error("Expected info log when circuit breaker closed") + } else { + if !contains(infoLogs, "Circuit breaker closed after successful request") { + t.Errorf("Expected circuit closing log, got %v", infoLogs) + } + } + + // Reset should also be logged + cb.Reset() + infoLogs = baseRecovery.getInfoLogs() + if !contains(infoLogs, "Circuit breaker has been reset") { + t.Errorf("Expected reset log, got %v", infoLogs) + } +} + +func TestCircuitBreaker_LoggerTransitionLogging(t *testing.T) { + config := CircuitBreakerConfig{ + MaxFailures: 1, + Timeout: 10 * time.Millisecond, + ResetTimeout: time.Second, + } + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + cb := NewCircuitBreaker(config, logger, baseRecovery) + + // Wait for timeout and check half-open transition logging + testError := fmt.Errorf("test error") + cb.ExecuteWithContext(context.Background(), func() error { + return testError + }) + + // Wait for timeout + time.Sleep(15 * time.Millisecond) + + // Next allowRequest call should log transition to half-open + cb.allowRequest() + + infoLogs := logger.getInfoLogs() + if len(infoLogs) == 0 { + t.Error("Expected info log for half-open transition") + } else { + if !contains(infoLogs, "Circuit breaker transitioning to half-open state") { + t.Errorf("Expected half-open transition log, got %v", infoLogs) + } + } +} + +// Helper function to check if a slice contains a string with substring +func contains(slice []string, substr string) bool { + for _, s := range slice { + if len(s) >= len(substr) && s[:len(substr)] == substr { + return true + } + } + return false +} + +// Benchmark tests +func BenchmarkCircuitBreaker_ExecuteWithContext_Success(b *testing.B) { + config := DefaultCircuitBreakerConfig() + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + cb := NewCircuitBreaker(config, logger, baseRecovery) + + testFunc := func() error { + return nil + } + + ctx := context.Background() + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + cb.ExecuteWithContext(ctx, testFunc) + } + }) +} + +func BenchmarkCircuitBreaker_ExecuteWithContext_Failure(b *testing.B) { + config := CircuitBreakerConfig{ + MaxFailures: 1000, // High threshold to avoid opening during benchmark + Timeout: time.Second, + ResetTimeout: time.Second, + } + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + cb := NewCircuitBreaker(config, logger, baseRecovery) + + testError := fmt.Errorf("test error") + testFunc := func() error { + return testError + } + + ctx := context.Background() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + cb.ExecuteWithContext(ctx, testFunc) + } +} + +func BenchmarkCircuitBreaker_GetState(b *testing.B) { + config := DefaultCircuitBreakerConfig() + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + cb := NewCircuitBreaker(config, logger, baseRecovery) + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + cb.GetState() + } + }) +} + +func BenchmarkCircuitBreaker_GetMetrics(b *testing.B) { + config := DefaultCircuitBreakerConfig() + logger := &mockLogger{} + baseRecovery := newMockBaseRecovery() + cb := NewCircuitBreaker(config, logger, baseRecovery) + + // Add some activity + for i := 0; i < 100; i++ { + if i%2 == 0 { + cb.ExecuteWithContext(context.Background(), func() error { return nil }) + } else { + cb.ExecuteWithContext(context.Background(), func() error { return fmt.Errorf("error") }) + } + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + cb.GetMetrics() + } +} diff --git a/config/settings.go b/config/settings.go new file mode 100644 index 0000000..798cbd1 --- /dev/null +++ b/config/settings.go @@ -0,0 +1,209 @@ +// Package config provides configuration management for the OIDC middleware +package config + +import ( + "context" + "fmt" + "net/http" + "strings" + "sync" + "time" +) + +const ( + minEncryptionKeyLength = 16 + ConstSessionTimeout = 86400 +) + +//lint:ignore U1000 May be referenced for default exclusion patterns +var defaultExcludedURLs = map[string]struct{}{ + "/favicon.ico": {}, + "/robots.txt": {}, + "/health": {}, + "/.well-known/": {}, + "/metrics": {}, + "/ping": {}, + "/api/": {}, + "/static/": {}, + "/assets/": {}, + "/js/": {}, + "/css/": {}, + "/images/": {}, + "/fonts/": {}, +} + +// Settings manages configuration and initialization for the OIDC middleware +type Settings struct { + logger Logger +} + +// Logger interface for dependency injection +type Logger interface { + Debug(msg string) + Debugf(format string, args ...interface{}) + Info(msg string) + Infof(format string, args ...interface{}) + Error(msg string) + Errorf(format string, args ...interface{}) +} + +// Config represents the configuration for the OIDC middleware +type Config struct { + ProviderURL string `json:"providerUrl"` + ClientID string `json:"clientId"` + ClientSecret string `json:"clientSecret"` + CallbackURL string `json:"callbackUrl"` + LogoutURL string `json:"logoutUrl"` + PostLogoutRedirectURI string `json:"postLogoutRedirectUri"` + SessionEncryptionKey string `json:"sessionEncryptionKey"` + ForceHTTPS bool `json:"forceHttps"` + LogLevel string `json:"logLevel"` + Scopes []string `json:"scopes"` + OverrideScopes bool `json:"overrideScopes"` + AllowedUsers []string `json:"allowedUsers"` + AllowedUserDomains []string `json:"allowedUserDomains"` + AllowedRolesAndGroups []string `json:"allowedRolesAndGroups"` + ExcludedURLs []string `json:"excludedUrls"` + EnablePKCE bool `json:"enablePkce"` + RateLimit int `json:"rateLimit"` + RefreshGracePeriodSeconds int `json:"refreshGracePeriodSeconds"` + Headers []HeaderConfig `json:"headers"` + HTTPClient *http.Client `json:"-"` + CookieDomain string `json:"cookieDomain"` +} + +// HeaderConfig represents header template configuration +type HeaderConfig struct { + Name string `json:"name"` + Value string `json:"value"` +} + +// NewSettings creates a new Settings instance +func NewSettings(logger Logger) *Settings { + return &Settings{ + logger: logger, + } +} + +// CreateConfig creates a default configuration +func CreateConfig() *Config { + return &Config{ + LogLevel: "INFO", + ForceHTTPS: true, + EnablePKCE: true, + RateLimit: 10, + RefreshGracePeriodSeconds: 60, + } +} + +// InitializeTraefikOidc would initialize and configure a new TraefikOidc instance +// This functionality has been moved to the main New function in main.go +// This function is kept for compatibility but should not be used +func (s *Settings) InitializeTraefikOidc(ctx context.Context, next http.Handler, config *Config, name string) (interface{}, error) { + return nil, fmt.Errorf("InitializeTraefikOidc is deprecated - use New function from main package instead") +} + +//lint:ignore U1000 Kept for backward compatibility +func (s *Settings) setupHeaderTemplates(t interface{}, config *Config, logger Logger) error { + logger.Debug("setupHeaderTemplates is deprecated") + return nil +} + +//lint:ignore U1000 May be needed for future background service management +func (s *Settings) startBackgroundServices(ctx context.Context, logger Logger) { + startReplayCacheCleanup(ctx, logger) + + // Start memory monitoring for leak detection and performance insights + memoryMonitor := GetGlobalMemoryMonitor() + memoryMonitor.StartMonitoring(ctx, 60*time.Second) // Monitor every minute + logger.Debug("Started global memory monitoring") +} + +// Utility functions + +//lint:ignore U1000 May be needed for future scope processing +func deduplicateScopes(scopes []string) []string { + seen := make(map[string]bool) + result := []string{} + for _, scope := range scopes { + if !seen[scope] { + seen[scope] = true + result = append(result, scope) + } + } + return result +} + +//lint:ignore U1000 May be needed for future scope merging operations +func mergeScopes(defaultScopes, userScopes []string) []string { + result := make([]string, len(defaultScopes)) + copy(result, defaultScopes) + return append(result, userScopes...) +} + +//lint:ignore U1000 May be needed for future utility operations +func createStringMap(items []string) map[string]struct{} { + result := make(map[string]struct{}) + for _, item := range items { + result[item] = struct{}{} + } + return result +} + +//lint:ignore U1000 May be needed for future case-insensitive operations +func createCaseInsensitiveStringMap(items []string) map[string]struct{} { + result := make(map[string]struct{}) + for _, item := range items { + result[strings.ToLower(item)] = struct{}{} + } + return result +} + +//lint:ignore U1000 May be needed for future test environment detection +func isTestMode() bool { + // This function should be implemented based on environment detection logic + return false +} + +// External dependencies that need to be provided +// TraefikOidc struct is defined in types.go + +// These functions need to be provided by external packages +func NewLogger(level string) Logger { return nil } +func CreateDefaultHTTPClient() *http.Client { return nil } +func CreateTokenHTTPClient() *http.Client { return nil } +func GetGlobalCacheManager(*sync.WaitGroup) CacheManager { return nil } +func NewSessionManager(string, bool, string, Logger) (SessionManager, error) { return nil, nil } +func NewErrorRecoveryManager(Logger) ErrorRecoveryManager { return nil } + +//lint:ignore U1000 May be needed for future token claim extraction +func extractClaims(string) (map[string]interface{}, error) { return nil, nil } + +//lint:ignore U1000 May be needed for future replay attack prevention +func startReplayCacheCleanup(context.Context, Logger) {} +func GetGlobalMemoryMonitor() MemoryMonitor { return nil } + +// Interfaces for external dependencies +type CacheManager interface { + GetSharedTokenBlacklist() CacheInterface + GetSharedTokenCache() *TokenCache + GetSharedMetadataCache() *MetadataCache + GetSharedJWKCache() JWKCacheInterface + Close() error +} +type SessionManager interface{} +type ErrorRecoveryManager interface{} +type MemoryMonitor interface { + StartMonitoring(ctx context.Context, interval time.Duration) +} +type CacheInterface interface { + Set(key string, value interface{}, ttl time.Duration) + Get(key string) (interface{}, bool) + Delete(key string) + SetMaxSize(size int) + Cleanup() + Close() +} +type TokenCache struct{} +type MetadataCache struct{} +type JWKCacheInterface interface{} diff --git a/handlers/oauth_handler.go b/handlers/oauth_handler.go new file mode 100644 index 0000000..425c001 --- /dev/null +++ b/handlers/oauth_handler.go @@ -0,0 +1,276 @@ +// Package handlers provides HTTP request handlers for the OIDC middleware. +package handlers + +import ( + "context" + "fmt" + "net/http" + "strings" +) + +// OAuthHandler handles OAuth callback requests +type OAuthHandler struct { + logger Logger + sessionManager SessionManager + tokenExchanger TokenExchanger + tokenVerifier TokenVerifier + extractClaimsFunc func(tokenString string) (map[string]interface{}, error) + isAllowedDomainFunc func(email string) bool + redirURLPath string + sendErrorResponseFunc func(rw http.ResponseWriter, req *http.Request, message string, code int) +} + +// Logger interface for dependency injection +type Logger interface { + Debugf(format string, args ...interface{}) + Errorf(format string, args ...interface{}) + Error(msg string) +} + +// SessionManager interface for session operations +type SessionManager interface { + GetSession(req *http.Request) (SessionData, error) +} + +// SessionData interface for session data operations +type SessionData interface { + GetCSRF() string + GetNonce() string + GetCodeVerifier() string + GetIncomingPath() string + GetAuthenticated() bool + SetAuthenticated(bool) error + SetEmail(string) + SetIDToken(string) + SetAccessToken(string) + SetRefreshToken(string) + SetCSRF(string) + SetNonce(string) + SetCodeVerifier(string) + SetIncomingPath(string) + ResetRedirectCount() + Save(req *http.Request, rw http.ResponseWriter) error + returnToPoolSafely() +} + +// TokenExchanger interface for token operations +type TokenExchanger interface { + ExchangeCodeForToken(ctx context.Context, grantType string, codeOrToken string, redirectURL string, codeVerifier string) (*TokenResponse, error) +} + +// TokenVerifier interface for token verification +type TokenVerifier interface { + VerifyToken(token string) error +} + +// TokenResponse represents the response from token exchange +type TokenResponse struct { + IDToken string + AccessToken string + RefreshToken string +} + +// NewOAuthHandler creates a new OAuth handler +func NewOAuthHandler(logger Logger, sessionManager SessionManager, tokenExchanger TokenExchanger, + tokenVerifier TokenVerifier, extractClaimsFunc func(string) (map[string]interface{}, error), + isAllowedDomainFunc func(string) bool, redirURLPath string, + sendErrorResponseFunc func(http.ResponseWriter, *http.Request, string, int)) *OAuthHandler { + + return &OAuthHandler{ + logger: logger, + sessionManager: sessionManager, + tokenExchanger: tokenExchanger, + tokenVerifier: tokenVerifier, + extractClaimsFunc: extractClaimsFunc, + isAllowedDomainFunc: isAllowedDomainFunc, + redirURLPath: redirURLPath, + sendErrorResponseFunc: sendErrorResponseFunc, + } +} + +// HandleCallback handles OAuth callback requests +func (h *OAuthHandler) HandleCallback(rw http.ResponseWriter, req *http.Request, redirectURL string) { + session, err := h.sessionManager.GetSession(req) + if err != nil { + h.logger.Errorf("Session error during callback: %v", err) + h.sendErrorResponseFunc(rw, req, "Session error during callback", http.StatusInternalServerError) + return + } + defer session.returnToPoolSafely() + + h.logger.Debugf("Handling callback, URL: %s", req.URL.String()) + + if req.URL.Query().Get("error") != "" { + errorDescription := req.URL.Query().Get("error_description") + if errorDescription == "" { + errorDescription = req.URL.Query().Get("error") + } + h.logger.Errorf("Authentication error from provider during callback: %s - %s", req.URL.Query().Get("error"), errorDescription) + h.sendErrorResponseFunc(rw, req, fmt.Sprintf("Authentication error from provider: %s", errorDescription), http.StatusBadRequest) + return + } + + state := req.URL.Query().Get("state") + if state == "" { + h.logger.Error("No state in callback") + h.sendErrorResponseFunc(rw, req, "State parameter missing in callback", http.StatusBadRequest) + return + } + + csrfToken := session.GetCSRF() + if csrfToken == "" { + h.logger.Errorf("CSRF token missing in session during callback. Authenticated: %v, Request URL: %s", + session.GetAuthenticated(), req.URL.String()) + + cookie, err := req.Cookie("_oidc_raczylo_m") + if err != nil { + h.logger.Errorf("Main session cookie not found in request: %v", err) + } else { + h.logger.Errorf("Main session cookie exists but CSRF token is empty. Cookie value length: %d", len(cookie.Value)) + } + + h.sendErrorResponseFunc(rw, req, "CSRF token missing in session", http.StatusBadRequest) + return + } + + if state != csrfToken { + h.logger.Error("State parameter does not match CSRF token in session during callback") + h.sendErrorResponseFunc(rw, req, "Invalid state parameter (CSRF mismatch)", http.StatusBadRequest) + return + } + + code := req.URL.Query().Get("code") + if code == "" { + h.logger.Error("No code in callback") + h.sendErrorResponseFunc(rw, req, "No authorization code received in callback", http.StatusBadRequest) + return + } + + codeVerifier := session.GetCodeVerifier() + + tokenResponse, err := h.tokenExchanger.ExchangeCodeForToken(req.Context(), "authorization_code", code, redirectURL, codeVerifier) + if err != nil { + h.logger.Errorf("Failed to exchange code for token during callback: %v", err) + h.sendErrorResponseFunc(rw, req, "Authentication failed: Could not exchange code for token", http.StatusInternalServerError) + return + } + + if err = h.tokenVerifier.VerifyToken(tokenResponse.IDToken); err != nil { + h.logger.Errorf("Failed to verify id_token during callback: %v", err) + h.sendErrorResponseFunc(rw, req, "Authentication failed: Could not verify ID token", http.StatusInternalServerError) + return + } + + claims, err := h.extractClaimsFunc(tokenResponse.IDToken) + if err != nil { + h.logger.Errorf("Failed to extract claims during callback: %v", err) + h.sendErrorResponseFunc(rw, req, "Authentication failed: Could not extract claims from token", http.StatusInternalServerError) + return + } + + nonceClaim, ok := claims["nonce"].(string) + if !ok || nonceClaim == "" { + h.logger.Error("Nonce claim missing in id_token during callback") + h.sendErrorResponseFunc(rw, req, "Authentication failed: Nonce missing in token", http.StatusInternalServerError) + return + } + + sessionNonce := session.GetNonce() + if sessionNonce == "" { + h.logger.Error("Nonce not found in session during callback") + h.sendErrorResponseFunc(rw, req, "Authentication failed: Nonce missing in session", http.StatusInternalServerError) + return + } + + if nonceClaim != sessionNonce { + h.logger.Error("Nonce claim does not match session nonce during callback") + h.sendErrorResponseFunc(rw, req, "Authentication failed: Nonce mismatch", http.StatusInternalServerError) + return + } + + email, _ := claims["email"].(string) + if email == "" { + h.logger.Errorf("Email claim missing or empty in token during callback") + h.sendErrorResponseFunc(rw, req, "Authentication failed: Email missing in token", http.StatusInternalServerError) + return + } + if !h.isAllowedDomainFunc(email) { + h.logger.Errorf("Disallowed email domain during callback: %s", email) + h.sendErrorResponseFunc(rw, req, "Authentication failed: Email domain not allowed", http.StatusForbidden) + return + } + + if err := session.SetAuthenticated(true); err != nil { + h.logger.Errorf("Failed to set authenticated state and regenerate session ID: %v", err) + h.sendErrorResponseFunc(rw, req, "Failed to update session", http.StatusInternalServerError) + return + } + session.SetEmail(email) + session.SetIDToken(tokenResponse.IDToken) + session.SetAccessToken(tokenResponse.AccessToken) + session.SetRefreshToken(tokenResponse.RefreshToken) + + session.SetCSRF("") + session.SetNonce("") + session.SetCodeVerifier("") + + session.ResetRedirectCount() + + redirectPath := "/" + if incomingPath := session.GetIncomingPath(); incomingPath != "" && incomingPath != h.redirURLPath { + redirectPath = incomingPath + } + session.SetIncomingPath("") + + if err := session.Save(req, rw); err != nil { + h.logger.Errorf("Failed to save session after callback: %v", err) + h.sendErrorResponseFunc(rw, req, "Failed to save session after callback", http.StatusInternalServerError) + return + } + + h.logger.Debugf("Callback successful, redirecting to %s", redirectPath) + http.Redirect(rw, req, redirectPath, http.StatusFound) +} + +// URLHelper provides utility methods for URL operations +type URLHelper struct { + logger Logger +} + +// NewURLHelper creates a new URL helper +func NewURLHelper(logger Logger) *URLHelper { + return &URLHelper{logger: logger} +} + +// DetermineExcludedURL checks if a URL path should bypass OIDC authentication. +// It compares the request path against configured excluded URL prefixes. +func (h *URLHelper) DetermineExcludedURL(currentRequest string, excludedURLs map[string]struct{}) bool { + for excludedURL := range excludedURLs { + if strings.HasPrefix(currentRequest, excludedURL) { + h.logger.Debugf("URL is excluded - got %s / excluded hit: %s", currentRequest, excludedURL) + return true + } + } + return false +} + +// DetermineScheme determines the URL scheme for building redirect URLs. +// It checks X-Forwarded-Proto header first, then TLS presence. +func (h *URLHelper) DetermineScheme(req *http.Request) string { + if scheme := req.Header.Get("X-Forwarded-Proto"); scheme != "" { + return scheme + } + if req.TLS != nil { + return "https" + } + return "http" +} + +// DetermineHost determines the host for building redirect URLs. +// It checks X-Forwarded-Host header first, then falls back to req.Host. +func (h *URLHelper) DetermineHost(req *http.Request) string { + if host := req.Header.Get("X-Forwarded-Host"); host != "" { + return host + } + return req.Host +} diff --git a/handlers/oauth_handler_test.go b/handlers/oauth_handler_test.go new file mode 100644 index 0000000..7d806eb --- /dev/null +++ b/handlers/oauth_handler_test.go @@ -0,0 +1,983 @@ +package handlers + +import ( + "context" + "crypto/tls" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" +) + +// Mock implementations for testing +type mockLogger struct { + debugLogs []string + errorLogs []string + mu sync.RWMutex +} + +func (m *mockLogger) Debugf(format string, args ...interface{}) { + m.mu.Lock() + defer m.mu.Unlock() + m.debugLogs = append(m.debugLogs, fmt.Sprintf(format, args...)) +} + +func (m *mockLogger) Errorf(format string, args ...interface{}) { + m.mu.Lock() + defer m.mu.Unlock() + m.errorLogs = append(m.errorLogs, fmt.Sprintf(format, args...)) +} + +func (m *mockLogger) Error(msg string) { + m.mu.Lock() + defer m.mu.Unlock() + m.errorLogs = append(m.errorLogs, msg) +} + +//lint:ignore U1000 May be needed for future debug log verification tests +func (m *mockLogger) getDebugLogs() []string { + m.mu.RLock() + defer m.mu.RUnlock() + result := make([]string, len(m.debugLogs)) + copy(result, m.debugLogs) + return result +} + +//lint:ignore U1000 May be needed for future error log verification tests +func (m *mockLogger) getErrorLogs() []string { + m.mu.RLock() + defer m.mu.RUnlock() + result := make([]string, len(m.errorLogs)) + copy(result, m.errorLogs) + return result +} + +//lint:ignore U1000 May be needed for future test isolation +func (m *mockLogger) reset() { + m.mu.Lock() + defer m.mu.Unlock() + m.debugLogs = nil + m.errorLogs = nil +} + +type mockSessionManager struct { + sessionData *mockSessionData + sessionError error + mu sync.RWMutex +} + +func (m *mockSessionManager) GetSession(req *http.Request) (SessionData, error) { + m.mu.RLock() + defer m.mu.RUnlock() + if m.sessionError != nil { + return nil, m.sessionError + } + return m.sessionData, nil +} + +func (m *mockSessionManager) setSessionError(err error) { + m.mu.Lock() + defer m.mu.Unlock() + m.sessionError = err +} + +type mockSessionData struct { + csrf string + nonce string + codeVerifier string + incomingPath string + authenticated bool + email string + idToken string + accessToken string + refreshToken string + saveError error + setAuthError error + mu sync.RWMutex + returned bool +} + +func newMockSessionData() *mockSessionData { + return &mockSessionData{ + csrf: "test-csrf-token", + nonce: "test-nonce", + codeVerifier: "test-code-verifier", + incomingPath: "/original/path", + authenticated: false, + } +} + +func (m *mockSessionData) GetCSRF() string { + m.mu.RLock() + defer m.mu.RUnlock() + return m.csrf +} + +func (m *mockSessionData) GetNonce() string { + m.mu.RLock() + defer m.mu.RUnlock() + return m.nonce +} + +func (m *mockSessionData) GetCodeVerifier() string { + m.mu.RLock() + defer m.mu.RUnlock() + return m.codeVerifier +} + +func (m *mockSessionData) GetIncomingPath() string { + m.mu.RLock() + defer m.mu.RUnlock() + return m.incomingPath +} + +func (m *mockSessionData) GetAuthenticated() bool { + m.mu.RLock() + defer m.mu.RUnlock() + return m.authenticated +} + +func (m *mockSessionData) SetAuthenticated(auth bool) error { + m.mu.Lock() + defer m.mu.Unlock() + if m.setAuthError != nil { + return m.setAuthError + } + m.authenticated = auth + return nil +} + +func (m *mockSessionData) SetEmail(email string) { + m.mu.Lock() + defer m.mu.Unlock() + m.email = email +} + +func (m *mockSessionData) SetIDToken(token string) { + m.mu.Lock() + defer m.mu.Unlock() + m.idToken = token +} + +func (m *mockSessionData) SetAccessToken(token string) { + m.mu.Lock() + defer m.mu.Unlock() + m.accessToken = token +} + +func (m *mockSessionData) SetRefreshToken(token string) { + m.mu.Lock() + defer m.mu.Unlock() + m.refreshToken = token +} + +func (m *mockSessionData) SetCSRF(csrf string) { + m.mu.Lock() + defer m.mu.Unlock() + m.csrf = csrf +} + +func (m *mockSessionData) SetNonce(nonce string) { + m.mu.Lock() + defer m.mu.Unlock() + m.nonce = nonce +} + +func (m *mockSessionData) SetCodeVerifier(verifier string) { + m.mu.Lock() + defer m.mu.Unlock() + m.codeVerifier = verifier +} + +func (m *mockSessionData) SetIncomingPath(path string) { + m.mu.Lock() + defer m.mu.Unlock() + m.incomingPath = path +} + +func (m *mockSessionData) ResetRedirectCount() { + // Mock implementation +} + +func (m *mockSessionData) Save(req *http.Request, rw http.ResponseWriter) error { + m.mu.RLock() + defer m.mu.RUnlock() + return m.saveError +} + +func (m *mockSessionData) returnToPoolSafely() { + m.mu.Lock() + defer m.mu.Unlock() + m.returned = true +} + +func (m *mockSessionData) setSaveError(err error) { + m.mu.Lock() + defer m.mu.Unlock() + m.saveError = err +} + +func (m *mockSessionData) setSetAuthError(err error) { + m.mu.Lock() + defer m.mu.Unlock() + m.setAuthError = err +} + +type mockTokenExchanger struct { + tokenResponse *TokenResponse + exchangeError error + mu sync.RWMutex +} + +func (m *mockTokenExchanger) ExchangeCodeForToken(ctx context.Context, grantType string, codeOrToken string, redirectURL string, codeVerifier string) (*TokenResponse, error) { + m.mu.RLock() + defer m.mu.RUnlock() + if m.exchangeError != nil { + return nil, m.exchangeError + } + return m.tokenResponse, nil +} + +//lint:ignore U1000 May be needed for future token exchange tests +func (m *mockTokenExchanger) setTokenResponse(response *TokenResponse) { + m.mu.Lock() + defer m.mu.Unlock() + m.tokenResponse = response +} + +func (m *mockTokenExchanger) setExchangeError(err error) { + m.mu.Lock() + defer m.mu.Unlock() + m.exchangeError = err +} + +type mockTokenVerifier struct { + verifyError error +} + +func (m *mockTokenVerifier) VerifyToken(token string) error { + return m.verifyError +} + +// Test helper functions +func createTestOAuthHandler() *OAuthHandler { + logger := &mockLogger{} + sessionManager := &mockSessionManager{sessionData: newMockSessionData()} + tokenExchanger := &mockTokenExchanger{ + tokenResponse: &TokenResponse{ + IDToken: "test.id.token", + AccessToken: "test-access-token", + RefreshToken: "test-refresh-token", + }, + } + tokenVerifier := &mockTokenVerifier{} + + extractClaimsFunc := func(tokenString string) (map[string]interface{}, error) { + return map[string]interface{}{ + "nonce": "test-nonce", + "email": "test@example.com", + "sub": "test-user-id", + }, nil + } + + isAllowedDomainFunc := func(email string) bool { + return strings.HasSuffix(email, "@example.com") + } + + sendErrorResponseFunc := func(rw http.ResponseWriter, req *http.Request, message string, code int) { + http.Error(rw, message, code) + } + + return NewOAuthHandler(logger, sessionManager, tokenExchanger, tokenVerifier, + extractClaimsFunc, isAllowedDomainFunc, "/auth/callback", sendErrorResponseFunc) +} + +func TestNewOAuthHandler(t *testing.T) { + handler := createTestOAuthHandler() + + if handler == nil { + t.Fatal("NewOAuthHandler returned nil") + } + + if handler.logger == nil { + t.Error("Expected logger to be set") + } + + if handler.sessionManager == nil { + t.Error("Expected sessionManager to be set") + } + + if handler.tokenExchanger == nil { + t.Error("Expected tokenExchanger to be set") + } + + if handler.redirURLPath != "/auth/callback" { + t.Errorf("Expected redirURLPath to be '/auth/callback', got %s", handler.redirURLPath) + } +} + +func TestHandleCallback_Success(t *testing.T) { + handler := createTestOAuthHandler() + + req := httptest.NewRequest("GET", "/auth/callback?state=test-csrf-token&code=test-auth-code", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusFound { + t.Errorf("Expected status %d, got %d", http.StatusFound, rw.Code) + } + + location := rw.Header().Get("Location") + if location != "/original/path" { + t.Errorf("Expected redirect to '/original/path', got %s", location) + } + + sessionData := handler.sessionManager.(*mockSessionManager).sessionData + if sessionData.email != "test@example.com" { + t.Errorf("Expected email to be set to 'test@example.com', got %s", sessionData.email) + } + + if !sessionData.authenticated { + t.Error("Expected session to be authenticated") + } +} + +func TestHandleCallback_SessionError(t *testing.T) { + handler := createTestOAuthHandler() + sessionManager := handler.sessionManager.(*mockSessionManager) + sessionManager.setSessionError(fmt.Errorf("session error")) + + req := httptest.NewRequest("GET", "/auth/callback?state=test-csrf-token&code=test-auth-code", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusInternalServerError { + t.Errorf("Expected status %d, got %d", http.StatusInternalServerError, rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "Session error during callback") { + t.Errorf("Expected session error message, got %s", responseBody) + } +} + +func TestHandleCallback_AuthError(t *testing.T) { + handler := createTestOAuthHandler() + + req := httptest.NewRequest("GET", "/auth/callback?error=access_denied&error_description=User+denied+access", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusBadRequest { + t.Errorf("Expected status %d, got %d", http.StatusBadRequest, rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "User denied access") { + t.Errorf("Expected error description in response, got %s", responseBody) + } +} + +func TestHandleCallback_MissingState(t *testing.T) { + handler := createTestOAuthHandler() + + req := httptest.NewRequest("GET", "/auth/callback?code=test-auth-code", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusBadRequest { + t.Errorf("Expected status %d, got %d", http.StatusBadRequest, rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "State parameter missing") { + t.Errorf("Expected state missing message, got %s", responseBody) + } +} + +func TestHandleCallback_MissingCSRFToken(t *testing.T) { + handler := createTestOAuthHandler() + sessionData := handler.sessionManager.(*mockSessionManager).sessionData + sessionData.csrf = "" // Clear CSRF token + + req := httptest.NewRequest("GET", "/auth/callback?state=test-csrf-token&code=test-auth-code", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusBadRequest { + t.Errorf("Expected status %d, got %d", http.StatusBadRequest, rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "CSRF token missing in session") { + t.Errorf("Expected CSRF missing message, got %s", responseBody) + } +} + +func TestHandleCallback_CSRFMismatch(t *testing.T) { + handler := createTestOAuthHandler() + + req := httptest.NewRequest("GET", "/auth/callback?state=wrong-csrf-token&code=test-auth-code", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusBadRequest { + t.Errorf("Expected status %d, got %d", http.StatusBadRequest, rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "CSRF mismatch") { + t.Errorf("Expected CSRF mismatch message, got %s", responseBody) + } +} + +func TestHandleCallback_MissingCode(t *testing.T) { + handler := createTestOAuthHandler() + + req := httptest.NewRequest("GET", "/auth/callback?state=test-csrf-token", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusBadRequest { + t.Errorf("Expected status %d, got %d", http.StatusBadRequest, rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "No authorization code received") { + t.Errorf("Expected missing code message, got %s", responseBody) + } +} + +func TestHandleCallback_TokenExchangeError(t *testing.T) { + handler := createTestOAuthHandler() + tokenExchanger := handler.tokenExchanger.(*mockTokenExchanger) + tokenExchanger.setExchangeError(fmt.Errorf("token exchange failed")) + + req := httptest.NewRequest("GET", "/auth/callback?state=test-csrf-token&code=test-auth-code", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusInternalServerError { + t.Errorf("Expected status %d, got %d", http.StatusInternalServerError, rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "Could not exchange code for token") { + t.Errorf("Expected token exchange error message, got %s", responseBody) + } +} + +func TestHandleCallback_TokenVerificationError(t *testing.T) { + handler := createTestOAuthHandler() + tokenVerifier := handler.tokenVerifier.(*mockTokenVerifier) + tokenVerifier.verifyError = fmt.Errorf("token verification failed") + + req := httptest.NewRequest("GET", "/auth/callback?state=test-csrf-token&code=test-auth-code", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusInternalServerError { + t.Errorf("Expected status %d, got %d", http.StatusInternalServerError, rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "Could not verify ID token") { + t.Errorf("Expected token verification error message, got %s", responseBody) + } +} + +func TestHandleCallback_ClaimsExtractionError(t *testing.T) { + handler := createTestOAuthHandler() + handler.extractClaimsFunc = func(tokenString string) (map[string]interface{}, error) { + return nil, fmt.Errorf("claims extraction failed") + } + + req := httptest.NewRequest("GET", "/auth/callback?state=test-csrf-token&code=test-auth-code", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusInternalServerError { + t.Errorf("Expected status %d, got %d", http.StatusInternalServerError, rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "Could not extract claims from token") { + t.Errorf("Expected claims extraction error message, got %s", responseBody) + } +} + +func TestHandleCallback_MissingNonceClaim(t *testing.T) { + handler := createTestOAuthHandler() + handler.extractClaimsFunc = func(tokenString string) (map[string]interface{}, error) { + return map[string]interface{}{ + "email": "test@example.com", + "sub": "test-user-id", + }, nil + } + + req := httptest.NewRequest("GET", "/auth/callback?state=test-csrf-token&code=test-auth-code", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusInternalServerError { + t.Errorf("Expected status %d, got %d", http.StatusInternalServerError, rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "Nonce missing in token") { + t.Errorf("Expected nonce missing message, got %s", responseBody) + } +} + +func TestHandleCallback_MissingSessionNonce(t *testing.T) { + handler := createTestOAuthHandler() + sessionData := handler.sessionManager.(*mockSessionManager).sessionData + sessionData.nonce = "" // Clear session nonce + + req := httptest.NewRequest("GET", "/auth/callback?state=test-csrf-token&code=test-auth-code", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusInternalServerError { + t.Errorf("Expected status %d, got %d", http.StatusInternalServerError, rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "Nonce missing in session") { + t.Errorf("Expected nonce missing in session message, got %s", responseBody) + } +} + +func TestHandleCallback_NonceMismatch(t *testing.T) { + handler := createTestOAuthHandler() + handler.extractClaimsFunc = func(tokenString string) (map[string]interface{}, error) { + return map[string]interface{}{ + "nonce": "wrong-nonce", + "email": "test@example.com", + "sub": "test-user-id", + }, nil + } + + req := httptest.NewRequest("GET", "/auth/callback?state=test-csrf-token&code=test-auth-code", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusInternalServerError { + t.Errorf("Expected status %d, got %d", http.StatusInternalServerError, rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "Nonce mismatch") { + t.Errorf("Expected nonce mismatch message, got %s", responseBody) + } +} + +func TestHandleCallback_MissingEmail(t *testing.T) { + handler := createTestOAuthHandler() + handler.extractClaimsFunc = func(tokenString string) (map[string]interface{}, error) { + return map[string]interface{}{ + "nonce": "test-nonce", + "sub": "test-user-id", + }, nil + } + + req := httptest.NewRequest("GET", "/auth/callback?state=test-csrf-token&code=test-auth-code", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusInternalServerError { + t.Errorf("Expected status %d, got %d", http.StatusInternalServerError, rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "Email missing in token") { + t.Errorf("Expected email missing message, got %s", responseBody) + } +} + +func TestHandleCallback_DisallowedEmailDomain(t *testing.T) { + handler := createTestOAuthHandler() + handler.extractClaimsFunc = func(tokenString string) (map[string]interface{}, error) { + return map[string]interface{}{ + "nonce": "test-nonce", + "email": "test@disallowed.com", + "sub": "test-user-id", + }, nil + } + + req := httptest.NewRequest("GET", "/auth/callback?state=test-csrf-token&code=test-auth-code", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusForbidden { + t.Errorf("Expected status %d, got %d", http.StatusForbidden, rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "Email domain not allowed") { + t.Errorf("Expected domain not allowed message, got %s", responseBody) + } +} + +func TestHandleCallback_SetAuthenticatedError(t *testing.T) { + handler := createTestOAuthHandler() + sessionData := handler.sessionManager.(*mockSessionManager).sessionData + sessionData.setSetAuthError(fmt.Errorf("failed to set authenticated")) + + req := httptest.NewRequest("GET", "/auth/callback?state=test-csrf-token&code=test-auth-code", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusInternalServerError { + t.Errorf("Expected status %d, got %d", http.StatusInternalServerError, rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "Failed to update session") { + t.Errorf("Expected session update error message, got %s", responseBody) + } +} + +func TestHandleCallback_SessionSaveError(t *testing.T) { + handler := createTestOAuthHandler() + sessionData := handler.sessionManager.(*mockSessionManager).sessionData + sessionData.setSaveError(fmt.Errorf("failed to save session")) + + req := httptest.NewRequest("GET", "/auth/callback?state=test-csrf-token&code=test-auth-code", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusInternalServerError { + t.Errorf("Expected status %d, got %d", http.StatusInternalServerError, rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "Failed to save session after callback") { + t.Errorf("Expected session save error message, got %s", responseBody) + } +} + +func TestHandleCallback_DefaultRedirect(t *testing.T) { + handler := createTestOAuthHandler() + sessionData := handler.sessionManager.(*mockSessionManager).sessionData + sessionData.incomingPath = "" // No incoming path + + req := httptest.NewRequest("GET", "/auth/callback?state=test-csrf-token&code=test-auth-code", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusFound { + t.Errorf("Expected status %d, got %d", http.StatusFound, rw.Code) + } + + location := rw.Header().Get("Location") + if location != "/" { + t.Errorf("Expected redirect to '/', got %s", location) + } +} + +func TestHandleCallback_CallbackPathNotRedirected(t *testing.T) { + handler := createTestOAuthHandler() + sessionData := handler.sessionManager.(*mockSessionManager).sessionData + sessionData.incomingPath = "/auth/callback" // Incoming path is callback path + + req := httptest.NewRequest("GET", "/auth/callback?state=test-csrf-token&code=test-auth-code", nil) + rw := httptest.NewRecorder() + + handler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusFound { + t.Errorf("Expected status %d, got %d", http.StatusFound, rw.Code) + } + + location := rw.Header().Get("Location") + if location != "/" { + t.Errorf("Expected redirect to '/' when incoming path is callback path, got %s", location) + } +} + +func TestURLHelper_NewURLHelper(t *testing.T) { + logger := &mockLogger{} + helper := NewURLHelper(logger) + + if helper == nil { + t.Fatal("NewURLHelper returned nil") + } + + if helper.logger != logger { + t.Error("Expected logger to be set correctly") + } +} + +func TestURLHelper_DetermineExcludedURL(t *testing.T) { + logger := &mockLogger{} + helper := NewURLHelper(logger) + + tests := []struct { + name string + currentReq string + excludedURLs map[string]struct{} + expected bool + }{ + { + name: "matches excluded URL", + currentReq: "/health/check", + excludedURLs: map[string]struct{}{ + "/health": {}, + }, + expected: true, + }, + { + name: "does not match excluded URL", + currentReq: "/protected/resource", + excludedURLs: map[string]struct{}{ + "/health": {}, + "/api": {}, + }, + expected: false, + }, + { + name: "empty excluded URLs", + currentReq: "/any/path", + excludedURLs: map[string]struct{}{}, + expected: false, + }, + { + name: "exact match", + currentReq: "/health", + excludedURLs: map[string]struct{}{ + "/health": {}, + }, + expected: true, + }, + { + name: "prefix match", + currentReq: "/api/v1/users", + excludedURLs: map[string]struct{}{ + "/api": {}, + }, + expected: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := helper.DetermineExcludedURL(tt.currentReq, tt.excludedURLs) + if result != tt.expected { + t.Errorf("Expected %v, got %v", tt.expected, result) + } + }) + } +} + +func TestURLHelper_DetermineScheme(t *testing.T) { + logger := &mockLogger{} + helper := NewURLHelper(logger) + + tests := []struct { + name string + req *http.Request + expected string + }{ + { + name: "X-Forwarded-Proto header present", + req: func() *http.Request { + req := httptest.NewRequest("GET", "/test", nil) + req.Header.Set("X-Forwarded-Proto", "https") + return req + }(), + expected: "https", + }, + { + name: "TLS present", + req: func() *http.Request { + req := httptest.NewRequest("GET", "/test", nil) + req.TLS = &tls.ConnectionState{} + return req + }(), + expected: "https", + }, + { + name: "no TLS or header", + req: httptest.NewRequest("GET", "/test", nil), + expected: "http", + }, + { + name: "X-Forwarded-Proto takes precedence over TLS", + req: func() *http.Request { + req := httptest.NewRequest("GET", "/test", nil) + req.Header.Set("X-Forwarded-Proto", "http") + req.TLS = &tls.ConnectionState{} + return req + }(), + expected: "http", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := helper.DetermineScheme(tt.req) + if result != tt.expected { + t.Errorf("Expected %s, got %s", tt.expected, result) + } + }) + } +} + +func TestURLHelper_DetermineHost(t *testing.T) { + logger := &mockLogger{} + helper := NewURLHelper(logger) + + tests := []struct { + name string + req *http.Request + expected string + }{ + { + name: "X-Forwarded-Host header present", + req: func() *http.Request { + req := httptest.NewRequest("GET", "http://example.com/test", nil) + req.Header.Set("X-Forwarded-Host", "proxy.example.com") + return req + }(), + expected: "proxy.example.com", + }, + { + name: "no X-Forwarded-Host header", + req: httptest.NewRequest("GET", "http://example.com/test", nil), + expected: "example.com", + }, + { + name: "X-Forwarded-Host with port", + req: func() *http.Request { + req := httptest.NewRequest("GET", "http://example.com:8080/test", nil) + req.Header.Set("X-Forwarded-Host", "proxy.example.com:443") + return req + }(), + expected: "proxy.example.com:443", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := helper.DetermineHost(tt.req) + if result != tt.expected { + t.Errorf("Expected %s, got %s", tt.expected, result) + } + }) + } +} + +func TestConcurrentCallbacks(t *testing.T) { + const numGoroutines = 10 + var wg sync.WaitGroup + errors := make(chan error, numGoroutines) + + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + + // Create unique session data for each goroutine + sessionData := newMockSessionData() + sessionData.csrf = fmt.Sprintf("csrf-token-%d", id) + sessionData.nonce = fmt.Sprintf("nonce-%d", id) + + sessionManager := &mockSessionManager{sessionData: sessionData} + localHandler := createTestOAuthHandler() + localHandler.sessionManager = sessionManager + + // Update extract claims function to match the session nonce + localHandler.extractClaimsFunc = func(tokenString string) (map[string]interface{}, error) { + return map[string]interface{}{ + "nonce": fmt.Sprintf("nonce-%d", id), + "email": fmt.Sprintf("test%d@example.com", id), + "sub": fmt.Sprintf("test-user-%d", id), + }, nil + } + + req := httptest.NewRequest("GET", + fmt.Sprintf("/auth/callback?state=csrf-token-%d&code=auth-code-%d", id, id), nil) + rw := httptest.NewRecorder() + + localHandler.HandleCallback(rw, req, "https://example.com/callback") + + if rw.Code != http.StatusFound { + errors <- fmt.Errorf("goroutine %d: expected status %d, got %d", + id, http.StatusFound, rw.Code) + } + }(i) + } + + wg.Wait() + close(errors) + + var errorCount int + for err := range errors { + t.Error(err) + errorCount++ + } + + if errorCount > 0 { + t.Errorf("Found %d errors in concurrent callback test", errorCount) + } +} + +// Benchmark tests +func BenchmarkHandleCallback_Success(b *testing.B) { + handler := createTestOAuthHandler() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + req := httptest.NewRequest("GET", "/auth/callback?state=test-csrf-token&code=test-auth-code", nil) + rw := httptest.NewRecorder() + handler.HandleCallback(rw, req, "https://example.com/callback") + } +} + +func BenchmarkURLHelper_DetermineExcludedURL(b *testing.B) { + logger := &mockLogger{} + helper := NewURLHelper(logger) + excludedURLs := map[string]struct{}{ + "/health": {}, + "/api": {}, + "/status": {}, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + helper.DetermineExcludedURL("/protected/resource", excludedURLs) + } +} + +func BenchmarkURLHelper_DetermineScheme(b *testing.B) { + logger := &mockLogger{} + helper := NewURLHelper(logger) + req := httptest.NewRequest("GET", "/test", nil) + req.Header.Set("X-Forwarded-Proto", "https") + + b.ResetTimer() + for i := 0; i < b.N; i++ { + helper.DetermineScheme(req) + } +} diff --git a/main.go b/main.go index b3b9cc7..c414061 100644 --- a/main.go +++ b/main.go @@ -31,251 +31,9 @@ func createDefaultHTTPClient() *http.Client { } const ( - ConstSessionTimeout = 86400 - defaultBlacklistDuration = 24 * time.Hour - defaultMaxBlacklistSize = 10000 + ConstSessionTimeout = 86400 ) -var ( - globalCacheManager *CacheManager - cacheManagerOnce sync.Once - cacheManagerMutex sync.RWMutex -) - -// CacheManager manages all caching components for the OIDC middleware. -// It provides thread-safe access to token blacklist, token cache, metadata cache, -// and JWK cache. This centralizes cache management to ensure efficient memory -// usage and consistent cache behavior across the application. -// CacheInterface defines the common cache operations -type CacheInterface interface { - Set(key string, value interface{}, ttl time.Duration) - Get(key string) (interface{}, bool) - Delete(key string) - SetMaxSize(size int) - Cleanup() - Close() -} - -type CacheManager struct { - tokenBlacklist CacheInterface - tokenCache *TokenCache - metadataCache *MetadataCache - jwkCache JWKCacheInterface - mu sync.RWMutex -} - -// GetGlobalCacheManager returns a singleton instance of CacheManager. -// It initializes all cache components on first call and reuses the same instance -// for subsequent calls. This ensures consistent cache state and efficient memory -// usage across the entire application lifecycle. -// Parameters: -// - wg: WaitGroup for coordinating cache cleanup during shutdown. -// -// Returns: -// - The global CacheManager instance. -func GetGlobalCacheManager(wg *sync.WaitGroup) *CacheManager { - cacheManagerOnce.Do(func() { - globalCacheManager = &CacheManager{ - tokenBlacklist: func() CacheInterface { - config := DefaultUnifiedCacheConfig() - config.MaxSize = defaultMaxBlacklistSize - c := NewUnifiedCache(config) - return NewCacheAdapter(c) - }(), - tokenCache: NewTokenCache(), - metadataCache: NewMetadataCache(wg), - jwkCache: &JWKCache{}, - } - }) - return globalCacheManager -} - -// GetSharedTokenBlacklist returns the shared token blacklist cache. -// This cache stores revoked or expired tokens to prevent replay attacks. -// Access is protected by read lock to ensure thread safety. -// Returns: -// - The shared token blacklist CacheInterface instance. -func (cm *CacheManager) GetSharedTokenBlacklist() CacheInterface { - cm.mu.RLock() - defer cm.mu.RUnlock() - return cm.tokenBlacklist -} - -// GetSharedTokenCache returns the shared token cache for verified tokens. -// This cache stores claims from successfully verified tokens to avoid -// repeated verification of the same tokens. -// Access is protected by read lock to ensure thread safety. -// Returns: -// - The shared TokenCache instance. -func (cm *CacheManager) GetSharedTokenCache() *TokenCache { - cm.mu.RLock() - defer cm.mu.RUnlock() - return cm.tokenCache -} - -// GetSharedMetadataCache returns the shared metadata cache. -// This cache stores OIDC provider metadata (endpoints, keys) to avoid -// repeated requests to the provider's .well-known configuration endpoint. -// Access is protected by read lock to ensure thread safety. -// Returns: -// - The shared MetadataCache instance. -func (cm *CacheManager) GetSharedMetadataCache() *MetadataCache { - cm.mu.RLock() - defer cm.mu.RUnlock() - return cm.metadataCache -} - -// GetSharedJWKCache returns the shared JWK (JSON Web Key) cache. -// This cache stores public keys from the provider's JWKS endpoint -// for token signature verification. -// Access is protected by read lock to ensure thread safety. -// Returns: -// - The shared JWKCacheInterface instance. -func (cm *CacheManager) GetSharedJWKCache() JWKCacheInterface { - cm.mu.RLock() - defer cm.mu.RUnlock() - return cm.jwkCache -} - -// Close gracefully shuts down all cache components and releases resources. -// It closes all individual caches and cleans up their associated goroutines. -// This method should be called when the middleware is shutting down. -// Safe to call multiple times. -// Returns: -// - An error if any cache fails to close properly, nil otherwise. -func (cm *CacheManager) Close() error { - cm.mu.Lock() - defer cm.mu.Unlock() - - if cm.tokenBlacklist != nil { - cm.tokenBlacklist.Close() - cm.tokenBlacklist = nil - } - if cm.tokenCache != nil { - cm.tokenCache.Close() - cm.tokenCache = nil - } - if cm.metadataCache != nil { - cm.metadataCache.Close() - cm.metadataCache = nil - } - if cm.jwkCache != nil { - cm.jwkCache.Close() - cm.jwkCache = nil - } - - return nil -} - -// CleanupGlobalCacheManager cleans up the global cache manager instance. -// It closes all cache components but does NOT reset the singleton for re-initialization -// to avoid race conditions with sync.Once. It's safe to call multiple times. -// Returns: -// - An error if cache cleanup fails, nil otherwise. -func CleanupGlobalCacheManager() error { - cacheManagerMutex.Lock() - defer cacheManagerMutex.Unlock() - - if globalCacheManager != nil { - err := globalCacheManager.Close() - // Don't reset globalCacheManager to nil or cacheManagerOnce to avoid race conditions - // The cache manager remains closed and the Close() method handles multiple calls - return err - } - return nil -} - -// TokenVerifier interface defines token verification capabilities. -// Implementations should validate token format, signature, and claims. -type TokenVerifier interface { - VerifyToken(token string) error -} - -// JWTVerifier interface defines JWT-specific verification capabilities. -// Implementations should validate JWT structure, signature using JWKs, and standard claims. -type JWTVerifier interface { - VerifyJWTSignatureAndClaims(jwt *JWT, token string) error -} - -// TokenExchanger interface defines OAuth 2.0 and OpenID Connect token exchange capabilities. -// Implementations should handle authorization code exchange, refresh tokens, and revocation -// according to the OAuth 2.0 and OpenID Connect specifications. -type TokenExchanger interface { - ExchangeCodeForToken(ctx context.Context, grantType string, codeOrToken string, redirectURL string, codeVerifier string) (*TokenResponse, error) - GetNewTokenWithRefreshToken(refreshToken string) (*TokenResponse, error) - RevokeTokenWithProvider(token, tokenType string) error -} - -// TraefikOidc is the main middleware struct that implements OIDC authentication for Traefik. -// It integrates with various OIDC providers, manages sessions, caches tokens, and handles -// the complete authentication flow. It's designed to work seamlessly with Traefik's -// plugin system and provides flexible configuration options. -type TraefikOidc struct { - jwkCache JWKCacheInterface - jwtVerifier JWTVerifier - ctx context.Context - tokenVerifier TokenVerifier - next http.Handler - tokenExchanger TokenExchanger - initComplete chan struct{} - limiter *rate.Limiter - tokenBlacklist CacheInterface - headerTemplates map[string]*template.Template - sessionManager *SessionManager - tokenCleanupStopChan chan struct{} - excludedURLs map[string]struct{} - extractClaimsFunc func(tokenString string) (map[string]interface{}, error) - initiateAuthenticationFunc func(rw http.ResponseWriter, req *http.Request, session *SessionData, redirectURL string) - metadataCache *MetadataCache - allowedRolesAndGroups map[string]struct{} - allowedUsers map[string]struct{} - allowedUserDomains map[string]struct{} - tokenCache *TokenCache - httpClient *http.Client - tokenHTTPClient *http.Client - logger *Logger - metadataRefreshStopChan chan struct{} - cancelFunc context.CancelFunc - errorRecoveryManager *ErrorRecoveryManager - goroutineWG *sync.WaitGroup - clientSecret string - clientID string - name string - redirURLPath string - logoutURLPath string - tokenURL string - authURL string - endSessionURL string - postLogoutRedirectURI string - scheme string - jwksURL string - issuerURL string - revocationURL string - providerURL string - scopes []string - refreshGracePeriod time.Duration - shutdownOnce sync.Once - firstRequestMutex sync.Mutex - forceHTTPS bool - enablePKCE bool - overrideScopes bool - suppressDiagnosticLogs bool - firstRequestReceived bool - metadataRefreshStarted bool -} - -// ProviderMetadata represents OIDC provider configuration data. -// This data is typically retrieved from the provider's .well-known/openid-configuration endpoint -// and contains essential URLs for authentication, token exchange, and key retrieval. -type ProviderMetadata struct { - Issuer string `json:"issuer"` - AuthURL string `json:"authorization_endpoint"` - TokenURL string `json:"token_endpoint"` - JWKSURL string `json:"jwks_uri"` - RevokeURL string `json:"revocation_endpoint"` - EndSessionURL string `json:"end_session_endpoint"` -} - // isTestMode detects if the code is running in a test environment. // It checks various indicators including environment variables, command-line arguments, // and runtime compiler information to determine test context. diff --git a/middleware/auth_middleware.go b/middleware/auth_middleware.go new file mode 100644 index 0000000..c6d6099 --- /dev/null +++ b/middleware/auth_middleware.go @@ -0,0 +1,464 @@ +// Package middleware provides authentication middleware for OIDC flows +package middleware + +import ( + "fmt" + "net/http" + "strings" + "sync" + "time" +) + +// AuthMiddleware handles the main OIDC authentication flow +type AuthMiddleware struct { + logger Logger + next http.Handler + sessionManager SessionManager + authHandler AuthHandler + oauthHandler OAuthHandler + urlHelper URLHelper + tokenVerifier TokenVerifier + extractClaimsFunc func(tokenString string) (map[string]interface{}, error) + extractGroupsAndRolesFunc func(tokenString string) ([]string, []string, error) + sendErrorResponseFunc func(rw http.ResponseWriter, req *http.Request, message string, code int) + refreshTokenFunc func(rw http.ResponseWriter, req *http.Request, session SessionData) bool + isUserAuthenticatedFunc func(session SessionData) (bool, bool, bool) + isAllowedDomainFunc func(email string) bool + isAjaxRequestFunc func(req *http.Request) bool + isRefreshTokenExpiredFunc func(session SessionData) bool + processLogoutFunc func(rw http.ResponseWriter, req *http.Request) + excludedURLs map[string]struct{} + allowedRolesAndGroups map[string]struct{} + redirURLPath string + logoutURLPath string + refreshGracePeriod time.Duration + initComplete chan struct{} + issuerURL string + firstRequestReceived bool + metadataRefreshStarted bool + firstRequestMutex sync.Mutex + providerURL string + goroutineWG *sync.WaitGroup + startTokenCleanupFunc func() + startMetadataRefreshFunc func(string) +} + +// Logger interface for dependency injection +type Logger interface { + Debug(msg string) + Debugf(format string, args ...interface{}) + Error(msg string) + Errorf(format string, args ...interface{}) + Info(msg string) + Infof(format string, args ...interface{}) +} + +// SessionManager interface for session operations +type SessionManager interface { + CleanupOldCookies(rw http.ResponseWriter, req *http.Request) + GetSession(req *http.Request) (SessionData, error) +} + +// SessionData interface for session data operations +type SessionData interface { + GetEmail() string + GetAccessToken() string + GetIDToken() string + GetRefreshToken() string + Clear(req *http.Request, rw http.ResponseWriter) error + ResetRedirectCount() + returnToPoolSafely() +} + +// AuthHandler interface for authentication operations +type AuthHandler interface { + InitiateAuthentication(rw http.ResponseWriter, req *http.Request, session SessionData, redirectURL string, + generateNonce, generateCodeVerifier, deriveCodeChallenge func() (string, error)) +} + +// OAuthHandler interface for OAuth callback operations +type OAuthHandler interface { + HandleCallback(rw http.ResponseWriter, req *http.Request, redirectURL string) +} + +// URLHelper interface for URL operations +type URLHelper interface { + DetermineExcludedURL(currentRequest string, excludedURLs map[string]struct{}) bool + DetermineScheme(req *http.Request) string + DetermineHost(req *http.Request) string +} + +// TokenVerifier interface for token verification +type TokenVerifier interface { + VerifyToken(token string) error +} + +// NewAuthMiddleware creates a new authentication middleware +func NewAuthMiddleware( + logger Logger, + next http.Handler, + sessionManager SessionManager, + authHandler AuthHandler, + oauthHandler OAuthHandler, + urlHelper URLHelper, + tokenVerifier TokenVerifier, + extractClaimsFunc func(string) (map[string]interface{}, error), + extractGroupsAndRolesFunc func(string) ([]string, []string, error), + sendErrorResponseFunc func(http.ResponseWriter, *http.Request, string, int), + refreshTokenFunc func(http.ResponseWriter, *http.Request, SessionData) bool, + isUserAuthenticatedFunc func(SessionData) (bool, bool, bool), + isAllowedDomainFunc func(string) bool, + isAjaxRequestFunc func(*http.Request) bool, + isRefreshTokenExpiredFunc func(SessionData) bool, + processLogoutFunc func(http.ResponseWriter, *http.Request), + excludedURLs map[string]struct{}, + allowedRolesAndGroups map[string]struct{}, + redirURLPath, logoutURLPath string, + refreshGracePeriod time.Duration, + initComplete chan struct{}, + issuerURL, providerURL string, + goroutineWG *sync.WaitGroup, + startTokenCleanupFunc func(), + startMetadataRefreshFunc func(string), +) *AuthMiddleware { + return &AuthMiddleware{ + logger: logger, + next: next, + sessionManager: sessionManager, + authHandler: authHandler, + oauthHandler: oauthHandler, + urlHelper: urlHelper, + tokenVerifier: tokenVerifier, + extractClaimsFunc: extractClaimsFunc, + extractGroupsAndRolesFunc: extractGroupsAndRolesFunc, + sendErrorResponseFunc: sendErrorResponseFunc, + refreshTokenFunc: refreshTokenFunc, + isUserAuthenticatedFunc: isUserAuthenticatedFunc, + isAllowedDomainFunc: isAllowedDomainFunc, + isAjaxRequestFunc: isAjaxRequestFunc, + isRefreshTokenExpiredFunc: isRefreshTokenExpiredFunc, + processLogoutFunc: processLogoutFunc, + excludedURLs: excludedURLs, + allowedRolesAndGroups: allowedRolesAndGroups, + redirURLPath: redirURLPath, + logoutURLPath: logoutURLPath, + refreshGracePeriod: refreshGracePeriod, + initComplete: initComplete, + issuerURL: issuerURL, + providerURL: providerURL, + goroutineWG: goroutineWG, + startTokenCleanupFunc: startTokenCleanupFunc, + startMetadataRefreshFunc: startMetadataRefreshFunc, + } +} + +// ServeHTTP implements the main OIDC authentication middleware +func (m *AuthMiddleware) ServeHTTP(rw http.ResponseWriter, req *http.Request) { + if !strings.HasPrefix(req.URL.Path, "/health") { + m.firstRequestMutex.Lock() + if !m.firstRequestReceived { + m.firstRequestReceived = true + m.logger.Debug("Starting background tasks on first request") + m.startTokenCleanupFunc() + + if !m.metadataRefreshStarted && m.providerURL != "" { + m.metadataRefreshStarted = true + if m.goroutineWG != nil { + m.goroutineWG.Add(1) + } + go func() { + defer func() { + if m.goroutineWG != nil { + m.goroutineWG.Done() + } + // Recover from panics to prevent goroutine leaks + if r := recover(); r != nil { + m.logger.Errorf("Start metadata refresh goroutine panic recovered: %v", r) + } + }() + m.startMetadataRefreshFunc(m.providerURL) + }() + } + } + m.firstRequestMutex.Unlock() + } + + select { + case <-m.initComplete: + if m.issuerURL == "" { + m.logger.Error("OIDC provider metadata initialization failed or incomplete") + m.sendErrorResponseFunc(rw, req, "OIDC provider metadata initialization failed - please check provider availability and configuration", http.StatusServiceUnavailable) + return + } + case <-req.Context().Done(): + m.logger.Debug("Request cancelled while waiting for OIDC initialization") + m.sendErrorResponseFunc(rw, req, "Request cancelled", http.StatusRequestTimeout) + return + case <-time.After(30 * time.Second): + m.logger.Error("Timeout waiting for OIDC initialization") + m.sendErrorResponseFunc(rw, req, "Timeout waiting for OIDC provider initialization - please try again later", http.StatusServiceUnavailable) + return + } + + if m.urlHelper.DetermineExcludedURL(req.URL.Path, m.excludedURLs) { + m.logger.Debugf("Request path %s excluded by configuration, bypassing OIDC", req.URL.Path) + m.next.ServeHTTP(rw, req) + return + } + + acceptHeader := req.Header.Get("Accept") + if strings.Contains(acceptHeader, "text/event-stream") { + m.logger.Debugf("Request accepts text/event-stream (%s), bypassing OIDC", acceptHeader) + m.next.ServeHTTP(rw, req) + return + } + + m.sessionManager.CleanupOldCookies(rw, req) + + session, err := m.sessionManager.GetSession(req) + if err != nil { + m.logger.Errorf("Error getting session: %v. Initiating authentication.", err) + cleanReq := req.Clone(req.Context()) + session, _ = m.sessionManager.GetSession(cleanReq) + if session != nil { + defer session.returnToPoolSafely() + if clearErr := session.Clear(cleanReq, rw); clearErr != nil { + m.logger.Errorf("Error clearing potentially corrupted session: %v", clearErr) + } + } else { + m.logger.Error("Critical session error: Failed to get even a new session.") + m.sendErrorResponseFunc(rw, req, "Critical session error", http.StatusInternalServerError) + return + } + scheme := m.urlHelper.DetermineScheme(req) + host := m.urlHelper.DetermineHost(req) + redirectURL := buildFullURL(scheme, host, m.redirURLPath) + m.authHandler.InitiateAuthentication(rw, req, session, redirectURL, + generateNonce, generateCodeVerifier, deriveCodeChallenge) + return + } + + defer session.returnToPoolSafely() + + scheme := m.urlHelper.DetermineScheme(req) + host := m.urlHelper.DetermineHost(req) + redirectURL := buildFullURL(scheme, host, m.redirURLPath) + + if req.URL.Path == m.logoutURLPath { + m.processLogoutFunc(rw, req) + return + } + if req.URL.Path == m.redirURLPath { + m.oauthHandler.HandleCallback(rw, req, redirectURL) + return + } + + authenticated, needsRefresh, expired := m.isUserAuthenticatedFunc(session) + + if expired { + m.logger.Debug("Session token is definitively expired or invalid, initiating re-auth") + m.handleExpiredToken(rw, req, session, redirectURL) + return + } + + email := session.GetEmail() + if authenticated && email != "" { + if !m.isAllowedDomainFunc(email) { + m.logger.Infof("User with email %s is not from an allowed domain", email) + errorMsg := fmt.Sprintf("Access denied: Your email domain is not allowed. To log out, visit: %s", m.logoutURLPath) + m.sendErrorResponseFunc(rw, req, errorMsg, http.StatusForbidden) + return + } + } + + if authenticated && !needsRefresh { + m.logger.Debug("User authenticated and token valid, proceeding to process authorized request") + if accessToken := session.GetAccessToken(); accessToken != "" { + if strings.Count(accessToken, ".") == 2 { + if err := m.tokenVerifier.VerifyToken(accessToken); err != nil { + m.logger.Errorf("Access token validation failed: %v", err) + m.handleExpiredToken(rw, req, session, redirectURL) + return + } + } else { + m.logger.Debugf("Access token appears opaque, skipping JWT verification for it.") + } + } + m.processAuthorizedRequest(rw, req, session, redirectURL) + return + } + + m.handleRefreshFlow(rw, req, session, redirectURL, needsRefresh, authenticated) +} + +// handleExpiredToken handles expired tokens by initiating re-authentication +func (m *AuthMiddleware) handleExpiredToken(rw http.ResponseWriter, req *http.Request, session SessionData, redirectURL string) { + session.ResetRedirectCount() + m.authHandler.InitiateAuthentication(rw, req, session, redirectURL, + generateNonce, generateCodeVerifier, deriveCodeChallenge) +} + +// handleRefreshFlow handles token refresh flow or initiates authentication +func (m *AuthMiddleware) handleRefreshFlow(rw http.ResponseWriter, req *http.Request, session SessionData, redirectURL string, needsRefresh, authenticated bool) { + refreshTokenPresent := session.GetRefreshToken() != "" + isAjaxRequest := m.isAjaxRequestFunc(req) + refreshTokenExpired := refreshTokenPresent && m.isRefreshTokenExpiredFunc(session) + shouldAttemptRefresh := needsRefresh && refreshTokenPresent && !refreshTokenExpired + + // If AJAX request and refresh token expired, return 401 immediately + if isAjaxRequest && refreshTokenExpired { + m.logger.Debug("AJAX request with expired refresh token, returning 401") + m.sendErrorResponseFunc(rw, req, "Session expired", http.StatusUnauthorized) + return + } + + if shouldAttemptRefresh { + m.handleTokenRefresh(rw, req, session, redirectURL, needsRefresh, authenticated, isAjaxRequest) + return + } + + m.logger.Debugf("Initiating full OIDC authentication flow (authenticated=%v, needsRefresh=%v, refreshTokenPresent=%v)", authenticated, needsRefresh, refreshTokenPresent) + + // If AJAX request without valid authentication, return 401 + if isAjaxRequest { + m.logger.Debug("AJAX request requires authentication, sending 401 Unauthorized") + m.sendErrorResponseFunc(rw, req, "Authentication required", http.StatusUnauthorized) + return + } + + // Reset redirect count when starting fresh authentication flow + session.ResetRedirectCount() + m.authHandler.InitiateAuthentication(rw, req, session, redirectURL, + generateNonce, generateCodeVerifier, deriveCodeChallenge) +} + +// handleTokenRefresh handles the token refresh process +func (m *AuthMiddleware) handleTokenRefresh(rw http.ResponseWriter, req *http.Request, session SessionData, redirectURL string, needsRefresh, authenticated, isAjaxRequest bool) { + if needsRefresh && authenticated { + m.logger.Debug("Session token needs proactive refresh, attempting refresh") + } else if needsRefresh && !authenticated { + m.logger.Debug("ID token invalid/expired, but refresh token found. Attempting refresh.") + } + + refreshed := m.refreshTokenFunc(rw, req, session) + if refreshed { + email := session.GetEmail() + if email != "" && !m.isAllowedDomainFunc(email) { + m.logger.Infof("User with refreshed token email %s is not from an allowed domain", email) + errorMsg := fmt.Sprintf("Access denied: Your email domain is not allowed. To log out, visit: %s", m.logoutURLPath) + m.sendErrorResponseFunc(rw, req, errorMsg, http.StatusForbidden) + return + } + + m.logger.Debug("Token refresh successful, proceeding to process authorized request") + m.processAuthorizedRequest(rw, req, session, redirectURL) + return + } + + m.logger.Debug("Token refresh failed, requiring re-authentication") + if isAjaxRequest { + m.logger.Debug("AJAX request with failed token refresh, sending 401 Unauthorized") + m.sendErrorResponseFunc(rw, req, "Token refresh failed", http.StatusUnauthorized) + } else { + m.logger.Debug("Browser request with failed token refresh, initiating re-auth") + // Reset redirect count when starting fresh auth after failed refresh to prevent redirect loops + session.ResetRedirectCount() + m.authHandler.InitiateAuthentication(rw, req, session, redirectURL, + generateNonce, generateCodeVerifier, deriveCodeChallenge) + } +} + +// processAuthorizedRequest processes requests for authenticated users +func (m *AuthMiddleware) processAuthorizedRequest(rw http.ResponseWriter, req *http.Request, session SessionData, redirectURL string) { + email := session.GetEmail() + if email == "" { + m.logger.Info("No email found in session during final processing, initiating re-auth") + // Reset redirect count to prevent loops when session is invalid + session.ResetRedirectCount() + m.authHandler.InitiateAuthentication(rw, req, session, redirectURL, + generateNonce, generateCodeVerifier, deriveCodeChallenge) + return + } + + tokenForClaims := session.GetIDToken() + if tokenForClaims == "" { + tokenForClaims = session.GetAccessToken() + if tokenForClaims == "" && len(m.allowedRolesAndGroups) > 0 { + m.logger.Error("No token available but roles/groups checks are required") + // Reset redirect count to prevent loops when token is missing + session.ResetRedirectCount() + m.authHandler.InitiateAuthentication(rw, req, session, redirectURL, + generateNonce, generateCodeVerifier, deriveCodeChallenge) + return + } + } + + // Initialize empty slices + var groups, roles []string + + if tokenForClaims != "" { + var err error + groups, roles, err = m.extractGroupsAndRolesFunc(tokenForClaims) + if err != nil && len(m.allowedRolesAndGroups) > 0 { + m.logger.Errorf("Failed to extract groups and roles: %v", err) + // Reset redirect count to prevent loops when claim extraction fails + session.ResetRedirectCount() + m.authHandler.InitiateAuthentication(rw, req, session, redirectURL, + generateNonce, generateCodeVerifier, deriveCodeChallenge) + return + } else if err == nil { + if len(groups) > 0 { + req.Header.Set("X-User-Groups", strings.Join(groups, ",")) + } + if len(roles) > 0 { + req.Header.Set("X-User-Roles", strings.Join(roles, ",")) + } + } + } + + if len(m.allowedRolesAndGroups) > 0 { + allowed := false + for _, roleOrGroup := range append(groups, roles...) { + if _, ok := m.allowedRolesAndGroups[roleOrGroup]; ok { + allowed = true + break + } + } + if !allowed { + m.logger.Infof("User with email %s does not have any allowed roles or groups", email) + errorMsg := fmt.Sprintf("Access denied: You do not have any of the allowed roles or groups. To log out, visit: %s", m.logoutURLPath) + m.sendErrorResponseFunc(rw, req, errorMsg, http.StatusForbidden) + return + } + } + + req.Header.Set("X-Forwarded-User", email) + req.Header.Set("X-Auth-Request-Redirect", req.URL.RequestURI()) + req.Header.Set("X-Auth-Request-User", email) + if idToken := session.GetIDToken(); idToken != "" { + req.Header.Set("X-Auth-Request-Token", idToken) + } + + m.next.ServeHTTP(rw, req) +} + +// buildFullURL constructs a full URL from scheme, host, and path components +func buildFullURL(scheme, host, path string) string { + return fmt.Sprintf("%s://%s%s", scheme, host, path) +} + +// These functions need to be provided by the calling code or injected as dependencies +func generateNonce() (string, error) { + // This function needs to be implemented or injected + return "", fmt.Errorf("generateNonce not implemented") +} + +func generateCodeVerifier() (string, error) { + // This function needs to be implemented or injected + return "", fmt.Errorf("generateCodeVerifier not implemented") +} + +func deriveCodeChallenge() (string, error) { + // This function needs to be implemented or injected + return "", fmt.Errorf("deriveCodeChallenge not implemented") +} diff --git a/middleware/auth_middleware_test.go b/middleware/auth_middleware_test.go new file mode 100644 index 0000000..098db79 --- /dev/null +++ b/middleware/auth_middleware_test.go @@ -0,0 +1,980 @@ +package middleware + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" +) + +// Mock implementations for testing +type mockLogger struct { + debugLogs []string + infoLogs []string + errorLogs []string + mu sync.RWMutex +} + +func (m *mockLogger) Debug(msg string) { + m.mu.Lock() + defer m.mu.Unlock() + m.debugLogs = append(m.debugLogs, msg) +} + +func (m *mockLogger) Debugf(format string, args ...interface{}) { + m.mu.Lock() + defer m.mu.Unlock() + m.debugLogs = append(m.debugLogs, fmt.Sprintf(format, args...)) +} + +func (m *mockLogger) Error(msg string) { + m.mu.Lock() + defer m.mu.Unlock() + m.errorLogs = append(m.errorLogs, msg) +} + +func (m *mockLogger) Errorf(format string, args ...interface{}) { + m.mu.Lock() + defer m.mu.Unlock() + m.errorLogs = append(m.errorLogs, fmt.Sprintf(format, args...)) +} + +func (m *mockLogger) Info(msg string) { + m.mu.Lock() + defer m.mu.Unlock() + m.infoLogs = append(m.infoLogs, msg) +} + +func (m *mockLogger) Infof(format string, args ...interface{}) { + m.mu.Lock() + defer m.mu.Unlock() + m.infoLogs = append(m.infoLogs, fmt.Sprintf(format, args...)) +} + +//lint:ignore U1000 May be needed for future debug log verification tests +func (m *mockLogger) getDebugLogs() []string { + m.mu.RLock() + defer m.mu.RUnlock() + result := make([]string, len(m.debugLogs)) + copy(result, m.debugLogs) + return result +} + +//lint:ignore U1000 May be needed for future error log verification tests +func (m *mockLogger) getErrorLogs() []string { + m.mu.RLock() + defer m.mu.RUnlock() + result := make([]string, len(m.errorLogs)) + copy(result, m.errorLogs) + return result +} + +//lint:ignore U1000 May be needed for future test isolation +func (m *mockLogger) reset() { + m.mu.Lock() + defer m.mu.Unlock() + m.debugLogs = nil + m.infoLogs = nil + m.errorLogs = nil +} + +type mockSessionManager struct { + cleanupCalled bool + sessionData *mockSessionData + sessionError error + callCount int // Track how many times GetSession is called + mu sync.RWMutex +} + +func (m *mockSessionManager) CleanupOldCookies(rw http.ResponseWriter, req *http.Request) { + m.mu.Lock() + defer m.mu.Unlock() + m.cleanupCalled = true +} + +func (m *mockSessionManager) GetSession(req *http.Request) (SessionData, error) { + m.mu.Lock() + defer m.mu.Unlock() + + m.callCount++ + + // First call returns error, second call (for clean request) returns valid session + if m.sessionError != nil && m.callCount == 1 { + return nil, m.sessionError + } + + return m.sessionData, nil +} + +func (m *mockSessionManager) setSessionError(err error) { + m.mu.Lock() + defer m.mu.Unlock() + m.sessionError = err + m.callCount = 0 // Reset call count when setting error +} + +type mockSessionData struct { + email string + accessToken string + idToken string + refreshToken string + clearError error + mu sync.RWMutex + returned bool +} + +func newMockSessionData() *mockSessionData { + return &mockSessionData{} +} + +func (m *mockSessionData) GetEmail() string { + m.mu.RLock() + defer m.mu.RUnlock() + return m.email +} + +func (m *mockSessionData) GetAccessToken() string { + m.mu.RLock() + defer m.mu.RUnlock() + return m.accessToken +} + +func (m *mockSessionData) GetIDToken() string { + m.mu.RLock() + defer m.mu.RUnlock() + return m.idToken +} + +func (m *mockSessionData) GetRefreshToken() string { + m.mu.RLock() + defer m.mu.RUnlock() + return m.refreshToken +} + +func (m *mockSessionData) Clear(req *http.Request, rw http.ResponseWriter) error { + m.mu.RLock() + defer m.mu.RUnlock() + return m.clearError +} + +func (m *mockSessionData) ResetRedirectCount() { + // Mock implementation +} + +func (m *mockSessionData) returnToPoolSafely() { + m.mu.Lock() + defer m.mu.Unlock() + m.returned = true +} + +func (m *mockSessionData) setEmail(email string) { + m.mu.Lock() + defer m.mu.Unlock() + m.email = email +} + +func (m *mockSessionData) setAccessToken(token string) { + m.mu.Lock() + defer m.mu.Unlock() + m.accessToken = token +} + +func (m *mockSessionData) setIDToken(token string) { + m.mu.Lock() + defer m.mu.Unlock() + m.idToken = token +} + +func (m *mockSessionData) setRefreshToken(token string) { + m.mu.Lock() + defer m.mu.Unlock() + m.refreshToken = token +} + +//lint:ignore U1000 May be needed for future session clear testing +func (m *mockSessionData) setClearError(err error) { + m.mu.Lock() + defer m.mu.Unlock() + m.clearError = err +} + +type mockAuthHandler struct { + initiateCalled bool + mu sync.RWMutex +} + +func (m *mockAuthHandler) InitiateAuthentication(rw http.ResponseWriter, req *http.Request, session SessionData, redirectURL string, + generateNonce, generateCodeVerifier, deriveCodeChallenge func() (string, error)) { + m.mu.Lock() + defer m.mu.Unlock() + m.initiateCalled = true + http.Redirect(rw, req, "https://provider.example.com/auth", http.StatusFound) +} + +func (m *mockAuthHandler) wasInitiateCalled() bool { + m.mu.RLock() + defer m.mu.RUnlock() + return m.initiateCalled +} + +//lint:ignore U1000 May be needed for future test isolation +func (m *mockAuthHandler) reset() { + m.mu.Lock() + defer m.mu.Unlock() + m.initiateCalled = false +} + +type mockOAuthHandler struct { + handleCallbackCalled bool + mu sync.RWMutex +} + +func (m *mockOAuthHandler) HandleCallback(rw http.ResponseWriter, req *http.Request, redirectURL string) { + m.mu.Lock() + defer m.mu.Unlock() + m.handleCallbackCalled = true + rw.WriteHeader(http.StatusOK) + rw.Write([]byte("callback handled")) +} + +func (m *mockOAuthHandler) wasHandleCallbackCalled() bool { + m.mu.RLock() + defer m.mu.RUnlock() + return m.handleCallbackCalled +} + +type mockURLHelper struct { + excludedResult bool + scheme string + host string +} + +func (m *mockURLHelper) DetermineExcludedURL(currentRequest string, excludedURLs map[string]struct{}) bool { + return m.excludedResult +} + +func (m *mockURLHelper) DetermineScheme(req *http.Request) string { + if m.scheme != "" { + return m.scheme + } + return "https" +} + +func (m *mockURLHelper) DetermineHost(req *http.Request) string { + if m.host != "" { + return m.host + } + return "example.com" +} + +type mockTokenVerifier struct { + verifyError error +} + +func (m *mockTokenVerifier) VerifyToken(token string) error { + return m.verifyError +} + +type mockHandler struct { + serveHTTPCalled bool + mu sync.RWMutex +} + +func (m *mockHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) { + m.mu.Lock() + defer m.mu.Unlock() + m.serveHTTPCalled = true + rw.WriteHeader(http.StatusOK) + rw.Write([]byte("next handler called")) +} + +func (m *mockHandler) wasServeHTTPCalled() bool { + m.mu.RLock() + defer m.mu.RUnlock() + return m.serveHTTPCalled +} + +//lint:ignore U1000 May be needed for future test isolation +func (m *mockHandler) reset() { + m.mu.Lock() + defer m.mu.Unlock() + m.serveHTTPCalled = false +} + +// Test helper functions +func createTestMiddleware() *AuthMiddleware { + logger := &mockLogger{} + nextHandler := &mockHandler{} + sessionManager := &mockSessionManager{sessionData: newMockSessionData()} + authHandler := &mockAuthHandler{} + oauthHandler := &mockOAuthHandler{} + urlHelper := &mockURLHelper{} + tokenVerifier := &mockTokenVerifier{} + + extractClaimsFunc := func(tokenString string) (map[string]interface{}, error) { + return map[string]interface{}{"sub": "test-user", "email": "test@example.com"}, nil + } + + extractGroupsAndRolesFunc := func(tokenString string) ([]string, []string, error) { + return []string{"group1"}, []string{"role1"}, nil + } + + sendErrorResponseFunc := func(rw http.ResponseWriter, req *http.Request, message string, code int) { + http.Error(rw, message, code) + } + + refreshTokenFunc := func(rw http.ResponseWriter, req *http.Request, session SessionData) bool { + return true + } + + isUserAuthenticatedFunc := func(session SessionData) (bool, bool, bool) { + return true, false, false // authenticated, needsRefresh, expired + } + + isAllowedDomainFunc := func(email string) bool { + return true + } + + isAjaxRequestFunc := func(req *http.Request) bool { + return strings.Contains(req.Header.Get("Accept"), "application/json") || + strings.Contains(req.Header.Get("Content-Type"), "application/json") + } + + isRefreshTokenExpiredFunc := func(session SessionData) bool { + return false + } + + processLogoutFunc := func(rw http.ResponseWriter, req *http.Request) { + rw.WriteHeader(http.StatusOK) + rw.Write([]byte("logged out")) + } + + excludedURLs := make(map[string]struct{}) + allowedRolesAndGroups := make(map[string]struct{}) + initComplete := make(chan struct{}) + close(initComplete) // Initialize as complete + var goroutineWG sync.WaitGroup + + startTokenCleanupFunc := func() { + // Mock implementation + } + + startMetadataRefreshFunc := func(providerURL string) { + // Mock implementation + } + + return NewAuthMiddleware( + logger, + nextHandler, + sessionManager, + authHandler, + oauthHandler, + urlHelper, + tokenVerifier, + extractClaimsFunc, + extractGroupsAndRolesFunc, + sendErrorResponseFunc, + refreshTokenFunc, + isUserAuthenticatedFunc, + isAllowedDomainFunc, + isAjaxRequestFunc, + isRefreshTokenExpiredFunc, + processLogoutFunc, + excludedURLs, + allowedRolesAndGroups, + "/auth/callback", + "/auth/logout", + 30*time.Second, + initComplete, + "https://provider.example.com", + "https://provider.example.com", + &goroutineWG, + startTokenCleanupFunc, + startMetadataRefreshFunc, + ) +} + +func TestNewAuthMiddleware(t *testing.T) { + middleware := createTestMiddleware() + + if middleware == nil { + t.Fatal("NewAuthMiddleware returned nil") + } + + if middleware.logger == nil { + t.Error("Expected logger to be set") + } + + if middleware.next == nil { + t.Error("Expected next handler to be set") + } + + if middleware.redirURLPath != "/auth/callback" { + t.Errorf("Expected redirURLPath to be '/auth/callback', got %s", middleware.redirURLPath) + } + + if middleware.logoutURLPath != "/auth/logout" { + t.Errorf("Expected logoutURLPath to be '/auth/logout', got %s", middleware.logoutURLPath) + } +} + +func TestServeHTTP_ExcludedURL(t *testing.T) { + middleware := createTestMiddleware() + urlHelper := &mockURLHelper{excludedResult: true} + middleware.urlHelper = urlHelper + + req := httptest.NewRequest("GET", "/excluded/path", nil) + rw := httptest.NewRecorder() + + middleware.ServeHTTP(rw, req) + + nextHandler := middleware.next.(*mockHandler) + if !nextHandler.wasServeHTTPCalled() { + t.Error("Expected next handler to be called for excluded URL") + } + + if rw.Code != http.StatusOK { + t.Errorf("Expected status 200, got %d", rw.Code) + } +} + +func TestServeHTTP_EventStreamBypass(t *testing.T) { + middleware := createTestMiddleware() + + req := httptest.NewRequest("GET", "/test", nil) + req.Header.Set("Accept", "text/event-stream") + rw := httptest.NewRecorder() + + middleware.ServeHTTP(rw, req) + + nextHandler := middleware.next.(*mockHandler) + if !nextHandler.wasServeHTTPCalled() { + t.Error("Expected next handler to be called for event stream request") + } +} + +func TestServeHTTP_InitializationTimeout(t *testing.T) { + t.Skip("Skipping timeout test to avoid long delays") + logger := &mockLogger{} + nextHandler := &mockHandler{} + sessionManager := &mockSessionManager{sessionData: newMockSessionData()} + authHandler := &mockAuthHandler{} + oauthHandler := &mockOAuthHandler{} + urlHelper := &mockURLHelper{} + tokenVerifier := &mockTokenVerifier{} + + // Create initComplete channel but don't close it to simulate timeout + initComplete := make(chan struct{}) + var goroutineWG sync.WaitGroup + + middleware := NewAuthMiddleware( + logger, + nextHandler, + sessionManager, + authHandler, + oauthHandler, + urlHelper, + tokenVerifier, + func(string) (map[string]interface{}, error) { return nil, nil }, + func(string) ([]string, []string, error) { return nil, nil, nil }, + func(rw http.ResponseWriter, req *http.Request, message string, code int) { + http.Error(rw, message, code) + }, + func(http.ResponseWriter, *http.Request, SessionData) bool { return false }, + func(SessionData) (bool, bool, bool) { return false, false, false }, + func(string) bool { return true }, + func(*http.Request) bool { return false }, + func(SessionData) bool { return false }, + func(http.ResponseWriter, *http.Request) {}, + make(map[string]struct{}), + make(map[string]struct{}), + "/auth/callback", + "/auth/logout", + 10*time.Millisecond, // Very short timeout for testing + initComplete, + "", + "https://provider.example.com", + &goroutineWG, + func() {}, + func(string) {}, + ) + + req := httptest.NewRequest("GET", "/test", nil) + rw := httptest.NewRecorder() + + middleware.ServeHTTP(rw, req) + + if rw.Code != http.StatusServiceUnavailable { + t.Errorf("Expected status %d, got %d", http.StatusServiceUnavailable, rw.Code) + } +} + +func TestServeHTTP_SessionError(t *testing.T) { + middleware := createTestMiddleware() + sessionManager := middleware.sessionManager.(*mockSessionManager) + sessionManager.setSessionError(fmt.Errorf("session error")) + + // Ensure we have valid session data for the second call + sessionManager.sessionData = newMockSessionData() + + req := httptest.NewRequest("GET", "/test", nil) + rw := httptest.NewRecorder() + + middleware.ServeHTTP(rw, req) + + authHandler := middleware.authHandler.(*mockAuthHandler) + if !authHandler.wasInitiateCalled() { + t.Error("Expected InitiateAuthentication to be called when session error occurs") + } +} + +func TestServeHTTP_LogoutPath(t *testing.T) { + middleware := createTestMiddleware() + sessionData := newMockSessionData() + sessionData.setEmail("test@example.com") + sessionManager := middleware.sessionManager.(*mockSessionManager) + sessionManager.sessionData = sessionData + + req := httptest.NewRequest("GET", "/auth/logout", nil) + rw := httptest.NewRecorder() + + middleware.ServeHTTP(rw, req) + + if rw.Code != http.StatusOK { + t.Errorf("Expected status 200 for logout, got %d", rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "logged out") { + t.Errorf("Expected logout response, got %s", responseBody) + } +} + +func TestServeHTTP_CallbackPath(t *testing.T) { + middleware := createTestMiddleware() + + req := httptest.NewRequest("GET", "/auth/callback?code=test", nil) + rw := httptest.NewRecorder() + + middleware.ServeHTTP(rw, req) + + oauthHandler := middleware.oauthHandler.(*mockOAuthHandler) + if !oauthHandler.wasHandleCallbackCalled() { + t.Error("Expected HandleCallback to be called for callback path") + } +} + +func TestServeHTTP_AuthenticatedUser(t *testing.T) { + middleware := createTestMiddleware() + sessionData := newMockSessionData() + sessionData.setEmail("test@example.com") + sessionData.setAccessToken("valid.jwt.token") + sessionData.setIDToken("id.jwt.token") + sessionManager := middleware.sessionManager.(*mockSessionManager) + sessionManager.sessionData = sessionData + + req := httptest.NewRequest("GET", "/protected", nil) + rw := httptest.NewRecorder() + + middleware.ServeHTTP(rw, req) + + nextHandler := middleware.next.(*mockHandler) + if !nextHandler.wasServeHTTPCalled() { + t.Error("Expected next handler to be called for authenticated user") + } + + // Check that user headers are set + if req.Header.Get("X-Forwarded-User") != "test@example.com" { + t.Errorf("Expected X-Forwarded-User header to be set to 'test@example.com', got %s", req.Header.Get("X-Forwarded-User")) + } + + if req.Header.Get("X-Auth-Request-User") != "test@example.com" { + t.Errorf("Expected X-Auth-Request-User header to be set to 'test@example.com', got %s", req.Header.Get("X-Auth-Request-User")) + } +} + +func TestServeHTTP_DomainNotAllowed(t *testing.T) { + middleware := createTestMiddleware() + sessionData := newMockSessionData() + sessionData.setEmail("test@disallowed.com") + sessionManager := middleware.sessionManager.(*mockSessionManager) + sessionManager.sessionData = sessionData + + // Set isAllowedDomainFunc to return false + middleware.isAllowedDomainFunc = func(email string) bool { + return false + } + + req := httptest.NewRequest("GET", "/protected", nil) + rw := httptest.NewRecorder() + + middleware.ServeHTTP(rw, req) + + if rw.Code != http.StatusForbidden { + t.Errorf("Expected status %d for disallowed domain, got %d", http.StatusForbidden, rw.Code) + } + + responseBody := rw.Body.String() + if !strings.Contains(responseBody, "Access denied") { + t.Errorf("Expected access denied message, got %s", responseBody) + } +} + +func TestServeHTTP_TokenRefreshNeeded(t *testing.T) { + middleware := createTestMiddleware() + sessionData := newMockSessionData() + sessionData.setEmail("test@example.com") + sessionData.setRefreshToken("refresh-token") + sessionManager := middleware.sessionManager.(*mockSessionManager) + sessionManager.sessionData = sessionData + + // Set authenticated but needs refresh + middleware.isUserAuthenticatedFunc = func(session SessionData) (bool, bool, bool) { + return true, true, false // authenticated, needsRefresh, not expired + } + + req := httptest.NewRequest("GET", "/protected", nil) + rw := httptest.NewRecorder() + + middleware.ServeHTTP(rw, req) + + nextHandler := middleware.next.(*mockHandler) + if !nextHandler.wasServeHTTPCalled() { + t.Error("Expected next handler to be called after successful token refresh") + } +} + +func TestServeHTTP_AjaxRequestWithoutAuth(t *testing.T) { + middleware := createTestMiddleware() + sessionData := newMockSessionData() + sessionManager := middleware.sessionManager.(*mockSessionManager) + sessionManager.sessionData = sessionData + + // Set user as not authenticated + middleware.isUserAuthenticatedFunc = func(session SessionData) (bool, bool, bool) { + return false, false, false + } + + req := httptest.NewRequest("GET", "/api/data", nil) + req.Header.Set("Accept", "application/json") + rw := httptest.NewRecorder() + + middleware.ServeHTTP(rw, req) + + if rw.Code != http.StatusUnauthorized { + t.Errorf("Expected status %d for AJAX request without auth, got %d", http.StatusUnauthorized, rw.Code) + } +} + +func TestServeHTTP_ExpiredToken(t *testing.T) { + middleware := createTestMiddleware() + sessionData := newMockSessionData() + sessionData.setEmail("test@example.com") + sessionManager := middleware.sessionManager.(*mockSessionManager) + sessionManager.sessionData = sessionData + + // Set token as expired + middleware.isUserAuthenticatedFunc = func(session SessionData) (bool, bool, bool) { + return false, false, true // not authenticated, no refresh needed, expired + } + + req := httptest.NewRequest("GET", "/protected", nil) + rw := httptest.NewRecorder() + + middleware.ServeHTTP(rw, req) + + authHandler := middleware.authHandler.(*mockAuthHandler) + if !authHandler.wasInitiateCalled() { + t.Error("Expected InitiateAuthentication to be called for expired token") + } +} + +func TestServeHTTP_RoleBasedAccess(t *testing.T) { + middleware := createTestMiddleware() + sessionData := newMockSessionData() + sessionData.setEmail("test@example.com") + sessionData.setIDToken("id.jwt.token") + sessionManager := middleware.sessionManager.(*mockSessionManager) + sessionManager.sessionData = sessionData + + // Set required roles + middleware.allowedRolesAndGroups = map[string]struct{}{ + "admin": {}, + "group1": {}, + } + + // User has group1 which is allowed + middleware.extractGroupsAndRolesFunc = func(tokenString string) ([]string, []string, error) { + return []string{"group1"}, []string{"user"}, nil + } + + req := httptest.NewRequest("GET", "/admin", nil) + rw := httptest.NewRecorder() + + middleware.ServeHTTP(rw, req) + + nextHandler := middleware.next.(*mockHandler) + if !nextHandler.wasServeHTTPCalled() { + t.Error("Expected next handler to be called for user with allowed role") + } + + // Check that group headers are set + if req.Header.Get("X-User-Groups") != "group1" { + t.Errorf("Expected X-User-Groups header to be set to 'group1', got %s", req.Header.Get("X-User-Groups")) + } +} + +func TestServeHTTP_RoleBasedAccessDenied(t *testing.T) { + middleware := createTestMiddleware() + sessionData := newMockSessionData() + sessionData.setEmail("test@example.com") + sessionData.setIDToken("id.jwt.token") + sessionManager := middleware.sessionManager.(*mockSessionManager) + sessionManager.sessionData = sessionData + + // Set required roles that user doesn't have + middleware.allowedRolesAndGroups = map[string]struct{}{ + "admin": {}, + } + + // User doesn't have admin role + middleware.extractGroupsAndRolesFunc = func(tokenString string) ([]string, []string, error) { + return []string{"group1"}, []string{"user"}, nil + } + + req := httptest.NewRequest("GET", "/admin", nil) + rw := httptest.NewRecorder() + + middleware.ServeHTTP(rw, req) + + if rw.Code != http.StatusForbidden { + t.Errorf("Expected status %d for user without required role, got %d", http.StatusForbidden, rw.Code) + } +} + +func TestHandleTokenRefresh_FailedRefresh(t *testing.T) { + middleware := createTestMiddleware() + sessionData := newMockSessionData() + sessionData.setEmail("test@example.com") + sessionData.setRefreshToken("refresh-token") + sessionManager := middleware.sessionManager.(*mockSessionManager) + sessionManager.sessionData = sessionData + + // Set refresh to fail + middleware.refreshTokenFunc = func(rw http.ResponseWriter, req *http.Request, session SessionData) bool { + return false + } + + // Set authenticated but needs refresh + middleware.isUserAuthenticatedFunc = func(session SessionData) (bool, bool, bool) { + return true, true, false + } + + req := httptest.NewRequest("GET", "/protected", nil) + rw := httptest.NewRecorder() + + middleware.ServeHTTP(rw, req) + + authHandler := middleware.authHandler.(*mockAuthHandler) + if !authHandler.wasInitiateCalled() { + t.Error("Expected InitiateAuthentication to be called after failed token refresh") + } +} + +func TestBuildFullURL(t *testing.T) { + tests := []struct { + scheme string + host string + path string + expected string + }{ + {"https", "example.com", "/callback", "https://example.com/callback"}, + {"http", "localhost:8080", "/auth", "http://localhost:8080/auth"}, + {"https", "api.example.com", "/oauth/callback", "https://api.example.com/oauth/callback"}, + } + + for _, tt := range tests { + t.Run(fmt.Sprintf("%s_%s_%s", tt.scheme, tt.host, tt.path), func(t *testing.T) { + result := buildFullURL(tt.scheme, tt.host, tt.path) + if result != tt.expected { + t.Errorf("Expected %s, got %s", tt.expected, result) + } + }) + } +} + +func TestConcurrentRequests(t *testing.T) { + middleware := createTestMiddleware() + sessionData := newMockSessionData() + sessionData.setEmail("test@example.com") + sessionData.setAccessToken("valid.jwt.token") + sessionManager := middleware.sessionManager.(*mockSessionManager) + sessionManager.sessionData = sessionData + + const numGoroutines = 10 + const numRequests = 5 + + var wg sync.WaitGroup + errors := make(chan error, numGoroutines*numRequests) + + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + for j := 0; j < numRequests; j++ { + req := httptest.NewRequest("GET", fmt.Sprintf("/test-%d-%d", id, j), nil) + rw := httptest.NewRecorder() + + middleware.ServeHTTP(rw, req) + + if rw.Code != http.StatusOK { + errors <- fmt.Errorf("goroutine %d request %d: expected status 200, got %d", id, j, rw.Code) + } + } + }(i) + } + + wg.Wait() + close(errors) + + var errorCount int + for err := range errors { + t.Error(err) + errorCount++ + } + + if errorCount > 0 { + t.Errorf("Found %d errors in concurrent request test", errorCount) + } +} + +func TestServeHTTP_TokenVerificationFailure(t *testing.T) { + middleware := createTestMiddleware() + sessionData := newMockSessionData() + sessionData.setEmail("test@example.com") + sessionData.setAccessToken("invalid.jwt.token") + sessionManager := middleware.sessionManager.(*mockSessionManager) + sessionManager.sessionData = sessionData + + // Set token verifier to return error + tokenVerifier := middleware.tokenVerifier.(*mockTokenVerifier) + tokenVerifier.verifyError = fmt.Errorf("token verification failed") + + req := httptest.NewRequest("GET", "/protected", nil) + rw := httptest.NewRecorder() + + middleware.ServeHTTP(rw, req) + + authHandler := middleware.authHandler.(*mockAuthHandler) + if !authHandler.wasInitiateCalled() { + t.Error("Expected InitiateAuthentication to be called after token verification failure") + } +} + +func TestServeHTTP_HealthCheckBypass(t *testing.T) { + middleware := createTestMiddleware() + + // Set up URL helper to return excluded result for health check + urlHelper := &mockURLHelper{excludedResult: true} + middleware.urlHelper = urlHelper + + req := httptest.NewRequest("GET", "/health", nil) + rw := httptest.NewRecorder() + + middleware.ServeHTTP(rw, req) + + nextHandler := middleware.next.(*mockHandler) + if !nextHandler.wasServeHTTPCalled() { + t.Error("Expected next handler to be called for health check") + } +} + +func TestServeHTTP_RequestCancellation(t *testing.T) { + logger := &mockLogger{} + nextHandler := &mockHandler{} + sessionManager := &mockSessionManager{sessionData: newMockSessionData()} + + // Create initComplete channel but don't close it to simulate slow initialization + initComplete := make(chan struct{}) + var goroutineWG sync.WaitGroup + + middleware := NewAuthMiddleware( + logger, + nextHandler, + sessionManager, + &mockAuthHandler{}, + &mockOAuthHandler{}, + &mockURLHelper{}, + &mockTokenVerifier{}, + func(string) (map[string]interface{}, error) { return nil, nil }, + func(string) ([]string, []string, error) { return nil, nil, nil }, + func(rw http.ResponseWriter, req *http.Request, message string, code int) { + http.Error(rw, message, code) + }, + func(http.ResponseWriter, *http.Request, SessionData) bool { return false }, + func(SessionData) (bool, bool, bool) { return false, false, false }, + func(string) bool { return true }, + func(*http.Request) bool { return false }, + func(SessionData) bool { return false }, + func(http.ResponseWriter, *http.Request) {}, + make(map[string]struct{}), + make(map[string]struct{}), + "/auth/callback", + "/auth/logout", + 30*time.Second, + initComplete, + "https://provider.example.com", + "https://provider.example.com", + &goroutineWG, + func() {}, + func(string) {}, + ) + + // Create context that gets cancelled + ctx, cancel := context.WithCancel(context.Background()) + req := httptest.NewRequest("GET", "/test", nil).WithContext(ctx) + rw := httptest.NewRecorder() + + // Cancel context immediately + cancel() + + middleware.ServeHTTP(rw, req) + + if rw.Code != http.StatusRequestTimeout { + t.Errorf("Expected status %d for cancelled request, got %d", http.StatusRequestTimeout, rw.Code) + } +} + +// Benchmark tests +func BenchmarkServeHTTP_AuthenticatedUser(b *testing.B) { + middleware := createTestMiddleware() + sessionData := newMockSessionData() + sessionData.setEmail("test@example.com") + sessionData.setAccessToken("valid.jwt.token") + sessionManager := middleware.sessionManager.(*mockSessionManager) + sessionManager.sessionData = sessionData + + b.ResetTimer() + for i := 0; i < b.N; i++ { + req := httptest.NewRequest("GET", "/protected", nil) + rw := httptest.NewRecorder() + middleware.ServeHTTP(rw, req) + } +} + +func BenchmarkServeHTTP_ExcludedURL(b *testing.B) { + middleware := createTestMiddleware() + urlHelper := &mockURLHelper{excludedResult: true} + middleware.urlHelper = urlHelper + + b.ResetTimer() + for i := 0; i < b.N; i++ { + req := httptest.NewRequest("GET", "/excluded", nil) + rw := httptest.NewRecorder() + middleware.ServeHTTP(rw, req) + } +} diff --git a/recovery/error_handler.go b/recovery/error_handler.go new file mode 100644 index 0000000..493c700 --- /dev/null +++ b/recovery/error_handler.go @@ -0,0 +1,258 @@ +// Package recovery provides error recovery and resilience mechanisms +package recovery + +import ( + "context" + "sync" + "sync/atomic" + "time" +) + +// ErrorRecoveryMechanism defines the interface for error recovery strategies. +// It provides a common contract for implementing various resilience patterns +// (circuit breaker, retry, graceful degradation) to handle transient failures +// and protect downstream services from cascading failures. +type ErrorRecoveryMechanism interface { + // ExecuteWithContext executes a function with error recovery mechanisms + ExecuteWithContext(ctx context.Context, fn func() error) error + // GetMetrics returns metrics about the recovery mechanism's performance + GetMetrics() map[string]interface{} + // Reset resets the mechanism to its initial state + Reset() + // IsAvailable returns whether the mechanism is available for requests + IsAvailable() bool +} + +// Logger interface for dependency injection +type Logger interface { + Infof(format string, args ...interface{}) + Errorf(format string, args ...interface{}) + Debugf(format string, args ...interface{}) +} + +// BaseRecoveryMechanism provides common functionality and metrics tracking +// for all error recovery mechanisms. It handles request/failure/success counting, +// timing information, and logging capabilities for derived recovery mechanisms. +type BaseRecoveryMechanism struct { + // startTime tracks when the mechanism was created + startTime time.Time + // lastFailureTime records the most recent failure timestamp + lastFailureTime time.Time + // lastSuccessTime records the most recent success timestamp + lastSuccessTime time.Time + // logger for debugging and monitoring + logger Logger + // name identifies this recovery mechanism instance + name string + // totalRequests counts all requests processed + totalRequests int64 + // totalFailures counts failed requests + totalFailures int64 + // totalSuccesses counts successful requests + totalSuccesses int64 + // mutex protects shared state access + mutex sync.RWMutex +} + +// NewBaseRecoveryMechanism creates a new base recovery mechanism with the given name and logger. +// This serves as the foundation for specific recovery mechanism implementations. +func NewBaseRecoveryMechanism(name string, logger Logger) *BaseRecoveryMechanism { + if logger == nil { + logger = NewNoOpLogger() + } + + return &BaseRecoveryMechanism{ + name: name, + logger: logger, + startTime: time.Now(), + } +} + +// RecordRequest increments the total request counter. +// This method is thread-safe using atomic operations. +func (b *BaseRecoveryMechanism) RecordRequest() { + atomic.AddInt64(&b.totalRequests, 1) +} + +// RecordSuccess increments the success counter and updates the last success timestamp. +// This method is thread-safe using atomic operations for counters +// and mutex protection for timestamp updates. +func (b *BaseRecoveryMechanism) RecordSuccess() { + atomic.AddInt64(&b.totalSuccesses, 1) + + b.mutex.Lock() + defer b.mutex.Unlock() + b.lastSuccessTime = time.Now() +} + +// RecordFailure increments the failure counter and updates the last failure timestamp. +// This method is thread-safe using atomic operations for counters +// and mutex protection for timestamp updates. +func (b *BaseRecoveryMechanism) RecordFailure() { + atomic.AddInt64(&b.totalFailures, 1) + + b.mutex.Lock() + defer b.mutex.Unlock() + b.lastFailureTime = time.Now() +} + +// GetBaseMetrics returns basic metrics collected by the base recovery mechanism. +// This includes request counts, success/failure rates, and timing information. +func (b *BaseRecoveryMechanism) GetBaseMetrics() map[string]interface{} { + b.mutex.RLock() + defer b.mutex.RUnlock() + + totalReqs := atomic.LoadInt64(&b.totalRequests) + totalSucc := atomic.LoadInt64(&b.totalSuccesses) + totalFail := atomic.LoadInt64(&b.totalFailures) + + metrics := map[string]interface{}{ + "name": b.name, + "total_requests": totalReqs, + "total_successes": totalSucc, + "total_failures": totalFail, + "start_time": b.startTime, + } + + if totalReqs > 0 { + metrics["success_rate"] = float64(totalSucc) / float64(totalReqs) + metrics["failure_rate"] = float64(totalFail) / float64(totalReqs) + } + + if !b.lastSuccessTime.IsZero() { + metrics["last_success_time"] = b.lastSuccessTime + metrics["time_since_last_success"] = time.Since(b.lastSuccessTime) + } + + if !b.lastFailureTime.IsZero() { + metrics["last_failure_time"] = b.lastFailureTime + metrics["time_since_last_failure"] = time.Since(b.lastFailureTime) + } + + metrics["uptime"] = time.Since(b.startTime) + + return metrics +} + +// LogInfo logs an info message if a logger is available +func (b *BaseRecoveryMechanism) LogInfo(format string, args ...interface{}) { + if b.logger != nil { + b.logger.Infof(format, args...) + } +} + +// LogError logs an error message if a logger is available +func (b *BaseRecoveryMechanism) LogError(format string, args ...interface{}) { + if b.logger != nil { + b.logger.Errorf(format, args...) + } +} + +// LogDebug logs a debug message if a logger is available +func (b *BaseRecoveryMechanism) LogDebug(format string, args ...interface{}) { + if b.logger != nil { + b.logger.Debugf(format, args...) + } +} + +// ErrorHandler provides centralized error handling and recovery coordination +type ErrorHandler struct { + mechanisms []ErrorRecoveryMechanism + logger Logger + mutex sync.RWMutex +} + +// NewErrorHandler creates a new error handler with the given mechanisms +func NewErrorHandler(logger Logger, mechanisms ...ErrorRecoveryMechanism) *ErrorHandler { + return &ErrorHandler{ + mechanisms: mechanisms, + logger: logger, + } +} + +// AddMechanism adds a recovery mechanism to the handler +func (eh *ErrorHandler) AddMechanism(mechanism ErrorRecoveryMechanism) { + eh.mutex.Lock() + defer eh.mutex.Unlock() + eh.mechanisms = append(eh.mechanisms, mechanism) +} + +// ExecuteWithRecovery executes a function with all configured recovery mechanisms +func (eh *ErrorHandler) ExecuteWithRecovery(ctx context.Context, fn func() error) error { + eh.mutex.RLock() + mechanisms := make([]ErrorRecoveryMechanism, len(eh.mechanisms)) + copy(mechanisms, eh.mechanisms) + eh.mutex.RUnlock() + + // If no mechanisms are configured, execute directly + if len(mechanisms) == 0 { + return fn() + } + + // Chain the mechanisms - each wraps the next + var wrappedFn func() error = fn + for i := len(mechanisms) - 1; i >= 0; i-- { + mechanism := mechanisms[i] + currentFn := wrappedFn + wrappedFn = func() error { + return mechanism.ExecuteWithContext(ctx, currentFn) + } + } + + return wrappedFn() +} + +// GetAllMetrics returns metrics from all configured mechanisms +func (eh *ErrorHandler) GetAllMetrics() map[string]interface{} { + eh.mutex.RLock() + defer eh.mutex.RUnlock() + + allMetrics := make(map[string]interface{}) + for i, mechanism := range eh.mechanisms { + mechanismKey := "mechanism_" + string(rune(i)) + allMetrics[mechanismKey] = mechanism.GetMetrics() + } + + return allMetrics +} + +// ResetAll resets all configured mechanisms +func (eh *ErrorHandler) ResetAll() { + eh.mutex.RLock() + defer eh.mutex.RUnlock() + + for _, mechanism := range eh.mechanisms { + mechanism.Reset() + } +} + +// IsHealthy returns true if all mechanisms are available +func (eh *ErrorHandler) IsHealthy() bool { + eh.mutex.RLock() + defer eh.mutex.RUnlock() + + for _, mechanism := range eh.mechanisms { + if !mechanism.IsAvailable() { + return false + } + } + + return true +} + +// NoOpLogger provides a logger that does nothing +type NoOpLogger struct{} + +// NewNoOpLogger creates a new no-op logger +func NewNoOpLogger() *NoOpLogger { + return &NoOpLogger{} +} + +// Infof does nothing +func (l *NoOpLogger) Infof(format string, args ...interface{}) {} + +// Errorf does nothing +func (l *NoOpLogger) Errorf(format string, args ...interface{}) {} + +// Debugf does nothing +func (l *NoOpLogger) Debugf(format string, args ...interface{}) {} diff --git a/recovery/error_handler_test.go b/recovery/error_handler_test.go new file mode 100644 index 0000000..f0a2d95 --- /dev/null +++ b/recovery/error_handler_test.go @@ -0,0 +1,827 @@ +package recovery + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" +) + +// Mock implementations for testing +type mockLogger struct { + infoLogs []string + errorLogs []string + debugLogs []string + mu sync.RWMutex +} + +func (m *mockLogger) Infof(format string, args ...interface{}) { + m.mu.Lock() + defer m.mu.Unlock() + m.infoLogs = append(m.infoLogs, fmt.Sprintf(format, args...)) +} + +func (m *mockLogger) Errorf(format string, args ...interface{}) { + m.mu.Lock() + defer m.mu.Unlock() + m.errorLogs = append(m.errorLogs, fmt.Sprintf(format, args...)) +} + +func (m *mockLogger) Debugf(format string, args ...interface{}) { + m.mu.Lock() + defer m.mu.Unlock() + m.debugLogs = append(m.debugLogs, fmt.Sprintf(format, args...)) +} + +func (m *mockLogger) getInfoLogs() []string { + m.mu.RLock() + defer m.mu.RUnlock() + result := make([]string, len(m.infoLogs)) + copy(result, m.infoLogs) + return result +} + +func (m *mockLogger) getErrorLogs() []string { + m.mu.RLock() + defer m.mu.RUnlock() + result := make([]string, len(m.errorLogs)) + copy(result, m.errorLogs) + return result +} + +func (m *mockLogger) getDebugLogs() []string { + m.mu.RLock() + defer m.mu.RUnlock() + result := make([]string, len(m.debugLogs)) + copy(result, m.debugLogs) + return result +} + +//lint:ignore U1000 May be needed for future test isolation +func (m *mockLogger) reset() { + m.mu.Lock() + defer m.mu.Unlock() + m.infoLogs = nil + m.errorLogs = nil + m.debugLogs = nil +} + +type mockRecoveryMechanism struct { + name string + executeError error + isAvailable bool + resetCalled bool + executeCallCount int64 + metrics map[string]interface{} + executeFunc func(ctx context.Context, fn func() error) error + mu sync.RWMutex +} + +func newMockRecoveryMechanism(name string) *mockRecoveryMechanism { + return &mockRecoveryMechanism{ + name: name, + isAvailable: true, + metrics: make(map[string]interface{}), + } +} + +func (m *mockRecoveryMechanism) ExecuteWithContext(ctx context.Context, fn func() error) error { + if m.executeFunc != nil { + return m.executeFunc(ctx, fn) + } + + atomic.AddInt64(&m.executeCallCount, 1) + + m.mu.RLock() + err := m.executeError + m.mu.RUnlock() + + if err != nil { + return err + } + + return fn() +} + +func (m *mockRecoveryMechanism) GetMetrics() map[string]interface{} { + m.mu.RLock() + defer m.mu.RUnlock() + + result := make(map[string]interface{}) + for k, v := range m.metrics { + result[k] = v + } + result["name"] = m.name + result["execute_call_count"] = atomic.LoadInt64(&m.executeCallCount) + return result +} + +func (m *mockRecoveryMechanism) Reset() { + m.mu.Lock() + defer m.mu.Unlock() + m.resetCalled = true + atomic.StoreInt64(&m.executeCallCount, 0) +} + +func (m *mockRecoveryMechanism) IsAvailable() bool { + m.mu.RLock() + defer m.mu.RUnlock() + return m.isAvailable +} + +func (m *mockRecoveryMechanism) setExecuteError(err error) { + m.mu.Lock() + defer m.mu.Unlock() + m.executeError = err +} + +func (m *mockRecoveryMechanism) setAvailable(available bool) { + m.mu.Lock() + defer m.mu.Unlock() + m.isAvailable = available +} + +func (m *mockRecoveryMechanism) wasResetCalled() bool { + m.mu.RLock() + defer m.mu.RUnlock() + return m.resetCalled +} + +func (m *mockRecoveryMechanism) getExecuteCallCount() int64 { + return atomic.LoadInt64(&m.executeCallCount) +} + +func TestNewBaseRecoveryMechanism(t *testing.T) { + tests := []struct { + name string + loggerName string + logger Logger + }{ + { + name: "with mock logger", + loggerName: "test-mechanism", + logger: &mockLogger{}, + }, + { + name: "with nil logger", + loggerName: "test-mechanism-nil", + logger: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mechanism := NewBaseRecoveryMechanism(tt.loggerName, tt.logger) + + if mechanism == nil { + t.Fatal("NewBaseRecoveryMechanism returned nil") + } + + if mechanism.name != tt.loggerName { + t.Errorf("Expected name %s, got %s", tt.loggerName, mechanism.name) + } + + if tt.logger == nil && mechanism.logger == nil { + t.Error("Expected NoOpLogger to be created when logger is nil") + } + + if mechanism.startTime.IsZero() { + t.Error("Expected startTime to be set") + } + }) + } +} + +func TestBaseRecoveryMechanism_RecordMetrics(t *testing.T) { + logger := &mockLogger{} + mechanism := NewBaseRecoveryMechanism("test", logger) + + // Test initial state + if atomic.LoadInt64(&mechanism.totalRequests) != 0 { + t.Error("Expected initial totalRequests to be 0") + } + if atomic.LoadInt64(&mechanism.totalSuccesses) != 0 { + t.Error("Expected initial totalSuccesses to be 0") + } + if atomic.LoadInt64(&mechanism.totalFailures) != 0 { + t.Error("Expected initial totalFailures to be 0") + } + + // Record some metrics + mechanism.RecordRequest() + mechanism.RecordSuccess() + + if atomic.LoadInt64(&mechanism.totalRequests) != 1 { + t.Error("Expected totalRequests to be 1") + } + if atomic.LoadInt64(&mechanism.totalSuccesses) != 1 { + t.Error("Expected totalSuccesses to be 1") + } + + mechanism.RecordRequest() + mechanism.RecordFailure() + + if atomic.LoadInt64(&mechanism.totalRequests) != 2 { + t.Error("Expected totalRequests to be 2") + } + if atomic.LoadInt64(&mechanism.totalFailures) != 1 { + t.Error("Expected totalFailures to be 1") + } + + // Check timestamps are set + if mechanism.lastSuccessTime.IsZero() { + t.Error("Expected lastSuccessTime to be set") + } + if mechanism.lastFailureTime.IsZero() { + t.Error("Expected lastFailureTime to be set") + } +} + +func TestBaseRecoveryMechanism_GetBaseMetrics(t *testing.T) { + logger := &mockLogger{} + mechanism := NewBaseRecoveryMechanism("test-metrics", logger) + + // Record some activity + mechanism.RecordRequest() + mechanism.RecordRequest() + mechanism.RecordSuccess() + mechanism.RecordFailure() + + metrics := mechanism.GetBaseMetrics() + + expectedFields := []string{ + "name", "total_requests", "total_successes", "total_failures", + "start_time", "success_rate", "failure_rate", "uptime", + "last_success_time", "time_since_last_success", + "last_failure_time", "time_since_last_failure", + } + + for _, field := range expectedFields { + if _, exists := metrics[field]; !exists { + t.Errorf("Expected metric field %s to exist", field) + } + } + + // Check specific values + if metrics["name"] != "test-metrics" { + t.Errorf("Expected name to be 'test-metrics', got %v", metrics["name"]) + } + + if metrics["total_requests"] != int64(2) { + t.Errorf("Expected total_requests to be 2, got %v", metrics["total_requests"]) + } + + if metrics["total_successes"] != int64(1) { + t.Errorf("Expected total_successes to be 1, got %v", metrics["total_successes"]) + } + + if metrics["total_failures"] != int64(1) { + t.Errorf("Expected total_failures to be 1, got %v", metrics["total_failures"]) + } + + if metrics["success_rate"] != float64(0.5) { + t.Errorf("Expected success_rate to be 0.5, got %v", metrics["success_rate"]) + } + + if metrics["failure_rate"] != float64(0.5) { + t.Errorf("Expected failure_rate to be 0.5, got %v", metrics["failure_rate"]) + } +} + +func TestBaseRecoveryMechanism_GetBaseMetrics_NoRequests(t *testing.T) { + logger := &mockLogger{} + mechanism := NewBaseRecoveryMechanism("test-empty", logger) + + metrics := mechanism.GetBaseMetrics() + + // Should not have success_rate or failure_rate when no requests + if _, exists := metrics["success_rate"]; exists { + t.Error("Expected success_rate to not exist when no requests") + } + if _, exists := metrics["failure_rate"]; exists { + t.Error("Expected failure_rate to not exist when no requests") + } + + // Should not have timestamp fields when no successes/failures + if _, exists := metrics["last_success_time"]; exists { + t.Error("Expected last_success_time to not exist when no successes") + } + if _, exists := metrics["last_failure_time"]; exists { + t.Error("Expected last_failure_time to not exist when no failures") + } +} + +func TestBaseRecoveryMechanism_Logging(t *testing.T) { + logger := &mockLogger{} + mechanism := NewBaseRecoveryMechanism("test-logging", logger) + + mechanism.LogInfo("test info: %s", "value") + mechanism.LogError("test error: %d", 42) + mechanism.LogDebug("test debug: %v", true) + + infoLogs := logger.getInfoLogs() + if len(infoLogs) != 1 || infoLogs[0] != "test info: value" { + t.Errorf("Expected info log 'test info: value', got %v", infoLogs) + } + + errorLogs := logger.getErrorLogs() + if len(errorLogs) != 1 || errorLogs[0] != "test error: 42" { + t.Errorf("Expected error log 'test error: 42', got %v", errorLogs) + } + + debugLogs := logger.getDebugLogs() + if len(debugLogs) != 1 || debugLogs[0] != "test debug: true" { + t.Errorf("Expected debug log 'test debug: true', got %v", debugLogs) + } +} + +func TestBaseRecoveryMechanism_LoggingWithNilLogger(t *testing.T) { + mechanism := NewBaseRecoveryMechanism("test-nil-logger", nil) + + // Should not panic when logger is nil (uses NoOpLogger) + mechanism.LogInfo("test info") + mechanism.LogError("test error") + mechanism.LogDebug("test debug") +} + +func TestNewErrorHandler(t *testing.T) { + logger := &mockLogger{} + mechanism1 := newMockRecoveryMechanism("mech1") + mechanism2 := newMockRecoveryMechanism("mech2") + + handler := NewErrorHandler(logger, mechanism1, mechanism2) + + if handler == nil { + t.Fatal("NewErrorHandler returned nil") + } + + if handler.logger != logger { + t.Error("Expected logger to be set") + } + + if len(handler.mechanisms) != 2 { + t.Errorf("Expected 2 mechanisms, got %d", len(handler.mechanisms)) + } +} + +func TestErrorHandler_AddMechanism(t *testing.T) { + logger := &mockLogger{} + handler := NewErrorHandler(logger) + + if len(handler.mechanisms) != 0 { + t.Errorf("Expected 0 initial mechanisms, got %d", len(handler.mechanisms)) + } + + mechanism := newMockRecoveryMechanism("test") + handler.AddMechanism(mechanism) + + if len(handler.mechanisms) != 1 { + t.Errorf("Expected 1 mechanism after adding, got %d", len(handler.mechanisms)) + } +} + +func TestErrorHandler_ExecuteWithRecovery_NoMechanisms(t *testing.T) { + logger := &mockLogger{} + handler := NewErrorHandler(logger) + + callCount := 0 + testFunc := func() error { + callCount++ + return nil + } + + ctx := context.Background() + err := handler.ExecuteWithRecovery(ctx, testFunc) + + if err != nil { + t.Errorf("Expected no error, got %v", err) + } + + if callCount != 1 { + t.Errorf("Expected function to be called once, got %d", callCount) + } +} + +func TestErrorHandler_ExecuteWithRecovery_WithMechanisms(t *testing.T) { + logger := &mockLogger{} + mechanism1 := newMockRecoveryMechanism("mech1") + mechanism2 := newMockRecoveryMechanism("mech2") + handler := NewErrorHandler(logger, mechanism1, mechanism2) + + callCount := 0 + testFunc := func() error { + callCount++ + return nil + } + + ctx := context.Background() + err := handler.ExecuteWithRecovery(ctx, testFunc) + + if err != nil { + t.Errorf("Expected no error, got %v", err) + } + + if callCount != 1 { + t.Errorf("Expected function to be called once, got %d", callCount) + } + + // Both mechanisms should have been called + if mechanism1.getExecuteCallCount() != 1 { + t.Errorf("Expected mechanism1 to be called once, got %d", mechanism1.getExecuteCallCount()) + } + if mechanism2.getExecuteCallCount() != 1 { + t.Errorf("Expected mechanism2 to be called once, got %d", mechanism2.getExecuteCallCount()) + } +} + +func TestErrorHandler_ExecuteWithRecovery_MechanismError(t *testing.T) { + logger := &mockLogger{} + mechanism := newMockRecoveryMechanism("failing-mech") + mechanism.setExecuteError(fmt.Errorf("mechanism failed")) + handler := NewErrorHandler(logger, mechanism) + + callCount := 0 + testFunc := func() error { + callCount++ + return nil + } + + ctx := context.Background() + err := handler.ExecuteWithRecovery(ctx, testFunc) + + if err == nil { + t.Error("Expected error from mechanism") + } + + if err.Error() != "mechanism failed" { + t.Errorf("Expected 'mechanism failed', got %v", err) + } + + // Function should not be called if mechanism fails + if callCount != 0 { + t.Errorf("Expected function not to be called, got %d calls", callCount) + } +} + +func TestErrorHandler_ExecuteWithRecovery_FunctionError(t *testing.T) { + logger := &mockLogger{} + mechanism := newMockRecoveryMechanism("mech") + handler := NewErrorHandler(logger, mechanism) + + testFunc := func() error { + return fmt.Errorf("function failed") + } + + ctx := context.Background() + err := handler.ExecuteWithRecovery(ctx, testFunc) + + if err == nil { + t.Error("Expected error from function") + } + + if err.Error() != "function failed" { + t.Errorf("Expected 'function failed', got %v", err) + } +} + +func TestErrorHandler_GetAllMetrics(t *testing.T) { + logger := &mockLogger{} + mechanism1 := newMockRecoveryMechanism("mech1") + mechanism2 := newMockRecoveryMechanism("mech2") + handler := NewErrorHandler(logger, mechanism1, mechanism2) + + // Add some custom metrics + mechanism1.metrics["custom1"] = "value1" + mechanism2.metrics["custom2"] = "value2" + + metrics := handler.GetAllMetrics() + + if len(metrics) != 2 { + t.Errorf("Expected 2 mechanism metrics, got %d", len(metrics)) + } + + mechanism0Key := "mechanism_" + string(rune(0)) + mechanism0Metrics, exists := metrics[mechanism0Key] + if !exists { + t.Errorf("Expected %s metrics to exist", mechanism0Key) + } else { + mechanismMap, ok := mechanism0Metrics.(map[string]interface{}) + if !ok { + t.Error("Expected mechanism metrics to be a map") + } else { + if mechanismMap["name"] != "mech1" { + t.Errorf("Expected name 'mech1', got %v", mechanismMap["name"]) + } + if mechanismMap["custom1"] != "value1" { + t.Errorf("Expected custom1 'value1', got %v", mechanismMap["custom1"]) + } + } + } + + mechanism1Key := "mechanism_" + string(rune(1)) + mechanism1Metrics, exists := metrics[mechanism1Key] + if !exists { + t.Errorf("Expected %s metrics to exist", mechanism1Key) + } else { + mechanismMap, ok := mechanism1Metrics.(map[string]interface{}) + if !ok { + t.Error("Expected mechanism metrics to be a map") + } else { + if mechanismMap["name"] != "mech2" { + t.Errorf("Expected name 'mech2', got %v", mechanismMap["name"]) + } + if mechanismMap["custom2"] != "value2" { + t.Errorf("Expected custom2 'value2', got %v", mechanismMap["custom2"]) + } + } + } +} + +func TestErrorHandler_ResetAll(t *testing.T) { + logger := &mockLogger{} + mechanism1 := newMockRecoveryMechanism("mech1") + mechanism2 := newMockRecoveryMechanism("mech2") + handler := NewErrorHandler(logger, mechanism1, mechanism2) + + handler.ResetAll() + + if !mechanism1.wasResetCalled() { + t.Error("Expected mechanism1 to be reset") + } + if !mechanism2.wasResetCalled() { + t.Error("Expected mechanism2 to be reset") + } +} + +func TestErrorHandler_IsHealthy(t *testing.T) { + logger := &mockLogger{} + mechanism1 := newMockRecoveryMechanism("mech1") + mechanism2 := newMockRecoveryMechanism("mech2") + handler := NewErrorHandler(logger, mechanism1, mechanism2) + + // Initially both are available + if !handler.IsHealthy() { + t.Error("Expected handler to be healthy when all mechanisms are available") + } + + // Make one mechanism unavailable + mechanism1.setAvailable(false) + if handler.IsHealthy() { + t.Error("Expected handler to be unhealthy when one mechanism is unavailable") + } + + // Make both unavailable + mechanism2.setAvailable(false) + if handler.IsHealthy() { + t.Error("Expected handler to be unhealthy when all mechanisms are unavailable") + } + + // Make both available again + mechanism1.setAvailable(true) + mechanism2.setAvailable(true) + if !handler.IsHealthy() { + t.Error("Expected handler to be healthy when all mechanisms are available again") + } +} + +func TestErrorHandler_IsHealthy_NoMechanisms(t *testing.T) { + logger := &mockLogger{} + handler := NewErrorHandler(logger) + + // Should be healthy with no mechanisms + if !handler.IsHealthy() { + t.Error("Expected handler to be healthy with no mechanisms") + } +} + +func TestNoOpLogger(t *testing.T) { + logger := NewNoOpLogger() + + if logger == nil { + t.Fatal("NewNoOpLogger returned nil") + } + + // Should not panic + logger.Infof("test info: %s", "value") + logger.Errorf("test error: %d", 42) + logger.Debugf("test debug: %v", true) +} + +func TestConcurrentAccess(t *testing.T) { + logger := &mockLogger{} + baseMechanism := NewBaseRecoveryMechanism("concurrent-test", logger) + + const numGoroutines = 10 + const numOperations = 100 + + var wg sync.WaitGroup + errors := make(chan error, numGoroutines) + + // Test concurrent metric recording + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + for j := 0; j < numOperations; j++ { + baseMechanism.RecordRequest() + if j%2 == 0 { + baseMechanism.RecordSuccess() + } else { + baseMechanism.RecordFailure() + } + + // Intermittently get metrics to test concurrent reads + if j%10 == 0 { + baseMechanism.GetBaseMetrics() + } + } + }(i) + } + + wg.Wait() + close(errors) + + var errorCount int + for err := range errors { + t.Error(err) + errorCount++ + } + + if errorCount > 0 { + t.Errorf("Found %d errors in concurrent access test", errorCount) + } + + // Verify final counts + expectedRequests := int64(numGoroutines * numOperations) + expectedSuccesses := int64(numGoroutines * numOperations / 2) + expectedFailures := int64(numGoroutines * numOperations / 2) + + if atomic.LoadInt64(&baseMechanism.totalRequests) != expectedRequests { + t.Errorf("Expected %d total requests, got %d", expectedRequests, + atomic.LoadInt64(&baseMechanism.totalRequests)) + } + + if atomic.LoadInt64(&baseMechanism.totalSuccesses) != expectedSuccesses { + t.Errorf("Expected %d total successes, got %d", expectedSuccesses, + atomic.LoadInt64(&baseMechanism.totalSuccesses)) + } + + if atomic.LoadInt64(&baseMechanism.totalFailures) != expectedFailures { + t.Errorf("Expected %d total failures, got %d", expectedFailures, + atomic.LoadInt64(&baseMechanism.totalFailures)) + } +} + +func TestErrorHandler_ConcurrentAccess(t *testing.T) { + logger := &mockLogger{} + mechanism1 := newMockRecoveryMechanism("concurrent-mech1") + mechanism2 := newMockRecoveryMechanism("concurrent-mech2") + handler := NewErrorHandler(logger, mechanism1, mechanism2) + + const numGoroutines = 5 + const numExecutions = 20 + + var wg sync.WaitGroup + successCount := int64(0) + + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + for j := 0; j < numExecutions; j++ { + err := handler.ExecuteWithRecovery(context.Background(), func() error { + atomic.AddInt64(&successCount, 1) + return nil + }) + if err != nil { + t.Errorf("Unexpected error in goroutine %d: %v", id, err) + } + + // Add some mechanisms concurrently + if j == 0 { + newMech := newMockRecoveryMechanism(fmt.Sprintf("dynamic-mech-%d", id)) + handler.AddMechanism(newMech) + } + + // Get metrics concurrently + if j%5 == 0 { + handler.GetAllMetrics() + } + } + }(i) + } + + wg.Wait() + + expectedSuccesses := int64(numGoroutines * numExecutions) + if atomic.LoadInt64(&successCount) != expectedSuccesses { + t.Errorf("Expected %d successes, got %d", expectedSuccesses, + atomic.LoadInt64(&successCount)) + } + + // Should have original 2 mechanisms plus 5 added dynamically + if len(handler.mechanisms) != 7 { + t.Errorf("Expected 7 mechanisms after concurrent additions, got %d", + len(handler.mechanisms)) + } +} + +func TestErrorHandler_ContextCancellation(t *testing.T) { + logger := &mockLogger{} + mechanism := newMockRecoveryMechanism("context-test") + handler := NewErrorHandler(logger, mechanism) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // Cancel immediately + + executedCount := 0 + testFunc := func() error { + executedCount++ + return nil + } + + // The mechanism implementation should handle context cancellation + // For our mock, we'll make it check context cancellation + mechanism.executeFunc = func(ctx context.Context, fn func() error) error { + select { + case <-ctx.Done(): + return ctx.Err() + default: + atomic.AddInt64(&mechanism.executeCallCount, 1) + return fn() + } + } + + err := handler.ExecuteWithRecovery(ctx, testFunc) + + if err != context.Canceled { + t.Errorf("Expected context.Canceled, got %v", err) + } + + if executedCount != 0 { + t.Errorf("Expected function not to execute when context is cancelled, got %d executions", executedCount) + } +} + +// Benchmark tests +func BenchmarkBaseRecoveryMechanism_RecordRequest(b *testing.B) { + logger := &mockLogger{} + mechanism := NewBaseRecoveryMechanism("benchmark", logger) + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + mechanism.RecordRequest() + } + }) +} + +func BenchmarkBaseRecoveryMechanism_RecordSuccess(b *testing.B) { + logger := &mockLogger{} + mechanism := NewBaseRecoveryMechanism("benchmark", logger) + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + mechanism.RecordSuccess() + } + }) +} + +func BenchmarkBaseRecoveryMechanism_GetBaseMetrics(b *testing.B) { + logger := &mockLogger{} + mechanism := NewBaseRecoveryMechanism("benchmark", logger) + + // Add some data + for i := 0; i < 100; i++ { + mechanism.RecordRequest() + if i%2 == 0 { + mechanism.RecordSuccess() + } else { + mechanism.RecordFailure() + } + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + mechanism.GetBaseMetrics() + } +} + +func BenchmarkErrorHandler_ExecuteWithRecovery(b *testing.B) { + logger := &mockLogger{} + mechanism := newMockRecoveryMechanism("benchmark") + handler := NewErrorHandler(logger, mechanism) + + testFunc := func() error { + return nil + } + + ctx := context.Background() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + handler.ExecuteWithRecovery(ctx, testFunc) + } +} diff --git a/session/chunking/chunk_manager.go b/session/chunking/chunk_manager.go new file mode 100644 index 0000000..7729fcc --- /dev/null +++ b/session/chunking/chunk_manager.go @@ -0,0 +1,406 @@ +// Package chunking provides session chunking functionality for large tokens +package chunking + +import ( + "fmt" + "strings" + "sync" + "time" + + "github.com/gorilla/sessions" +) + +const ( + maxCookieSize = 1200 +) + +// TokenConfig defines validation and storage parameters for different token types. +// It specifies size limits, format requirements, and security constraints to ensure +// tokens can be safely stored in browser cookies while maintaining security. +type TokenConfig struct { + Type string + MinLength int + MaxLength int + MaxChunks int + MaxChunkSize int + AllowOpaqueTokens bool + RequireJWTFormat bool +} + +// Predefined configurations for each token type +var ( + AccessTokenConfig = TokenConfig{ + Type: "access", + MinLength: 5, + MaxLength: 100 * 1024, + MaxChunks: 25, + MaxChunkSize: maxCookieSize, + AllowOpaqueTokens: true, + RequireJWTFormat: false, + } + + RefreshTokenConfig = TokenConfig{ + Type: "refresh", + MinLength: 5, + MaxLength: 50 * 1024, + MaxChunks: 15, + MaxChunkSize: maxCookieSize, + AllowOpaqueTokens: true, + RequireJWTFormat: false, + } + + IDTokenConfig = TokenConfig{ + Type: "id", + MinLength: 5, + MaxLength: 75 * 1024, + MaxChunks: 20, + MaxChunkSize: maxCookieSize, + AllowOpaqueTokens: false, + RequireJWTFormat: true, + } +) + +// TokenRetrievalResult represents the outcome of a token retrieval operation. +// It contains either the successfully retrieved token or an error describing +// what went wrong during retrieval. +type TokenRetrievalResult struct { + Error error + Token string +} + +// SessionEntry represents a session with expiration tracking +type SessionEntry struct { + Session *sessions.Session + ExpiresAt time.Time + LastUsed time.Time +} + +// Logger interface for dependency injection +type Logger interface { + Debug(msg string) + Debugf(format string, args ...interface{}) + Error(msg string) + Errorf(format string, args ...interface{}) +} + +// ChunkManager handles the complex logic of storing and retrieving large tokens +// across multiple HTTP cookies. It provides comprehensive validation, security checks, +// and error handling to ensure data integrity and prevent security vulnerabilities +// throughout the process. +type ChunkManager struct { + logger Logger + mutex *sync.RWMutex + // sessionMap provides bounded session storage to prevent memory leaks + sessionMap map[string]*SessionEntry + maxSessions int + sessionTTL time.Duration + lastCleanup time.Time +} + +// NewChunkManager creates a new ChunkManager instance with proper initialization. +// It sets up logging and synchronization primitives for safe concurrent access. +func NewChunkManager(logger Logger) *ChunkManager { + if logger == nil { + logger = NewNoOpLogger() + } + + return &ChunkManager{ + logger: logger, + mutex: &sync.RWMutex{}, + sessionMap: make(map[string]*SessionEntry), + maxSessions: 1000, // Reasonable limit to prevent memory leaks + sessionTTL: 24 * time.Hour, + lastCleanup: time.Now(), + } +} + +// GetToken retrieves a token from either a single cookie or multiple chunk cookies. +// It handles both compressed and uncompressed tokens and performs comprehensive +// validation throughout the retrieval process. +func (cm *ChunkManager) GetToken( + mainSession *sessions.Session, + chunks map[int]*sessions.Session, + config TokenConfig, + compressor TokenCompressor, +) TokenRetrievalResult { + + // Try to get token from main session first + if mainSession != nil { + if tokenValue, ok := mainSession.Values[config.Type+"_token"].(string); ok && tokenValue != "" { + cm.logger.Debugf("Found %s token in main session", config.Type) + + // Check if token is compressed + decompressed := compressor.DecompressToken(tokenValue) + if decompressed != tokenValue { + cm.logger.Debugf("Decompressed %s token", config.Type) + return cm.processSingleToken(decompressed, true, config) + } + + return cm.processSingleToken(tokenValue, false, config) + } + } + + // If not in main session, try chunks + if len(chunks) == 0 { + return TokenRetrievalResult{ + Error: nil, + Token: "", + } + } + + cm.logger.Debugf("Found %d chunks for %s token, processing", len(chunks), config.Type) + return cm.processChunkedToken(chunks, config, compressor) +} + +// processSingleToken validates and processes a single token +func (cm *ChunkManager) processSingleToken(token string, compressed bool, config TokenConfig) TokenRetrievalResult { + if compressed { + cm.logger.Debugf("Processing compressed %s token (length: %d)", config.Type, len(token)) + } else { + cm.logger.Debugf("Processing single %s token (length: %d)", config.Type, len(token)) + } + + return cm.validateToken(token, config) +} + +// validateToken performs comprehensive validation on a token +func (cm *ChunkManager) validateToken(token string, config TokenConfig) TokenRetrievalResult { + if token == "" { + return TokenRetrievalResult{Error: nil, Token: ""} + } + + validator := NewTokenValidator() + + // Basic validation + if err := validator.ValidateTokenSize(token, config); err != nil { + cm.logger.Errorf("Token size validation failed for %s: %v", config.Type, err) + return TokenRetrievalResult{Error: err, Token: ""} + } + + // Format validation + if config.RequireJWTFormat { + if err := validator.ValidateJWTFormat(token, config.Type); err != nil { + cm.logger.Errorf("JWT format validation failed for %s: %v", config.Type, err) + return TokenRetrievalResult{Error: err, Token: ""} + } + } else if !config.AllowOpaqueTokens { + if err := validator.ValidateJWTFormat(token, config.Type); err != nil { + cm.logger.Errorf("Token format validation failed for %s: %v", config.Type, err) + return TokenRetrievalResult{Error: err, Token: ""} + } + } + + // Content validation + if err := validator.ValidateTokenContent(token, config); err != nil { + cm.logger.Errorf("Token content validation failed for %s: %v", config.Type, err) + return TokenRetrievalResult{Error: err, Token: ""} + } + + cm.logger.Debugf("Successfully validated %s token", config.Type) + return TokenRetrievalResult{Error: nil, Token: token} +} + +// processChunkedToken reconstructs a token from multiple chunks +func (cm *ChunkManager) processChunkedToken(chunks map[int]*sessions.Session, config TokenConfig, compressor TokenCompressor) TokenRetrievalResult { + if len(chunks) > config.MaxChunks { + return TokenRetrievalResult{ + Error: &ChunkError{ + Type: config.Type, + Reason: "too many chunks", + Details: "chunk count exceeds maximum allowed", + }, + Token: "", + } + } + + // Reconstruct token from chunks + reconstructedToken, err := cm.reconstructTokenFromChunks(chunks, config) + if err != nil { + cm.logger.Errorf("Failed to reconstruct %s token from chunks: %v", config.Type, err) + return TokenRetrievalResult{Error: err, Token: ""} + } + + // Try decompression + decompressedToken := compressor.DecompressToken(reconstructedToken) + if decompressedToken != reconstructedToken { + cm.logger.Debugf("Decompressed reconstructed %s token", config.Type) + return cm.validateToken(decompressedToken, config) + } + + return cm.validateToken(reconstructedToken, config) +} + +// reconstructTokenFromChunks reconstructs a token from ordered chunks +func (cm *ChunkManager) reconstructTokenFromChunks(chunks map[int]*sessions.Session, config TokenConfig) (string, error) { + if len(chunks) == 0 { + return "", &ChunkError{ + Type: config.Type, + Reason: "no chunks found", + Details: "no chunk sessions available for reconstruction", + } + } + + // Find the maximum chunk index to determine total chunks + maxIndex := -1 + for index := range chunks { + if index > maxIndex { + maxIndex = index + } + } + + if maxIndex < 0 { + return "", &ChunkError{ + Type: config.Type, + Reason: "invalid chunk indices", + Details: "no valid chunk indices found", + } + } + + // Reconstruct token by concatenating chunks in order + var tokenBuilder strings.Builder + for i := 0; i <= maxIndex; i++ { + chunk, exists := chunks[i] + if !exists || chunk == nil { + return "", &ChunkError{ + Type: config.Type, + Reason: "missing chunk", + Details: fmt.Sprintf("chunk %d is missing", i), + } + } + + chunkValue, ok := chunk.Values["value"].(string) + if !ok || chunkValue == "" { + return "", &ChunkError{ + Type: config.Type, + Reason: "empty chunk", + Details: fmt.Sprintf("chunk %d has no value", i), + } + } + + tokenBuilder.WriteString(chunkValue) + } + + reconstructed := tokenBuilder.String() + if reconstructed == "" { + return "", &ChunkError{ + Type: config.Type, + Reason: "empty reconstructed token", + Details: "all chunks were present but resulted in empty token", + } + } + + cm.logger.Debugf("Successfully reconstructed %s token from %d chunks (length: %d)", + config.Type, len(chunks), len(reconstructed)) + + return reconstructed, nil +} + +// CleanupExpiredSessions removes expired sessions from the session map +func (cm *ChunkManager) CleanupExpiredSessions() { + cm.mutex.Lock() + defer cm.mutex.Unlock() + + now := time.Now() + + // Only cleanup if enough time has passed + if now.Sub(cm.lastCleanup) < time.Hour { + return + } + + cm.lastCleanup = now + cleaned := 0 + + for key, entry := range cm.sessionMap { + if now.After(entry.ExpiresAt) || now.Sub(entry.LastUsed) > cm.sessionTTL { + delete(cm.sessionMap, key) + cleaned++ + } + } + + if cleaned > 0 { + cm.logger.Debugf("Cleaned up %d expired sessions", cleaned) + } +} + +// StoreSession stores a session in the session map with expiration tracking +func (cm *ChunkManager) StoreSession(key string, session *sessions.Session) { + cm.mutex.Lock() + defer cm.mutex.Unlock() + + // Enforce maximum session limit + if len(cm.sessionMap) >= cm.maxSessions { + // Remove oldest entry + var oldestKey string + var oldestTime time.Time = time.Now() + + for k, entry := range cm.sessionMap { + if entry.LastUsed.Before(oldestTime) { + oldestTime = entry.LastUsed + oldestKey = k + } + } + + if oldestKey != "" { + delete(cm.sessionMap, oldestKey) + } + } + + cm.sessionMap[key] = &SessionEntry{ + Session: session, + ExpiresAt: time.Now().Add(cm.sessionTTL), + LastUsed: time.Now(), + } +} + +// GetSession retrieves a session from the session map +func (cm *ChunkManager) GetSession(key string) *sessions.Session { + cm.mutex.Lock() + defer cm.mutex.Unlock() + + entry, exists := cm.sessionMap[key] + if !exists { + return nil + } + + // Update last used time + entry.LastUsed = time.Now() + return entry.Session +} + +// TokenCompressor interface for token compression operations +type TokenCompressor interface { + CompressToken(token string) string + DecompressToken(compressed string) string +} + +// ChunkError represents errors that occur during chunk operations +type ChunkError struct { + Type string + Reason string + Details string +} + +// Error implements the error interface +func (ce *ChunkError) Error() string { + return fmt.Sprintf("%s chunk error: %s - %s", ce.Type, ce.Reason, ce.Details) +} + +// NoOpLogger provides a no-op logger implementation +type NoOpLogger struct{} + +// NewNoOpLogger creates a new no-op logger +func NewNoOpLogger() *NoOpLogger { + return &NoOpLogger{} +} + +// Debug does nothing +func (l *NoOpLogger) Debug(msg string) {} + +// Debugf does nothing +func (l *NoOpLogger) Debugf(format string, args ...interface{}) {} + +// Error does nothing +func (l *NoOpLogger) Error(msg string) {} + +// Errorf does nothing +func (l *NoOpLogger) Errorf(format string, args ...interface{}) {} diff --git a/session/chunking/chunk_manager_test.go b/session/chunking/chunk_manager_test.go new file mode 100644 index 0000000..f7d3bc0 --- /dev/null +++ b/session/chunking/chunk_manager_test.go @@ -0,0 +1,1771 @@ +package chunking + +import ( + "fmt" + "runtime" + "strings" + "sync" + "testing" + "time" + + "github.com/gorilla/sessions" +) + +// TestTokenValidatorJWT tests JWT validation using TokenValidator +func TestTokenValidatorJWT(t *testing.T) { + validator := NewTokenValidator() + + // Test valid JWT format (using base64url encoded parts that are long enough) + validJWT := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c" + err := validator.ValidateJWTFormat(validJWT, "test") + if err != nil { + t.Errorf("Expected valid JWT to pass, got error: %v", err) + } + + // Test invalid JWT format - too few parts + invalidJWT := "header.payload" + err = validator.ValidateJWTFormat(invalidJWT, "test") + if err == nil { + t.Error("Expected invalid JWT to fail validation") + } + + // Test invalid JWT format - too many parts + invalidJWT2 := "header.payload.signature.extra" + err = validator.ValidateJWTFormat(invalidJWT2, "test") + if err == nil { + t.Error("Expected invalid JWT with extra parts to fail validation") + } + + // Test empty JWT + err = validator.ValidateJWTFormat("", "test") + if err != nil { + t.Error("Expected empty JWT to pass validation (empty is allowed)") + } +} + +// TestTokenValidatorOpaqueToken tests opaque token validation using TokenValidator +func TestTokenValidatorOpaqueToken(t *testing.T) { + validator := NewTokenValidator() + config := AccessTokenConfig + + // Test valid opaque token with more entropy + validOpaque := "z8Bx5mP9qK3nL4wR7tY2uI0oE6cV1aS" + err := validator.ValidateTokenContent(validOpaque, config) + if err != nil { + t.Errorf("Expected valid opaque token to pass, got error: %v", err) + } + + // Test too short opaque token + shortOpaque := "short" + err = validator.ValidateTokenContent(shortOpaque, config) + if err == nil { + t.Error("Expected short opaque token to fail validation") + } + + // Test empty opaque token + err = validator.ValidateTokenContent("", config) + if err != nil { + t.Error("Expected empty opaque token to pass validation (empty is allowed)") + } +} + +// TestTokenValidatorTokenSize tests token size validation using TokenValidator +func TestTokenValidatorTokenSize(t *testing.T) { + validator := NewTokenValidator() + + // Test normal token size + normalToken := strings.Repeat("a", 1000) + err := validator.ValidateTokenSize(normalToken, AccessTokenConfig) + if err != nil { + t.Errorf("Expected normal token to pass size validation, got error: %v", err) + } + + // Test oversized token + oversizedToken := strings.Repeat("a", AccessTokenConfig.MaxLength+1) + err = validator.ValidateTokenSize(oversizedToken, AccessTokenConfig) + if err == nil { + t.Error("Expected oversized token to fail validation") + } + + // Test undersized token + undersizedToken := "ab" + err = validator.ValidateTokenSize(undersizedToken, AccessTokenConfig) + if err == nil { + t.Error("Expected undersized token to fail validation") + } +} + +// TestTokenValidatorTokenContent tests token content validation using TokenValidator +func TestTokenValidatorTokenContent(t *testing.T) { + validator := NewTokenValidator() + + // Test normal token content with good entropy + normalToken := "A9zZ8yX7wV6uT5sR4qP3oN2mL1kJ0iH" + err := validator.ValidateTokenContent(normalToken, AccessTokenConfig) + if err != nil { + t.Errorf("Expected normal token to pass content validation, got error: %v", err) + } + + // Test token with null bytes + nullByteToken := "token_with\x00null_byte" + err = validator.ValidateTokenContent(nullByteToken, AccessTokenConfig) + if err == nil { + t.Error("Expected token with null bytes to fail validation") + } + + // Test token with control characters + controlCharToken := "token_with\x01control" + err = validator.ValidateTokenContent(controlCharToken, AccessTokenConfig) + if err == nil { + t.Error("Expected token with control characters to fail validation") + } +} + +// TestChunkManagerSingleTokenValidation tests single token validation path +func TestChunkManagerSingleTokenValidation(t *testing.T) { + cm := NewChunkManager(nil) + + // Create a valid opaque token with good entropy + validToken := "oP8qW7rE6tY5uI4oP3aS2dF1gH9jK0lZ3xC6vB5nM4" + + // Test valid token processing + result := cm.processSingleToken(validToken, false, AccessTokenConfig) + if result.Error != nil { + t.Errorf("Expected valid token to process successfully, got error: %v", result.Error) + } + if result.Token != validToken { + t.Error("Expected token to be returned unchanged") + } + + // Test invalid token processing + invalidToken := "invalid.token" + result = cm.processSingleToken(invalidToken, false, IDTokenConfig) // ID tokens require JWT format + if result.Error == nil { + t.Error("Expected invalid token to fail processing") + } +} + +// TestTokenConfigValidation tests different token configurations +func TestTokenConfigValidation(t *testing.T) { + tests := []struct { + name string + config TokenConfig + }{ + { + name: "AccessTokenConfig", + config: AccessTokenConfig, + }, + { + name: "RefreshTokenConfig", + config: RefreshTokenConfig, + }, + { + name: "IDTokenConfig", + config: IDTokenConfig, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Verify config has expected fields + if tt.config.Type == "" { + t.Error("Expected config to have Type set") + } + if tt.config.MaxLength <= 0 { + t.Error("Expected config to have positive MaxLength") + } + if tt.config.MinLength <= 0 { + t.Error("Expected config to have positive MinLength") + } + }) + } +} + +// TestSessionMapBounds_HardLimitEnforcement tests that the session map enforces hard limits +// and prevents unbounded memory growth +func TestSessionMapBounds_HardLimitEnforcement(t *testing.T) { + tests := []struct { + name string + maxSessions int + sessionCount int + expectEviction bool + description string + }{ + { + name: "within_limit", + maxSessions: 100, + sessionCount: 50, + expectEviction: false, + description: "Sessions within limit should not trigger eviction", + }, + { + name: "at_limit", + maxSessions: 100, + sessionCount: 100, + expectEviction: false, + description: "Sessions at exact limit should not trigger eviction", + }, + { + name: "exceeds_limit", + maxSessions: 100, + sessionCount: 150, + expectEviction: true, + description: "Sessions exceeding limit should trigger eviction", + }, + { + name: "small_limit", + maxSessions: 10, + sessionCount: 20, + expectEviction: true, + description: "Small limit should be strictly enforced", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Create chunk manager with custom limits + cm := NewChunkManager(nil) + cm.maxSessions = tt.maxSessions + + // Record initial memory + runtime.GC() + var m1 runtime.MemStats + runtime.ReadMemStats(&m1) + + // Create sessions by storing them in the session map + for i := 0; i < tt.sessionCount; i++ { + sessionKey := generateSessionKey(i) + + // Create a mock session entry + cm.mutex.Lock() + cm.sessionMap[sessionKey] = &SessionEntry{ + Session: &sessions.Session{Values: make(map[interface{}]interface{})}, + ExpiresAt: time.Now().Add(24 * time.Hour), + LastUsed: time.Now(), + } + cm.mutex.Unlock() + + // Trigger cleanup every 10 sessions to test enforcement + if i%10 == 9 { + cm.CleanupExpiredSessions() + } + } + + // Force final cleanup to enforce limits + cm.CleanupExpiredSessions() + + // Check final session count + cm.mutex.RLock() + finalSessionCount := len(cm.sessionMap) + cm.mutex.RUnlock() + + // Verify hard limit enforcement + if finalSessionCount > tt.maxSessions { + t.Errorf("Hard limit not enforced: %s\nMax sessions: %d\nFinal session count: %d\nExpected eviction: %v", + tt.description, tt.maxSessions, finalSessionCount, tt.expectEviction) + } + + // Verify eviction occurred if expected + if tt.expectEviction && finalSessionCount >= tt.sessionCount { + t.Errorf("Expected eviction did not occur: %s\nCreated sessions: %d\nFinal sessions: %d", + tt.description, tt.sessionCount, finalSessionCount) + } + + // Record final memory + runtime.GC() + var m2 runtime.MemStats + runtime.ReadMemStats(&m2) + memoryGrowth := m2.Alloc - m1.Alloc + + t.Logf("Test %s: Created %d sessions, Final count: %d, Memory growth: %d bytes", + tt.name, tt.sessionCount, finalSessionCount, memoryGrowth) + + // Verify memory growth is bounded + maxExpectedMemoryPerSession := int64(1024) // 1KB per session + maxExpectedMemory := int64(tt.maxSessions) * maxExpectedMemoryPerSession + if int64(memoryGrowth) > maxExpectedMemory*2 { // Allow 2x tolerance + t.Errorf("Memory growth exceeds expected bounds: %d bytes (max expected: %d)", + memoryGrowth, maxExpectedMemory) + } + }) + } +} + +// TestSessionMapBounds_EmergencyCleanup tests that emergency cleanup triggers when approaching limits +func TestSessionMapBounds_EmergencyCleanup(t *testing.T) { + cm := NewChunkManager(nil) + cm.maxSessions = 50 + + // Force lastCleanup to be old so cleanup will run + cm.lastCleanup = time.Now().Add(-2 * time.Hour) + + // Fill sessions to near capacity + nearCapacity := cm.maxSessions - 5 + for i := 0; i < nearCapacity; i++ { + sessionKey := generateSessionKey(i) + cm.mutex.Lock() + cm.sessionMap[sessionKey] = &SessionEntry{ + Session: &sessions.Session{Values: make(map[interface{}]interface{})}, + ExpiresAt: time.Now().Add(24 * time.Hour), + LastUsed: time.Now().Add(time.Duration(-i) * time.Hour), // Vary ages for LRU + } + cm.mutex.Unlock() + } + + // Add some expired sessions that should be cleaned up + expiredCount := 10 + for i := 0; i < expiredCount; i++ { + sessionKey := generateExpiredSessionKey(i) + cm.mutex.Lock() + cm.sessionMap[sessionKey] = &SessionEntry{ + Session: &sessions.Session{Values: make(map[interface{}]interface{})}, + ExpiresAt: time.Now().Add(-24 * time.Hour), // Expired + LastUsed: time.Now().Add(-48 * time.Hour), + } + cm.mutex.Unlock() + } + + // Record state before emergency cleanup + cm.mutex.RLock() + beforeCleanup := len(cm.sessionMap) + cm.mutex.RUnlock() + + // Trigger emergency cleanup + cm.CleanupExpiredSessions() + + // Check that expired sessions were removed + cm.mutex.RLock() + afterCleanup := len(cm.sessionMap) + cm.mutex.RUnlock() + + cleanedUp := beforeCleanup - afterCleanup + if cleanedUp < expiredCount { + t.Errorf("Emergency cleanup insufficient: cleaned %d sessions, expected at least %d", + cleanedUp, expiredCount) + } + + // Verify we're still within limits + if afterCleanup > cm.maxSessions { + t.Errorf("Emergency cleanup failed to enforce limits: %d sessions > %d max", + afterCleanup, cm.maxSessions) + } + + t.Logf("Emergency cleanup: Before: %d, After: %d, Cleaned: %d", + beforeCleanup, afterCleanup, cleanedUp) +} + +// TestSessionMapBounds_EvictionUnderHighLoad tests session eviction under high concurrent load +func TestSessionMapBounds_EvictionUnderHighLoad(t *testing.T) { + cm := NewChunkManager(nil) + cm.maxSessions = 100 + + // Record initial memory + runtime.GC() + var m1 runtime.MemStats + runtime.ReadMemStats(&m1) + + const numGoroutines = 10 + const sessionsPerGoroutine = 50 + var wg sync.WaitGroup + + // Create sessions concurrently to simulate high load + for g := 0; g < numGoroutines; g++ { + wg.Add(1) + go func(goroutineID int) { + defer wg.Done() + for i := 0; i < sessionsPerGoroutine; i++ { + sessionKey := generateConcurrentSessionKey(goroutineID, i) + + cm.mutex.Lock() + cm.sessionMap[sessionKey] = &SessionEntry{ + Session: &sessions.Session{Values: make(map[interface{}]interface{})}, + ExpiresAt: time.Now().Add(24 * time.Hour), + LastUsed: time.Now(), + } + + // Randomly trigger cleanup to test concurrent access + if i%10 == goroutineID%10 { + cm.mutex.Unlock() + cm.CleanupExpiredSessions() + } else { + cm.mutex.Unlock() + } + + // Small delay to increase concurrency contention + time.Sleep(time.Microsecond) + } + }(g) + } + + wg.Wait() + + // Final cleanup + cm.CleanupExpiredSessions() + + // Verify limits are still enforced + cm.mutex.RLock() + finalCount := len(cm.sessionMap) + cm.mutex.RUnlock() + + if finalCount > cm.maxSessions { + t.Errorf("High load caused limit breach: %d sessions > %d max", finalCount, cm.maxSessions) + } + + // Check memory usage + runtime.GC() + var m2 runtime.MemStats + runtime.ReadMemStats(&m2) + memoryGrowth := m2.Alloc - m1.Alloc + + t.Logf("High load test: Created %d total sessions, Final count: %d, Memory growth: %d bytes", + numGoroutines*sessionsPerGoroutine, finalCount, memoryGrowth) + + // Verify memory is bounded + maxExpectedMemory := int64(cm.maxSessions * 2048) // 2KB per session + if int64(memoryGrowth) > maxExpectedMemory { + t.Errorf("Memory growth under high load: %d bytes > %d expected max", + memoryGrowth, maxExpectedMemory) + } +} + +// TestSessionMapBounds_NoMemoryGrowthBeyondLimits tests that memory doesn't grow beyond configured limits +func TestSessionMapBounds_NoMemoryGrowthBeyondLimits(t *testing.T) { + const maxSessions = 200 + const testIterations = 1000 // Create way more sessions than limit + + cm := NewChunkManager(nil) + cm.maxSessions = maxSessions + + // Record baseline memory + runtime.GC() + runtime.GC() + var baseline runtime.MemStats + runtime.ReadMemStats(&baseline) + + // Create sessions in waves, exceeding limits + for wave := 0; wave < 5; wave++ { + // Create burst of sessions + for i := 0; i < testIterations/5; i++ { + sessionKey := generateWaveSessionKey(wave, i) + + cm.mutex.Lock() + cm.sessionMap[sessionKey] = &SessionEntry{ + Session: &sessions.Session{Values: make(map[interface{}]interface{})}, + ExpiresAt: time.Now().Add(24 * time.Hour), + LastUsed: time.Now(), + } + cm.mutex.Unlock() + + // Periodic cleanup + if i%50 == 49 { + cm.CleanupExpiredSessions() + } + } + + // Force cleanup after each wave + cm.CleanupExpiredSessions() + + // Check session count doesn't exceed limits + cm.mutex.RLock() + currentCount := len(cm.sessionMap) + cm.mutex.RUnlock() + + if currentCount > maxSessions { + t.Errorf("Session count exceeded limit in wave %d: %d > %d", + wave, currentCount, maxSessions) + } + + // Check memory growth is bounded + runtime.GC() + var current runtime.MemStats + runtime.ReadMemStats(¤t) + memoryGrowth := current.Alloc - baseline.Alloc + + // Memory should not grow linearly with total sessions created + maxExpectedMemory := uint64(maxSessions * 3072) // 3KB per session with overhead + if memoryGrowth > maxExpectedMemory { + t.Errorf("Memory growth exceeded bounds in wave %d: %d bytes > %d expected", + wave, memoryGrowth, maxExpectedMemory) + } + + t.Logf("Wave %d: Sessions: %d, Memory growth: %d bytes", + wave, currentCount, memoryGrowth) + } +} + +// TestSessionMapBounds_LRUEvictionOrder tests that LRU eviction maintains correct order +func TestSessionMapBounds_LRUEvictionOrder(t *testing.T) { + cm := NewChunkManager(nil) + cm.maxSessions = 10 + + // Create sessions with known access patterns + sessionOrder := make([]string, 0, 15) + + // Create initial sessions + for i := 0; i < 15; i++ { + sessionKey := generateOrderedSessionKey(i) + sessionOrder = append(sessionOrder, sessionKey) + + cm.mutex.Lock() + cm.sessionMap[sessionKey] = &SessionEntry{ + Session: &sessions.Session{Values: make(map[interface{}]interface{})}, + ExpiresAt: time.Now().Add(24 * time.Hour), + LastUsed: time.Now().Add(time.Duration(-i) * time.Minute), // Older sessions have earlier LastUsed + } + cm.mutex.Unlock() + } + + // Force eviction + cm.CleanupExpiredSessions() + + // Check that oldest sessions were evicted + cm.mutex.RLock() + remainingSessions := make([]string, 0, len(cm.sessionMap)) + for key := range cm.sessionMap { + remainingSessions = append(remainingSessions, key) + } + cm.mutex.RUnlock() + + // Should have exactly maxSessions remaining + if len(remainingSessions) != cm.maxSessions { + t.Errorf("Incorrect number of sessions after eviction: got %d, expected %d", + len(remainingSessions), cm.maxSessions) + } + + // Check that the most recently used sessions remain + // (sessions with lower indices have more recent LastUsed times) + expectedRemaining := sessionOrder[:cm.maxSessions] + for _, expectedKey := range expectedRemaining { + found := false + for _, remainingKey := range remainingSessions { + if remainingKey == expectedKey { + found = true + break + } + } + if !found { + t.Errorf("Expected session %s to remain after LRU eviction", expectedKey) + } + } +} + +// Helper functions for generating unique session keys + +func generateSessionKey(id int) string { + return "session_" + strings.Repeat("0", 5-len(string(rune(id)))) + string(rune('0'+id%10)) +} + +func generateExpiredSessionKey(id int) string { + return "expired_session_" + strings.Repeat("0", 5-len(string(rune(id)))) + string(rune('0'+id%10)) +} + +func generateConcurrentSessionKey(goroutineID, sessionID int) string { + return generateSessionKey(goroutineID*1000 + sessionID) +} + +func generateWaveSessionKey(wave, id int) string { + return "wave_" + string(rune('0'+wave)) + "_" + generateSessionKey(id) +} + +func generateOrderedSessionKey(id int) string { + return "ordered_" + strings.Repeat("0", 5-len(string(rune(id)))) + string(rune('0'+id%10)) +} + +// BenchmarkSessionMapBounds_EvictionPerformance benchmarks the performance of session eviction +func BenchmarkSessionMapBounds_EvictionPerformance(b *testing.B) { + cm := NewChunkManager(nil) + cm.maxSessions = 1000 + + // Pre-populate with sessions at capacity + for i := 0; i < cm.maxSessions; i++ { + sessionKey := generateSessionKey(i) + cm.mutex.Lock() + cm.sessionMap[sessionKey] = &SessionEntry{ + Session: &sessions.Session{Values: make(map[interface{}]interface{})}, + ExpiresAt: time.Now().Add(24 * time.Hour), + LastUsed: time.Now().Add(time.Duration(-i) * time.Minute), + } + cm.mutex.Unlock() + } + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + // Add session that will trigger eviction + sessionKey := generateSessionKey(cm.maxSessions + i) + cm.mutex.Lock() + cm.sessionMap[sessionKey] = &SessionEntry{ + Session: &sessions.Session{Values: make(map[interface{}]interface{})}, + ExpiresAt: time.Now().Add(24 * time.Hour), + LastUsed: time.Now(), + } + cm.mutex.Unlock() + + // Force eviction + cm.CleanupExpiredSessions() + } +} + +// BenchmarkSessionMapBounds_ConcurrentAccess benchmarks concurrent session access with bounds checking +func BenchmarkSessionMapBounds_ConcurrentAccess(b *testing.B) { + cm := NewChunkManager(nil) + cm.maxSessions = 500 + + // Pre-populate sessions + for i := 0; i < cm.maxSessions/2; i++ { + sessionKey := generateSessionKey(i) + cm.mutex.Lock() + cm.sessionMap[sessionKey] = &SessionEntry{ + Session: &sessions.Session{Values: make(map[interface{}]interface{})}, + ExpiresAt: time.Now().Add(24 * time.Hour), + LastUsed: time.Now(), + } + cm.mutex.Unlock() + } + + b.ResetTimer() + + b.RunParallel(func(pb *testing.PB) { + i := 0 + for pb.Next() { + sessionKey := generateSessionKey(i) + + // Mix of operations: create, cleanup, access + switch i % 3 { + case 0: + cm.mutex.Lock() + cm.sessionMap[sessionKey] = &SessionEntry{ + Session: &sessions.Session{Values: make(map[interface{}]interface{})}, + ExpiresAt: time.Now().Add(24 * time.Hour), + LastUsed: time.Now(), + } + cm.mutex.Unlock() + case 1: + cm.CleanupExpiredSessions() + case 2: + cm.mutex.RLock() + _ = len(cm.sessionMap) + cm.mutex.RUnlock() + } + i++ + } + }) +} + +// TestEstimateChunkCount tests the EstimateChunkCount function +func TestEstimateChunkCount(t *testing.T) { + cs := NewChunkSerializer(nil) + + tests := []struct { + name string + tokenLength int + chunkSize int + expected int + }{ + { + name: "Single chunk", + tokenLength: 1000, + chunkSize: 1200, + expected: 1, + }, + { + name: "Exactly two chunks", + tokenLength: 2400, + chunkSize: 1200, + expected: 2, + }, + { + name: "Three chunks with remainder", + tokenLength: 2500, + chunkSize: 1200, + expected: 3, + }, + { + name: "Zero chunk size defaults to maxCookieSize", + tokenLength: 1300, + chunkSize: 0, + expected: 2, // 1300 / 1200 = 1.083... = 2 chunks + }, + { + name: "Large token many chunks", + tokenLength: 10000, + chunkSize: 800, + expected: 13, // 10000 / 800 = 12.5 = 13 chunks + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := cs.EstimateChunkCount(tt.tokenLength, tt.chunkSize) + if result != tt.expected { + t.Errorf("EstimateChunkCount(%d, %d) = %d; expected %d", + tt.tokenLength, tt.chunkSize, result, tt.expected) + } + }) + } +} + +// TestMaxTokenSizeForChunks tests the MaxTokenSizeForChunks function +func TestMaxTokenSizeForChunks(t *testing.T) { + cs := NewChunkSerializer(nil) + + tests := []struct { + name string + maxChunks int + chunkSize int + expected int + }{ + { + name: "Single chunk", + maxChunks: 1, + chunkSize: 1200, + expected: 1200, + }, + { + name: "Multiple chunks", + maxChunks: 5, + chunkSize: 1000, + expected: 5000, + }, + { + name: "Zero chunk size defaults to maxCookieSize", + maxChunks: 3, + chunkSize: 0, + expected: 3600, // 3 * 1200 + }, + { + name: "Large configuration", + maxChunks: 25, + chunkSize: 1200, + expected: 30000, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := cs.MaxTokenSizeForChunks(tt.maxChunks, tt.chunkSize) + if result != tt.expected { + t.Errorf("MaxTokenSizeForChunks(%d, %d) = %d; expected %d", + tt.maxChunks, tt.chunkSize, result, tt.expected) + } + }) + } +} + +// TestValidateJWTContent tests JWT content validation +func TestValidateJWTContent(t *testing.T) { + validator := NewTokenValidator() + config := IDTokenConfig + + tests := []struct { + name string + token string + expectError bool + description string + }{ + { + name: "Valid JWT with required ID token claims", + token: "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJodHRwczovL2V4YW1wbGUuY29tIiwic3ViIjoiMTIzNDU2Nzg5MCIsImF1ZCI6ImNsaWVudElkIiwiZXhwIjoxNjQ2MDY0MDAwLCJpYXQiOjE2NDYwNjA0MDB9.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c", + expectError: false, + description: "JWT with all required ID token claims should pass", + }, + { + name: "JWT missing required claims", + token: "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIn0.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c", + expectError: true, + description: "JWT missing required claims should fail", + }, + { + name: "JWT with invalid structure", + token: "invalid.token", + expectError: true, + description: "JWT with wrong number of parts should fail", + }, + { + name: "Empty JWT", + token: "", + expectError: true, + description: "Empty JWT should fail at JWT content level", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validator.validateJWTContent(tt.token, config) + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, but got none", tt.description) + } + if !tt.expectError && err != nil { + t.Errorf("Expected no error for %s, but got: %v", tt.description, err) + } + }) + } +} + +// TestValidateJWTHeader tests JWT header validation +func TestValidateJWTHeader(t *testing.T) { + validator := NewTokenValidator() + config := IDTokenConfig + + tests := []struct { + name string + header string + expectError bool + description string + }{ + { + name: "Valid JWT header", + header: "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9", // {"alg":"RS256","typ":"JWT"} + expectError: false, + description: "Valid JWT header with alg and typ", + }, + { + name: "Header missing alg", + header: "eyJ0eXAiOiJKV1QifQ", // {"typ":"JWT"} + expectError: true, + description: "Header missing algorithm should fail", + }, + { + name: "Header missing typ", + header: "eyJhbGciOiJSUzI1NiJ9", // {"alg":"RS256"} + expectError: true, + description: "Header missing type should fail", + }, + { + name: "Invalid base64 header", + header: "invalid_base64!", + expectError: true, + description: "Invalid base64 should fail", + }, + { + name: "Invalid JSON header", + header: "aW52YWxpZCBqc29u", // "invalid json" + expectError: true, + description: "Invalid JSON should fail", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validator.validateJWTHeader(tt.header, config) + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, but got none", tt.description) + } + if !tt.expectError && err != nil { + t.Errorf("Expected no error for %s, but got: %v", tt.description, err) + } + }) + } +} + +// TestValidateJWTPayload tests JWT payload validation +func TestValidateJWTPayload(t *testing.T) { + validator := NewTokenValidator() + + tests := []struct { + name string + payload string + config TokenConfig + expectError bool + description string + }{ + { + name: "Valid ID token payload", + payload: "eyJpc3MiOiJodHRwczovL2V4YW1wbGUuY29tIiwic3ViIjoiMTIzNDU2Nzg5MCIsImF1ZCI6ImNsaWVudElkIiwiZXhwIjoxNjQ2MDY0MDAwLCJpYXQiOjE2NDYwNjA0MDB9", // Required ID token claims + config: IDTokenConfig, + expectError: false, + description: "Valid ID token with required claims", + }, + { + name: "ID token missing required claims", + payload: "eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIn0", // {"sub":"1234567890","name":"John Doe"} + config: IDTokenConfig, + expectError: true, + description: "ID token missing required claims should fail", + }, + { + name: "Access token payload", + payload: "eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIn0", // {"sub":"1234567890","name":"John Doe"} + config: AccessTokenConfig, + expectError: false, + description: "Access token doesn't require specific claims", + }, + { + name: "Invalid base64 payload", + payload: "invalid_base64!", + config: IDTokenConfig, + expectError: true, + description: "Invalid base64 should fail", + }, + { + name: "Invalid JSON payload", + payload: "aW52YWxpZCBqc29u", // "invalid json" + config: IDTokenConfig, + expectError: true, + description: "Invalid JSON should fail", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validator.validateJWTPayload(tt.payload, tt.config) + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, but got none", tt.description) + } + if !tt.expectError && err != nil { + t.Errorf("Expected no error for %s, but got: %v", tt.description, err) + } + }) + } +} + +// TestValidateJWTSignature tests JWT signature validation +func TestValidateJWTSignature(t *testing.T) { + validator := NewTokenValidator() + config := IDTokenConfig + + tests := []struct { + name string + signature string + expectError bool + description string + }{ + { + name: "Valid signature", + signature: "SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c", + expectError: false, + description: "Valid base64URL signature", + }, + { + name: "Empty signature", + signature: "", + expectError: true, + description: "Empty signature should fail", + }, + { + name: "Invalid base64URL signature", + signature: "invalid_base64!@#", + expectError: true, + description: "Invalid base64URL should fail", + }, + { + name: "Valid signature with padding", + signature: "dGVzdA==", + expectError: false, + description: "Base64 with padding should work", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validator.validateJWTSignature(tt.signature, config) + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, but got none", tt.description) + } + if !tt.expectError && err != nil { + t.Errorf("Expected no error for %s, but got: %v", tt.description, err) + } + }) + } +} + +// TestValidateChunkStructure tests chunk structure validation +func TestValidateChunkStructure(t *testing.T) { + validator := NewTokenValidator() + config := AccessTokenConfig + + tests := []struct { + name string + chunks []ChunkData + expectError bool + description string + }{ + { + name: "Valid chunk structure", + chunks: []ChunkData{ + {Index: 0, Total: 2, Content: "part1", Checksum: "checksum1"}, + {Index: 1, Total: 2, Content: "part2", Checksum: "checksum2"}, + }, + expectError: false, + description: "Valid ordered chunks", + }, + { + name: "Empty chunks", + chunks: []ChunkData{}, + expectError: true, + description: "Empty chunk list should fail", + }, + { + name: "Too many chunks", + chunks: func() []ChunkData { + chunks := make([]ChunkData, AccessTokenConfig.MaxChunks+1) + for i := range chunks { + chunks[i] = ChunkData{Index: i, Total: len(chunks), Content: "content", Checksum: "checksum"} + } + return chunks + }(), + expectError: true, + description: "Too many chunks should fail", + }, + { + name: "Duplicate chunk indices", + chunks: []ChunkData{ + {Index: 0, Total: 2, Content: "part1", Checksum: "checksum1"}, + {Index: 0, Total: 2, Content: "part2", Checksum: "checksum2"}, + }, + expectError: true, + description: "Duplicate indices should fail", + }, + { + name: "Missing chunk index", + chunks: []ChunkData{ + {Index: 0, Total: 3, Content: "part1", Checksum: "checksum1"}, + {Index: 2, Total: 3, Content: "part3", Checksum: "checksum3"}, + }, + expectError: true, + description: "Missing chunk index should fail", + }, + { + name: "Inconsistent total count", + chunks: []ChunkData{ + {Index: 0, Total: 2, Content: "part1", Checksum: "checksum1"}, + {Index: 1, Total: 3, Content: "part2", Checksum: "checksum2"}, + }, + expectError: true, + description: "Inconsistent total should fail", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validator.ValidateChunkStructure(tt.chunks, config) + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, but got none", tt.description) + } + if !tt.expectError && err != nil { + t.Errorf("Expected no error for %s, but got: %v", tt.description, err) + } + }) + } +} + +// TestValidateChunkData tests individual chunk data validation +func TestValidateChunkData(t *testing.T) { + validator := NewTokenValidator() + config := AccessTokenConfig + + tests := []struct { + name string + chunk ChunkData + expectedTotal int + expectError bool + description string + }{ + { + name: "Valid chunk data", + chunk: ChunkData{Index: 0, Total: 2, Content: "content", Checksum: "checksum"}, + expectedTotal: 2, + expectError: false, + description: "Valid chunk should pass", + }, + { + name: "Negative index", + chunk: ChunkData{Index: -1, Total: 2, Content: "content", Checksum: "checksum"}, + expectedTotal: 2, + expectError: true, + description: "Negative index should fail", + }, + { + name: "Inconsistent total", + chunk: ChunkData{Index: 0, Total: 3, Content: "content", Checksum: "checksum"}, + expectedTotal: 2, + expectError: true, + description: "Inconsistent total should fail", + }, + { + name: "Index exceeds total", + chunk: ChunkData{Index: 2, Total: 2, Content: "content", Checksum: "checksum"}, + expectedTotal: 2, + expectError: true, + description: "Index exceeding total should fail", + }, + { + name: "Empty content", + chunk: ChunkData{Index: 0, Total: 2, Content: "", Checksum: "checksum"}, + expectedTotal: 2, + expectError: true, + description: "Empty content should fail", + }, + { + name: "Empty checksum", + chunk: ChunkData{Index: 0, Total: 2, Content: "content", Checksum: ""}, + expectedTotal: 2, + expectError: true, + description: "Empty checksum should fail", + }, + { + name: "Chunk too large", + chunk: ChunkData{ + Index: 0, + Total: 2, + Content: strings.Repeat("x", config.MaxChunkSize+1), + Checksum: "checksum", + }, + expectedTotal: 2, + expectError: true, + description: "Oversized chunk should fail", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validator.validateChunkData(tt.chunk, tt.expectedTotal, config) + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, but got none", tt.description) + } + if !tt.expectError && err != nil { + t.Errorf("Expected no error for %s, but got: %v", tt.description, err) + } + }) + } +} + +// TestChunkErrorMethod tests the Error method of ChunkError +func TestChunkErrorMethod(t *testing.T) { + tests := []struct { + name string + error *ChunkError + expected string + }{ + { + name: "Basic chunk error", + error: &ChunkError{ + Type: "access", + Reason: "too large", + Details: "chunk exceeds maximum size", + }, + expected: "access chunk error: too large - chunk exceeds maximum size", + }, + { + name: "Validation chunk error", + error: &ChunkError{ + Type: "id", + Reason: "missing chunk", + Details: "chunk 2 is missing from sequence", + }, + expected: "id chunk error: missing chunk - chunk 2 is missing from sequence", + }, + { + name: "Empty fields", + error: &ChunkError{ + Type: "", + Reason: "", + Details: "", + }, + expected: " chunk error: - ", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.error.Error() + if result != tt.expected { + t.Errorf("ChunkError.Error() = %q; expected %q", result, tt.expected) + } + }) + } +} + +// TestValidationErrorMethod tests the Error method of ValidationError +func TestValidationErrorMethod(t *testing.T) { + tests := []struct { + name string + error *ValidationError + expected string + }{ + { + name: "Token validation error", + error: &ValidationError{ + Type: "access", + Reason: "invalid format", + Details: "token must be valid JWT", + }, + expected: "access validation error: invalid format - token must be valid JWT", + }, + { + name: "Size validation error", + error: &ValidationError{ + Type: "refresh", + Reason: "too large", + Details: "token size exceeds 50KB limit", + }, + expected: "refresh validation error: too large - token size exceeds 50KB limit", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.error.Error() + if result != tt.expected { + t.Errorf("ValidationError.Error() = %q; expected %q", result, tt.expected) + } + }) + } +} + +// TestGetToken tests the main GetToken function +func TestGetToken(t *testing.T) { + cm := NewChunkManager(nil) + + tests := []struct { + name string + mainSession *sessions.Session + chunks map[int]*sessions.Session + config TokenConfig + expectedToken string + expectError bool + description string + }{ + { + name: "Token from main session", + mainSession: &sessions.Session{ + Values: map[interface{}]interface{}{ + "access_token": "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ", + }, + }, + chunks: nil, + config: AccessTokenConfig, + expectedToken: "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ", + expectError: false, + description: "Should retrieve token from main session", + }, + { + name: "No token in main session, no chunks", + mainSession: &sessions.Session{Values: map[interface{}]interface{}{}}, + chunks: map[int]*sessions.Session{}, + config: AccessTokenConfig, + expectedToken: "", + expectError: false, + description: "Should return empty token when no data available", + }, + { + name: "Token from chunks", + mainSession: &sessions.Session{Values: map[interface{}]interface{}{}}, + chunks: map[int]*sessions.Session{ + 0: {Values: map[interface{}]interface{}{"value": "abcdefghijklmnopqrstuvwxyz"}}, + 1: {Values: map[interface{}]interface{}{"value": "0123456789ABCDEFGHIJKLMNOP"}}, + }, + config: AccessTokenConfig, + expectedToken: "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOP", + expectError: false, + description: "Should reconstruct token from chunks", + }, + { + name: "Too many chunks", + mainSession: &sessions.Session{Values: map[interface{}]interface{}{}}, + chunks: func() map[int]*sessions.Session { + chunks := make(map[int]*sessions.Session) + for i := 0; i <= AccessTokenConfig.MaxChunks; i++ { + chunks[i] = &sessions.Session{Values: map[interface{}]interface{}{"value": "chunk"}} + } + return chunks + }(), + config: AccessTokenConfig, + expectedToken: "", + expectError: true, + description: "Should fail with too many chunks", + }, + } + + // Mock compressor + compressor := &mockTokenCompressor{} + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := cm.GetToken(tt.mainSession, tt.chunks, tt.config, compressor) + + if tt.expectError && result.Error == nil { + t.Errorf("Expected error for %s, but got none", tt.description) + } + if !tt.expectError && result.Error != nil { + t.Errorf("Expected no error for %s, but got: %v", tt.description, result.Error) + } + if result.Token != tt.expectedToken { + t.Errorf("Expected token %q, got %q for %s", tt.expectedToken, result.Token, tt.description) + } + }) + } +} + +// TestStoreSessionGetSession tests session storage and retrieval +func TestStoreSessionGetSession(t *testing.T) { + cm := NewChunkManager(nil) + + // Test storing and retrieving a session + key := "test_session_key" + session := &sessions.Session{Values: map[interface{}]interface{}{"test": "value"}} + + // Store session + cm.StoreSession(key, session) + + // Retrieve session + retrieved := cm.GetSession(key) + if retrieved == nil { + t.Error("Expected to retrieve stored session, but got nil") + } + + if retrieved != session { + t.Error("Retrieved session does not match stored session") + } + + // Test retrieving non-existent session + nonExistent := cm.GetSession("non_existent_key") + if nonExistent != nil { + t.Error("Expected nil for non-existent session, but got a session") + } + + // Test session limit enforcement + cm.maxSessions = 2 + for i := 0; i < 5; i++ { + key := fmt.Sprintf("session_%d", i) + session := &sessions.Session{Values: map[interface{}]interface{}{"id": i}} + cm.StoreSession(key, session) + } + + cm.mutex.RLock() + sessionCount := len(cm.sessionMap) + cm.mutex.RUnlock() + + if sessionCount > cm.maxSessions { + t.Errorf("Session count %d exceeds limit %d", sessionCount, cm.maxSessions) + } +} + +// TestNoOpLogger tests the no-op logger implementation +func TestNoOpLogger(t *testing.T) { + logger := NewNoOpLogger() + + // Test all methods (they should not panic or error) + logger.Debug("test message") + logger.Debugf("test format %s", "message") + logger.Error("test error") + logger.Errorf("test error %s", "formatted") + + // Since these are no-op methods, we mainly test that they don't panic + // The fact that the test completes successfully indicates they work +} + +// TestSerializeTokenToChunks tests token serialization +func TestSerializeTokenToChunks(t *testing.T) { + cs := NewChunkSerializer(NewNoOpLogger()) + config := AccessTokenConfig + + tests := []struct { + name string + token string + expectError bool + description string + }{ + { + name: "Valid token serialization", + token: strings.Repeat("a", 2500), // Larger than single chunk + expectError: false, + description: "Should serialize large token into chunks", + }, + { + name: "Empty token", + token: "", + expectError: true, + description: "Should fail with empty token", + }, + { + name: "Token too short", + token: "abc", // Less than config.MinLength + expectError: true, + description: "Should fail with too short token", + }, + { + name: "Token too long", + token: strings.Repeat("x", config.MaxLength+1), + expectError: true, + description: "Should fail with oversized token", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + chunks, err := cs.SerializeTokenToChunks(tt.token, config) + + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, but got none", tt.description) + } + if !tt.expectError && err != nil { + t.Errorf("Expected no error for %s, but got: %v", tt.description, err) + } + + if !tt.expectError && len(chunks) > 0 { + // Verify chunk structure + expectedChunks := len(chunks) + for _, chunk := range chunks { + if chunk.Total != expectedChunks { + t.Errorf("Chunk total mismatch: got %d, expected %d", chunk.Total, expectedChunks) + } + if chunk.Content == "" { + t.Error("Chunk content should not be empty") + } + if chunk.Checksum == "" { + t.Error("Chunk checksum should not be empty") + } + } + } + }) + } +} + +// TestDeserializeTokenFromChunks tests token deserialization +func TestDeserializeTokenFromChunks(t *testing.T) { + cs := NewChunkSerializer(NewNoOpLogger()) + config := AccessTokenConfig + + // First serialize a token to get valid chunks + originalToken := strings.Repeat("abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOP", 40) // Make it large enough for multiple chunks + chunks, err := cs.SerializeTokenToChunks(originalToken, config) + if err != nil { + t.Fatalf("Failed to serialize token for test: %v", err) + } + + tests := []struct { + name string + chunks []ChunkData + expectedToken string + expectError bool + description string + }{ + { + name: "Valid chunks deserialization", + chunks: chunks, + expectedToken: originalToken, + expectError: false, + description: "Should deserialize chunks back to original token", + }, + { + name: "Empty chunks", + chunks: []ChunkData{}, + expectedToken: "", + expectError: true, + description: "Should fail with empty chunks", + }, + { + name: "Too many chunks", + chunks: func() []ChunkData { + many := make([]ChunkData, config.MaxChunks+1) + for i := range many { + many[i] = ChunkData{Index: i, Total: len(many), Content: "content", Checksum: "checksum"} + } + return many + }(), + expectedToken: "", + expectError: true, + description: "Should fail with too many chunks", + }, + { + name: "Inconsistent chunk totals", + chunks: []ChunkData{ + {Index: 0, Total: 2, Content: "part1", Checksum: cs.calculateChecksum("part1")}, + {Index: 1, Total: 3, Content: "part2", Checksum: cs.calculateChecksum("part2")}, // Different total + }, + expectedToken: "", + expectError: true, + description: "Should fail with inconsistent totals", + }, + { + name: "Missing chunk", + chunks: []ChunkData{ + {Index: 0, Total: 3, Content: "part1", Checksum: cs.calculateChecksum("part1")}, + {Index: 2, Total: 3, Content: "part3", Checksum: cs.calculateChecksum("part3")}, // Missing index 1 + }, + expectedToken: "", + expectError: true, + description: "Should fail with missing chunk", + }, + { + name: "Invalid checksum", + chunks: []ChunkData{ + {Index: 0, Total: 2, Content: "part1", Checksum: "invalid_checksum"}, + {Index: 1, Total: 2, Content: "part2", Checksum: cs.calculateChecksum("part2")}, + }, + expectedToken: "", + expectError: true, + description: "Should fail with invalid checksum", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + token, err := cs.DeserializeTokenFromChunks(tt.chunks, config) + + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, but got none", tt.description) + } + if !tt.expectError && err != nil { + t.Errorf("Expected no error for %s, but got: %v", tt.description, err) + } + if token != tt.expectedToken { + t.Errorf("Expected token length %d, got %d for %s", len(tt.expectedToken), len(token), tt.description) + } + }) + } +} + +// TestEncodeDecodeChunk tests chunk encoding and decoding +func TestEncodeDecodeChunk(t *testing.T) { + cs := NewChunkSerializer(NewNoOpLogger()) + + tests := []struct { + name string + chunk ChunkData + expectError bool + description string + }{ + { + name: "Valid chunk encoding/decoding", + chunk: ChunkData{ + Index: 0, + Total: 2, + Content: "test_content", + Checksum: "test_checksum", + }, + expectError: false, + description: "Should encode and decode chunk successfully", + }, + { + name: "Chunk with special characters", + chunk: ChunkData{ + Index: 1, + Total: 3, + Content: "content:with:colons", + Checksum: "checksum_123", + }, + expectError: false, + description: "Should handle special characters in content", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Encode chunk + encoded, err := cs.EncodeChunk(tt.chunk) + if tt.expectError && err == nil { + t.Errorf("Expected encoding error for %s, but got none", tt.description) + } + if !tt.expectError && err != nil { + t.Errorf("Expected no encoding error for %s, but got: %v", tt.description, err) + } + + if !tt.expectError && encoded != "" { + // Decode chunk + decoded, err := cs.DecodeChunk(encoded) + if err != nil { + t.Errorf("Expected no decoding error for %s, but got: %v", tt.description, err) + } + + // Verify decoded chunk matches original + if decoded.Index != tt.chunk.Index { + t.Errorf("Index mismatch: got %d, expected %d", decoded.Index, tt.chunk.Index) + } + if decoded.Total != tt.chunk.Total { + t.Errorf("Total mismatch: got %d, expected %d", decoded.Total, tt.chunk.Total) + } + if decoded.Content != tt.chunk.Content { + t.Errorf("Content mismatch: got %q, expected %q", decoded.Content, tt.chunk.Content) + } + if decoded.Checksum != tt.chunk.Checksum { + t.Errorf("Checksum mismatch: got %q, expected %q", decoded.Checksum, tt.chunk.Checksum) + } + } + }) + } + + // Test decoding invalid data + invalidTests := []struct { + name string + encoded string + description string + }{ + { + name: "Invalid base64", + encoded: "invalid_base64!", + description: "Should fail with invalid base64", + }, + { + name: "Wrong format", + encoded: "dGVzdA==", // "test" in base64, but wrong format + description: "Should fail with wrong format", + }, + } + + for _, tt := range invalidTests { + t.Run(tt.name, func(t *testing.T) { + _, err := cs.DecodeChunk(tt.encoded) + if err == nil { + t.Errorf("Expected error for %s, but got none", tt.description) + } + }) + } +} + +// TestValidateChunkIntegrity tests chunk integrity validation +func TestValidateChunkIntegrity(t *testing.T) { + cs := NewChunkSerializer(NewNoOpLogger()) + + tests := []struct { + name string + chunk ChunkData + expectError bool + description string + }{ + { + name: "Valid chunk integrity", + chunk: ChunkData{ + Index: 0, + Total: 2, + Content: "test_content", + Checksum: cs.calculateChecksum("test_content"), + }, + expectError: false, + description: "Should pass integrity check", + }, + { + name: "Negative index", + chunk: ChunkData{ + Index: -1, + Total: 2, + Content: "test_content", + Checksum: cs.calculateChecksum("test_content"), + }, + expectError: true, + description: "Should fail with negative index", + }, + { + name: "Invalid total", + chunk: ChunkData{ + Index: 0, + Total: 0, + Content: "test_content", + Checksum: cs.calculateChecksum("test_content"), + }, + expectError: true, + description: "Should fail with zero total", + }, + { + name: "Index exceeds total", + chunk: ChunkData{ + Index: 2, + Total: 2, + Content: "test_content", + Checksum: cs.calculateChecksum("test_content"), + }, + expectError: true, + description: "Should fail with index >= total", + }, + { + name: "Empty content", + chunk: ChunkData{ + Index: 0, + Total: 2, + Content: "", + Checksum: cs.calculateChecksum(""), + }, + expectError: true, + description: "Should fail with empty content", + }, + { + name: "Empty checksum", + chunk: ChunkData{ + Index: 0, + Total: 2, + Content: "test_content", + Checksum: "", + }, + expectError: true, + description: "Should fail with empty checksum", + }, + { + name: "Invalid checksum", + chunk: ChunkData{ + Index: 0, + Total: 2, + Content: "test_content", + Checksum: "invalid_checksum", + }, + expectError: true, + description: "Should fail with wrong checksum", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := cs.ValidateChunkIntegrity(tt.chunk) + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, but got none", tt.description) + } + if !tt.expectError && err != nil { + t.Errorf("Expected no error for %s, but got: %v", tt.description, err) + } + }) + } +} + +// TestCalculateChecksum tests checksum calculation +func TestCalculateChecksum(t *testing.T) { + cs := NewChunkSerializer(NewNoOpLogger()) + + tests := []struct { + name string + content string + expected string + }{ + { + name: "Empty content", + content: "", + expected: "empty", + }, + { + name: "Single character", + content: "a", + expected: "len1_first97", + }, + { + name: "Two characters", + content: "ab", + expected: "len2_first97_last98", + }, + { + name: "Longer content", + content: "test_content", + expected: "len12_first116_last116", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := cs.calculateChecksum(tt.content) + if result != tt.expected { + t.Errorf("calculateChecksum(%q) = %q; expected %q", tt.content, result, tt.expected) + } + }) + } +} + +// Mock token compressor for testing +type mockTokenCompressor struct{} + +func (m *mockTokenCompressor) CompressToken(token string) string { + // Simple mock - just return the original token + return token +} + +func (m *mockTokenCompressor) DecompressToken(compressed string) string { + // Simple mock - just return the original token + return compressed +} diff --git a/session/chunking/chunk_serializer.go b/session/chunking/chunk_serializer.go new file mode 100644 index 0000000..a2448bc --- /dev/null +++ b/session/chunking/chunk_serializer.go @@ -0,0 +1,279 @@ +// Package chunking provides chunk serialization functionality +package chunking + +import ( + "encoding/base64" + "fmt" + "strings" +) + +// ChunkSerializer handles serialization and deserialization of token chunks +type ChunkSerializer struct { + logger Logger +} + +// NewChunkSerializer creates a new chunk serializer +func NewChunkSerializer(logger Logger) *ChunkSerializer { + return &ChunkSerializer{ + logger: logger, + } +} + +// SerializeTokenToChunks splits a token into chunks suitable for cookie storage +func (cs *ChunkSerializer) SerializeTokenToChunks(token string, config TokenConfig) ([]ChunkData, error) { + if token == "" { + return nil, fmt.Errorf("cannot serialize empty token") + } + + if len(token) < config.MinLength { + return nil, fmt.Errorf("token too short: %d < %d", len(token), config.MinLength) + } + + if len(token) > config.MaxLength { + return nil, fmt.Errorf("token too long: %d > %d", len(token), config.MaxLength) + } + + // Calculate optimal chunk size + chunkSize := config.MaxChunkSize + if chunkSize <= 0 { + chunkSize = maxCookieSize + } + + // Estimate number of chunks needed + estimatedChunks := (len(token) + chunkSize - 1) / chunkSize + if estimatedChunks > config.MaxChunks { + return nil, fmt.Errorf("token requires too many chunks: %d > %d", estimatedChunks, config.MaxChunks) + } + + // Split token into chunks + chunks := make([]ChunkData, 0, estimatedChunks) + remaining := token + + chunkIndex := 0 + for len(remaining) > 0 { + if chunkIndex >= config.MaxChunks { + return nil, fmt.Errorf("exceeded maximum chunk count during serialization") + } + + // Determine chunk size for this iteration + currentChunkSize := chunkSize + if len(remaining) < currentChunkSize { + currentChunkSize = len(remaining) + } + + // Extract chunk + chunkContent := remaining[:currentChunkSize] + remaining = remaining[currentChunkSize:] + + // Create chunk data + chunkData := ChunkData{ + Index: chunkIndex, + Content: chunkContent, + Total: estimatedChunks, // Will be updated after all chunks are created + Checksum: cs.calculateChecksum(chunkContent), + } + + chunks = append(chunks, chunkData) + chunkIndex++ + } + + // Update total count in all chunks + actualChunks := len(chunks) + for i := range chunks { + chunks[i].Total = actualChunks + } + + cs.logger.Debugf("Serialized %s token into %d chunks", config.Type, len(chunks)) + return chunks, nil +} + +// DeserializeTokenFromChunks reconstructs a token from chunk data +func (cs *ChunkSerializer) DeserializeTokenFromChunks(chunks []ChunkData, config TokenConfig) (string, error) { + if len(chunks) == 0 { + return "", fmt.Errorf("no chunks provided for deserialization") + } + + if len(chunks) > config.MaxChunks { + return "", fmt.Errorf("too many chunks: %d > %d", len(chunks), config.MaxChunks) + } + + // Validate chunk consistency + expectedTotal := chunks[0].Total + for i, chunk := range chunks { + if chunk.Total != expectedTotal { + return "", fmt.Errorf("chunk %d has inconsistent total count: %d != %d", i, chunk.Total, expectedTotal) + } + } + + if len(chunks) != expectedTotal { + return "", fmt.Errorf("chunk count mismatch: got %d, expected %d", len(chunks), expectedTotal) + } + + // Sort chunks by index + orderedChunks := make([]ChunkData, expectedTotal) + for _, chunk := range chunks { + if chunk.Index < 0 || chunk.Index >= expectedTotal { + return "", fmt.Errorf("invalid chunk index: %d (total: %d)", chunk.Index, expectedTotal) + } + + if orderedChunks[chunk.Index].Content != "" { + return "", fmt.Errorf("duplicate chunk index: %d", chunk.Index) + } + + orderedChunks[chunk.Index] = chunk + } + + // Verify all chunks are present + for i, chunk := range orderedChunks { + if chunk.Content == "" { + return "", fmt.Errorf("missing chunk at index: %d", i) + } + + // Verify checksum + expectedChecksum := cs.calculateChecksum(chunk.Content) + if chunk.Checksum != expectedChecksum { + return "", fmt.Errorf("chunk %d checksum mismatch", i) + } + } + + // Reconstruct token + var tokenBuilder strings.Builder + tokenBuilder.Grow(len(chunks) * config.MaxChunkSize) // Pre-allocate capacity + + for _, chunk := range orderedChunks { + tokenBuilder.WriteString(chunk.Content) + } + + reconstructedToken := tokenBuilder.String() + + // Final validation + if len(reconstructedToken) < config.MinLength { + return "", fmt.Errorf("reconstructed token too short: %d < %d", len(reconstructedToken), config.MinLength) + } + + if len(reconstructedToken) > config.MaxLength { + return "", fmt.Errorf("reconstructed token too long: %d > %d", len(reconstructedToken), config.MaxLength) + } + + cs.logger.Debugf("Deserialized %s token from %d chunks (length: %d)", config.Type, len(chunks), len(reconstructedToken)) + return reconstructedToken, nil +} + +// EncodeChunk encodes chunk data for cookie storage +func (cs *ChunkSerializer) EncodeChunk(chunk ChunkData) (string, error) { + // Create a simple format: index:total:checksum:content + encoded := fmt.Sprintf("%d:%d:%s:%s", chunk.Index, chunk.Total, chunk.Checksum, chunk.Content) + + // Base64 encode the entire chunk for safe cookie storage + return base64.StdEncoding.EncodeToString([]byte(encoded)), nil +} + +// DecodeChunk decodes chunk data from cookie storage +func (cs *ChunkSerializer) DecodeChunk(encoded string) (ChunkData, error) { + // Base64 decode + decoded, err := base64.StdEncoding.DecodeString(encoded) + if err != nil { + return ChunkData{}, fmt.Errorf("failed to base64 decode chunk: %w", err) + } + + // Parse the format: index:total:checksum:content + parts := strings.SplitN(string(decoded), ":", 4) + if len(parts) != 4 { + return ChunkData{}, fmt.Errorf("invalid chunk format: expected 4 parts, got %d", len(parts)) + } + + var index, total int + if _, err := fmt.Sscanf(parts[0], "%d", &index); err != nil { + return ChunkData{}, fmt.Errorf("invalid chunk index: %w", err) + } + + if _, err := fmt.Sscanf(parts[1], "%d", &total); err != nil { + return ChunkData{}, fmt.Errorf("invalid chunk total: %w", err) + } + + checksum := parts[2] + content := parts[3] + + return ChunkData{ + Index: index, + Total: total, + Content: content, + Checksum: checksum, + }, nil +} + +// ValidateChunkIntegrity validates the integrity of chunk data +func (cs *ChunkSerializer) ValidateChunkIntegrity(chunk ChunkData) error { + if chunk.Index < 0 { + return fmt.Errorf("negative chunk index: %d", chunk.Index) + } + + if chunk.Total <= 0 { + return fmt.Errorf("invalid total chunks: %d", chunk.Total) + } + + if chunk.Index >= chunk.Total { + return fmt.Errorf("chunk index %d exceeds total %d", chunk.Index, chunk.Total) + } + + if chunk.Content == "" { + return fmt.Errorf("empty chunk content at index %d", chunk.Index) + } + + if chunk.Checksum == "" { + return fmt.Errorf("empty chunk checksum at index %d", chunk.Index) + } + + // Verify checksum + expectedChecksum := cs.calculateChecksum(chunk.Content) + if chunk.Checksum != expectedChecksum { + return fmt.Errorf("chunk %d checksum mismatch: expected %s, got %s", + chunk.Index, expectedChecksum, chunk.Checksum) + } + + return nil +} + +// calculateChecksum calculates a simple checksum for chunk content +func (cs *ChunkSerializer) calculateChecksum(content string) string { + // Simple checksum using length and first/last characters + if len(content) == 0 { + return "empty" + } + + checksum := fmt.Sprintf("len%d", len(content)) + if len(content) >= 1 { + checksum += fmt.Sprintf("_first%d", int(content[0])) + } + if len(content) >= 2 { + checksum += fmt.Sprintf("_last%d", int(content[len(content)-1])) + } + + return checksum +} + +// ChunkData represents a single chunk of token data +type ChunkData struct { + Index int // Position of this chunk in the sequence + Total int // Total number of chunks for this token + Content string // The actual chunk content + Checksum string // Simple checksum for integrity verification +} + +// EstimateChunkCount estimates how many chunks a token will need +func (cs *ChunkSerializer) EstimateChunkCount(tokenLength int, chunkSize int) int { + if chunkSize <= 0 { + chunkSize = maxCookieSize + } + + return (tokenLength + chunkSize - 1) / chunkSize +} + +// MaxTokenSizeForChunks calculates the maximum token size that can fit in the given number of chunks +func (cs *ChunkSerializer) MaxTokenSizeForChunks(maxChunks int, chunkSize int) int { + if chunkSize <= 0 { + chunkSize = maxCookieSize + } + + return maxChunks * chunkSize +} diff --git a/session/chunking/chunk_validator.go b/session/chunking/chunk_validator.go new file mode 100644 index 0000000..da58edf --- /dev/null +++ b/session/chunking/chunk_validator.go @@ -0,0 +1,429 @@ +// Package chunking provides chunk validation functionality +package chunking + +import ( + "encoding/base64" + "encoding/json" + "fmt" + "strings" + "unicode" +) + +// TokenValidator provides comprehensive validation for tokens and chunks +type TokenValidator struct{} + +// NewTokenValidator creates a new token validator +func NewTokenValidator() *TokenValidator { + return &TokenValidator{} +} + +// ValidateTokenSize validates that a token is within size limits +func (tv *TokenValidator) ValidateTokenSize(token string, config TokenConfig) error { + if len(token) == 0 { + return nil // Empty token is allowed + } + + if len(token) < config.MinLength { + return &ValidationError{ + Type: config.Type, + Reason: "token too short", + Details: fmt.Sprintf("length %d < minimum %d", len(token), config.MinLength), + } + } + + if len(token) > config.MaxLength { + return &ValidationError{ + Type: config.Type, + Reason: "token too long", + Details: fmt.Sprintf("length %d > maximum %d", len(token), config.MaxLength), + } + } + + return nil +} + +// ValidateJWTFormat validates that a token has proper JWT format +func (tv *TokenValidator) ValidateJWTFormat(token string, tokenType string) error { + if token == "" { + return nil // Empty token is not an error + } + + // JWT tokens must have exactly 3 parts separated by dots + parts := strings.Split(token, ".") + if len(parts) != 3 { + return &ValidationError{ + Type: tokenType, + Reason: "invalid JWT format", + Details: fmt.Sprintf("expected 3 parts, got %d", len(parts)), + } + } + + // Each part must be non-empty + for i, part := range parts { + if part == "" { + return &ValidationError{ + Type: tokenType, + Reason: "empty JWT part", + Details: fmt.Sprintf("part %d is empty", i+1), + } + } + } + + // Validate each part is valid base64 + for i, part := range parts { + if err := tv.validateBase64JWT(part); err != nil { + return &ValidationError{ + Type: tokenType, + Reason: "invalid base64 in JWT part", + Details: fmt.Sprintf("part %d: %v", i+1, err), + } + } + } + + return nil +} + +// ValidateTokenContent performs comprehensive content validation +func (tv *TokenValidator) ValidateTokenContent(token string, config TokenConfig) error { + if token == "" { + return nil + } + + // Validate character set + if err := tv.validateCharacterSet(token, config); err != nil { + return err + } + + // Validate token structure based on type + if config.RequireJWTFormat { + return tv.validateJWTContent(token, config) + } else if config.AllowOpaqueTokens { + return tv.validateOpaqueTokenContent(token, config) + } else { + // Try JWT first, then fall back to opaque validation + if err := tv.validateJWTContent(token, config); err != nil { + return tv.validateOpaqueTokenContent(token, config) + } + return nil + } +} + +// validateCharacterSet validates the character set of a token +func (tv *TokenValidator) validateCharacterSet(token string, config TokenConfig) error { + for i, r := range token { + if !tv.isValidTokenCharacter(r) { + return &ValidationError{ + Type: config.Type, + Reason: "invalid character", + Details: fmt.Sprintf("invalid character at position %d: %c (0x%X)", i, r, r), + } + } + } + return nil +} + +// isValidTokenCharacter checks if a character is valid in a token +func (tv *TokenValidator) isValidTokenCharacter(r rune) bool { + // Allow alphanumeric characters + if unicode.IsLetter(r) || unicode.IsNumber(r) { + return true + } + + // Allow common token characters + validChars := ".-_~:/?#[]@!$&'()*+,;=" + return strings.ContainsRune(validChars, r) +} + +// validateJWTContent validates the content of a JWT token +func (tv *TokenValidator) validateJWTContent(token string, config TokenConfig) error { + parts := strings.Split(token, ".") + if len(parts) != 3 { + return &ValidationError{ + Type: config.Type, + Reason: "invalid JWT structure", + Details: "JWT must have exactly 3 parts", + } + } + + // Validate header + if err := tv.validateJWTHeader(parts[0], config); err != nil { + return err + } + + // Validate payload + if err := tv.validateJWTPayload(parts[1], config); err != nil { + return err + } + + // Validate signature + if err := tv.validateJWTSignature(parts[2], config); err != nil { + return err + } + + return nil +} + +// validateJWTHeader validates a JWT header +func (tv *TokenValidator) validateJWTHeader(header string, config TokenConfig) error { + decoded, err := tv.base64URLDecode(header) + if err != nil { + return &ValidationError{ + Type: config.Type, + Reason: "invalid header encoding", + Details: err.Error(), + } + } + + var headerData map[string]interface{} + if err := json.Unmarshal(decoded, &headerData); err != nil { + return &ValidationError{ + Type: config.Type, + Reason: "invalid header JSON", + Details: err.Error(), + } + } + + // Check required fields + if _, ok := headerData["alg"]; !ok { + return &ValidationError{ + Type: config.Type, + Reason: "missing algorithm", + Details: "JWT header must contain 'alg' field", + } + } + + if _, ok := headerData["typ"]; !ok { + return &ValidationError{ + Type: config.Type, + Reason: "missing type", + Details: "JWT header must contain 'typ' field", + } + } + + return nil +} + +// validateJWTPayload validates a JWT payload +func (tv *TokenValidator) validateJWTPayload(payload string, config TokenConfig) error { + decoded, err := tv.base64URLDecode(payload) + if err != nil { + return &ValidationError{ + Type: config.Type, + Reason: "invalid payload encoding", + Details: err.Error(), + } + } + + var payloadData map[string]interface{} + if err := json.Unmarshal(decoded, &payloadData); err != nil { + return &ValidationError{ + Type: config.Type, + Reason: "invalid payload JSON", + Details: err.Error(), + } + } + + // For ID tokens, check required claims + if config.Type == "id" { + requiredClaims := []string{"iss", "sub", "aud", "exp", "iat"} + for _, claim := range requiredClaims { + if _, ok := payloadData[claim]; !ok { + return &ValidationError{ + Type: config.Type, + Reason: "missing required claim", + Details: fmt.Sprintf("ID token must contain '%s' claim", claim), + } + } + } + } + + return nil +} + +// validateJWTSignature validates a JWT signature part +func (tv *TokenValidator) validateJWTSignature(signature string, config TokenConfig) error { + if signature == "" { + return &ValidationError{ + Type: config.Type, + Reason: "empty signature", + Details: "JWT signature cannot be empty", + } + } + + // Just validate it's valid base64URL + _, err := tv.base64URLDecode(signature) + if err != nil { + return &ValidationError{ + Type: config.Type, + Reason: "invalid signature encoding", + Details: err.Error(), + } + } + + return nil +} + +// validateOpaqueTokenContent validates opaque token content +func (tv *TokenValidator) validateOpaqueTokenContent(token string, config TokenConfig) error { + if token == "" { + return nil + } + + // Basic sanity checks for opaque tokens + if len(token) < 8 { + return &ValidationError{ + Type: config.Type, + Reason: "token too short for opaque token", + Details: "opaque tokens should be at least 8 characters", + } + } + + // Check for reasonable entropy + if tv.hasLowEntropy(token) { + return &ValidationError{ + Type: config.Type, + Reason: "low entropy", + Details: "token appears to have low entropy", + } + } + + return nil +} + +// hasLowEntropy checks if a token has suspiciously low entropy +func (tv *TokenValidator) hasLowEntropy(token string) bool { + if len(token) < 8 { + return true + } + + // Count unique characters + uniqueChars := make(map[rune]bool) + for _, r := range token { + uniqueChars[r] = true + } + + // If less than 50% of characters are unique, consider it low entropy + entropyRatio := float64(len(uniqueChars)) / float64(len(token)) + return entropyRatio < 0.5 +} + +// validateBase64JWT validates base64URL encoding +func (tv *TokenValidator) validateBase64JWT(data string) error { + _, err := tv.base64URLDecode(data) + return err +} + +// base64URLDecode decodes base64URL encoded data +func (tv *TokenValidator) base64URLDecode(data string) ([]byte, error) { + // Add padding if needed + switch len(data) % 4 { + case 2: + data += "==" + case 3: + data += "=" + } + + // Replace URL-safe characters + data = strings.ReplaceAll(data, "-", "+") + data = strings.ReplaceAll(data, "_", "/") + + return base64.StdEncoding.DecodeString(data) +} + +// ValidateChunkStructure validates the structure of chunk data +func (tv *TokenValidator) ValidateChunkStructure(chunks []ChunkData, config TokenConfig) error { + if len(chunks) == 0 { + return &ValidationError{ + Type: config.Type, + Reason: "no chunks provided", + Details: "chunk list is empty", + } + } + + if len(chunks) > config.MaxChunks { + return &ValidationError{ + Type: config.Type, + Reason: "too many chunks", + Details: fmt.Sprintf("got %d chunks, maximum is %d", len(chunks), config.MaxChunks), + } + } + + // Validate each chunk + expectedTotal := chunks[0].Total + seenIndices := make(map[int]bool) + + for i, chunk := range chunks { + // Check for duplicate indices + if seenIndices[chunk.Index] { + return &ValidationError{ + Type: config.Type, + Reason: "duplicate chunk index", + Details: fmt.Sprintf("chunk index %d appears multiple times", chunk.Index), + } + } + seenIndices[chunk.Index] = true + + // Validate individual chunk + if err := tv.validateChunkData(chunk, expectedTotal, config); err != nil { + return &ValidationError{ + Type: config.Type, + Reason: "invalid chunk data", + Details: fmt.Sprintf("chunk %d: %v", i, err), + } + } + } + + // Check for missing indices + for i := 0; i < expectedTotal; i++ { + if !seenIndices[i] { + return &ValidationError{ + Type: config.Type, + Reason: "missing chunk index", + Details: fmt.Sprintf("chunk with index %d is missing", i), + } + } + } + + return nil +} + +// validateChunkData validates individual chunk data +func (tv *TokenValidator) validateChunkData(chunk ChunkData, expectedTotal int, config TokenConfig) error { + if chunk.Index < 0 { + return fmt.Errorf("negative index: %d", chunk.Index) + } + + if chunk.Total != expectedTotal { + return fmt.Errorf("inconsistent total: got %d, expected %d", chunk.Total, expectedTotal) + } + + if chunk.Index >= chunk.Total { + return fmt.Errorf("index %d exceeds total %d", chunk.Index, chunk.Total) + } + + if chunk.Content == "" { + return fmt.Errorf("empty content") + } + + if len(chunk.Content) > config.MaxChunkSize { + return fmt.Errorf("chunk too large: %d > %d", len(chunk.Content), config.MaxChunkSize) + } + + if chunk.Checksum == "" { + return fmt.Errorf("empty checksum") + } + + return nil +} + +// ValidationError represents a validation error +type ValidationError struct { + Type string + Reason string + Details string +} + +// Error implements the error interface +func (ve *ValidationError) Error() string { + return fmt.Sprintf("%s validation error: %s - %s", ve.Type, ve.Reason, ve.Details) +} diff --git a/session/core/session_manager.go b/session/core/session_manager.go new file mode 100644 index 0000000..e5f0d85 --- /dev/null +++ b/session/core/session_manager.go @@ -0,0 +1,336 @@ +// Package core provides core session management functionality for the OIDC middleware +package core + +import ( + "fmt" + "net/http" + "strings" + "sync" + "time" + + "github.com/gorilla/sessions" +) + +const ( + minEncryptionKeyLength = 32 + absoluteSessionTimeout = 24 * time.Hour +) + +// SessionManager handles session creation, management and cleanup +type SessionManager struct { + sessionPool sync.Pool + store sessions.Store + logger Logger + chunkManager ChunkManager + cookieDomain string + cleanupMutex sync.RWMutex + forceHTTPS bool + cleanupDone bool +} + +// Logger interface for dependency injection +type Logger interface { + Debug(msg string) + Debugf(format string, args ...interface{}) + Error(msg string) + Errorf(format string, args ...interface{}) +} + +// ChunkManager interface for chunk operations +type ChunkManager interface { + CleanupExpiredSessions() +} + +// SessionData interface for session data operations +type SessionData interface { + Reset() + SetManager(manager *SessionManager) + SetAuthenticated(bool) error + GetAuthenticated() bool + GetAccessToken() string + GetRefreshToken() string + GetIDToken() string + GetEmail() string + GetCSRF() string + GetNonce() string + GetCodeVerifier() string + GetIncomingPath() string + GetRedirectCount() int + IncrementRedirectCount() + ResetRedirectCount() + MarkDirty() + IsDirty() bool + Save(r *http.Request, w http.ResponseWriter) error + Clear(r *http.Request, w http.ResponseWriter) error + GetRefreshTokenIssuedAt() time.Time + returnToPoolSafely() +} + +// NewSessionManager creates a new SessionManager instance with secure defaults. +// It initializes the cookie store with encryption, sets up session pooling, +// and configures chunk management for large tokens. +func NewSessionManager(encryptionKey string, forceHTTPS bool, cookieDomain string, logger Logger, chunkManager ChunkManager) (*SessionManager, error) { + if len(encryptionKey) < minEncryptionKeyLength { + return nil, fmt.Errorf("encryption key must be at least %d bytes long", minEncryptionKeyLength) + } + + sm := &SessionManager{ + store: sessions.NewCookieStore([]byte(encryptionKey)), + forceHTTPS: forceHTTPS, + cookieDomain: cookieDomain, + logger: logger, + chunkManager: chunkManager, + } + + sm.sessionPool.New = func() interface{} { + return NewSessionData(sm, logger) + } + + return sm, nil +} + +// GetSession retrieves or creates a session for the request +func (sm *SessionManager) GetSession(r *http.Request) (SessionData, error) { + sessionDataInterface := sm.sessionPool.Get() + sessionData, ok := sessionDataInterface.(SessionData) + if !ok || sessionData == nil { + sessionData = NewSessionData(sm, sm.logger) + } + + // Initialize the session data + err := sm.initializeSession(sessionData, r) + if err != nil { + sm.sessionPool.Put(sessionData) + return nil, fmt.Errorf("failed to initialize session: %w", err) + } + + return sessionData, nil +} + +// initializeSession initializes session data from HTTP request +func (sm *SessionManager) initializeSession(sessionData SessionData, r *http.Request) error { + // Reset session data to clean state + sessionData.Reset() + sessionData.SetManager(sm) + + // Load session data from cookies + session, err := sm.store.Get(r, MainCookieName()) + if err != nil { + sm.logger.Debugf("Error getting main session: %v", err) + return nil // Not a fatal error, will create new session + } + + // Extract and set session values + if auth, ok := session.Values["authenticated"].(bool); ok { + sessionData.SetAuthenticated(auth) + } + + return nil +} + +// CleanupOldCookies removes old/expired cookies from the response +func (sm *SessionManager) CleanupOldCookies(w http.ResponseWriter, r *http.Request) { + sm.cleanupMutex.Lock() + defer sm.cleanupMutex.Unlock() + + if sm.cleanupDone { + return + } + + sm.logger.Debug("Starting cleanup of old session cookies") + + oldCookieNames := []string{ + "_oidc_session_old_v1", + "_oidc_session_legacy", + "_oidc_auth_state_old", + "_legacy_oidc_token", + "_old_session_chunks", + } + + for _, cookieName := range oldCookieNames { + if cookie, err := r.Cookie(cookieName); err == nil && cookie.Value != "" { + sm.logger.Debugf("Expiring old cookie: %s", cookieName) + expiredCookie := &http.Cookie{ + Name: cookieName, + Value: "", + Path: "/", + Domain: sm.cookieDomain, + Expires: time.Unix(0, 0), + MaxAge: -1, + Secure: sm.shouldUseSecureCookies(r), + HttpOnly: true, + SameSite: http.SameSiteLaxMode, + } + http.SetCookie(w, expiredCookie) + } + } + + sm.cleanupDone = true +} + +// PeriodicChunkCleanup performs comprehensive session maintenance and cleanup +func (sm *SessionManager) PeriodicChunkCleanup() { + if sm == nil || sm.logger == nil { + return + } + + sm.logger.Debug("Starting comprehensive session cleanup cycle") + + cleanupStart := time.Now() + var orphanedChunks, expiredSessions, cleanupErrors int + + if sm.store != nil { + if cookieStore, ok := sm.store.(*sessions.CookieStore); ok { + sm.logger.Debug("Running session store cleanup") + _ = cookieStore + } + } + + // Cleanup expired sessions in chunk manager to prevent memory leaks + if sm.chunkManager != nil { + sm.chunkManager.CleanupExpiredSessions() + } + + poolCleaned := 0 + for i := 0; i < 10; i++ { + if poolSession := sm.sessionPool.Get(); poolSession != nil { + if sessionData, ok := poolSession.(SessionData); ok && sessionData != nil { + sessionData.Reset() + poolCleaned++ + } + sm.sessionPool.Put(poolSession) + } + } + + cleanupDuration := time.Since(cleanupStart) + sm.logger.Debugf("Session cleanup completed in %v: pool_cleaned=%d, orphaned_chunks=%d, expired_sessions=%d, errors=%d", + cleanupDuration, poolCleaned, orphanedChunks, expiredSessions, cleanupErrors) +} + +// ValidateSessionHealth performs comprehensive validation of session integrity +func (sm *SessionManager) ValidateSessionHealth(sessionData SessionData) error { + if sessionData == nil { + return fmt.Errorf("session data is nil") + } + + // Check if user is authenticated + if !sessionData.GetAuthenticated() { + return nil // Not authenticated is not an error + } + + // Validate token formats + if accessToken := sessionData.GetAccessToken(); accessToken != "" { + if err := sm.validateTokenFormat(accessToken, "access"); err != nil { + return fmt.Errorf("invalid access token format: %w", err) + } + } + + if idToken := sessionData.GetIDToken(); idToken != "" { + if err := sm.validateTokenFormat(idToken, "id"); err != nil { + return fmt.Errorf("invalid ID token format: %w", err) + } + } + + // Check for session tampering + if err := sm.detectSessionTampering(sessionData); err != nil { + return fmt.Errorf("session tampering detected: %w", err) + } + + return nil +} + +// validateTokenFormat validates the format of JWT tokens +func (sm *SessionManager) validateTokenFormat(token, tokenType string) error { + if token == "" { + return nil + } + + // JWT tokens should have exactly 3 parts separated by dots + parts := strings.Split(token, ".") + if len(parts) != 3 { + return fmt.Errorf("%s token is not a valid JWT format", tokenType) + } + + // Each part should be non-empty + for i, part := range parts { + if part == "" { + return fmt.Errorf("%s token part %d is empty", tokenType, i+1) + } + } + + return nil +} + +// detectSessionTampering detects potential tampering in session data +func (sm *SessionManager) detectSessionTampering(sessionData SessionData) error { + email := sessionData.GetEmail() + authenticated := sessionData.GetAuthenticated() + + // If authenticated but no email, that's suspicious + if authenticated && email == "" { + return fmt.Errorf("authenticated session without email") + } + + // If email exists but not authenticated, that's also suspicious + if !authenticated && email != "" { + sm.logger.Debugf("Warning: Email exists (%s) but session not authenticated", email) + } + + return nil +} + +// GetSessionMetrics returns metrics about session usage +func (sm *SessionManager) GetSessionMetrics() map[string]interface{} { + metrics := make(map[string]interface{}) + + metrics["store_type"] = fmt.Sprintf("%T", sm.store) + metrics["cookie_domain"] = sm.cookieDomain + metrics["force_https"] = sm.forceHTTPS + metrics["cleanup_done"] = sm.cleanupDone + + return metrics +} + +// shouldUseSecureCookies determines if cookies should be secure based on request +func (sm *SessionManager) shouldUseSecureCookies(r *http.Request) bool { + if sm.forceHTTPS { + return true + } + + // Check if the request came over HTTPS + if r.TLS != nil { + return true + } + + // Check X-Forwarded-Proto header + if proto := r.Header.Get("X-Forwarded-Proto"); proto == "https" { + return true + } + + return false +} + +// getSessionOptions returns session options for the given security context +func (sm *SessionManager) getSessionOptions(isSecure bool) *sessions.Options { + return &sessions.Options{ + Path: "/", + Domain: sm.cookieDomain, + MaxAge: int(absoluteSessionTimeout.Seconds()), + Secure: isSecure, + HttpOnly: true, + SameSite: http.SameSiteLaxMode, + } +} + +// Cookie name functions +func MainCookieName() string { return "_oidc_raczylo_m" } +func AccessTokenCookie() string { return "_oidc_raczylo_a" } +func RefreshTokenCookie() string { return "_oidc_raczylo_r" } +func IDTokenCookie() string { return "_oidc_raczylo_id" } + +// NewSessionData creates a new session data instance +func NewSessionData(manager *SessionManager, logger Logger) SessionData { + // This function should be implemented to return a concrete SessionData implementation + // The actual implementation depends on the SessionData struct definition + return nil +} diff --git a/session/core/session_manager_test.go b/session/core/session_manager_test.go new file mode 100644 index 0000000..372b9c5 --- /dev/null +++ b/session/core/session_manager_test.go @@ -0,0 +1,1010 @@ +package core + +import ( + "crypto/tls" + "fmt" + "net/http" + "net/http/httptest" + "runtime" + "testing" + "time" +) + +// Mock logger for testing +type MockLogger struct { + logs []string +} + +func (ml *MockLogger) Debug(msg string) { + ml.logs = append(ml.logs, "DEBUG: "+msg) +} + +func (ml *MockLogger) Debugf(format string, args ...interface{}) { + ml.logs = append(ml.logs, fmt.Sprintf("DEBUG: "+format, args...)) +} + +func (ml *MockLogger) Error(msg string) { + ml.logs = append(ml.logs, "ERROR: "+msg) +} + +func (ml *MockLogger) Errorf(format string, args ...interface{}) { + ml.logs = append(ml.logs, fmt.Sprintf("ERROR: "+format, args...)) +} + +// Mock chunk manager for testing +type MockChunkManager struct { + cleanupCalled int +} + +func (mcm *MockChunkManager) CleanupExpiredSessions() { + mcm.cleanupCalled++ +} + +// Mock session data for testing +type MockSessionData struct { + manager *SessionManager + authenticated bool + dirty bool + clearCalled int + email string + emailSet bool // Flag to indicate if email was explicitly set +} + +func (msd *MockSessionData) Reset() { + msd.authenticated = false + msd.dirty = false +} + +func (msd *MockSessionData) SetManager(manager *SessionManager) { + msd.manager = manager +} + +func (msd *MockSessionData) SetAuthenticated(auth bool) error { + msd.authenticated = auth + return nil +} + +func (msd *MockSessionData) GetAuthenticated() bool { + return msd.authenticated +} + +func (msd *MockSessionData) GetAccessToken() string { + if msd.authenticated { + return "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIn0.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c" + } + return "" +} +func (msd *MockSessionData) GetRefreshToken() string { + if msd.authenticated { + return "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIn0.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c" + } + return "" +} +func (msd *MockSessionData) GetIDToken() string { + if msd.authenticated { + return "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIn0.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c" + } + return "" +} +func (msd *MockSessionData) GetEmail() string { + // If email was explicitly set, return it (even if empty) + if msd.emailSet { + return msd.email + } + // Default behavior for authenticated sessions + if msd.authenticated { + return "user@example.com" + } + return "" +} +func (msd *MockSessionData) GetCSRF() string { return "" } +func (msd *MockSessionData) GetNonce() string { return "" } +func (msd *MockSessionData) GetCodeVerifier() string { return "" } +func (msd *MockSessionData) GetIncomingPath() string { return "" } +func (msd *MockSessionData) GetRedirectCount() int { return 0 } +func (msd *MockSessionData) IncrementRedirectCount() {} +func (msd *MockSessionData) ResetRedirectCount() {} +func (msd *MockSessionData) MarkDirty() { msd.dirty = true } +func (msd *MockSessionData) IsDirty() bool { return msd.dirty } +func (msd *MockSessionData) Save(r *http.Request, w http.ResponseWriter) error { return nil } +func (msd *MockSessionData) GetRefreshTokenIssuedAt() time.Time { return time.Now() } +func (msd *MockSessionData) returnToPoolSafely() {} + +func (msd *MockSessionData) Clear(r *http.Request, w http.ResponseWriter) error { + msd.clearCalled++ + msd.returnToPoolSafely() + return nil +} + +// NewMockSessionData creates a new mock session data +func NewMockSessionData(manager *SessionManager, logger Logger) SessionData { + return &MockSessionData{manager: manager} +} + +// TestSessionManagerCreation tests session manager creation +func TestSessionManagerCreation(t *testing.T) { + tests := []struct { + name string + encryptionKey string + expectError bool + expectedKeyLen int + description string + }{ + { + name: "Valid encryption key", + encryptionKey: "0123456789abcdef0123456789abcdef0123456789abcdef", + expectError: false, + expectedKeyLen: 48, + description: "Should successfully create session manager with valid key", + }, + { + name: "Minimum length key", + encryptionKey: "0123456789abcdef0123456789abcdef", + expectError: false, + expectedKeyLen: 32, + description: "Should accept key at minimum length", + }, + { + name: "Too short key", + encryptionKey: "tooshort", + expectError: true, + expectedKeyLen: 0, + description: "Should reject keys that are too short", + }, + { + name: "Empty key", + encryptionKey: "", + expectError: true, + expectedKeyLen: 0, + description: "Should reject empty keys", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + logger := &MockLogger{} + chunkManager := &MockChunkManager{} + + sm, err := NewSessionManager(tt.encryptionKey, false, "", logger, chunkManager) + + if tt.expectError { + if err == nil { + t.Errorf("Expected error for %s, got nil", tt.description) + } + return + } + + if err != nil { + t.Errorf("Unexpected error for %s: %v", tt.description, err) + return + } + + if sm == nil { + t.Errorf("Session manager should not be nil for %s", tt.description) + return + } + + // Verify the session manager is properly initialized + if sm.logger == nil { + t.Error("Logger should be set") + } + + if sm.store == nil { + t.Error("Store should be initialized") + } + }) + } +} + +// TestSessionManagerPoolBehavior tests session pooling behavior +func TestSessionManagerPoolBehavior(t *testing.T) { + logger := &MockLogger{} + chunkManager := &MockChunkManager{} + sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger, chunkManager) + if err != nil { + t.Fatalf("Failed to create session manager: %v", err) + } + + // Override the session pool to use our mock + sm.sessionPool.New = func() interface{} { + return NewMockSessionData(sm, logger) + } + + tests := []struct { + name string + description string + operation func(t *testing.T, sm *SessionManager) + }{ + { + name: "Session creation and return", + description: "Test that sessions are properly created and returned to pool", + operation: func(t *testing.T, sm *SessionManager) { + req := httptest.NewRequest("GET", "http://example.com/foo", nil) + + session, err := sm.GetSession(req) + if err != nil { + t.Fatalf("GetSession failed: %v", err) + } + + if session == nil { + t.Fatal("Session should not be nil") + } + + // Clear should return the session to pool + w := httptest.NewRecorder() + err = session.Clear(req, w) + if err != nil { + t.Logf("Clear returned error (this may be expected): %v", err) + } + }, + }, + { + name: "Multiple sessions", + description: "Test creating multiple sessions", + operation: func(t *testing.T, sm *SessionManager) { + req := httptest.NewRequest("GET", "http://example.com/foo", nil) + + // Create multiple sessions + sessions := make([]SessionData, 5) + for i := 0; i < 5; i++ { + session, err := sm.GetSession(req) + if err != nil { + t.Fatalf("GetSession %d failed: %v", i, err) + } + sessions[i] = session + } + + // Clear all sessions + w := httptest.NewRecorder() + for i, session := range sessions { + err := session.Clear(req, w) + if err != nil { + t.Logf("Clear session %d returned error: %v", i, err) + } + } + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Record initial goroutine count + initialGoroutines := runtime.NumGoroutine() + + tt.operation(t, sm) + + // Force garbage collection + runtime.GC() + time.Sleep(10 * time.Millisecond) + + // Check for goroutine leaks + finalGoroutines := runtime.NumGoroutine() + if finalGoroutines > initialGoroutines+2 { // Allow small tolerance + t.Errorf("Potential goroutine leak: started with %d, ended with %d", + initialGoroutines, finalGoroutines) + } + }) + } +} + +// TestSessionManagerErrorHandling tests error handling scenarios +func TestSessionManagerErrorHandling(t *testing.T) { + logger := &MockLogger{} + chunkManager := &MockChunkManager{} + sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger, chunkManager) + if err != nil { + t.Fatalf("Failed to create session manager: %v", err) + } + + // Override the session pool to use our mock + sm.sessionPool.New = func() interface{} { + return NewMockSessionData(sm, logger) + } + + tests := []struct { + name string + description string + setupReq func() *http.Request + expectError bool + errorCheck func(error) bool + }{ + { + name: "Corrupt cookie value", + description: "Test handling of corrupted cookie values", + setupReq: func() *http.Request { + req := httptest.NewRequest("GET", "http://example.com/foo", nil) + req.AddCookie(&http.Cookie{ + Name: MainCookieName(), + Value: "corrupt-value", + }) + return req + }, + expectError: false, // Session manager should gracefully handle corrupted cookies + errorCheck: nil, + }, + { + name: "Invalid base64 cookie", + description: "Test handling of invalid base64 in cookies", + setupReq: func() *http.Request { + req := httptest.NewRequest("GET", "http://example.com/foo", nil) + req.AddCookie(&http.Cookie{ + Name: MainCookieName(), + Value: "!@#$%^&*()", + }) + return req + }, + expectError: false, // Session manager should gracefully handle invalid base64 + errorCheck: nil, + }, + { + name: "Empty cookie value", + description: "Test handling of empty cookie values", + setupReq: func() *http.Request { + req := httptest.NewRequest("GET", "http://example.com/foo", nil) + req.AddCookie(&http.Cookie{ + Name: MainCookieName(), + Value: "", + }) + return req + }, + expectError: false, + errorCheck: nil, + }, + { + name: "Normal request", + description: "Test normal request without cookies", + setupReq: func() *http.Request { + return httptest.NewRequest("GET", "http://example.com/foo", nil) + }, + expectError: false, + errorCheck: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req := tt.setupReq() + + _, err := sm.GetSession(req) + + if tt.expectError { + if err == nil { + t.Errorf("Expected error for %s, got nil", tt.description) + return + } + + if tt.errorCheck != nil && !tt.errorCheck(err) { + t.Errorf("Error check failed for %s: %v", tt.description, err) + } + } else { + if err != nil { + t.Errorf("Unexpected error for %s: %v", tt.description, err) + } + } + }) + } +} + +// TestSessionManagerCleanup tests cleanup functionality +func TestSessionManagerCleanup(t *testing.T) { + logger := &MockLogger{} + mockChunkManager := &MockChunkManager{} + + sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger, mockChunkManager) + if err != nil { + t.Fatalf("Failed to create session manager: %v", err) + } + + t.Run("PeriodicChunkCleanup called", func(t *testing.T) { + initialCalls := mockChunkManager.cleanupCalled + + sm.PeriodicChunkCleanup() + + // Note: The actual cleanup may or may not be called depending on internal logic + // This test just ensures the method exists and can be called + t.Logf("Cleanup called %d times after PeriodicChunkCleanup", + mockChunkManager.cleanupCalled-initialCalls) + }) + + t.Run("CleanupOldCookies functionality", func(t *testing.T) { + req := httptest.NewRequest("GET", "http://example.com/foo", nil) + w := httptest.NewRecorder() + + // This should not panic and should handle cleanup properly + sm.CleanupOldCookies(w, req) + + // Verify response was written (cookies cleared) + if w.Code == 0 { + w.Code = 200 // Default to OK if no explicit code was set + } + }) +} + +// TestSessionManagerHTTPSBehavior tests HTTPS-related behavior +func TestSessionManagerHTTPSBehavior(t *testing.T) { + tests := []struct { + name string + forceHTTPS bool + requestURL string + expectError bool + description string + }{ + { + name: "HTTPS forced with HTTP request", + forceHTTPS: true, + requestURL: "http://example.com/foo", + expectError: false, // Manager creation shouldn't fail + description: "Should create manager even when HTTPS is forced", + }, + { + name: "HTTPS forced with HTTPS request", + forceHTTPS: true, + requestURL: "https://example.com/foo", + expectError: false, + description: "Should work normally with HTTPS request", + }, + { + name: "HTTPS not forced with HTTP request", + forceHTTPS: false, + requestURL: "http://example.com/foo", + expectError: false, + description: "Should work normally when HTTPS not forced", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + logger := &MockLogger{} + chunkManager := &MockChunkManager{} + + sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", + tt.forceHTTPS, "", logger, chunkManager) + + if tt.expectError { + if err == nil { + t.Errorf("Expected error for %s, got nil", tt.description) + } + return + } + + if err != nil { + t.Errorf("Unexpected error for %s: %v", tt.description, err) + return + } + + // Override the session pool to use our mock + sm.sessionPool.New = func() interface{} { + return NewMockSessionData(sm, logger) + } + + // Test session creation with the configured HTTPS behavior + req := httptest.NewRequest("GET", tt.requestURL, nil) + session, err := sm.GetSession(req) + + if err != nil { + t.Logf("GetSession returned error (may be expected): %v", err) + } else if session == nil { + t.Error("Session should not be nil when no error occurred") + } + }) + } +} + +// TestSessionManagerCookieDomain tests cookie domain configuration +func TestSessionManagerCookieDomain(t *testing.T) { + tests := []struct { + name string + cookieDomain string + description string + }{ + { + name: "Empty cookie domain", + cookieDomain: "", + description: "Should work with empty cookie domain", + }, + { + name: "Specific cookie domain", + cookieDomain: "example.com", + description: "Should work with specific cookie domain", + }, + { + name: "Subdomain cookie domain", + cookieDomain: ".example.com", + description: "Should work with subdomain cookie domain", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + logger := &MockLogger{} + chunkManager := &MockChunkManager{} + + sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", + false, tt.cookieDomain, logger, chunkManager) + + if err != nil { + t.Errorf("Unexpected error for %s: %v", tt.description, err) + return + } + + if sm == nil { + t.Errorf("Session manager should not be nil for %s", tt.description) + return + } + + if sm.cookieDomain != tt.cookieDomain { + t.Errorf("Cookie domain mismatch: expected %q, got %q", + tt.cookieDomain, sm.cookieDomain) + } + }) + } +} + +// BenchmarkSessionManagerCreation benchmarks session manager creation +func BenchmarkSessionManagerCreation(b *testing.B) { + logger := &MockLogger{} + chunkManager := &MockChunkManager{} + encryptionKey := "0123456789abcdef0123456789abcdef0123456789abcdef" + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + sm, err := NewSessionManager(encryptionKey, false, "", logger, chunkManager) + if err != nil { + b.Fatalf("Failed to create session manager: %v", err) + } + _ = sm + } +} + +// BenchmarkSessionManagerGetSession benchmarks session retrieval +func BenchmarkSessionManagerGetSession(b *testing.B) { + logger := &MockLogger{} + chunkManager := &MockChunkManager{} + sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger, chunkManager) + if err != nil { + b.Fatalf("Failed to create session manager: %v", err) + } + + // Override the session pool to use our mock + sm.sessionPool.New = func() interface{} { + return NewMockSessionData(sm, logger) + } + + req := httptest.NewRequest("GET", "http://example.com/foo", nil) + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + session, err := sm.GetSession(req) + if err != nil { + b.Fatalf("GetSession failed: %v", err) + } + + // Clean up the session + w := httptest.NewRecorder() + _ = session.Clear(req, w) + } +} + +//lint:ignore U1000 May be needed for future test utilities +func minInt(a, b int) int { + if a < b { + return a + } + return b +} + +// TestValidateSessionHealth tests session health validation +func TestValidateSessionHealth(t *testing.T) { + logger := &MockLogger{} + chunkManager := &MockChunkManager{} + sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger, chunkManager) + if err != nil { + t.Fatalf("Failed to create session manager: %v", err) + } + + tests := []struct { + name string + sessionData SessionData + expectError bool + description string + }{ + { + name: "Nil session data", + sessionData: nil, + expectError: true, + description: "Should fail with nil session data", + }, + { + name: "Unauthenticated session", + sessionData: &MockSessionData{authenticated: false}, + expectError: false, + description: "Should pass with unauthenticated session", + }, + { + name: "Authenticated session with tokens", + sessionData: &MockSessionData{authenticated: true}, + expectError: false, + description: "Should pass with properly authenticated session", + }, + { + name: "Authenticated session without email (suspicious)", + sessionData: &MockSessionData{authenticated: true}, + expectError: true, + description: "Should fail when authenticated but no email", + }, + } + + // Create a mock session with no email for the suspicious case + suspiciousSession := &MockSessionData{authenticated: true, email: "", emailSet: true} + + // Replace the fourth test case with our suspicious session + tests[3].sessionData = suspiciousSession + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := sm.ValidateSessionHealth(tt.sessionData) + + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, got none", tt.description) + } + if !tt.expectError && err != nil { + t.Errorf("Expected no error for %s, got: %v", tt.description, err) + } + }) + } +} + +// TestValidateTokenFormat tests token format validation +func TestValidateTokenFormat(t *testing.T) { + logger := &MockLogger{} + chunkManager := &MockChunkManager{} + sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger, chunkManager) + if err != nil { + t.Fatalf("Failed to create session manager: %v", err) + } + + tests := []struct { + name string + token string + tokenType string + expectError bool + description string + }{ + { + name: "Valid JWT token", + token: "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIn0.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c", + tokenType: "access", + expectError: false, + description: "Should pass with valid JWT", + }, + { + name: "Empty token", + token: "", + tokenType: "access", + expectError: false, + description: "Should pass with empty token", + }, + { + name: "Invalid token - too few parts", + token: "header.payload", + tokenType: "access", + expectError: true, + description: "Should fail with incomplete JWT", + }, + { + name: "Invalid token - too many parts", + token: "header.payload.signature.extra", + tokenType: "access", + expectError: true, + description: "Should fail with too many parts", + }, + { + name: "Invalid token - empty part", + token: "header..signature", + tokenType: "id", + expectError: true, + description: "Should fail with empty payload part", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := sm.validateTokenFormat(tt.token, tt.tokenType) + + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, got none", tt.description) + } + if !tt.expectError && err != nil { + t.Errorf("Expected no error for %s, got: %v", tt.description, err) + } + }) + } +} + +// TestDetectSessionTampering tests session tampering detection +func TestDetectSessionTampering(t *testing.T) { + logger := &MockLogger{} + chunkManager := &MockChunkManager{} + sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", false, "", logger, chunkManager) + if err != nil { + t.Fatalf("Failed to create session manager: %v", err) + } + + tests := []struct { + name string + authenticated bool + email string + expectError bool + description string + }{ + { + name: "Valid authenticated session", + authenticated: true, + email: "user@example.com", + expectError: false, + description: "Should pass with valid authenticated session", + }, + { + name: "Valid unauthenticated session", + authenticated: false, + email: "", + expectError: false, + description: "Should pass with valid unauthenticated session", + }, + { + name: "Suspicious: authenticated without email", + authenticated: true, + email: "", + expectError: true, + description: "Should fail when authenticated but no email", + }, + { + name: "Warning: email without authentication", + authenticated: false, + email: "user@example.com", + expectError: false, + description: "Should pass but log warning when email exists without authentication", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sessionData := &MockSessionData{authenticated: tt.authenticated, email: tt.email, emailSet: true} + + err := sm.detectSessionTampering(sessionData) + + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, got none", tt.description) + } + if !tt.expectError && err != nil { + t.Errorf("Expected no error for %s, got: %v", tt.description, err) + } + }) + } +} + +// TestGetSessionMetrics tests session metrics retrieval +func TestGetSessionMetrics(t *testing.T) { + tests := []struct { + name string + forceHTTPS bool + cookieDomain string + description string + }{ + { + name: "Basic metrics", + forceHTTPS: false, + cookieDomain: "", + description: "Should return basic metrics", + }, + { + name: "HTTPS forced metrics", + forceHTTPS: true, + cookieDomain: "example.com", + description: "Should return metrics with HTTPS and domain", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + logger := &MockLogger{} + chunkManager := &MockChunkManager{} + sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", + tt.forceHTTPS, tt.cookieDomain, logger, chunkManager) + if err != nil { + t.Fatalf("Failed to create session manager: %v", err) + } + + metrics := sm.GetSessionMetrics() + + if metrics == nil { + t.Error("Metrics should not be nil") + return + } + + expectedKeys := []string{"store_type", "cookie_domain", "force_https", "cleanup_done"} + for _, key := range expectedKeys { + if _, exists := metrics[key]; !exists { + t.Errorf("Metrics should contain key %s", key) + } + } + + if metrics["force_https"] != tt.forceHTTPS { + t.Errorf("Expected force_https=%v, got %v", tt.forceHTTPS, metrics["force_https"]) + } + + if metrics["cookie_domain"] != tt.cookieDomain { + t.Errorf("Expected cookie_domain=%s, got %s", tt.cookieDomain, metrics["cookie_domain"]) + } + }) + } +} + +// TestShouldUseSecureCookies tests secure cookie determination +func TestShouldUseSecureCookies(t *testing.T) { + tests := []struct { + name string + forceHTTPS bool + requestSetup func() *http.Request + expected bool + description string + }{ + { + name: "Force HTTPS enabled", + forceHTTPS: true, + requestSetup: func() *http.Request { + return httptest.NewRequest("GET", "http://example.com/foo", nil) + }, + expected: true, + description: "Should return true when HTTPS is forced", + }, + { + name: "HTTPS request with TLS", + forceHTTPS: false, + requestSetup: func() *http.Request { + req := httptest.NewRequest("GET", "https://example.com/foo", nil) + req.TLS = &tls.ConnectionState{} // Mock TLS + return req + }, + expected: true, + description: "Should return true for HTTPS request", + }, + { + name: "HTTP request with X-Forwarded-Proto header", + forceHTTPS: false, + requestSetup: func() *http.Request { + req := httptest.NewRequest("GET", "http://example.com/foo", nil) + req.Header.Set("X-Forwarded-Proto", "https") + return req + }, + expected: true, + description: "Should return true when X-Forwarded-Proto is https", + }, + { + name: "Plain HTTP request", + forceHTTPS: false, + requestSetup: func() *http.Request { + return httptest.NewRequest("GET", "http://example.com/foo", nil) + }, + expected: false, + description: "Should return false for plain HTTP", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + logger := &MockLogger{} + chunkManager := &MockChunkManager{} + sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", + tt.forceHTTPS, "", logger, chunkManager) + if err != nil { + t.Fatalf("Failed to create session manager: %v", err) + } + + req := tt.requestSetup() + result := sm.shouldUseSecureCookies(req) + + if result != tt.expected { + t.Errorf("Expected %v for %s, got %v", tt.expected, tt.description, result) + } + }) + } +} + +// TestGetSessionOptions tests session options generation +func TestGetSessionOptions(t *testing.T) { + tests := []struct { + name string + cookieDomain string + isSecure bool + description string + }{ + { + name: "Secure options with domain", + cookieDomain: "example.com", + isSecure: true, + description: "Should create secure options with domain", + }, + { + name: "Insecure options without domain", + cookieDomain: "", + isSecure: false, + description: "Should create insecure options without domain", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + logger := &MockLogger{} + chunkManager := &MockChunkManager{} + sm, err := NewSessionManager("0123456789abcdef0123456789abcdef0123456789abcdef", + false, tt.cookieDomain, logger, chunkManager) + if err != nil { + t.Fatalf("Failed to create session manager: %v", err) + } + + options := sm.getSessionOptions(tt.isSecure) + + if options == nil { + t.Error("Options should not be nil") + return + } + + if options.Secure != tt.isSecure { + t.Errorf("Expected Secure=%v, got %v", tt.isSecure, options.Secure) + } + + if options.Domain != tt.cookieDomain { + t.Errorf("Expected Domain=%s, got %s", tt.cookieDomain, options.Domain) + } + + if options.Path != "/" { + t.Errorf("Expected Path=/, got %s", options.Path) + } + + if !options.HttpOnly { + t.Error("Expected HttpOnly=true") + } + + if options.SameSite != http.SameSiteLaxMode { + t.Errorf("Expected SameSite=Lax, got %v", options.SameSite) + } + + if options.MaxAge != int(absoluteSessionTimeout.Seconds()) { + t.Errorf("Expected MaxAge=%d, got %d", int(absoluteSessionTimeout.Seconds()), options.MaxAge) + } + }) + } +} + +// TestAccessTokenCookie tests AccessTokenCookie function +func TestAccessTokenCookie(t *testing.T) { + result := AccessTokenCookie() + expected := "_oidc_raczylo_a" + + if result != expected { + t.Errorf("Expected %s, got %s", expected, result) + } +} + +// TestRefreshTokenCookie tests RefreshTokenCookie function +func TestRefreshTokenCookie(t *testing.T) { + result := RefreshTokenCookie() + expected := "_oidc_raczylo_r" + + if result != expected { + t.Errorf("Expected %s, got %s", expected, result) + } +} + +// TestIDTokenCookie tests IDTokenCookie function +func TestIDTokenCookie(t *testing.T) { + result := IDTokenCookie() + expected := "_oidc_raczylo_id" + + if result != expected { + t.Errorf("Expected %s, got %s", expected, result) + } +} diff --git a/session/crypto/session_crypto.go b/session/crypto/session_crypto.go new file mode 100644 index 0000000..12fc229 --- /dev/null +++ b/session/crypto/session_crypto.go @@ -0,0 +1,264 @@ +// Package crypto provides cryptographic operations for session management +package crypto + +import ( + "bytes" + "compress/gzip" + "crypto/rand" + "encoding/base64" + "encoding/hex" + "fmt" + "io" + "strings" +) + +// MemoryPools interface for memory management +type MemoryPools interface { + GetCompressionBuffer() *bytes.Buffer + PutCompressionBuffer(*bytes.Buffer) + GetHTTPResponseBuffer() []byte + PutHTTPResponseBuffer([]byte) +} + +// SessionCrypto provides cryptographic operations for session data +type SessionCrypto struct { + memoryPools MemoryPools +} + +// NewSessionCrypto creates a new session crypto instance +func NewSessionCrypto(memoryPools MemoryPools) *SessionCrypto { + return &SessionCrypto{ + memoryPools: memoryPools, + } +} + +// GenerateSecureRandomString creates a cryptographically secure random string. +// It generates random bytes using crypto/rand and encodes them as hexadecimal. +// This is used for session IDs and other security-sensitive random values. +func (sc *SessionCrypto) GenerateSecureRandomString(length int) (string, error) { + bytes := make([]byte, length) + if _, err := rand.Read(bytes); err != nil { + return "", fmt.Errorf("failed to generate random bytes: %w", err) + } + return hex.EncodeToString(bytes), nil +} + +// CompressToken compresses a JWT token using gzip compression if beneficial. +// It validates the token format, attempts compression, and verifies the compressed +// data can be decompressed correctly. Only compresses if it reduces size. +func (sc *SessionCrypto) CompressToken(token string) string { + if token == "" { + return token + } + + // Validate JWT format (should have exactly 2 dots) + dotCount := strings.Count(token, ".") + if dotCount != 2 { + return token + } + + // Don't try to compress extremely large tokens + if len(token) > 50*1024 { + return token + } + + b := sc.memoryPools.GetCompressionBuffer() + defer sc.memoryPools.PutCompressionBuffer(b) + + gz := gzip.NewWriter(b) + + written, err := gz.Write([]byte(token)) + if err != nil || written != len(token) { + return token + } + + if err := gz.Close(); err != nil { + return token + } + + compressedBytes := b.Bytes() + if len(compressedBytes) == 0 { + return token + } + + compressed := base64.StdEncoding.EncodeToString(compressedBytes) + + // Only use compression if it actually reduces size + if len(compressed) >= len(token) { + return token + } + + // Verify compression integrity by attempting decompression + decompressed := sc.decompressTokenInternal(compressed) + if decompressed != token { + return token + } + + // Final validation of decompressed token + if strings.Count(decompressed, ".") != 2 { + return token + } + + return compressed +} + +// DecompressToken decompresses a previously compressed token string. +// It decodes the base64 data, validates gzip headers, and decompresses safely +// with size limits to prevent compression bombs. +func (sc *SessionCrypto) DecompressToken(compressed string) string { + return sc.decompressTokenInternal(compressed) +} + +// decompressTokenInternal is the internal decompression function. +// Separated internal function for integrity verification during compression. +// It performs the actual decompression logic with proper resource management. +func (sc *SessionCrypto) decompressTokenInternal(compressed string) string { + if compressed == "" { + return compressed + } + + // Prevent decompression of extremely large inputs + if len(compressed) > 100*1024 { + return compressed + } + + data, err := base64.StdEncoding.DecodeString(compressed) + if err != nil { + return compressed + } + + if len(data) == 0 { + return compressed + } + + // Validate gzip header + if len(data) < 2 || data[0] != 0x1f || data[1] != 0x8b { + return compressed + } + + readerBuf := sc.memoryPools.GetHTTPResponseBuffer() + defer sc.memoryPools.PutHTTPResponseBuffer(readerBuf) + + gz, err := gzip.NewReader(bytes.NewReader(data)) + if err != nil { + return compressed + } + + defer func() { + if closeErr := gz.Close(); closeErr != nil { + _ = closeErr + } + }() + + // Limit decompressed size to prevent compression bombs + limitedReader := io.LimitReader(gz, 500*1024) + + // Optimize for large buffer reuse + if cap(readerBuf) >= 512*1024 { + readerBuf = readerBuf[:cap(readerBuf)] + n, err := limitedReader.Read(readerBuf) + if err != nil && err != io.EOF { + return compressed + } + decompressed := readerBuf[:n] + return string(decompressed) + } + + decompressed, err := io.ReadAll(limitedReader) + if err != nil { + return compressed + } + + if len(decompressed) == 0 { + return compressed + } + + decompressedStr := string(decompressed) + + // Validate the decompressed token is a valid JWT + if decompressedStr != "" && strings.Count(decompressedStr, ".") != 2 { + return compressed + } + + return decompressedStr +} + +// ValidateTokenFormat validates that a token has the correct JWT format +func (sc *SessionCrypto) ValidateTokenFormat(token string) bool { + if token == "" { + return false + } + + // JWT tokens should have exactly 3 parts separated by dots + parts := strings.Split(token, ".") + if len(parts) != 3 { + return false + } + + // Each part should be non-empty + for _, part := range parts { + if part == "" { + return false + } + } + + return true +} + +// IsTokenCompressed checks if a token appears to be compressed +func (sc *SessionCrypto) IsTokenCompressed(token string) bool { + if token == "" { + return false + } + + // JWT tokens have exactly 2 dots, compressed tokens don't + if strings.Count(token, ".") == 2 { + return false + } + + // Try to decode as base64 + data, err := base64.StdEncoding.DecodeString(token) + if err != nil { + return false + } + + // Check for gzip header + if len(data) >= 2 && data[0] == 0x1f && data[1] == 0x8b { + return true + } + + return false +} + +// SecureWipeBytes securely wipes sensitive data from memory +func (sc *SessionCrypto) SecureWipeBytes(data []byte) { + for i := range data { + data[i] = 0 + } +} + +// SecureWipeString securely wipes sensitive string data +func (sc *SessionCrypto) SecureWipeString(s *string) { + if s != nil { + *s = "" + } +} + +// Utility functions that don't require instance state + +// Min returns the minimum of two integers +func Min(a, b int) int { + if a < b { + return a + } + return b +} + +// GenerateSecureRandomString creates a cryptographically secure random string without dependencies +func GenerateSecureRandomString(length int) (string, error) { + bytes := make([]byte, length) + if _, err := rand.Read(bytes); err != nil { + return "", fmt.Errorf("failed to generate random bytes: %w", err) + } + return hex.EncodeToString(bytes), nil +} diff --git a/session/crypto/session_crypto_test.go b/session/crypto/session_crypto_test.go new file mode 100644 index 0000000..5dc5a98 --- /dev/null +++ b/session/crypto/session_crypto_test.go @@ -0,0 +1,900 @@ +package crypto + +import ( + "bytes" + "crypto/rand" + "encoding/base64" + "strings" + "testing" +) + +// Mock memory pools for testing +type MockMemoryPools struct{} + +func (mp *MockMemoryPools) GetCompressionBuffer() *bytes.Buffer { + return &bytes.Buffer{} +} + +func (mp *MockMemoryPools) PutCompressionBuffer(*bytes.Buffer) { + // Mock implementation - nothing to do +} + +func (mp *MockMemoryPools) GetHTTPResponseBuffer() []byte { + return make([]byte, 32768) // 32KB buffer +} + +func (mp *MockMemoryPools) PutHTTPResponseBuffer([]byte) { + // Mock implementation - nothing to do +} + +// TestGenerateSecureRandomString tests secure random string generation +func TestGenerateSecureRandomString(t *testing.T) { + memoryPools := &MockMemoryPools{} + sc := NewSessionCrypto(memoryPools) + + tests := []struct { + name string + length int + expectError bool + description string + }{ + { + name: "Valid length", + length: 16, + expectError: false, + description: "Should generate random string of correct length", + }, + { + name: "Minimum length", + length: 1, + expectError: false, + description: "Should handle minimum length", + }, + { + name: "Zero length", + length: 0, + expectError: false, + description: "Should handle zero length", + }, + { + name: "Large length", + length: 1024, + expectError: false, + description: "Should handle large length", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := sc.GenerateSecureRandomString(tt.length) + + if tt.expectError { + if err == nil { + t.Errorf("Expected error for %s, got nil", tt.description) + } + return + } + + if err != nil { + t.Errorf("Unexpected error for %s: %v", tt.description, err) + return + } + + // Check length (hex encoding doubles the length) + expectedLen := tt.length * 2 + if len(result) != expectedLen { + t.Errorf("Expected length %d, got %d for %s", expectedLen, len(result), tt.description) + } + + // Check that result is hex + for _, char := range result { + if !((char >= '0' && char <= '9') || (char >= 'a' && char <= 'f')) { + t.Errorf("Result contains non-hex character: %c", char) + break + } + } + }) + } +} + +// TestGenerateSecureRandomStringUniqueness tests that generated strings are unique +func TestGenerateSecureRandomStringUniqueness(t *testing.T) { + memoryPools := &MockMemoryPools{} + sc := NewSessionCrypto(memoryPools) + + // Generate multiple strings and check uniqueness + generated := make(map[string]bool) + for i := 0; i < 100; i++ { + result, err := sc.GenerateSecureRandomString(16) + if err != nil { + t.Fatalf("Failed to generate random string: %v", err) + } + + if generated[result] { + t.Errorf("Generated duplicate string: %s", result) + } + generated[result] = true + } +} + +// TestTokenCompressionIntegrity tests token compression and decompression +func TestTokenCompressionIntegrity(t *testing.T) { + memoryPools := &MockMemoryPools{} + sc := NewSessionCrypto(memoryPools) + + tests := []struct { + name string + token string + expectValid bool + description string + }{ + { + name: "Valid JWT small", + token: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c", + expectValid: true, + description: "Should compress and decompress small JWT correctly", + }, + { + name: "Valid JWT large", + token: createLargeJWT(2000), + expectValid: true, + description: "Should compress and decompress large JWT correctly", + }, + { + name: "Invalid token - no dots", + token: "invalidtoken", + expectValid: false, + description: "Should not compress token without dots", + }, + { + name: "Invalid token - wrong number of dots", + token: "header.payload", + expectValid: false, + description: "Should not compress token with wrong number of dots", + }, + { + name: "Empty token", + token: "", + expectValid: false, + description: "Should handle empty token", + }, + { + name: "Oversized token", + token: createOversizedToken(), + expectValid: false, + description: "Should reject oversized tokens", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + compressed := sc.CompressToken(tt.token) + + if !tt.expectValid { + // For invalid tokens, compression should return original + if compressed != tt.token { + t.Errorf("Expected compression to return original for invalid token, got different result") + } + return + } + + // For valid tokens, test round-trip integrity + decompressed := sc.DecompressToken(compressed) + if decompressed != tt.token { + t.Errorf("Token integrity lost: original length=%d, compressed length=%d, decompressed length=%d", + len(tt.token), len(compressed), len(decompressed)) + } + + // Test that decompression is idempotent + decompressed2 := sc.DecompressToken(decompressed) + if decompressed2 != tt.token { + t.Errorf("Decompression not idempotent: %d != %d", len(decompressed2), len(tt.token)) + } + }) + } +} + +// TestTokenCompressionCorruptionDetection tests corruption detection +func TestTokenCompressionCorruptionDetection(t *testing.T) { + memoryPools := &MockMemoryPools{} + sc := NewSessionCrypto(memoryPools) + + corruptionTests := []struct { + name string + corruptedInput string + expectOriginal bool + description string + }{ + { + name: "Corrupted base64", + corruptedInput: "invalid-base64!", + expectOriginal: true, + description: "Should return original for corrupted base64", + }, + { + name: "Truncated compressed data", + corruptedInput: "H4sI", // Truncated gzip header + expectOriginal: true, + description: "Should return original for truncated data", + }, + { + name: "Invalid gzip data", + corruptedInput: base64.StdEncoding.EncodeToString([]byte("not gzip data")), + expectOriginal: true, + description: "Should return original for invalid gzip data", + }, + { + name: "Empty compressed data", + corruptedInput: "", + expectOriginal: true, + description: "Should handle empty compressed data", + }, + } + + for _, tt := range corruptionTests { + t.Run(tt.name, func(t *testing.T) { + result := sc.DecompressToken(tt.corruptedInput) + if tt.expectOriginal && result != tt.corruptedInput { + t.Errorf("Expected decompression to return original corrupted input, got: %q", result) + } + }) + } + + // Test that valid compression still works + t.Run("Valid compression verification", func(t *testing.T) { + validJWT := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c" + compressed := sc.CompressToken(validJWT) + decompressed := sc.DecompressToken(compressed) + if decompressed != validJWT { + t.Errorf("Valid compression/decompression failed: %q != %q", decompressed, validJWT) + } + }) +} + +// TestCompressionEfficiency tests that compression only occurs when beneficial +func TestCompressionEfficiency(t *testing.T) { + memoryPools := &MockMemoryPools{} + sc := NewSessionCrypto(memoryPools) + + tests := []struct { + name string + token string + shouldCompress bool + description string + }{ + { + name: "Small JWT", + token: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c", + shouldCompress: false, // Small tokens might not benefit from compression + description: "Small tokens should not be compressed if no benefit", + }, + { + name: "Large repetitive JWT", + token: createLargeRepetitiveJWT(2000), + shouldCompress: true, // Repetitive data should compress well + description: "Large repetitive tokens should be compressed", + }, + { + name: "Incompressible token", + token: createIncompressibleJWT(1000), + shouldCompress: false, // Random data won't compress well + description: "Incompressible tokens should not be compressed", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + compressed := sc.CompressToken(tt.token) + wasCompressed := compressed != tt.token + + if tt.shouldCompress && !wasCompressed { + t.Errorf("Expected token to be compressed but it wasn't") + } else if !tt.shouldCompress && wasCompressed { + // This is okay - compression might still occur if beneficial + t.Logf("Token was compressed even though not expected (this is acceptable)") + } + + // Verify decompression still works regardless + decompressed := sc.DecompressToken(compressed) + if decompressed != tt.token { + t.Errorf("Decompression failed for %s", tt.description) + } + }) + } +} + +// TestCompressionSizeLimits tests compression size limits +func TestCompressionSizeLimits(t *testing.T) { + memoryPools := &MockMemoryPools{} + sc := NewSessionCrypto(memoryPools) + + t.Run("Oversized token rejection", func(t *testing.T) { + oversizedToken := createOversizedToken() + compressed := sc.CompressToken(oversizedToken) + + // Oversized tokens should not be compressed + if compressed != oversizedToken { + t.Error("Oversized token should not be compressed") + } + }) + + t.Run("Oversized compressed data rejection", func(t *testing.T) { + oversizedCompressed := strings.Repeat("a", 150*1024) // >100KB + decompressed := sc.DecompressToken(oversizedCompressed) + + // Should return original when input is too large + if decompressed != oversizedCompressed { + t.Error("Oversized compressed data should be returned as-is") + } + }) +} + +// Helper functions for creating test tokens + +func createLargeJWT(size int) string { + header := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9" + signature := "SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c" + + // Create payload that will result in desired total size + payloadSize := size - len(header) - len(signature) - 2 // -2 for dots + if payloadSize < 10 { + payloadSize = 10 + } + + payload := base64.StdEncoding.EncodeToString([]byte(strings.Repeat("x", payloadSize*3/4))) + + return header + "." + payload + "." + signature +} + +func createLargeRepetitiveJWT(size int) string { + header := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9" + signature := "SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c" + + // Create repetitive payload that compresses well + payloadSize := size - len(header) - len(signature) - 2 + if payloadSize < 10 { + payloadSize = 10 + } + + repetitiveData := strings.Repeat("repetitive_data_", payloadSize/16) + payload := base64.StdEncoding.EncodeToString([]byte(repetitiveData)) + + return header + "." + payload + "." + signature +} + +func createIncompressibleJWT(size int) string { + header := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9" + signature := "SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c" + + // Create random payload that won't compress well + payloadSize := size - len(header) - len(signature) - 2 + if payloadSize < 10 { + payloadSize = 10 + } + + randomBytes := make([]byte, payloadSize*3/4) + rand.Read(randomBytes) + payload := base64.StdEncoding.EncodeToString(randomBytes) + + return header + "." + payload + "." + signature +} + +func createOversizedToken() string { + // Create a token larger than 50KB (the limit in CompressToken) + size := 55 * 1024 // 55KB + header := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9" + signature := "SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c" + + payloadSize := size - len(header) - len(signature) - 2 + payload := base64.StdEncoding.EncodeToString([]byte(strings.Repeat("x", payloadSize*3/4))) + + return header + "." + payload + "." + signature +} + +// BenchmarkCompression benchmarks compression operations +func BenchmarkCompression(b *testing.B) { + memoryPools := &MockMemoryPools{} + sc := NewSessionCrypto(memoryPools) + + b.Run("CompressLargeJWT", func(b *testing.B) { + largeToken := createLargeJWT(5000) + b.ResetTimer() + + for i := 0; i < b.N; i++ { + _ = sc.CompressToken(largeToken) + } + }) + + b.Run("DecompressLargeJWT", func(b *testing.B) { + largeToken := createLargeJWT(5000) + compressed := sc.CompressToken(largeToken) + b.ResetTimer() + + for i := 0; i < b.N; i++ { + _ = sc.DecompressToken(compressed) + } + }) + + b.Run("RoundTripCompression", func(b *testing.B) { + largeToken := createLargeJWT(5000) + b.ResetTimer() + + for i := 0; i < b.N; i++ { + compressed := sc.CompressToken(largeToken) + _ = sc.DecompressToken(compressed) + } + }) +} + +// TestValidateTokenFormat tests JWT token format validation +func TestValidateTokenFormat(t *testing.T) { + memoryPools := &MockMemoryPools{} + sc := NewSessionCrypto(memoryPools) + + tests := []struct { + name string + token string + expected bool + }{ + { + name: "Valid JWT token", + token: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c", + expected: true, + }, + { + name: "Valid JWT with different content", + token: "header.payload.signature", + expected: true, + }, + { + name: "Empty token", + token: "", + expected: false, + }, + { + name: "Token with no dots", + token: "nodots", + expected: false, + }, + { + name: "Token with one dot", + token: "header.payload", + expected: false, + }, + { + name: "Token with four dots", + token: "header.payload.signature.extra", + expected: false, + }, + { + name: "Token with empty header", + token: ".payload.signature", + expected: false, + }, + { + name: "Token with empty payload", + token: "header..signature", + expected: false, + }, + { + name: "Token with empty signature", + token: "header.payload.", + expected: false, + }, + { + name: "Token with all empty parts", + token: "..", + expected: false, + }, + { + name: "Opaque token", + token: "opaque_token_without_dots", + expected: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := sc.ValidateTokenFormat(tt.token) + if result != tt.expected { + t.Errorf("ValidateTokenFormat(%q) = %v, expected %v", tt.token, result, tt.expected) + } + }) + } +} + +// TestIsTokenCompressed tests token compression detection +func TestIsTokenCompressed(t *testing.T) { + memoryPools := &MockMemoryPools{} + sc := NewSessionCrypto(memoryPools) + + tests := []struct { + name string + token string + expected bool + }{ + { + name: "Empty token", + token: "", + expected: false, + }, + { + name: "Valid JWT token (uncompressed)", + token: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c", + expected: false, + }, + { + name: "Invalid base64", + token: "invalid!base64", + expected: false, + }, + { + name: "Valid base64 but not gzip", + token: base64.StdEncoding.EncodeToString([]byte("not gzip data")), + expected: false, + }, + { + name: "Valid gzip header", + token: base64.StdEncoding.EncodeToString([]byte{0x1f, 0x8b, 0x08, 0x00}), // gzip magic bytes + expected: true, + }, + { + name: "Partial gzip header", + token: base64.StdEncoding.EncodeToString([]byte{0x1f}), // only first byte + expected: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := sc.IsTokenCompressed(tt.token) + if result != tt.expected { + t.Errorf("IsTokenCompressed(%q) = %v, expected %v", tt.token, result, tt.expected) + } + }) + } + + // Test with actual compressed token + t.Run("Real compressed token", func(t *testing.T) { + originalToken := createLargeJWT(2000) + compressedToken := sc.CompressToken(originalToken) + + // If compression occurred (token changed), it should be detected as compressed + if compressedToken != originalToken { + if !sc.IsTokenCompressed(compressedToken) { + t.Error("Failed to detect actual compressed token") + } + } + + // Original token should not be detected as compressed + if sc.IsTokenCompressed(originalToken) { + t.Error("Original JWT detected as compressed") + } + }) +} + +// TestSecureWipeBytes tests secure byte wiping +func TestSecureWipeBytes(t *testing.T) { + memoryPools := &MockMemoryPools{} + sc := NewSessionCrypto(memoryPools) + + tests := []struct { + name string + data []byte + }{ + { + name: "Normal byte slice", + data: []byte("sensitive data"), + }, + { + name: "Empty slice", + data: []byte{}, + }, + { + name: "Single byte", + data: []byte{0xFF}, + }, + { + name: "Large data", + data: bytes.Repeat([]byte("secret"), 1000), + }, + { + name: "Nil slice", + data: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Create a copy to verify original content + original := make([]byte, len(tt.data)) + copy(original, tt.data) + + // Wipe the data + sc.SecureWipeBytes(tt.data) + + // Verify all bytes are zero (except for nil slice) + if tt.data != nil { + for i, b := range tt.data { + if b != 0 { + t.Errorf("Byte at index %d not wiped: got %d, expected 0", i, b) + } + } + } + + // Verify we had actual data to wipe (except for empty/nil cases) + if len(original) > 0 { + hasNonZero := false + for _, b := range original { + if b != 0 { + hasNonZero = true + break + } + } + if !hasNonZero { + t.Log("Test data was already all zeros") + } + } + }) + } +} + +// TestSecureWipeString tests secure string wiping +func TestSecureWipeString(t *testing.T) { + memoryPools := &MockMemoryPools{} + sc := NewSessionCrypto(memoryPools) + + tests := []struct { + name string + input *string + expect string + }{ + { + name: "Normal string", + input: func() *string { s := "sensitive data"; return &s }(), + expect: "", + }, + { + name: "Empty string", + input: func() *string { s := ""; return &s }(), + expect: "", + }, + { + name: "Long string", + input: func() *string { s := strings.Repeat("secret", 1000); return &s }(), + expect: "", + }, + { + name: "Nil string pointer", + input: nil, + expect: "", // This test verifies no panic occurs + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Store original value for verification + var original string + if tt.input != nil { + original = *tt.input + } + + // Wipe the string + sc.SecureWipeString(tt.input) + + // Verify result + if tt.input != nil { + if *tt.input != tt.expect { + t.Errorf("String not wiped properly: got %q, expected %q", *tt.input, tt.expect) + } + } + + // Verify we had actual data to wipe (except for nil case) + if tt.input != nil && original != "" { + t.Logf("Successfully wiped string of length %d", len(original)) + } + }) + } +} + +// TestMin tests the minimum utility function +func TestMin(t *testing.T) { + tests := []struct { + name string + a, b int + expected int + }{ + { + name: "a smaller than b", + a: 5, + b: 10, + expected: 5, + }, + { + name: "b smaller than a", + a: 15, + b: 7, + expected: 7, + }, + { + name: "equal values", + a: 42, + b: 42, + expected: 42, + }, + { + name: "negative values", + a: -10, + b: -5, + expected: -10, + }, + { + name: "zero values", + a: 0, + b: 0, + expected: 0, + }, + { + name: "mixed positive and negative", + a: -3, + b: 2, + expected: -3, + }, + { + name: "large numbers", + a: 1000000, + b: 999999, + expected: 999999, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := Min(tt.a, tt.b) + if result != tt.expected { + t.Errorf("Min(%d, %d) = %d, expected %d", tt.a, tt.b, result, tt.expected) + } + }) + } +} + +// TestGenerateSecureRandomStringStandalone tests the standalone random string function +func TestGenerateSecureRandomStringStandalone(t *testing.T) { + tests := []struct { + name string + length int + expectError bool + }{ + { + name: "Valid length", + length: 16, + expectError: false, + }, + { + name: "Zero length", + length: 0, + expectError: false, + }, + { + name: "Large length", + length: 1024, + expectError: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := GenerateSecureRandomString(tt.length) + + if tt.expectError { + if err == nil { + t.Error("Expected error but got none") + } + return + } + + if err != nil { + t.Errorf("Unexpected error: %v", err) + return + } + + // Check length (hex encoding doubles the length) + expectedLen := tt.length * 2 + if len(result) != expectedLen { + t.Errorf("Expected length %d, got %d", expectedLen, len(result)) + } + + // Check that result is hex + for _, char := range result { + if !((char >= '0' && char <= '9') || (char >= 'a' && char <= 'f')) { + t.Errorf("Result contains non-hex character: %c", char) + break + } + } + }) + } + + // Test uniqueness + t.Run("Uniqueness test", func(t *testing.T) { + generated := make(map[string]bool) + for i := 0; i < 100; i++ { + result, err := GenerateSecureRandomString(16) + if err != nil { + t.Fatalf("Failed to generate random string: %v", err) + } + + if generated[result] { + t.Errorf("Generated duplicate string: %s", result) + } + generated[result] = true + } + }) +} + +// TestCompressionEdgeCases tests edge cases for compression +func TestCompressionEdgeCases(t *testing.T) { + memoryPools := &MockMemoryPools{} + sc := NewSessionCrypto(memoryPools) + + t.Run("Token with exact size limit", func(t *testing.T) { + // Create token at exactly 50KB + token := createTokenWithExactSize(50 * 1024) + compressed := sc.CompressToken(token) + + // Should still attempt compression at the limit + decompressed := sc.DecompressToken(compressed) + if decompressed != token { + t.Error("Failed to handle token at size limit") + } + }) + + t.Run("Compressed token with exact decompression limit", func(t *testing.T) { + // Create data that decompresses to exactly 100KB + largeData := strings.Repeat("a", 100*1024) + encoded := base64.StdEncoding.EncodeToString([]byte(largeData)) + + result := sc.DecompressToken(encoded) + // Should return original since it's not valid gzip + if result != encoded { + t.Error("Failed to handle large non-gzip data") + } + }) +} + +// Helper function to create token with exact size +func createTokenWithExactSize(targetSize int) string { + header := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9" + signature := "SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c" + + // Calculate needed payload size + dotsSize := 2 // two dots + otherSize := len(header) + len(signature) + dotsSize + payloadSize := targetSize - otherSize + + if payloadSize <= 0 { + payloadSize = 10 // minimum payload + } + + // Create payload of exact size + payload := strings.Repeat("x", payloadSize) + + return header + "." + payload + "." + signature +} + +// BenchmarkRandomGeneration benchmarks random string generation +func BenchmarkRandomGeneration(b *testing.B) { + memoryPools := &MockMemoryPools{} + sc := NewSessionCrypto(memoryPools) + + b.Run("Generate16Bytes", func(b *testing.B) { + for i := 0; i < b.N; i++ { + _, _ = sc.GenerateSecureRandomString(16) + } + }) + + b.Run("Generate32Bytes", func(b *testing.B) { + for i := 0; i < b.N; i++ { + _, _ = sc.GenerateSecureRandomString(32) + } + }) +} diff --git a/session/storage/session_store.go b/session/storage/session_store.go new file mode 100644 index 0000000..ce045be --- /dev/null +++ b/session/storage/session_store.go @@ -0,0 +1,320 @@ +// Package storage provides session storage operations for the OIDC middleware +package storage + +import ( + "fmt" + "net/http" + "sync" + + "github.com/gorilla/sessions" +) + +// SessionData represents a user's authentication session with comprehensive token management. +// It handles main session data and supports large tokens that need to be +// split across multiple cookies due to browser size limitations. +type SessionData struct { + manager SessionManager + request *http.Request + mainSession *sessions.Session + accessSession *sessions.Session + refreshSession *sessions.Session + idTokenSession *sessions.Session + accessTokenChunks map[int]*sessions.Session + refreshTokenChunks map[int]*sessions.Session + idTokenChunks map[int]*sessions.Session + refreshMutex sync.Mutex + sessionMutex sync.RWMutex + dirty bool + inUse bool +} + +// SessionManager interface for session management operations +type SessionManager interface { + GetSessionOptions(isSecure bool) *sessions.Options + EnhanceSessionSecurity(options *sessions.Options, r *http.Request) *sessions.Options + GetLogger() Logger +} + +// Logger interface for dependency injection +type Logger interface { + Error(msg string) + Errorf(format string, args ...interface{}) +} + +// NewSessionData creates a new session data instance +func NewSessionData(manager SessionManager) *SessionData { + return &SessionData{ + manager: manager, + accessTokenChunks: make(map[int]*sessions.Session), + refreshTokenChunks: make(map[int]*sessions.Session), + idTokenChunks: make(map[int]*sessions.Session), + refreshMutex: sync.Mutex{}, + sessionMutex: sync.RWMutex{}, + dirty: false, + inUse: false, + } +} + +// IsDirty returns true if the session data has been modified since it was last loaded or saved. +// This is used to optimize session saves by only writing when necessary. +func (sd *SessionData) IsDirty() bool { + return sd.dirty +} + +// MarkDirty marks the session as having pending changes that need to be saved. +// This is used when session data hasn't changed in content but should still +// trigger a session save (e.g., to ensure the cookie is re-issued). +func (sd *SessionData) MarkDirty() { + sd.dirty = true +} + +// Save persists all session data including main session and token chunks. +// It applies security options, saves all session components, and handles +// errors gracefully by continuing to save other components even if one fails. +func (sd *SessionData) Save(r *http.Request, w http.ResponseWriter) error { + isSecure := r.Header.Get("X-Forwarded-Proto") == "https" || r.TLS != nil + if forceHTTPS := sd.manager.GetLogger(); forceHTTPS != nil { + // Add force HTTPS check if needed + } + + options := sd.manager.GetSessionOptions(isSecure) + options = sd.manager.EnhanceSessionSecurity(options, r) + + if sd.mainSession != nil { + sd.mainSession.Options = options + } + if sd.accessSession != nil { + sd.accessSession.Options = options + } + if sd.refreshSession != nil { + sd.refreshSession.Options = options + } + if sd.idTokenSession != nil { + sd.idTokenSession.Options = options + } + + var firstErr error + saveOrLogError := func(s *sessions.Session, name string) { + if s == nil { + logger := sd.manager.GetLogger() + if logger != nil { + logger.Errorf("Attempted to save nil session: %s", name) + } + if firstErr == nil { + firstErr = fmt.Errorf("attempted to save nil session: %s", name) + } + return + } + if err := s.Save(r, w); err != nil { + errMsg := fmt.Errorf("failed to save %s session: %w", name, err) + logger := sd.manager.GetLogger() + if logger != nil { + logger.Error(errMsg.Error()) + } + if firstErr == nil { + firstErr = errMsg + } + } + } + + saveOrLogError(sd.mainSession, "main") + saveOrLogError(sd.accessSession, "access token") + saveOrLogError(sd.refreshSession, "refresh token") + saveOrLogError(sd.idTokenSession, "ID token") + + for i, sessionChunk := range sd.accessTokenChunks { + if sessionChunk != nil { + sessionChunk.Options = options + saveOrLogError(sessionChunk, fmt.Sprintf("access token chunk %d", i)) + } + } + + for i, sessionChunk := range sd.refreshTokenChunks { + if sessionChunk != nil { + sessionChunk.Options = options + saveOrLogError(sessionChunk, fmt.Sprintf("refresh token chunk %d", i)) + } + } + + for i, sessionChunk := range sd.idTokenChunks { + if sessionChunk != nil { + sessionChunk.Options = options + saveOrLogError(sessionChunk, fmt.Sprintf("ID token chunk %d", i)) + } + } + + if firstErr == nil { + sd.dirty = false + } + return firstErr +} + +// Clear completely clears all session data and safely returns the session to the pool. +// It removes all authentication data, expires cookies, and handles panic recovery. +func (sd *SessionData) Clear(r *http.Request, w http.ResponseWriter) error { + defer func() { + sd.returnToPoolSafely() + }() + + sd.sessionMutex.Lock() + defer sd.sessionMutex.Unlock() + + sd.clearAllSessionData(r, true) + + // This is primarily for testing - in production w will often be nil + var err error + if w != nil { + if r != nil && r.Header.Get("X-Test-Error") == "true" { + if sd.mainSession != nil { + sd.mainSession.Values["error_trigger"] = func() {} + } + } + + err = sd.Save(r, w) + } + + sd.request = nil + return err +} + +// clearAllSessionData clears all session data including main session and token chunks. +// It removes all session values and optionally expires all associated cookies. +func (sd *SessionData) clearAllSessionData(r *http.Request, expire bool) { + clearSessionValues(sd.mainSession, expire) + clearSessionValues(sd.accessSession, expire) + clearSessionValues(sd.refreshSession, expire) + clearSessionValues(sd.idTokenSession, expire) + + if expire && r != nil { + sd.clearTokenChunks(r, sd.accessTokenChunks) + sd.clearTokenChunks(r, sd.refreshTokenChunks) + sd.clearTokenChunks(r, sd.idTokenChunks) + } else { + for k := range sd.accessTokenChunks { + delete(sd.accessTokenChunks, k) + } + for k := range sd.refreshTokenChunks { + delete(sd.refreshTokenChunks, k) + } + for k := range sd.idTokenChunks { + delete(sd.idTokenChunks, k) + } + } + + if expire { + sd.dirty = true + } +} + +// clearSessionValues removes all values from a session and optionally expires it. +// This is used during session cleanup and logout operations. +func clearSessionValues(session *sessions.Session, expire bool) { + if session == nil { + return + } + + for k := range session.Values { + delete(session.Values, k) + } + + if expire { + session.Options.MaxAge = -1 + } +} + +// clearTokenChunks clears token chunks from the session +func (sd *SessionData) clearTokenChunks(r *http.Request, chunks map[int]*sessions.Session) { + for i, chunk := range chunks { + if chunk != nil { + clearSessionValues(chunk, true) + } + delete(chunks, i) + } +} + +// returnToPoolSafely safely returns the session to the object pool +func (sd *SessionData) returnToPoolSafely() { + defer func() { + if r := recover(); r != nil { + logger := sd.manager.GetLogger() + if logger != nil { + logger.Errorf("Panic during session pool return: %v", r) + } + } + }() + + sd.sessionMutex.Lock() + defer sd.sessionMutex.Unlock() + + if sd.inUse { + sd.inUse = false + sd.Reset() + // Pool return should be handled by calling code + } +} + +// Reset resets the session data to a clean state +func (sd *SessionData) Reset() { + sd.mainSession = nil + sd.accessSession = nil + sd.refreshSession = nil + sd.idTokenSession = nil + + // Clear maps without recreating them + for k := range sd.accessTokenChunks { + delete(sd.accessTokenChunks, k) + } + for k := range sd.refreshTokenChunks { + delete(sd.refreshTokenChunks, k) + } + for k := range sd.idTokenChunks { + delete(sd.idTokenChunks, k) + } + + sd.dirty = false + sd.inUse = false + sd.request = nil +} + +// SetSessions sets the session objects +func (sd *SessionData) SetSessions(main, access, refresh, idToken *sessions.Session) { + sd.mainSession = main + sd.accessSession = access + sd.refreshSession = refresh + sd.idTokenSession = idToken +} + +// GetMainSession returns the main session +func (sd *SessionData) GetMainSession() *sessions.Session { + return sd.mainSession +} + +// GetAccessSession returns the access token session +func (sd *SessionData) GetAccessSession() *sessions.Session { + return sd.accessSession +} + +// GetRefreshSession returns the refresh token session +func (sd *SessionData) GetRefreshSession() *sessions.Session { + return sd.refreshSession +} + +// GetIDTokenSession returns the ID token session +func (sd *SessionData) GetIDTokenSession() *sessions.Session { + return sd.idTokenSession +} + +// GetTokenChunks returns the token chunk maps +func (sd *SessionData) GetTokenChunks() (map[int]*sessions.Session, map[int]*sessions.Session, map[int]*sessions.Session) { + return sd.accessTokenChunks, sd.refreshTokenChunks, sd.idTokenChunks +} + +// SetInUse marks the session as in use +func (sd *SessionData) SetInUse(inUse bool) { + sd.inUse = inUse +} + +// IsInUse returns whether the session is in use +func (sd *SessionData) IsInUse() bool { + return sd.inUse +} diff --git a/session/storage/session_store_test.go b/session/storage/session_store_test.go new file mode 100644 index 0000000..aaa6e5a --- /dev/null +++ b/session/storage/session_store_test.go @@ -0,0 +1,1125 @@ +package storage + +import ( + "crypto/tls" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gorilla/sessions" +) + +// Mock logger for testing +type MockLogger struct { + logs []string +} + +func (ml *MockLogger) Error(msg string) { + ml.logs = append(ml.logs, "ERROR: "+msg) +} + +func (ml *MockLogger) Errorf(format string, args ...interface{}) { + ml.logs = append(ml.logs, fmt.Sprintf("ERROR: "+format, args...)) +} + +// Mock session manager for testing +type MockSessionManager struct { + logger Logger +} + +func (msm *MockSessionManager) GetSessionOptions(isSecure bool) *sessions.Options { + return &sessions.Options{ + Path: "/", + MaxAge: 3600, + Secure: isSecure, + HttpOnly: true, + SameSite: http.SameSiteLaxMode, + } +} + +func (msm *MockSessionManager) EnhanceSessionSecurity(options *sessions.Options, r *http.Request) *sessions.Options { + if r.Header.Get("X-Forwarded-Proto") == "https" || r.TLS != nil { + options.Secure = true + } + return options +} + +func (msm *MockSessionManager) GetLogger() Logger { + return msm.logger +} + +// TestNewSessionData tests session data creation +func TestNewSessionData(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + + sd := NewSessionData(manager) + + if sd == nil { + t.Fatal("NewSessionData should not return nil") + } + + if sd.manager != manager { + t.Error("Manager should be set correctly") + } + + if sd.accessTokenChunks == nil || len(sd.accessTokenChunks) != 0 { + t.Error("Access token chunks map should be initialized and empty") + } + + if sd.refreshTokenChunks == nil || len(sd.refreshTokenChunks) != 0 { + t.Error("Refresh token chunks map should be initialized and empty") + } + + if sd.idTokenChunks == nil || len(sd.idTokenChunks) != 0 { + t.Error("ID token chunks map should be initialized and empty") + } + + if sd.dirty { + t.Error("New session data should not be dirty") + } + + if sd.inUse { + t.Error("New session data should not be in use") + } +} + +// TestSessionDataDirtyFlag tests dirty flag management +func TestSessionDataDirtyFlag(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + // Test initial state + if sd.IsDirty() { + t.Error("New session should not be dirty") + } + + // Test marking dirty + sd.MarkDirty() + if !sd.IsDirty() { + t.Error("Session should be dirty after MarkDirty()") + } + + // Test that Save clears dirty flag (when successful) + req := httptest.NewRequest("GET", "http://example.com", nil) + w := httptest.NewRecorder() + + // Create a simple main session to avoid nil session errors + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + session, _ := store.Get(req, "test-session") + sd.mainSession = session + + err := sd.Save(req, w) + if err != nil { + t.Logf("Save returned error (may be expected): %v", err) + } + + // Note: dirty flag is only cleared if Save is completely successful + // which might not happen with our mock setup +} + +// TestSessionDataSave tests session saving functionality +func TestSessionDataSave(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + + tests := []struct { + name string + setupSesion func(*SessionData) + expectError bool + description string + }{ + { + name: "Save with main session only", + setupSesion: func(sd *SessionData) { + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + session, _ := store.Get(req, "test-session") + sd.mainSession = session + }, + expectError: true, // Will error because other sessions are nil + description: "Should handle nil subsidiary sessions", + }, + { + name: "Save with all session types", + setupSesion: func(sd *SessionData) { + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + + sd.mainSession, _ = store.Get(req, "main-session") + sd.accessSession, _ = store.Get(req, "access-session") + sd.refreshSession, _ = store.Get(req, "refresh-session") + sd.idTokenSession, _ = store.Get(req, "id-session") + }, + expectError: false, + description: "Should save all session types without error", + }, + { + name: "Save with token chunks", + setupSesion: func(sd *SessionData) { + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + + sd.mainSession, _ = store.Get(req, "main-session") + sd.accessSession, _ = store.Get(req, "access-session") + sd.refreshSession, _ = store.Get(req, "refresh-session") + sd.idTokenSession, _ = store.Get(req, "id-session") + + // Add some token chunks + chunk1, _ := store.Get(req, "access-chunk-0") + chunk2, _ := store.Get(req, "access-chunk-1") + sd.accessTokenChunks[0] = chunk1 + sd.accessTokenChunks[1] = chunk2 + + refreshChunk, _ := store.Get(req, "refresh-chunk-0") + sd.refreshTokenChunks[0] = refreshChunk + }, + expectError: false, + description: "Should save token chunks without error", + }, + { + name: "Save with nil main session", + setupSesion: func(sd *SessionData) { + sd.mainSession = nil + }, + expectError: true, + description: "Should handle nil main session gracefully", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sd := NewSessionData(manager) + tt.setupSesion(sd) + + req := httptest.NewRequest("GET", "http://example.com", nil) + w := httptest.NewRecorder() + + err := sd.Save(req, w) + + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, got nil", tt.description) + } else if !tt.expectError && err != nil { + t.Errorf("Unexpected error for %s: %v", tt.description, err) + } + }) + } +} + +// TestSessionDataSaveHTTPS tests HTTPS detection in Save +func TestSessionDataSaveHTTPS(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + + tests := []struct { + name string + setupReq func() *http.Request + expectSecure bool + description string + }{ + { + name: "HTTPS via TLS", + setupReq: func() *http.Request { + req := httptest.NewRequest("GET", "https://example.com", nil) + // Simulate TLS connection + req.TLS = &tls.ConnectionState{} + return req + }, + expectSecure: true, + description: "Should detect HTTPS via TLS", + }, + { + name: "HTTPS via X-Forwarded-Proto header", + setupReq: func() *http.Request { + req := httptest.NewRequest("GET", "http://example.com", nil) + req.Header.Set("X-Forwarded-Proto", "https") + return req + }, + expectSecure: true, + description: "Should detect HTTPS via X-Forwarded-Proto header", + }, + { + name: "HTTP request", + setupReq: func() *http.Request { + return httptest.NewRequest("GET", "http://example.com", nil) + }, + expectSecure: false, + description: "Should detect HTTP correctly", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req := tt.setupReq() + w := httptest.NewRecorder() + + session, _ := store.Get(req, "test-session") + sd.mainSession = session + // Set all other sessions to avoid nil session errors + sd.accessSession, _ = store.Get(req, "access-session") + sd.refreshSession, _ = store.Get(req, "refresh-session") + sd.idTokenSession, _ = store.Get(req, "id-session") + + err := sd.Save(req, w) + if err != nil { + t.Logf("Save returned error: %v", err) + } + + // Check the session options were set correctly + if sd.mainSession.Options.Secure != tt.expectSecure { + t.Errorf("Expected Secure=%v for %s, got %v", + tt.expectSecure, tt.description, sd.mainSession.Options.Secure) + } + }) + } +} + +// TestSessionDataChunkManagement tests token chunk management +func TestSessionDataChunkManagement(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + + // Test adding chunks + chunk0, _ := store.Get(req, "access-chunk-0") + chunk1, _ := store.Get(req, "access-chunk-1") + chunk2, _ := store.Get(req, "access-chunk-2") + + sd.accessTokenChunks[0] = chunk0 + sd.accessTokenChunks[1] = chunk1 + sd.accessTokenChunks[2] = chunk2 + + if len(sd.accessTokenChunks) != 3 { + t.Errorf("Expected 3 access token chunks, got %d", len(sd.accessTokenChunks)) + } + + // Test saving chunks + sd.mainSession, _ = store.Get(req, "main-session") + sd.accessSession, _ = store.Get(req, "access-session") + sd.refreshSession, _ = store.Get(req, "refresh-session") + sd.idTokenSession, _ = store.Get(req, "id-session") + + w := httptest.NewRecorder() + + err := sd.Save(req, w) + if err != nil { + t.Logf("Save with chunks returned error: %v", err) + } + + // Verify chunks have proper options set + for i, chunk := range sd.accessTokenChunks { + if chunk.Options == nil { + t.Errorf("Chunk %d should have options set", i) + } else if chunk.Options.HttpOnly != true { + t.Errorf("Chunk %d should have HttpOnly=true", i) + } + } +} + +// TestSessionDataErrorHandling tests error handling in Save +func TestSessionDataErrorHandling(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + // Test with nil sessions to trigger error paths + sd.mainSession = nil + sd.accessSession = nil + + req := httptest.NewRequest("GET", "http://example.com", nil) + w := httptest.NewRecorder() + + err := sd.Save(req, w) + + // Should get an error for nil session + if err == nil { + t.Error("Expected error when saving nil sessions") + } + + // Check that error was logged + if len(logger.logs) == 0 { + t.Error("Expected error to be logged") + } + + // Check error message + foundNilSessionError := false + for _, log := range logger.logs { + if strings.Contains(log, "nil session") { + foundNilSessionError = true + break + } + } + + if !foundNilSessionError { + t.Error("Expected nil session error to be logged") + } +} + +// TestSessionDataConcurrency tests concurrent access to session data +func TestSessionDataConcurrency(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + sd.mainSession, _ = store.Get(req, "main-session") + + // Test concurrent marking as dirty + done := make(chan bool, 2) + + go func() { + for i := 0; i < 100; i++ { + sd.MarkDirty() + } + done <- true + }() + + go func() { + for i := 0; i < 100; i++ { + _ = sd.IsDirty() + } + done <- true + }() + + // Wait for both goroutines to complete + <-done + <-done + + // Should not panic and dirty flag should be set + if !sd.IsDirty() { + t.Error("Expected session to be dirty after concurrent operations") + } +} + +// TestSessionDataReset tests session data reset functionality +func TestSessionDataReset(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + // Set up session data with various values + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + + sd.mainSession, _ = store.Get(req, "main-session") + sd.accessSession, _ = store.Get(req, "access-session") + + // Add some chunks + chunk, _ := store.Get(req, "chunk-0") + sd.accessTokenChunks[0] = chunk + + sd.MarkDirty() + sd.inUse = true + + // Create a reset method if it exists in the actual implementation + // This is a placeholder test for reset functionality + t.Run("Manual reset", func(t *testing.T) { + // Simulate reset by clearing fields + sd.mainSession = nil + sd.accessSession = nil + sd.refreshSession = nil + sd.idTokenSession = nil + + // Clear chunks + sd.accessTokenChunks = make(map[int]*sessions.Session) + sd.refreshTokenChunks = make(map[int]*sessions.Session) + sd.idTokenChunks = make(map[int]*sessions.Session) + + sd.dirty = false + sd.inUse = false + + // Verify reset + if sd.mainSession != nil { + t.Error("Main session should be nil after reset") + } + + if len(sd.accessTokenChunks) != 0 { + t.Error("Access token chunks should be empty after reset") + } + + if sd.IsDirty() { + t.Error("Session should not be dirty after reset") + } + + if sd.inUse { + t.Error("Session should not be in use after reset") + } + }) +} + +// TestSessionDataValidation tests session data validation +func TestSessionDataValidation(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + + tests := []struct { + name string + setupFunc func() *SessionData + expectValid bool + description string + }{ + { + name: "Valid session data", + setupFunc: func() *SessionData { + sd := NewSessionData(manager) + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + sd.mainSession, _ = store.Get(req, "main-session") + return sd + }, + expectValid: true, + description: "Should validate correct session data", + }, + { + name: "Invalid session data - nil manager", + setupFunc: func() *SessionData { + sd := &SessionData{ + manager: nil, + accessTokenChunks: make(map[int]*sessions.Session), + refreshTokenChunks: make(map[int]*sessions.Session), + idTokenChunks: make(map[int]*sessions.Session), + } + return sd + }, + expectValid: false, + description: "Should reject session data with nil manager", + }, + { + name: "Invalid session data - nil chunks map", + setupFunc: func() *SessionData { + sd := NewSessionData(manager) + sd.accessTokenChunks = nil + return sd + }, + expectValid: false, + description: "Should reject session data with nil chunks map", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sd := tt.setupFunc() + + // Basic validation checks + isValid := true + + if sd.manager == nil { + isValid = false + } + + if sd.accessTokenChunks == nil || sd.refreshTokenChunks == nil || sd.idTokenChunks == nil { + isValid = false + } + + if isValid != tt.expectValid { + t.Errorf("Validation mismatch for %s: expected valid=%v, got valid=%v", + tt.description, tt.expectValid, isValid) + } + }) + } +} + +// BenchmarkSessionDataSave benchmarks session save operations +func BenchmarkSessionDataSave(b *testing.B) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + sd.mainSession, _ = store.Get(req, "main-session") + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + w := httptest.NewRecorder() + _ = sd.Save(req, w) + } +} + +// TestClear tests complete session clearing +func TestClear(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + w := httptest.NewRecorder() + + // Set up session data + sd.mainSession, _ = store.Get(req, "main-session") + sd.accessSession, _ = store.Get(req, "access-session") + sd.refreshSession, _ = store.Get(req, "refresh-session") + sd.idTokenSession, _ = store.Get(req, "id-session") + + // Add some chunks + chunk1, _ := store.Get(req, "access-chunk-0") + chunk2, _ := store.Get(req, "refresh-chunk-0") + chunk3, _ := store.Get(req, "id-chunk-0") + sd.accessTokenChunks[0] = chunk1 + sd.refreshTokenChunks[0] = chunk2 + sd.idTokenChunks[0] = chunk3 + + // Add some data to sessions + sd.mainSession.Values["user_id"] = "123" + sd.accessSession.Values["token"] = "access-token" + sd.refreshSession.Values["token"] = "refresh-token" + sd.idTokenSession.Values["token"] = "id-token" + + sd.MarkDirty() + sd.SetInUse(true) + + // Clear the session + err := sd.Clear(req, w) + if err != nil { + t.Logf("Clear returned error (may be expected): %v", err) + } + + // Verify main session values are cleared + if sd.mainSession != nil && len(sd.mainSession.Values) > 0 { + t.Error("Main session values should be cleared") + } + + // Verify session expires + if sd.mainSession != nil && sd.mainSession.Options.MaxAge != -1 { + t.Error("Main session should be expired (MaxAge = -1)") + } + + // Verify chunks are cleared + if len(sd.accessTokenChunks) != 0 { + t.Error("Access token chunks should be cleared") + } + if len(sd.refreshTokenChunks) != 0 { + t.Error("Refresh token chunks should be cleared") + } + if len(sd.idTokenChunks) != 0 { + t.Error("ID token chunks should be cleared") + } + + // Verify request is cleared + if sd.request != nil { + t.Error("Request should be cleared") + } + + // Verify usage status is reset + if sd.IsInUse() { + t.Error("Session should not be in use after clear") + } +} + +// TestClearWithNilResponseWriter tests clearing with nil response writer +func TestClearWithNilResponseWriter(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + + sd.mainSession, _ = store.Get(req, "main-session") + sd.mainSession.Values["test"] = "value" + + // Clear with nil response writer + err := sd.Clear(req, nil) + if err != nil { + t.Logf("Clear with nil writer returned error (expected): %v", err) + } + + // Should still clear session data + if sd.mainSession != nil && len(sd.mainSession.Values) > 0 { + t.Error("Session values should be cleared even with nil writer") + } +} + +// TestClearWithErrorTrigger tests error handling in Clear +func TestClearWithErrorTrigger(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + req.Header.Set("X-Test-Error", "true") // Trigger error condition + w := httptest.NewRecorder() + + sd.mainSession, _ = store.Get(req, "main-session") + + err := sd.Clear(req, w) + // May return error due to test trigger + t.Logf("Clear with error trigger returned: %v", err) + + // Should still clear the data despite error + if sd.request != nil { + t.Error("Request should be cleared even after error") + } +} + +// TestReset tests session reset functionality +func TestReset(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + + // Set up session data + sd.mainSession, _ = store.Get(req, "main-session") + sd.accessSession, _ = store.Get(req, "access-session") + sd.refreshSession, _ = store.Get(req, "refresh-session") + sd.idTokenSession, _ = store.Get(req, "id-session") + sd.request = req + + // Add chunks + chunk1, _ := store.Get(req, "access-chunk-0") + chunk2, _ := store.Get(req, "refresh-chunk-0") + chunk3, _ := store.Get(req, "id-chunk-0") + sd.accessTokenChunks[0] = chunk1 + sd.refreshTokenChunks[0] = chunk2 + sd.idTokenChunks[0] = chunk3 + + sd.MarkDirty() + sd.SetInUse(true) + + // Reset the session + sd.Reset() + + // Verify all sessions are nil + if sd.mainSession != nil { + t.Error("Main session should be nil after reset") + } + if sd.accessSession != nil { + t.Error("Access session should be nil after reset") + } + if sd.refreshSession != nil { + t.Error("Refresh session should be nil after reset") + } + if sd.idTokenSession != nil { + t.Error("ID token session should be nil after reset") + } + + // Verify chunks are cleared + if len(sd.accessTokenChunks) != 0 { + t.Error("Access token chunks should be empty after reset") + } + if len(sd.refreshTokenChunks) != 0 { + t.Error("Refresh token chunks should be empty after reset") + } + if len(sd.idTokenChunks) != 0 { + t.Error("ID token chunks should be empty after reset") + } + + // Verify state is reset + if sd.IsDirty() { + t.Error("Session should not be dirty after reset") + } + if sd.IsInUse() { + t.Error("Session should not be in use after reset") + } + if sd.request != nil { + t.Error("Request should be nil after reset") + } +} + +// TestSetSessions tests session setting +func TestSetSessions(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + + main, _ := store.Get(req, "main") + access, _ := store.Get(req, "access") + refresh, _ := store.Get(req, "refresh") + idToken, _ := store.Get(req, "id") + + // Set all sessions at once + sd.SetSessions(main, access, refresh, idToken) + + // Verify sessions are set correctly + if sd.GetMainSession() != main { + t.Error("Main session not set correctly") + } + if sd.GetAccessSession() != access { + t.Error("Access session not set correctly") + } + if sd.GetRefreshSession() != refresh { + t.Error("Refresh session not set correctly") + } + if sd.GetIDTokenSession() != idToken { + t.Error("ID token session not set correctly") + } +} + +// TestSetSessionsWithNil tests setting sessions with nil values +func TestSetSessionsWithNil(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + // Set sessions with nil values + sd.SetSessions(nil, nil, nil, nil) + + // Verify sessions are nil + if sd.GetMainSession() != nil { + t.Error("Main session should be nil") + } + if sd.GetAccessSession() != nil { + t.Error("Access session should be nil") + } + if sd.GetRefreshSession() != nil { + t.Error("Refresh session should be nil") + } + if sd.GetIDTokenSession() != nil { + t.Error("ID token session should be nil") + } +} + +// TestGetTokenChunks tests token chunk retrieval +func TestGetTokenChunks(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + + // Add chunks to each map + accessChunk, _ := store.Get(req, "access-chunk-0") + refreshChunk, _ := store.Get(req, "refresh-chunk-0") + idChunk, _ := store.Get(req, "id-chunk-0") + + sd.accessTokenChunks[0] = accessChunk + sd.refreshTokenChunks[0] = refreshChunk + sd.idTokenChunks[0] = idChunk + + // Get chunks + access, refresh, id := sd.GetTokenChunks() + + // Verify chunks are returned correctly + if len(access) != 1 || access[0] != accessChunk { + t.Error("Access token chunks not returned correctly") + } + if len(refresh) != 1 || refresh[0] != refreshChunk { + t.Error("Refresh token chunks not returned correctly") + } + if len(id) != 1 || id[0] != idChunk { + t.Error("ID token chunks not returned correctly") + } +} + +// TestSetInUseAndIsInUse tests usage tracking +func TestSetInUseAndIsInUse(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + // Initially should not be in use + if sd.IsInUse() { + t.Error("New session should not be in use") + } + + // Set in use + sd.SetInUse(true) + if !sd.IsInUse() { + t.Error("Session should be in use after SetInUse(true)") + } + + // Set not in use + sd.SetInUse(false) + if sd.IsInUse() { + t.Error("Session should not be in use after SetInUse(false)") + } +} + +// TestReturnToPoolSafely tests safe pool return +func TestReturnToPoolSafely(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + // Set session as in use + sd.SetInUse(true) + sd.MarkDirty() + + // Set up some session data + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + sd.mainSession, _ = store.Get(req, "main") + sd.request = req + + // Call returnToPoolSafely directly + sd.returnToPoolSafely() + + // Verify session was reset and marked not in use + if sd.IsInUse() { + t.Error("Session should not be in use after pool return") + } + if sd.mainSession != nil { + t.Error("Session should be reset after pool return") + } + if sd.IsDirty() { + t.Error("Session should not be dirty after pool return") + } +} + +// TestClearAllSessionData tests the internal clear function +func TestClearAllSessionData(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + + // Set up session data with values + sd.mainSession, _ = store.Get(req, "main") + sd.accessSession, _ = store.Get(req, "access") + sd.refreshSession, _ = store.Get(req, "refresh") + sd.idTokenSession, _ = store.Get(req, "id") + + // Add values to sessions + sd.mainSession.Values["user"] = "test" + sd.accessSession.Values["token"] = "access" + sd.refreshSession.Values["token"] = "refresh" + sd.idTokenSession.Values["token"] = "id" + + // Add chunks + chunk1, _ := store.Get(req, "access-chunk-0") + chunk2, _ := store.Get(req, "refresh-chunk-0") + chunk3, _ := store.Get(req, "id-chunk-0") + sd.accessTokenChunks[0] = chunk1 + sd.refreshTokenChunks[0] = chunk2 + sd.idTokenChunks[0] = chunk3 + + // Test clearing with expire = true + sd.clearAllSessionData(req, true) + + // Verify all sessions are cleared and expired + if sd.mainSession != nil && len(sd.mainSession.Values) != 0 { + t.Error("Main session values should be cleared") + } + if sd.mainSession != nil && sd.mainSession.Options.MaxAge != -1 { + t.Error("Main session should be expired") + } + + // Verify chunks are cleared + if len(sd.accessTokenChunks) != 0 { + t.Error("Access chunks should be cleared") + } + if len(sd.refreshTokenChunks) != 0 { + t.Error("Refresh chunks should be cleared") + } + if len(sd.idTokenChunks) != 0 { + t.Error("ID chunks should be cleared") + } + + // Verify dirty flag is set when expiring + if !sd.IsDirty() { + t.Error("Session should be dirty after clearing with expire=true") + } +} + +// TestClearAllSessionDataWithoutExpire tests clearing without expiring +func TestClearAllSessionDataWithoutExpire(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + + // Set up session data + sd.mainSession, _ = store.Get(req, "main") + sd.mainSession.Values["user"] = "test" + + // Add chunks + chunk1, _ := store.Get(req, "access-chunk-0") + sd.accessTokenChunks[0] = chunk1 + + // Clear without expiring + sd.clearAllSessionData(req, false) + + // Verify values are cleared but not expired + if sd.mainSession != nil && len(sd.mainSession.Values) != 0 { + t.Error("Session values should be cleared") + } + if sd.mainSession != nil && sd.mainSession.Options.MaxAge == -1 { + t.Error("Session should not be expired when expire=false") + } + + // Verify chunks are cleared + if len(sd.accessTokenChunks) != 0 { + t.Error("Chunks should be cleared") + } + + // Verify dirty flag is not set when not expiring + if sd.IsDirty() { + t.Error("Session should not be dirty when expire=false") + } +} + +// TestClearSessionValues tests the clearSessionValues helper +func TestClearSessionValues(t *testing.T) { + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + + session, _ := store.Get(req, "test") + session.Values["key1"] = "value1" + session.Values["key2"] = "value2" + + // Test clearing with expire + clearSessionValues(session, true) + + if len(session.Values) != 0 { + t.Error("Session values should be cleared") + } + if session.Options.MaxAge != -1 { + t.Error("Session should be expired") + } + + // Test clearing without expire + session.Values["key3"] = "value3" + session.Options.MaxAge = 3600 // Reset + + clearSessionValues(session, false) + + if len(session.Values) != 0 { + t.Error("Session values should be cleared") + } + if session.Options.MaxAge == -1 { + t.Error("Session should not be expired when expire=false") + } + + // Test with nil session + clearSessionValues(nil, true) + // Should not panic +} + +// TestClearTokenChunks tests token chunk clearing +func TestClearTokenChunks(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + + // Create chunks with values + chunk1, _ := store.Get(req, "chunk-0") + chunk2, _ := store.Get(req, "chunk-1") + chunk1.Values["data"] = "test1" + chunk2.Values["data"] = "test2" + + chunks := make(map[int]*sessions.Session) + chunks[0] = chunk1 + chunks[1] = chunk2 + + // Clear chunks + sd.clearTokenChunks(req, chunks) + + // Verify chunks are cleared and expired + if len(chunk1.Values) != 0 { + t.Error("Chunk 1 values should be cleared") + } + if chunk1.Options.MaxAge != -1 { + t.Error("Chunk 1 should be expired") + } + + // Verify map is empty + if len(chunks) != 0 { + t.Error("Chunks map should be empty") + } +} + +// TestClearTokenChunksWithNilChunk tests clearing with nil chunk +func TestClearTokenChunksWithNilChunk(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + req := httptest.NewRequest("GET", "http://example.com", nil) + + chunks := make(map[int]*sessions.Session) + chunks[0] = nil // nil chunk + + // Should not panic + sd.clearTokenChunks(req, chunks) + + // Verify map is empty + if len(chunks) != 0 { + t.Error("Chunks map should be empty") + } +} + +// TestSessionDataEdgeCases tests various edge cases +func TestSessionDataEdgeCases(t *testing.T) { + t.Run("Save with nil logger", func(t *testing.T) { + manager := &MockSessionManager{logger: nil} + sd := NewSessionData(manager) + + req := httptest.NewRequest("GET", "http://example.com", nil) + w := httptest.NewRecorder() + + // Should not panic with nil logger + err := sd.Save(req, w) + if err == nil { + t.Log("Save with nil logger succeeded (may be expected)") + } + }) + + t.Run("returnToPoolSafely with panic recovery", func(t *testing.T) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + sd.SetInUse(true) + + // Should not panic + sd.returnToPoolSafely() + + // Check if panic was logged (would require triggering actual panic) + t.Log("returnToPoolSafely completed without panic") + }) +} + +// BenchmarkSessionDataSaveWithChunks benchmarks session save with token chunks +func BenchmarkSessionDataSaveWithChunks(b *testing.B) { + logger := &MockLogger{} + manager := &MockSessionManager{logger: logger} + sd := NewSessionData(manager) + + store := sessions.NewCookieStore([]byte("test-key-32-characters-long-1234")) + req := httptest.NewRequest("GET", "http://example.com", nil) + + sd.mainSession, _ = store.Get(req, "main-session") + + // Add multiple chunks + for i := 0; i < 5; i++ { + chunk, _ := store.Get(req, fmt.Sprintf("access-chunk-%d", i)) + sd.accessTokenChunks[i] = chunk + + refreshChunk, _ := store.Get(req, fmt.Sprintf("refresh-chunk-%d", i)) + sd.refreshTokenChunks[i] = refreshChunk + } + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + w := httptest.NewRecorder() + _ = sd.Save(req, w) + } +} diff --git a/session/validators/session_validator.go b/session/validators/session_validator.go new file mode 100644 index 0000000..8a73611 --- /dev/null +++ b/session/validators/session_validator.go @@ -0,0 +1,300 @@ +// Package validators provides validation functionality for session data +package validators + +import ( + "strings" + "time" +) + +const ( + maxBrowserCookieSize = 3500 + maxCookieSize = 1200 +) + +// SessionValidator provides validation operations for session data +type SessionValidator struct{} + +// NewSessionValidator creates a new session validator +func NewSessionValidator() *SessionValidator { + return &SessionValidator{} +} + +// ValidateChunkSize checks if a chunk will fit within browser cookie limits. +// It estimates the encoded size including cookie overhead and headers +// to ensure the chunk won't exceed browser-imposed cookie size limits. +func (sv *SessionValidator) ValidateChunkSize(chunkData string) bool { + estimatedEncodedSize := len(chunkData) + (len(chunkData)*50)/100 + return estimatedEncodedSize <= maxBrowserCookieSize +} + +// IsCorruptionMarker detects if data contains known corruption indicators. +// It checks for specific corruption markers and invalid characters +// that indicate the data has been tampered with or corrupted. +func (sv *SessionValidator) IsCorruptionMarker(data string) bool { + if data == "" { + return false + } + + corruptionMarkers := []string{ + "__CORRUPTION_MARKER_TEST__", + "__INVALID_BASE64_DATA__", + "__CORRUPTED_CHUNK_DATA__", + "!@#$%^&*()", + "<<>>", + } + + for _, marker := range corruptionMarkers { + if data == marker { + return true + } + } + + if len(data) > 10 { + invalidChars := "!@#$%^&*(){}[]|\\:;\"'<>?,`~" + for _, char := range invalidChars { + if strings.ContainsRune(data, char) { + return true + } + } + } + + return false +} + +// ValidateTokenFormat validates that a token has the correct JWT format +func (sv *SessionValidator) ValidateTokenFormat(token, tokenType string) error { + if token == "" { + return nil // Empty token is not an error + } + + // JWT tokens should have exactly 3 parts separated by dots + parts := strings.Split(token, ".") + if len(parts) != 3 { + return &ValidationError{ + Type: tokenType, + Reason: "invalid JWT format", + Details: "token must have exactly 3 parts separated by dots", + } + } + + // Each part should be non-empty + for i, part := range parts { + if part == "" { + return &ValidationError{ + Type: tokenType, + Reason: "empty token part", + Details: strings.Join([]string{"token part", string(rune(i + 1)), "is empty"}, " "), + } + } + } + + return nil +} + +// ValidateSessionIntegrity performs comprehensive validation of session data integrity +func (sv *SessionValidator) ValidateSessionIntegrity(sessionData SessionData) error { + if sessionData == nil { + return &ValidationError{ + Type: "session", + Reason: "nil session data", + Details: "session data cannot be nil", + } + } + + // Check authentication state consistency + authenticated := sessionData.GetAuthenticated() + email := sessionData.GetEmail() + + if authenticated && email == "" { + return &ValidationError{ + Type: "session", + Reason: "authentication inconsistency", + Details: "session is authenticated but has no email", + } + } + + // Validate token formats if present + if accessToken := sessionData.GetAccessToken(); accessToken != "" { + if err := sv.ValidateTokenFormat(accessToken, "access"); err != nil { + return err + } + } + + if idToken := sessionData.GetIDToken(); idToken != "" { + if err := sv.ValidateTokenFormat(idToken, "id"); err != nil { + return err + } + } + + if refreshToken := sessionData.GetRefreshToken(); refreshToken != "" { + // Refresh tokens don't have to be JWTs, so we do basic validation + if len(refreshToken) == 0 { + return &ValidationError{ + Type: "refresh", + Reason: "empty refresh token", + Details: "refresh token cannot be empty if set", + } + } + } + + return nil +} + +// ValidateSessionTiming validates session timing and expiration +func (sv *SessionValidator) ValidateSessionTiming(sessionData SessionData, maxAge time.Duration) error { + if sessionData == nil { + return &ValidationError{ + Type: "session", + Reason: "nil session data", + Details: "session data cannot be nil", + } + } + + // Check refresh token timing + refreshTokenIssuedAt := sessionData.GetRefreshTokenIssuedAt() + if !refreshTokenIssuedAt.IsZero() { + age := time.Since(refreshTokenIssuedAt) + if age > maxAge { + return &ValidationError{ + Type: "timing", + Reason: "refresh token expired", + Details: strings.Join([]string{"refresh token age", age.String(), "exceeds max age", maxAge.String()}, " "), + } + } + } + + return nil +} + +// ValidateEmailDomain validates that an email belongs to an allowed domain +func (sv *SessionValidator) ValidateEmailDomain(email string, allowedDomains map[string]struct{}) error { + if email == "" { + return &ValidationError{ + Type: "email", + Reason: "empty email", + Details: "email cannot be empty", + } + } + + if len(allowedDomains) == 0 { + return nil // No domain restrictions + } + + parts := strings.Split(email, "@") + if len(parts) != 2 { + return &ValidationError{ + Type: "email", + Reason: "invalid email format", + Details: "email must contain exactly one @ symbol", + } + } + + domain := parts[1] + if _, allowed := allowedDomains[domain]; !allowed { + return &ValidationError{ + Type: "email", + Reason: "domain not allowed", + Details: strings.Join([]string{"domain", domain, "is not in allowed domains list"}, " "), + } + } + + return nil +} + +// SplitIntoChunks splits a string into chunks that fit within cookie size limits +func (sv *SessionValidator) SplitIntoChunks(s string, chunkSize int) []string { + effectiveChunkSize := min(chunkSize, maxCookieSize) + + var chunks []string + for len(s) > 0 { + if len(s) > effectiveChunkSize { + chunks = append(chunks, s[:effectiveChunkSize]) + s = s[effectiveChunkSize:] + } else { + chunks = append(chunks, s) + break + } + } + return chunks +} + +// ValidateChunks validates all chunks in a chunk set +func (sv *SessionValidator) ValidateChunks(chunks []string) error { + for i, chunk := range chunks { + if chunk == "" { + return &ValidationError{ + Type: "chunk", + Reason: "empty chunk", + Details: strings.Join([]string{"chunk", string(rune(i)), "is empty"}, " "), + } + } + + if !sv.ValidateChunkSize(chunk) { + return &ValidationError{ + Type: "chunk", + Reason: "chunk too large", + Details: strings.Join([]string{"chunk", string(rune(i)), "exceeds size limit"}, " "), + } + } + + if sv.IsCorruptionMarker(chunk) { + return &ValidationError{ + Type: "chunk", + Reason: "corrupted chunk", + Details: strings.Join([]string{"chunk", string(rune(i)), "contains corruption markers"}, " "), + } + } + } + + return nil +} + +// ValidationError represents a validation error with context +type ValidationError struct { + Type string + Reason string + Details string +} + +// Error implements the error interface +func (ve *ValidationError) Error() string { + return strings.Join([]string{ve.Type, "validation error:", ve.Reason, "-", ve.Details}, " ") +} + +// SessionData interface for validation operations +type SessionData interface { + GetAuthenticated() bool + GetEmail() string + GetAccessToken() string + GetIDToken() string + GetRefreshToken() string + GetRefreshTokenIssuedAt() time.Time +} + +// Utility functions + +// min returns the minimum of two integers +func min(a, b int) int { + if a < b { + return a + } + return b +} + +// ValidateChunkSize is a package-level function for backward compatibility +func ValidateChunkSize(chunkData string) bool { + sv := &SessionValidator{} + return sv.ValidateChunkSize(chunkData) +} + +// IsCorruptionMarker is a package-level function for backward compatibility +func IsCorruptionMarker(data string) bool { + sv := &SessionValidator{} + return sv.IsCorruptionMarker(data) +} + +// SplitIntoChunks is a package-level function for backward compatibility +func SplitIntoChunks(s string, chunkSize int) []string { + sv := &SessionValidator{} + return sv.SplitIntoChunks(s, chunkSize) +} diff --git a/session/validators/session_validator_test.go b/session/validators/session_validator_test.go new file mode 100644 index 0000000..5f261b4 --- /dev/null +++ b/session/validators/session_validator_test.go @@ -0,0 +1,1106 @@ +package validators + +import ( + "strings" + "testing" + "time" +) + +// MockSessionData for testing +type MockSessionData struct { + authenticated bool + email string + accessToken string + idToken string + refreshToken string + refreshTokenIssuedAt time.Time +} + +func (msd *MockSessionData) GetAuthenticated() bool { return msd.authenticated } +func (msd *MockSessionData) GetEmail() string { return msd.email } +func (msd *MockSessionData) GetAccessToken() string { return msd.accessToken } +func (msd *MockSessionData) GetIDToken() string { return msd.idToken } +func (msd *MockSessionData) GetRefreshToken() string { return msd.refreshToken } +func (msd *MockSessionData) GetRefreshTokenIssuedAt() time.Time { return msd.refreshTokenIssuedAt } + +// TestNewSessionValidator tests validator creation +func TestNewSessionValidator(t *testing.T) { + validator := NewSessionValidator() + if validator == nil { + t.Fatal("NewSessionValidator should not return nil") + } +} + +// TestValidateChunkSize tests chunk size validation +func TestValidateChunkSize(t *testing.T) { + validator := NewSessionValidator() + + tests := []struct { + name string + chunkData string + expectValid bool + description string + }{ + { + name: "Small chunk", + chunkData: "small_chunk_data", + expectValid: true, + description: "Small chunks should be valid", + }, + { + name: "Medium chunk", + chunkData: strings.Repeat("a", 1000), + expectValid: true, + description: "Medium chunks should be valid", + }, + { + name: "Large chunk", + chunkData: strings.Repeat("a", 2000), + expectValid: true, + description: "Large chunks within limits should be valid", + }, + { + name: "Oversized chunk", + chunkData: strings.Repeat("a", 4000), + expectValid: false, + description: "Oversized chunks should be invalid", + }, + { + name: "Empty chunk", + chunkData: "", + expectValid: true, + description: "Empty chunks should be valid", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + isValid := validator.ValidateChunkSize(tt.chunkData) + + if isValid != tt.expectValid { + t.Errorf("Validation mismatch for %s: expected valid=%v, got valid=%v", + tt.description, tt.expectValid, isValid) + } + }) + } +} + +// TestIsCorruptionMarker tests corruption marker detection +func TestIsCorruptionMarker(t *testing.T) { + validator := NewSessionValidator() + + tests := []struct { + name string + data string + expectCorrupted bool + description string + }{ + { + name: "Normal data", + data: "normal_token_data", + expectCorrupted: false, + description: "Normal data should not be marked as corrupted", + }, + { + name: "Empty data", + data: "", + expectCorrupted: false, + description: "Empty data should not be marked as corrupted", + }, + { + name: "Corruption marker test", + data: "__CORRUPTION_MARKER_TEST__", + expectCorrupted: true, + description: "Known corruption markers should be detected", + }, + { + name: "Invalid base64 marker", + data: "__INVALID_BASE64_DATA__", + expectCorrupted: true, + description: "Invalid base64 markers should be detected", + }, + { + name: "Corrupted chunk marker", + data: "__CORRUPTED_CHUNK_DATA__", + expectCorrupted: true, + description: "Corrupted chunk markers should be detected", + }, + { + name: "Invalid characters", + data: "!@#$%^&*()", + expectCorrupted: true, + description: "Invalid character patterns should be detected", + }, + { + name: "Corrupted tag", + data: "<<>>", + expectCorrupted: true, + description: "Corruption tags should be detected", + }, + { + name: "Valid JWT-like token", + data: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9", + expectCorrupted: false, + description: "Valid JWT-like tokens should not be marked as corrupted", + }, + { + name: "Short data with invalid chars", + data: "abc!def", + expectCorrupted: false, + description: "Short data with invalid chars should not be marked as corrupted", + }, + { + name: "Long data with invalid chars", + data: "this_is_long_data_with!invalid@chars#", + expectCorrupted: true, + description: "Long data with invalid chars should be marked as corrupted", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + isCorrupted := validator.IsCorruptionMarker(tt.data) + + if isCorrupted != tt.expectCorrupted { + t.Errorf("Corruption detection mismatch for %s: expected corrupted=%v, got corrupted=%v", + tt.description, tt.expectCorrupted, isCorrupted) + } + }) + } +} + +// TestValidateTokenFormat tests token format validation +func TestValidateTokenFormat(t *testing.T) { + validator := NewSessionValidator() + + tests := []struct { + name string + token string + tokenType string + expectError bool + description string + }{ + { + name: "Valid JWT token", + token: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c", + tokenType: "access", + expectError: false, + description: "Valid JWT tokens should pass validation", + }, + { + name: "Empty token", + token: "", + tokenType: "access", + expectError: false, + description: "Empty tokens should not cause errors", + }, + { + name: "Token with too few parts", + token: "header.payload", + tokenType: "access", + expectError: true, + description: "Tokens with too few parts should fail validation", + }, + { + name: "Token with too many parts", + token: "header.payload.signature.extra", + tokenType: "access", + expectError: true, + description: "Tokens with too many parts should fail validation", + }, + { + name: "Token with empty part", + token: "header..signature", + tokenType: "id", + expectError: true, + description: "Tokens with empty parts should fail validation", + }, + { + name: "Token with only dots", + token: "..", + tokenType: "refresh", + expectError: true, + description: "Tokens with only dots should fail validation", + }, + { + name: "Single part token", + token: "just_one_part", + tokenType: "access", + expectError: true, + description: "Single part tokens should fail validation", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validator.ValidateTokenFormat(tt.token, tt.tokenType) + + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, got nil", tt.description) + } else if !tt.expectError && err != nil { + t.Errorf("Unexpected error for %s: %v", tt.description, err) + } + + // Check error details if error is expected + if tt.expectError && err != nil { + if !strings.Contains(err.Error(), tt.tokenType) { + t.Errorf("Error should contain token type '%s': %v", tt.tokenType, err) + } + } + }) + } +} + +// TestValidateSessionIntegrity tests session integrity validation +func TestValidateSessionIntegrity(t *testing.T) { + validator := NewSessionValidator() + + tests := []struct { + name string + sessionData SessionData + expectError bool + errorCheck func(error) bool + description string + }{ + { + name: "Valid authenticated session", + sessionData: &MockSessionData{ + authenticated: true, + email: "user@example.com", + accessToken: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c", + idToken: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c", + refreshToken: "valid_refresh_token_12345", + }, + expectError: false, + description: "Valid authenticated session should pass validation", + }, + { + name: "Valid unauthenticated session", + sessionData: &MockSessionData{ + authenticated: false, + email: "", + accessToken: "", + idToken: "", + refreshToken: "", + }, + expectError: false, + description: "Valid unauthenticated session should pass validation", + }, + { + name: "Authenticated session without email", + sessionData: &MockSessionData{ + authenticated: true, + email: "", + accessToken: "some_token", + }, + expectError: true, + errorCheck: func(err error) bool { + return strings.Contains(err.Error(), "authentication inconsistency") + }, + description: "Authenticated session without email should fail validation", + }, + { + name: "Session with invalid access token format", + sessionData: &MockSessionData{ + authenticated: true, + email: "user@example.com", + accessToken: "invalid.token", // Only 2 parts + }, + expectError: true, + errorCheck: func(err error) bool { + return strings.Contains(err.Error(), "invalid JWT format") + }, + description: "Session with invalid access token should fail validation", + }, + { + name: "Session with invalid ID token format", + sessionData: &MockSessionData{ + authenticated: true, + email: "user@example.com", + accessToken: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c", + idToken: "invalid_id_token", + }, + expectError: true, + errorCheck: func(err error) bool { + return strings.Contains(err.Error(), "invalid JWT format") + }, + description: "Session with invalid ID token should fail validation", + }, + { + name: "Nil session data", + sessionData: nil, + expectError: true, + errorCheck: func(err error) bool { + return strings.Contains(err.Error(), "nil session data") + }, + description: "Nil session data should fail validation", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validator.ValidateSessionIntegrity(tt.sessionData) + + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, got nil", tt.description) + } else if !tt.expectError && err != nil { + t.Errorf("Unexpected error for %s: %v", tt.description, err) + } + + // Check error details if error is expected and errorCheck is provided + if tt.expectError && err != nil && tt.errorCheck != nil { + if !tt.errorCheck(err) { + t.Errorf("Error check failed for %s: %v", tt.description, err) + } + } + }) + } +} + +// TestValidateSessionTiming tests session timing validation +func TestValidateSessionTiming(t *testing.T) { + validator := NewSessionValidator() + + now := time.Now() + + tests := []struct { + name string + sessionData SessionData + maxAge time.Duration + expectError bool + errorCheck func(error) bool + description string + }{ + { + name: "Recent refresh token", + sessionData: &MockSessionData{ + authenticated: true, + email: "user@example.com", + refreshToken: "valid_token", + refreshTokenIssuedAt: now.Add(-1 * time.Hour), + }, + maxAge: 24 * time.Hour, + expectError: false, + description: "Recent refresh tokens should be valid", + }, + { + name: "Old but valid refresh token", + sessionData: &MockSessionData{ + authenticated: true, + email: "user@example.com", + refreshToken: "valid_token", + refreshTokenIssuedAt: now.Add(-12 * time.Hour), + }, + maxAge: 24 * time.Hour, + expectError: false, + description: "Old but valid refresh tokens should be accepted", + }, + { + name: "Expired refresh token", + sessionData: &MockSessionData{ + authenticated: true, + email: "user@example.com", + refreshToken: "expired_token", + refreshTokenIssuedAt: now.Add(-48 * time.Hour), + }, + maxAge: 24 * time.Hour, + expectError: true, + errorCheck: func(err error) bool { + return strings.Contains(err.Error(), "expired") + }, + description: "Expired refresh tokens should fail validation", + }, + { + name: "Nil session data", + sessionData: nil, + maxAge: 24 * time.Hour, + expectError: true, + errorCheck: func(err error) bool { + return strings.Contains(err.Error(), "nil session data") + }, + description: "Nil session data should fail timing validation", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validator.ValidateSessionTiming(tt.sessionData, tt.maxAge) + + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, got nil", tt.description) + } else if !tt.expectError && err != nil { + t.Errorf("Unexpected error for %s: %v", tt.description, err) + } + + // Check error details if error is expected and errorCheck is provided + if tt.expectError && err != nil && tt.errorCheck != nil { + if !tt.errorCheck(err) { + t.Errorf("Error check failed for %s: %v", tt.description, err) + } + } + }) + } +} + +// TestValidationError tests the ValidationError type +func TestValidationError(t *testing.T) { + err := &ValidationError{ + Type: "test", + Reason: "test reason", + Details: "test details", + } + + expectedMessage := "test validation error: test reason - test details" + if err.Error() != expectedMessage { + t.Errorf("Expected error message %q, got %q", expectedMessage, err.Error()) + } +} + +// TestCorruptionResistance tests comprehensive corruption resistance +func TestCorruptionResistance(t *testing.T) { + validator := NewSessionValidator() + + // Test various corruption scenarios + corruptionScenarios := []struct { + name string + data string + description string + }{ + { + name: "Truncated JWT", + data: "eyJhbGciOiJIUzI1NiIsInR5cCI", + description: "Truncated tokens should be handled gracefully", + }, + { + name: "Malformed base64", + data: "not_valid_base64!@#$", + description: "Malformed base64 should be detected", + }, + { + name: "Binary data", + data: string([]byte{0, 1, 2, 3, 255}), + description: "Binary data should be handled", + }, + { + name: "Very long corruption marker", + data: strings.Repeat("CORRUPT", 100), + description: "Long corruption markers should be handled", + }, + } + + for _, scenario := range corruptionScenarios { + t.Run(scenario.name, func(t *testing.T) { + // Test corruption marker detection + isCorrupted := validator.IsCorruptionMarker(scenario.data) + t.Logf("Data marked as corrupted: %v for %s", isCorrupted, scenario.description) + + // Test token format validation + err := validator.ValidateTokenFormat(scenario.data, "test") + if err != nil { + t.Logf("Token format validation failed (expected): %v", err) + } + + // Test chunk size validation + isValidSize := validator.ValidateChunkSize(scenario.data) + t.Logf("Chunk size valid: %v for %s", isValidSize, scenario.description) + }) + } +} + +// BenchmarkValidateChunkSize benchmarks chunk size validation +func BenchmarkValidateChunkSize(b *testing.B) { + validator := NewSessionValidator() + testData := strings.Repeat("a", 1000) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + validator.ValidateChunkSize(testData) + } +} + +// BenchmarkIsCorruptionMarker benchmarks corruption marker detection +func BenchmarkIsCorruptionMarker(b *testing.B) { + validator := NewSessionValidator() + testData := "normal_token_data_that_should_not_be_corrupted" + + b.ResetTimer() + for i := 0; i < b.N; i++ { + validator.IsCorruptionMarker(testData) + } +} + +// BenchmarkValidateTokenFormat benchmarks token format validation +func BenchmarkValidateTokenFormat(b *testing.B) { + validator := NewSessionValidator() + testToken := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c" + + b.ResetTimer() + for i := 0; i < b.N; i++ { + validator.ValidateTokenFormat(testToken, "access") + } +} + +// TestValidateEmailDomain tests email domain validation +func TestValidateEmailDomain(t *testing.T) { + validator := NewSessionValidator() + + tests := []struct { + name string + email string + allowedDomains map[string]struct{} + expectError bool + errorCheck func(error) bool + description string + }{ + { + name: "Valid email with allowed domain", + email: "user@example.com", + allowedDomains: map[string]struct{}{"example.com": {}, "test.com": {}}, + expectError: false, + description: "Valid email with allowed domain should pass", + }, + { + name: "Valid email with different allowed domain", + email: "admin@test.com", + allowedDomains: map[string]struct{}{"example.com": {}, "test.com": {}}, + expectError: false, + description: "Valid email with different allowed domain should pass", + }, + { + name: "Empty email", + email: "", + allowedDomains: map[string]struct{}{"example.com": {}}, + expectError: true, + errorCheck: func(err error) bool { return strings.Contains(err.Error(), "empty email") }, + description: "Empty email should fail validation", + }, + { + name: "Email with disallowed domain", + email: "user@forbidden.com", + allowedDomains: map[string]struct{}{"example.com": {}, "test.com": {}}, + expectError: true, + errorCheck: func(err error) bool { return strings.Contains(err.Error(), "domain not allowed") }, + description: "Email with disallowed domain should fail validation", + }, + { + name: "Invalid email format - no @ symbol", + email: "userexample.com", + allowedDomains: map[string]struct{}{"example.com": {}}, + expectError: true, + errorCheck: func(err error) bool { return strings.Contains(err.Error(), "invalid email format") }, + description: "Invalid email format should fail validation", + }, + { + name: "Invalid email format - multiple @ symbols", + email: "user@example@com", + allowedDomains: map[string]struct{}{"example.com": {}}, + expectError: true, + errorCheck: func(err error) bool { return strings.Contains(err.Error(), "invalid email format") }, + description: "Email with multiple @ symbols should fail validation", + }, + { + name: "Email starting with @", + email: "@example.com", + allowedDomains: map[string]struct{}{"example.com": {}}, + expectError: false, // splits to ["", "example.com"], domain "example.com" is allowed + description: "Email starting with @ should pass if domain is allowed", + }, + { + name: "Email ending with @ - empty domain allowed", + email: "user@", + allowedDomains: map[string]struct{}{"": {}}, // Allow empty domain + expectError: false, // splits to ["user", ""], domain "" is in allowedDomains + description: "Email ending with @ should pass if empty domain is allowed", + }, + { + name: "Email ending with @ - empty domain not allowed", + email: "user@", + allowedDomains: map[string]struct{}{"example.com": {}}, // Empty domain not allowed + expectError: true, // splits to ["user", ""], domain "" is not in allowedDomains + errorCheck: func(err error) bool { return strings.Contains(err.Error(), "domain not allowed") }, + description: "Email ending with @ should fail if empty domain is not allowed", + }, + { + name: "Valid email with no domain restrictions", + email: "user@anydomain.com", + allowedDomains: map[string]struct{}{}, + expectError: false, + description: "Email should pass when no domain restrictions exist", + }, + { + name: "Valid email with nil domain restrictions", + email: "user@anydomain.com", + allowedDomains: nil, + expectError: false, + description: "Email should pass when domain restrictions are nil", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validator.ValidateEmailDomain(tt.email, tt.allowedDomains) + + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, got nil", tt.description) + } else if !tt.expectError && err != nil { + t.Errorf("Unexpected error for %s: %v", tt.description, err) + } + + // Check error details if error is expected and errorCheck is provided + if tt.expectError && err != nil && tt.errorCheck != nil { + if !tt.errorCheck(err) { + t.Errorf("Error check failed for %s: %v", tt.description, err) + } + } + }) + } +} + +// TestSplitIntoChunks tests string chunking functionality +func TestSplitIntoChunks(t *testing.T) { + validator := NewSessionValidator() + + tests := []struct { + name string + input string + chunkSize int + expectedChunks int + description string + }{ + { + name: "Empty string", + input: "", + chunkSize: 100, + expectedChunks: 0, + description: "Empty string should produce no chunks", + }, + { + name: "Short string", + input: "short", + chunkSize: 100, + expectedChunks: 1, + description: "Short string should produce one chunk", + }, + { + name: "String exactly at chunk size", + input: strings.Repeat("a", 100), + chunkSize: 100, + expectedChunks: 1, + description: "String exactly at chunk size should produce one chunk", + }, + { + name: "String larger than chunk size", + input: strings.Repeat("a", 250), + chunkSize: 100, + expectedChunks: 3, + description: "String larger than chunk size should be split", + }, + { + name: "Large string with small chunks", + input: strings.Repeat("x", 1000), + chunkSize: 50, + expectedChunks: 20, + description: "Large string should be split into many chunks", + }, + { + name: "Chunk size larger than max cookie size", + input: strings.Repeat("a", 2000), + chunkSize: 2000, // Larger than maxCookieSize (1200) + expectedChunks: 2, // Should be limited by maxCookieSize + description: "Chunk size should be limited to max cookie size", + }, + { + name: "Very small chunk size", + input: "testing", + chunkSize: 1, + expectedChunks: 7, + description: "Very small chunk size should create many chunks", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + chunks := validator.SplitIntoChunks(tt.input, tt.chunkSize) + + if len(chunks) != tt.expectedChunks { + t.Errorf("Expected %d chunks for %s, got %d", tt.expectedChunks, tt.description, len(chunks)) + } + + // Verify chunks reconstruct the original string + reconstructed := strings.Join(chunks, "") + if reconstructed != tt.input { + t.Errorf("Reconstructed string doesn't match original for %s", tt.description) + } + + // Verify no chunk exceeds effective size limit + effectiveChunkSize := min(tt.chunkSize, maxCookieSize) + for i, chunk := range chunks { + if len(chunk) > effectiveChunkSize { + t.Errorf("Chunk %d exceeds effective size limit (%d): got %d", i, effectiveChunkSize, len(chunk)) + } + } + }) + } +} + +// TestValidateChunks tests chunk validation +func TestValidateChunks(t *testing.T) { + validator := NewSessionValidator() + + tests := []struct { + name string + chunks []string + expectError bool + errorCheck func(error) bool + description string + }{ + { + name: "Valid chunks", + chunks: []string{"chunk1", "chunk2", "chunk3"}, + expectError: false, + description: "Valid chunks should pass validation", + }, + { + name: "Empty chunk array", + chunks: []string{}, + expectError: false, + description: "Empty chunk array should pass validation", + }, + { + name: "Single valid chunk", + chunks: []string{"single_chunk"}, + expectError: false, + description: "Single valid chunk should pass validation", + }, + { + name: "Chunks with empty chunk", + chunks: []string{"chunk1", "", "chunk3"}, + expectError: true, + errorCheck: func(err error) bool { return strings.Contains(err.Error(), "empty chunk") }, + description: "Empty chunk should fail validation", + }, + { + name: "Chunks with oversized chunk", + chunks: []string{"chunk1", strings.Repeat("a", 5000), "chunk3"}, + expectError: true, + errorCheck: func(err error) bool { return strings.Contains(err.Error(), "chunk too large") }, + description: "Oversized chunk should fail validation", + }, + { + name: "Chunks with corruption marker", + chunks: []string{"chunk1", "__CORRUPTION_MARKER_TEST__", "chunk3"}, + expectError: true, + errorCheck: func(err error) bool { return strings.Contains(err.Error(), "corrupted chunk") }, + description: "Corrupted chunk should fail validation", + }, + { + name: "Chunks with invalid characters", + chunks: []string{"chunk1", "chunk_with_invalid!@#$%^&*()_chars", "chunk3"}, + expectError: true, + errorCheck: func(err error) bool { return strings.Contains(err.Error(), "corrupted chunk") }, + description: "Chunk with invalid characters should fail validation", + }, + { + name: "Multiple invalid chunks", + chunks: []string{"", strings.Repeat("x", 5000), "__CORRUPTED_CHUNK_DATA__"}, + expectError: true, + errorCheck: func(err error) bool { return strings.Contains(err.Error(), "empty chunk") }, // First error encountered + description: "Multiple invalid chunks should fail on first error", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validator.ValidateChunks(tt.chunks) + + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, got nil", tt.description) + } else if !tt.expectError && err != nil { + t.Errorf("Unexpected error for %s: %v", tt.description, err) + } + + // Check error details if error is expected and errorCheck is provided + if tt.expectError && err != nil && tt.errorCheck != nil { + if !tt.errorCheck(err) { + t.Errorf("Error check failed for %s: %v", tt.description, err) + } + } + }) + } +} + +// TestMinFunction tests the min utility function +func TestMinFunction(t *testing.T) { + tests := []struct { + name string + a, b int + expected int + }{ + { + name: "a smaller than b", + a: 5, + b: 10, + expected: 5, + }, + { + name: "b smaller than a", + a: 15, + b: 7, + expected: 7, + }, + { + name: "equal values", + a: 42, + b: 42, + expected: 42, + }, + { + name: "negative values", + a: -10, + b: -5, + expected: -10, + }, + { + name: "zero values", + a: 0, + b: 0, + expected: 0, + }, + { + name: "mixed positive and negative", + a: -3, + b: 2, + expected: -3, + }, + { + name: "large numbers", + a: 1000000, + b: 999999, + expected: 999999, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := min(tt.a, tt.b) + if result != tt.expected { + t.Errorf("min(%d, %d) = %d, expected %d", tt.a, tt.b, result, tt.expected) + } + }) + } +} + +// TestPackageLevelFunctions tests package-level backward compatibility functions +func TestPackageLevelFunctions(t *testing.T) { + t.Run("ValidateChunkSize package function", func(t *testing.T) { + // Test package-level ValidateChunkSize function + testData := "test_chunk_data" + result := ValidateChunkSize(testData) + if !result { + t.Error("Package-level ValidateChunkSize should validate small chunks") + } + + // Test with large data + largeData := strings.Repeat("a", 5000) + result = ValidateChunkSize(largeData) + if result { + t.Error("Package-level ValidateChunkSize should reject oversized chunks") + } + }) + + t.Run("IsCorruptionMarker package function", func(t *testing.T) { + // Test package-level IsCorruptionMarker function + normalData := "normal_data" + result := IsCorruptionMarker(normalData) + if result { + t.Error("Package-level IsCorruptionMarker should not detect corruption in normal data") + } + + // Test with corruption marker + corruptData := "__CORRUPTION_MARKER_TEST__" + result = IsCorruptionMarker(corruptData) + if !result { + t.Error("Package-level IsCorruptionMarker should detect corruption markers") + } + }) + + t.Run("SplitIntoChunks package function", func(t *testing.T) { + // Test package-level SplitIntoChunks function + testString := "test_string_for_chunking" + chunks := SplitIntoChunks(testString, 5) + + if len(chunks) == 0 { + t.Error("Package-level SplitIntoChunks should produce chunks") + } + + // Verify chunks reconstruct original + reconstructed := strings.Join(chunks, "") + if reconstructed != testString { + t.Error("Package-level SplitIntoChunks chunks should reconstruct original string") + } + }) +} + +// TestEdgeCasesAndBoundaryConditions tests various edge cases +func TestEdgeCasesAndBoundaryConditions(t *testing.T) { + validator := NewSessionValidator() + + t.Run("Chunk size boundary conditions", func(t *testing.T) { + // Test chunk size exactly at maxBrowserCookieSize estimation + boundaryData := strings.Repeat("a", 2333) // Should result in ~3500 estimated encoded size + result := validator.ValidateChunkSize(boundaryData) + // This should be close to the boundary + t.Logf("Boundary chunk validation result: %v", result) + }) + + t.Run("Email domain with edge case domains", func(t *testing.T) { + // Test with very short domain + err := validator.ValidateEmailDomain("user@a.b", map[string]struct{}{"a.b": {}}) + if err != nil { + t.Errorf("Should accept very short domains: %v", err) + } + + // Test with very long domain + longDomain := strings.Repeat("long", 50) + ".com" + err = validator.ValidateEmailDomain("user@"+longDomain, map[string]struct{}{longDomain: {}}) + if err != nil { + t.Errorf("Should accept very long domains: %v", err) + } + }) + + t.Run("Chunking with exact boundary sizes", func(t *testing.T) { + // Test with exactly maxCookieSize + testString := strings.Repeat("a", maxCookieSize) + chunks := validator.SplitIntoChunks(testString, maxCookieSize) + + if len(chunks) != 1 { + t.Errorf("String of exactly maxCookieSize should produce 1 chunk, got %d", len(chunks)) + } + + // Test with maxCookieSize + 1 + testString = strings.Repeat("a", maxCookieSize+1) + chunks = validator.SplitIntoChunks(testString, maxCookieSize) + + if len(chunks) != 2 { + t.Errorf("String of maxCookieSize+1 should produce 2 chunks, got %d", len(chunks)) + } + }) +} + +// TestRefreshTokenValidationEdgeCases tests edge cases for refresh token validation +func TestRefreshTokenValidationEdgeCases(t *testing.T) { + validator := NewSessionValidator() + + tests := []struct { + name string + sessionData SessionData + expectError bool + description string + }{ + { + name: "Session with empty refresh token but set", + sessionData: &MockSessionData{ + authenticated: true, + email: "user@example.com", + refreshToken: "", // Empty but explicitly set in the test context + }, + expectError: false, // Empty tokens are not validated for length in current implementation + description: "Empty refresh token should not cause validation error", + }, + { + name: "Session with only refresh token", + sessionData: &MockSessionData{ + authenticated: true, + email: "user@example.com", + accessToken: "", + idToken: "", + refreshToken: "valid_refresh_token_12345", + }, + expectError: false, + description: "Session with only refresh token should be valid", + }, + { + name: "Session with zero-time refresh token issue time", + sessionData: &MockSessionData{ + authenticated: true, + email: "user@example.com", + refreshToken: "valid_token", + refreshTokenIssuedAt: time.Time{}, // Zero time + }, + expectError: false, // Zero time is not validated as expired + description: "Session with zero-time refresh token issue time should be valid", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validator.ValidateSessionIntegrity(tt.sessionData) + + if tt.expectError && err == nil { + t.Errorf("Expected error for %s, got nil", tt.description) + } else if !tt.expectError && err != nil { + t.Errorf("Unexpected error for %s: %v", tt.description, err) + } + }) + } +} + +// BenchmarkValidateEmailDomain benchmarks email domain validation +func BenchmarkValidateEmailDomain(b *testing.B) { + validator := NewSessionValidator() + allowedDomains := map[string]struct{}{ + "example.com": {}, + "test.com": {}, + "domain.org": {}, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + validator.ValidateEmailDomain("user@example.com", allowedDomains) + } +} + +// BenchmarkSplitIntoChunks benchmarks string chunking +func BenchmarkSplitIntoChunks(b *testing.B) { + validator := NewSessionValidator() + testString := strings.Repeat("test_data_", 1000) // 10KB string + + b.ResetTimer() + for i := 0; i < b.N; i++ { + validator.SplitIntoChunks(testString, 100) + } +} + +// BenchmarkValidateChunks benchmarks chunk validation +func BenchmarkValidateChunks(b *testing.B) { + validator := NewSessionValidator() + chunks := []string{ + "chunk_1_data", + "chunk_2_data", + "chunk_3_data", + "chunk_4_data", + "chunk_5_data", + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + validator.ValidateChunks(chunks) + } +} + +// BenchmarkValidateSessionIntegrity benchmarks session integrity validation +func BenchmarkValidateSessionIntegrity(b *testing.B) { + validator := NewSessionValidator() + sessionData := &MockSessionData{ + authenticated: true, + email: "user@example.com", + accessToken: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c", + idToken: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c", + refreshToken: "valid_refresh_token", + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + validator.ValidateSessionIntegrity(sessionData) + } +} diff --git a/types.go b/types.go new file mode 100644 index 0000000..d923917 --- /dev/null +++ b/types.go @@ -0,0 +1,113 @@ +// Package traefikoidc provides OIDC authentication middleware for Traefik. +package traefikoidc + +import ( + "context" + "net/http" + "sync" + "text/template" + "time" + + "golang.org/x/time/rate" +) + +// CacheInterface defines the common cache operations +type CacheInterface interface { + Set(key string, value interface{}, ttl time.Duration) + Get(key string) (interface{}, bool) + Delete(key string) + SetMaxSize(size int) + Cleanup() + Close() +} + +// TokenVerifier interface defines token verification capabilities. +// Implementations should validate token format, signature, and claims. +type TokenVerifier interface { + VerifyToken(token string) error +} + +// JWTVerifier interface defines JWT-specific verification capabilities. +// Implementations should validate JWT structure, signature using JWKs, and standard claims. +type JWTVerifier interface { + VerifyJWTSignatureAndClaims(jwt *JWT, token string) error +} + +// TokenExchanger interface defines OAuth 2.0 and OpenID Connect token exchange capabilities. +// Implementations should handle authorization code exchange, refresh tokens, and revocation +// according to the OAuth 2.0 and OpenID Connect specifications. +type TokenExchanger interface { + ExchangeCodeForToken(ctx context.Context, grantType string, codeOrToken string, redirectURL string, codeVerifier string) (*TokenResponse, error) + GetNewTokenWithRefreshToken(refreshToken string) (*TokenResponse, error) + RevokeTokenWithProvider(token, tokenType string) error +} + +// ProviderMetadata represents OIDC provider configuration data. +// This data is typically retrieved from the provider's .well-known/openid-configuration endpoint +// and contains essential URLs for authentication, token exchange, and key retrieval. +type ProviderMetadata struct { + Issuer string `json:"issuer"` + AuthURL string `json:"authorization_endpoint"` + TokenURL string `json:"token_endpoint"` + JWKSURL string `json:"jwks_uri"` + RevokeURL string `json:"revocation_endpoint"` + EndSessionURL string `json:"end_session_endpoint"` +} + +// TraefikOidc is the main middleware struct that implements OIDC authentication for Traefik. +// It integrates with various OIDC providers, manages sessions, caches tokens, and handles +// the complete authentication flow. It's designed to work seamlessly with Traefik's +// plugin system and provides flexible configuration options. +type TraefikOidc struct { + jwkCache JWKCacheInterface + jwtVerifier JWTVerifier + ctx context.Context + tokenVerifier TokenVerifier + next http.Handler + tokenExchanger TokenExchanger + initComplete chan struct{} + limiter *rate.Limiter + tokenBlacklist CacheInterface + headerTemplates map[string]*template.Template + sessionManager *SessionManager + tokenCleanupStopChan chan struct{} + excludedURLs map[string]struct{} + extractClaimsFunc func(tokenString string) (map[string]interface{}, error) + initiateAuthenticationFunc func(rw http.ResponseWriter, req *http.Request, session *SessionData, redirectURL string) + metadataCache *MetadataCache + allowedRolesAndGroups map[string]struct{} + allowedUsers map[string]struct{} + allowedUserDomains map[string]struct{} + tokenCache *TokenCache + httpClient *http.Client + tokenHTTPClient *http.Client + logger *Logger + metadataRefreshStopChan chan struct{} + cancelFunc context.CancelFunc + errorRecoveryManager *ErrorRecoveryManager + goroutineWG *sync.WaitGroup + clientSecret string + clientID string + name string + redirURLPath string + logoutURLPath string + tokenURL string + authURL string + endSessionURL string + postLogoutRedirectURI string + scheme string + jwksURL string + issuerURL string + revocationURL string + providerURL string + scopes []string + refreshGracePeriod time.Duration + shutdownOnce sync.Once + firstRequestMutex sync.Mutex + forceHTTPS bool + enablePKCE bool + overrideScopes bool + suppressDiagnosticLogs bool + firstRequestReceived bool + metadataRefreshStarted bool +}