Improve documentation.

helpers.go

@@ -16,6 +16,13 @@ import (
 	"github.com/gorilla/sessions"
 )
 
+// newSessionOptions creates secure session cookie options.
+// Parameters:
+//   - isSecure: Whether to set the Secure flag on cookies
+// Returns session options configured for security with:
+//   - HttpOnly flag to prevent JavaScript access
+//   - SameSite=Lax for CSRF protection
+//   - Appropriate timeout and path settings
 func newSessionOptions(isSecure bool) *sessions.Options {
 	return &sessions.Options{
 		HttpOnly: true,
@@ -26,7 +33,10 @@ func newSessionOptions(isSecure bool) *sessions.Options {
 	}
 }
 
-// generateNonce generates a random nonce
+// generateNonce creates a cryptographically secure random nonce
+// for use in the OIDC authentication flow. The nonce is used to
+// prevent replay attacks by ensuring the token received matches
+// the authentication request.
 func generateNonce() (string, error) {
 	nonceBytes := make([]byte, 32)
 	_, err := rand.Read(nonceBytes)
@@ -36,7 +46,33 @@ func generateNonce() (string, error) {
 	return base64.URLEncoding.EncodeToString(nonceBytes), nil
 }
 
-// exchangeTokens exchanges a code or refresh token for tokens
+// TokenResponse represents the response from the OIDC token endpoint.
+// It contains the various tokens and metadata returned after successful
+// code exchange or token refresh operations.
+type TokenResponse struct {
+	// IDToken is the OIDC ID token containing user claims
+	IDToken string `json:"id_token"`
+
+	// AccessToken is the OAuth 2.0 access token for API access
+	AccessToken string `json:"access_token"`
+
+	// RefreshToken is the OAuth 2.0 refresh token for obtaining new tokens
+	RefreshToken string `json:"refresh_token"`
+
+	// ExpiresIn is the lifetime in seconds of the access token
+	ExpiresIn int `json:"expires_in"`
+
+	// TokenType is the type of token, typically "Bearer"
+	TokenType string `json:"token_type"`
+}
+
+// exchangeTokens performs the OAuth 2.0 token exchange with the OIDC provider.
+// It supports both authorization code and refresh token grant types.
+// Parameters:
+//   - ctx: Context for the HTTP request
+//   - grantType: The OAuth 2.0 grant type ("authorization_code" or "refresh_token")
+//   - codeOrToken: Either the authorization code or refresh token
+//   - redirectURL: The callback URL for authorization code grant
 func (t *TraefikOidc) exchangeTokens(ctx context.Context, grantType, codeOrToken, redirectURL string) (*TokenResponse, error) {
 	data := url.Values{
 		"grant_type":    {grantType},
@@ -76,16 +112,8 @@ func (t *TraefikOidc) exchangeTokens(ctx context.Context, grantType, codeOrToken
 	return &tokenResponse, nil
 }
 
-// TokenResponse represents the response from the token endpoint
-type TokenResponse struct {
-	IDToken      string `json:"id_token"`
-	AccessToken  string `json:"access_token"`
-	RefreshToken string `json:"refresh_token"`
-	ExpiresIn    int    `json:"expires_in"`
-	TokenType    string `json:"token_type"`
-}
-
-// getNewTokenWithRefreshToken refreshes the token using the refresh token
+// getNewTokenWithRefreshToken obtains new tokens using a refresh token.
+// This is used to refresh access tokens before they expire.
 func (t *TraefikOidc) getNewTokenWithRefreshToken(refreshToken string) (*TokenResponse, error) {
 	ctx := context.Background()
 	tokenResponse, err := t.exchangeTokens(ctx, "refresh_token", refreshToken, "")
@@ -94,24 +122,23 @@ func (t *TraefikOidc) getNewTokenWithRefreshToken(refreshToken string) (*TokenRe
 	}
 
 	t.logger.Debugf("Token response: %+v", tokenResponse)
-
 	return tokenResponse, nil
 }
 
-// handleExpiredToken handles the case when a token has expired
+// handleExpiredToken manages token expiration by clearing the session
+// and initiating a new authentication flow.
 func (t *TraefikOidc) handleExpiredToken(rw http.ResponseWriter, req *http.Request, session *SessionData, redirectURL string) {
-	// Clear the existing session
 	if err := session.Clear(req, rw); err != nil {
 		t.logger.Errorf("Failed to clear session: %v", err)
 		http.Error(rw, "Internal Server Error", http.StatusInternalServerError)
 		return
 	}
-
-	// Initialize new authentication
 	t.defaultInitiateAuthentication(rw, req, session, redirectURL)
 }
 
-// handleCallback handles the callback from the OIDC provider
+// handleCallback processes the authentication callback from the OIDC provider.
+// It validates the callback parameters, exchanges the authorization code for
+// tokens, verifies the tokens, and establishes the user's session.
 func (t *TraefikOidc) handleCallback(rw http.ResponseWriter, req *http.Request, redirectURL string) {
 	session, err := t.sessionManager.GetSession(req)
 	if err != nil {
@@ -122,7 +149,7 @@ func (t *TraefikOidc) handleCallback(rw http.ResponseWriter, req *http.Request,
 
 	t.logger.Debugf("Handling callback, URL: %s", req.URL.String())
 
-	// Check for errors in the query parameters
+	// Check for errors in the callback
 	if req.URL.Query().Get("error") != "" {
 		errorDescription := req.URL.Query().Get("error_description")
 		t.logger.Errorf("Authentication error: %s - %s", req.URL.Query().Get("error"), errorDescription)
@@ -130,7 +157,7 @@ func (t *TraefikOidc) handleCallback(rw http.ResponseWriter, req *http.Request,
 		return
 	}
 
-	// Validate state parameter matches the session's CSRF token
+	// Validate CSRF state
 	state := req.URL.Query().Get("state")
 	if state == "" {
 		t.logger.Error("No state in callback")
@@ -166,7 +193,7 @@ func (t *TraefikOidc) handleCallback(rw http.ResponseWriter, req *http.Request,
 		return
 	}
 
-	// Verify and process tokens
+	// Verify tokens and claims
 	if err := t.verifyToken(tokenResponse.IDToken); err != nil {
 		t.logger.Errorf("Failed to verify id_token: %v", err)
 		http.Error(rw, "Authentication failed", http.StatusInternalServerError)
@@ -180,7 +207,7 @@ func (t *TraefikOidc) handleCallback(rw http.ResponseWriter, req *http.Request,
 		return
 	}
 
-	// Verify nonce
+	// Verify nonce to prevent replay attacks
 	nonceClaim, ok := claims["nonce"].(string)
 	if !ok || nonceClaim == "" {
 		t.logger.Error("Nonce claim missing in id_token")
@@ -201,7 +228,7 @@ func (t *TraefikOidc) handleCallback(rw http.ResponseWriter, req *http.Request,
 		return
 	}
 
-	// Process email
+	// Validate user's email domain
 	email, _ := claims["email"].(string)
 	if email == "" || !t.isAllowedDomain(email) {
 		t.logger.Errorf("Invalid or disallowed email: %s", email)
@@ -209,13 +236,12 @@ func (t *TraefikOidc) handleCallback(rw http.ResponseWriter, req *http.Request,
 		return
 	}
 
-	// Update session with new values
+	// Update session with authentication data
 	session.SetAuthenticated(true)
 	session.SetEmail(email)
 	session.SetAccessToken(tokenResponse.IDToken)
 	session.SetRefreshToken(tokenResponse.RefreshToken)
 
-	// Save session
 	if err := session.Save(req, rw); err != nil {
 		t.logger.Errorf("Failed to save session: %v", err)
 		http.Error(rw, "Failed to save session", http.StatusInternalServerError)
@@ -231,7 +257,8 @@ func (t *TraefikOidc) handleCallback(rw http.ResponseWriter, req *http.Request,
 	http.Redirect(rw, req, redirectPath, http.StatusFound)
 }
 
-// extractClaims extracts claims from a JWT token
+// extractClaims parses a JWT token and extracts its claims.
+// It handles base64url decoding and JSON parsing of the token payload.
 func extractClaims(tokenString string) (map[string]interface{}, error) {
 	parts := strings.Split(tokenString, ".")
 	if len(parts) != 3 {
@@ -251,27 +278,32 @@ func extractClaims(tokenString string) (map[string]interface{}, error) {
 	return claims, nil
 }
 
-// TokenBlacklist maintains a blacklist of tokens
+// TokenBlacklist maintains a thread-safe list of revoked tokens.
+// It stores tokens with their expiration times and automatically
+// removes expired entries during cleanup operations.
 type TokenBlacklist struct {
+	// blacklist maps token IDs to their expiration times
 	blacklist map[string]time.Time
-	mutex     sync.RWMutex
+
+	// mutex protects concurrent access to the blacklist
+	mutex sync.RWMutex
 }
 
-// NewTokenBlacklist creates a new TokenBlacklist
+// NewTokenBlacklist creates a new TokenBlacklist instance.
 func NewTokenBlacklist() *TokenBlacklist {
 	return &TokenBlacklist{
 		blacklist: make(map[string]time.Time),
 	}
 }
 
-// Add adds a token to the blacklist
+// Add adds a token to the blacklist with an expiration time.
 func (tb *TokenBlacklist) Add(tokenID string, expiration time.Time) {
 	tb.mutex.Lock()
 	defer tb.mutex.Unlock()
 	tb.blacklist[tokenID] = expiration
 }
 
-// IsBlacklisted checks if a token is blacklisted
+// IsBlacklisted checks if a token is in the blacklist and not expired.
 func (tb *TokenBlacklist) IsBlacklisted(tokenID string) bool {
 	tb.mutex.RLock()
 	defer tb.mutex.RUnlock()
@@ -279,7 +311,7 @@ func (tb *TokenBlacklist) IsBlacklisted(tokenID string) bool {
 	return exists && time.Now().Before(expiration)
 }
 
-// Cleanup removes expired tokens from the blacklist
+// Cleanup removes expired tokens from the blacklist.
 func (tb *TokenBlacklist) Cleanup() {
 	tb.mutex.Lock()
 	defer tb.mutex.Unlock()
@@ -291,25 +323,29 @@ func (tb *TokenBlacklist) Cleanup() {
 	}
 }
 
-// TokenCache caches tokens
+// TokenCache provides a caching mechanism for validated tokens.
+// It stores token claims to avoid repeated validation of the
+// same token, improving performance for frequently used tokens.
 type TokenCache struct {
+	// cache is the underlying cache implementation
 	cache *Cache
 }
 
-// NewTokenCache creates a new TokenCache
+// NewTokenCache creates a new TokenCache instance.
 func NewTokenCache() *TokenCache {
 	return &TokenCache{
 		cache: NewCache(),
 	}
 }
 
-// Set sets a token in the cache
+// Set stores a token's claims in the cache with an expiration time.
 func (tc *TokenCache) Set(token string, claims map[string]interface{}, expiration time.Duration) {
 	token = "t-" + token
 	tc.cache.Set(token, claims, expiration)
 }
 
-// Get retrieves a token from the cache
+// Get retrieves a token's claims from the cache.
+// Returns the claims and a boolean indicating if the token was found.
 func (tc *TokenCache) Get(token string) (map[string]interface{}, bool) {
 	token = "t-" + token
 	value, found := tc.cache.Get(token)
@@ -320,18 +356,18 @@ func (tc *TokenCache) Get(token string) (map[string]interface{}, bool) {
 	return claims, ok
 }
 
-// Delete removes a token from the cache
+// Delete removes a token from the cache.
 func (tc *TokenCache) Delete(token string) {
 	token = "t-" + token
 	tc.cache.Delete(token)
 }
 
-// Cleanup cleans up expired tokens from the cache
+// Cleanup removes expired tokens from the cache.
 func (tc *TokenCache) Cleanup() {
 	tc.cache.Cleanup()
 }
 
-// exchangeCodeForToken exchanges the authorization code for tokens
+// exchangeCodeForToken exchanges an authorization code for tokens.
 func (t *TraefikOidc) exchangeCodeForToken(code string, redirectURL string) (*TokenResponse, error) {
 	ctx := context.Background()
 	tokenResponse, err := t.exchangeTokens(ctx, "authorization_code", code, redirectURL)
@@ -341,7 +377,8 @@ func (t *TraefikOidc) exchangeCodeForToken(code string, redirectURL string) (*To
 	return tokenResponse, nil
 }
 
-// createStringMap creates a map from a slice of strings
+// createStringMap creates a map from a slice of strings.
+// Used for efficient lookups in allowed domains and roles.
 func createStringMap(keys []string) map[string]struct{} {
 	result := make(map[string]struct{})
 	for _, key := range keys {
@@ -350,7 +387,9 @@ func createStringMap(keys []string) map[string]struct{} {
 	return result
 }
 
-// handleLogout handles the logout request
+// handleLogout manages the OIDC logout process.
+// It clears the session and redirects either to the OIDC provider's
+// end session endpoint (if available) or to the configured post-logout URL.
 func (t *TraefikOidc) handleLogout(rw http.ResponseWriter, req *http.Request) {
 	session, err := t.sessionManager.GetSession(req)
 	if err != nil {
@@ -359,22 +398,18 @@ func (t *TraefikOidc) handleLogout(rw http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	// Get the access token before clearing session
 	accessToken := session.GetAccessToken()
 
-	// Clear all session data
 	if err := session.Clear(req, rw); err != nil {
 		t.logger.Errorf("Error clearing session: %v", err)
 		http.Error(rw, "Session error", http.StatusInternalServerError)
 		return
 	}
 
-	// Get the base URL for redirects
 	host := t.determineHost(req)
 	scheme := t.determineScheme(req)
 	baseURL := fmt.Sprintf("%s://%s", scheme, host)
 
-	// Determine post logout redirect URI
 	postLogoutRedirectURI := t.postLogoutRedirectURI
 	if postLogoutRedirectURI == "" {
 		postLogoutRedirectURI = fmt.Sprintf("%s/", baseURL)
@@ -382,7 +417,6 @@ func (t *TraefikOidc) handleLogout(rw http.ResponseWriter, req *http.Request) {
 		postLogoutRedirectURI = fmt.Sprintf("%s%s", baseURL, postLogoutRedirectURI)
 	}
 
-	// If we have an end session endpoint and an access token, use OIDC end session
 	if t.endSessionURL != "" && accessToken != "" {
 		logoutURL, err := BuildLogoutURL(t.endSessionURL, accessToken, postLogoutRedirectURI)
 		if err != nil {
@@ -394,11 +428,14 @@ func (t *TraefikOidc) handleLogout(rw http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	// Otherwise, redirect to post logout URI
 	http.Redirect(rw, req, postLogoutRedirectURI, http.StatusFound)
 }
 
-// BuildLogoutURL constructs the OIDC end session URL
+// BuildLogoutURL constructs the OIDC end session URL with appropriate parameters.
+// Parameters:
+//   - endSessionURL: The OIDC provider's end session endpoint
+//   - idToken: The ID token to be invalidated
+//   - postLogoutRedirectURI: Where to redirect after logout completes
 func BuildLogoutURL(endSessionURL, idToken, postLogoutRedirectURI string) (string, error) {
 	u, err := url.Parse(endSessionURL)
 	if err != nil {
@@ -408,7 +445,6 @@ func BuildLogoutURL(endSessionURL, idToken, postLogoutRedirectURI string) (strin
 	q := u.Query()
 	q.Set("id_token_hint", idToken)
 	if postLogoutRedirectURI != "" {
-		// Ensure postLogoutRedirectURI is properly URL encoded
 		q.Set("post_logout_redirect_uri", postLogoutRedirectURI)
 	}
 	u.RawQuery = q.Encode()