Cleanup and optimise the code.

helpers.go

@@ -6,6 +6,7 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	"net/url"
 	"strings"
@@ -16,6 +17,7 @@ import (
 	"github.com/gorilla/sessions"
 )
 
+// generateNonce generates a random nonce
 func generateNonce() (string, error) {
 	nonceBytes := make([]byte, 32)
 	_, err := rand.Read(nonceBytes)
@@ -25,6 +27,7 @@ func generateNonce() (string, error) {
 	return base64.URLEncoding.EncodeToString(nonceBytes), nil
 }
 
+// buildFullURL constructs a full URL from scheme, host, and path
 func buildFullURL(scheme, host, path string) string {
 	if scheme == "" {
 		scheme = "http"
@@ -32,7 +35,8 @@ func buildFullURL(scheme, host, path string) string {
 	return fmt.Sprintf("%s://%s%s", scheme, host, path)
 }
 
-func (t *TraefikOidc) exchangeTokens(ctx context.Context, grantType, codeOrToken, redirectURL string) (map[string]interface{}, error) {
+// exchangeTokens exchanges a code or refresh token for tokens
+func (t *TraefikOidc) exchangeTokens(ctx context.Context, grantType, codeOrToken, redirectURL string) (*TokenResponse, error) {
 	data := url.Values{
 		"grant_type":    {grantType},
 		"client_id":     {t.clientID},
@@ -58,14 +62,20 @@ func (t *TraefikOidc) exchangeTokens(ctx context.Context, grantType, codeOrToken
 	}
 	defer resp.Body.Close()
 
-	var result map[string]interface{}
-	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+	if resp.StatusCode != http.StatusOK {
+		bodyBytes, _ := io.ReadAll(resp.Body)
+		return nil, fmt.Errorf("token endpoint returned status %d: %s", resp.StatusCode, string(bodyBytes))
+	}
+
+	var tokenResponse TokenResponse
+	if err := json.NewDecoder(resp.Body).Decode(&tokenResponse); err != nil {
 		return nil, fmt.Errorf("failed to decode token response: %w", err)
 	}
 
-	return result, nil
+	return &tokenResponse, nil
 }
 
+// TokenResponse represents the response from the token endpoint
 type TokenResponse struct {
 	IDToken      string `json:"id_token"`
 	AccessToken  string `json:"access_token"`
@@ -74,47 +84,20 @@ type TokenResponse struct {
 	TokenType    string `json:"token_type"`
 }
 
+// getNewTokenWithRefreshToken refreshes the token using the refresh token
 func (t *TraefikOidc) getNewTokenWithRefreshToken(refreshToken string) (*TokenResponse, error) {
 	ctx := context.Background()
-	result, err := t.exchangeTokens(ctx, "refresh_token", refreshToken, "")
+	tokenResponse, err := t.exchangeTokens(ctx, "refresh_token", refreshToken, "")
 	if err != nil {
 		return nil, fmt.Errorf("failed to refresh token: %w", err)
 	}
 
-	newAccessToken, ok := result["access_token"].(string)
-	if !ok || newAccessToken == "" {
-		return nil, fmt.Errorf("no access_token field in token response")
-	}
+	t.logger.Debugf("Token response: %+v", tokenResponse)
 
-	rawIDToken, ok := result["id_token"].(string)
-	if !ok || rawIDToken == "" {
-		return nil, fmt.Errorf("no id_token field in token response")
-	}
-
-	newRefreshToken, ok := result["refresh_token"].(string)
-	if !ok || newRefreshToken == "" {
-		return nil, fmt.Errorf("no refresh_token field in token response")
-	}
-
-	response := &TokenResponse{
-		IDToken:     rawIDToken,
-		AccessToken: newAccessToken,
-		ExpiresIn:   int(result["expires_in"].(float64)),
-		TokenType:   result["token_type"].(string),
-	}
-
-	// The refresh token might not be returned if it hasn't changed
-	if newRefreshToken != refreshToken {
-		response.RefreshToken = newRefreshToken
-	} else {
-		response.RefreshToken = refreshToken
-	}
-
-	t.logger.Debug("Token response: %+v", response)
-
-	return response, nil
+	return tokenResponse, nil
 }
 
+// handleLogout handles the user logout
 func (t *TraefikOidc) handleLogout(rw http.ResponseWriter, req *http.Request) {
 	session, err := t.store.Get(req, cookieName)
 	t.logger.Debugf("Logging out user")
@@ -123,28 +106,41 @@ func (t *TraefikOidc) handleLogout(rw http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	if idToken, ok := session.Values["id_token"].(string); ok {
-		err := t.RevokeTokenWithProvider(idToken)
-		if err != nil {
-			handleError(rw, "Failed to revoke token", http.StatusInternalServerError, t.logger)
-			return
+	// Revoke tokens if available
+	if refreshToken, ok := session.Values["refresh_token"].(string); ok && refreshToken != "" {
+		if err := t.RevokeTokenWithProvider(refreshToken, "refresh_token"); err != nil {
+			t.logger.Errorf("Failed to revoke refresh token: %v", err)
 		}
-		t.RevokeToken(idToken)
+		t.RevokeToken(refreshToken)
+	}
+	if accessToken, ok := session.Values["access_token"].(string); ok && accessToken != "" {
+		if err := t.RevokeTokenWithProvider(accessToken, "access_token"); err != nil {
+			t.logger.Errorf("Failed to revoke access token: %v", err)
+		}
+		t.RevokeToken(accessToken)
 	}
 
+	// Remove tokens from session
+	delete(session.Values, "id_token")
+	delete(session.Values, "refresh_token")
+	delete(session.Values, "access_token")
+	delete(session.Values, "authenticated")
+
+	// Set session options to delete the session
 	session.Options = defaultSessionOptions
-	// Clear the session
 	session.Options.MaxAge = -1
-	session.Values = make(map[interface{}]interface{})
-	err = session.Save(req, rw)
-	if err != nil {
+
+	if err := session.Save(req, rw); err != nil {
 		handleError(rw, "Failed to save session", http.StatusInternalServerError, t.logger)
 		return
 	}
 
-	http.Error(rw, "Logged out", http.StatusForbidden)
+	// Redirect or display logout message
+	rw.WriteHeader(http.StatusOK)
+	rw.Write([]byte("Logged out successfully"))
 }
 
+// handleExpiredToken handles the case when a token has expired
 func (t *TraefikOidc) handleExpiredToken(rw http.ResponseWriter, req *http.Request, session *sessions.Session) {
 	// Clear the existing session
 	session.Options.MaxAge = -1
@@ -169,6 +165,7 @@ func (t *TraefikOidc) handleExpiredToken(rw http.ResponseWriter, req *http.Reque
 	t.initiateAuthenticationFunc(rw, req, session, t.redirectURL)
 }
 
+// handleCallback handles the callback from the OIDC provider
 func (t *TraefikOidc) handleCallback(rw http.ResponseWriter, req *http.Request) {
 	session, err := t.store.Get(req, cookieName)
 	if err != nil {
@@ -179,6 +176,34 @@ 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
+	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)
+		http.Error(rw, fmt.Sprintf("Authentication error: %s", errorDescription), http.StatusBadRequest)
+		return
+	}
+
+	// Validate the state parameter matches the session's CSRF token
+	state := req.URL.Query().Get("state")
+	if state == "" {
+		t.logger.Error("No state in callback")
+		http.Error(rw, "State parameter missing in callback", http.StatusBadRequest)
+		return
+	}
+	csrfToken, ok := session.Values["csrf"].(string)
+	if !ok || csrfToken == "" {
+		t.logger.Error("CSRF token missing in session")
+		http.Error(rw, "CSRF token missing", http.StatusBadRequest)
+		return
+	}
+	if state != csrfToken {
+		t.logger.Error("State parameter does not match CSRF token in session")
+		http.Error(rw, "Invalid state parameter", http.StatusBadRequest)
+		return
+	}
+
+	// Proceed to exchange the code for tokens
 	code := req.URL.Query().Get("code")
 	if code == "" {
 		t.logger.Error("No code in callback")
@@ -186,20 +211,29 @@ func (t *TraefikOidc) handleCallback(rw http.ResponseWriter, req *http.Request)
 		return
 	}
 
-	token, err := t.exchangeCodeForTokenFunc(code)
+	tokenResponse, err := t.exchangeCodeForTokenFunc(code)
 	if err != nil {
 		t.logger.Errorf("Failed to exchange code for token: %v", err)
 		http.Error(rw, "Authentication failed", http.StatusInternalServerError)
 		return
 	}
 
-	idToken, ok := token["id_token"].(string)
-	if !ok || idToken == "" {
+	// Extract id_token
+	idToken := tokenResponse.IDToken
+	if idToken == "" {
 		t.logger.Error("No id_token in token response")
 		http.Error(rw, "Authentication failed", http.StatusInternalServerError)
 		return
 	}
 
+	// Verify the id_token
+	if err := t.verifyToken(idToken); err != nil {
+		t.logger.Errorf("Failed to verify id_token: %v", err)
+		http.Error(rw, "Authentication failed", http.StatusInternalServerError)
+		return
+	}
+
+	// Extract claims from id_token
 	claims, err := t.extractClaimsFunc(idToken)
 	if err != nil {
 		t.logger.Errorf("Failed to extract claims: %v", err)
@@ -207,6 +241,26 @@ func (t *TraefikOidc) handleCallback(rw http.ResponseWriter, req *http.Request)
 		return
 	}
 
+	// Verify the nonce claim matches the one stored in session
+	nonceClaim, ok := claims["nonce"].(string)
+	if !ok || nonceClaim == "" {
+		t.logger.Error("Nonce claim missing in id_token")
+		http.Error(rw, "Authentication failed", http.StatusInternalServerError)
+		return
+	}
+	sessionNonce, ok := session.Values["nonce"].(string)
+	if !ok || sessionNonce == "" {
+		t.logger.Error("Nonce not found in session")
+		http.Error(rw, "Authentication failed", http.StatusInternalServerError)
+		return
+	}
+	if nonceClaim != sessionNonce {
+		t.logger.Error("Nonce claim does not match session nonce")
+		http.Error(rw, "Authentication failed", http.StatusInternalServerError)
+		return
+	}
+
+	// Get the email from claims
 	email, _ := claims["email"].(string)
 	if email == "" || !t.isAllowedDomain(email) {
 		t.logger.Errorf("Invalid or disallowed email: %s", email)
@@ -214,11 +268,17 @@ func (t *TraefikOidc) handleCallback(rw http.ResponseWriter, req *http.Request)
 		return
 	}
 
+	// Store tokens and authentication status in session
 	session.Values["authenticated"] = true
 	session.Values["email"] = email
 	session.Values["id_token"] = idToken
+	session.Values["refresh_token"] = tokenResponse.RefreshToken
 	session.Options = defaultSessionOptions
 
+	// Remove CSRF and nonce from session
+	delete(session.Values, "csrf")
+	delete(session.Values, "nonce")
+
 	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)
@@ -226,16 +286,17 @@ func (t *TraefikOidc) handleCallback(rw http.ResponseWriter, req *http.Request)
 	}
 
 	t.logger.Debugf("Authentication successful. User email: %s", email)
-	http.Redirect(rw, req, func() string {
-		if path, ok := session.Values["incoming_path"].(string); ok {
-			t.logger.Debug("Redirecting to incoming path from original request: %s", path)
-			return path
-		}
-		t.logger.Debug("Redirecting to root path as no incoming path found")
-		return "/"
-	}(), http.StatusFound)
+
+	// Redirect to the original requested path or default to root
+	redirectPath := "/"
+	if path, ok := session.Values["incoming_path"].(string); ok && path != t.redirURLPath {
+		t.logger.Debugf("Redirecting to incoming path from original request: %s", path)
+		redirectPath = path
+	}
+	http.Redirect(rw, req, redirectPath, http.StatusFound)
 }
 
+// extractClaims extracts claims from a JWT token
 func extractClaims(tokenString string) (map[string]interface{}, error) {
 	parts := strings.Split(tokenString, ".")
 	if len(parts) != 3 {
@@ -255,23 +316,27 @@ func extractClaims(tokenString string) (map[string]interface{}, error) {
 	return claims, nil
 }
 
+// TokenBlacklist maintains a blacklist of tokens
 type TokenBlacklist struct {
 	blacklist map[string]time.Time
 	mutex     sync.RWMutex
 }
 
+// NewTokenBlacklist creates a new TokenBlacklist
 func NewTokenBlacklist() *TokenBlacklist {
 	return &TokenBlacklist{
 		blacklist: make(map[string]time.Time),
 	}
 }
 
+// Add adds a token to the blacklist
 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
 func (tb *TokenBlacklist) IsBlacklisted(tokenID string) bool {
 	tb.mutex.RLock()
 	defer tb.mutex.RUnlock()
@@ -279,6 +344,7 @@ func (tb *TokenBlacklist) IsBlacklisted(tokenID string) bool {
 	return exists && time.Now().Before(expiration)
 }
 
+// Cleanup removes expired tokens from the blacklist
 func (tb *TokenBlacklist) Cleanup() {
 	tb.mutex.Lock()
 	defer tb.mutex.Unlock()
@@ -290,26 +356,25 @@ func (tb *TokenBlacklist) Cleanup() {
 	}
 }
 
+// TokenCache caches tokens
 type TokenCache struct {
 	cache *Cache
 }
 
-type TokenInfo struct {
-	Token     string
-	ExpiresAt time.Time
-}
-
+// NewTokenCache creates a new TokenCache
 func NewTokenCache() *TokenCache {
 	return &TokenCache{
 		cache: NewCache(),
 	}
 }
 
+// Set sets a token in the cache
 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
 func (tc *TokenCache) Get(token string) (map[string]interface{}, bool) {
 	token = "t-" + token
 	value, found := tc.cache.Get(token)
@@ -320,37 +385,28 @@ func (tc *TokenCache) Get(token string) (map[string]interface{}, bool) {
 	return claims, ok
 }
 
+// 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
 func (tc *TokenCache) Cleanup() {
 	tc.cache.Cleanup()
 }
 
-func (t *TraefikOidc) exchangeCodeForToken(code string) (map[string]interface{}, error) {
-	data := url.Values{}
-	data.Set("grant_type", "authorization_code")
-	data.Set("client_id", t.clientID)
-	data.Set("client_secret", t.clientSecret)
-	data.Set("code", code)
-	data.Set("redirect_uri", t.redirectURL)
-
-	resp, err := t.httpClient.PostForm(t.tokenURL, data)
+// exchangeCodeForToken exchanges the authorization code for tokens
+func (t *TraefikOidc) exchangeCodeForToken(code string) (*TokenResponse, error) {
+	ctx := context.Background()
+	tokenResponse, err := t.exchangeTokens(ctx, "authorization_code", code, t.redirectURL)
 	if err != nil {
-		return nil, fmt.Errorf("failed to exchange token: %v", err)
+		return nil, fmt.Errorf("failed to exchange code for token: %w", err)
 	}
-	defer resp.Body.Close()
-
-	var result map[string]interface{}
-	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
-		return nil, fmt.Errorf("failed to decode token response: %v", err)
-	}
-
-	return result, nil
+	return tokenResponse, nil
 }
 
+// createStringMap creates a map from a slice of strings
 func createStringMap(keys []string) map[string]struct{} {
 	result := make(map[string]struct{})
 	for _, key := range keys {