aitrade/pkg/app/web/auth.go

package web

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"log/slog"
	"net/http"
	"time"

	"github.com/coreos/go-oidc/v3/oidc"
	"github.com/pheinrich/aitrade/pkg/config"
	"golang.org/x/oauth2"
)

type AuthMiddleware struct {
	enabled      bool
	provider     *oidc.Provider
	oauth2Config oauth2.Config
	verifier     *oidc.IDTokenVerifier
	logger       *slog.Logger
}

func NewAuthMiddleware(cfg *config.OIDCConfig, logger *slog.Logger) (*AuthMiddleware, error) {
	if !cfg.Enabled {
		logger.Info("OIDC authentication disabled")
		return &AuthMiddleware{
			enabled: false,
			logger:  logger,
		}, nil
	}

	ctx := context.Background()
	provider, err := oidc.NewProvider(ctx, cfg.Issuer)
	if err != nil {
		return nil, fmt.Errorf("failed to create OIDC provider: %w", err)
	}

	oauth2Config := oauth2.Config{
		ClientID:     cfg.ClientID,
		ClientSecret: cfg.ClientSecret,
		RedirectURL:  cfg.RedirectURL,
		Endpoint:     provider.Endpoint(),
		Scopes:       cfg.Scopes,
	}

	verifier := provider.Verifier(&oidc.Config{ClientID: cfg.ClientID})

	logger.Info("OIDC authentication enabled",
		slog.String("issuer", cfg.Issuer),
		slog.String("client_id", cfg.ClientID),
	)

	return &AuthMiddleware{
		enabled:      true,
		provider:     provider,
		oauth2Config: oauth2Config,
		verifier:     verifier,
		logger:       logger,
	}, nil
}

func (a *AuthMiddleware) Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// If OIDC is disabled, allow all requests
		if !a.enabled {
			next.ServeHTTP(w, r)
			return
		}

		// Check for session cookie
		cookie, err := r.Cookie("aitrade_session")
		if err != nil {
			// No session, redirect to login
			a.redirectToLogin(w, r)
			return
		}

		// Verify the ID token
		ctx := r.Context()
		_, err = a.verifier.Verify(ctx, cookie.Value)
		if err != nil {
			a.logger.Warn("invalid session token", slog.Any("error", err))
			a.redirectToLogin(w, r)
			return
		}

		// Token is valid, proceed
		next.ServeHTTP(w, r)
	})
}

func (a *AuthMiddleware) redirectToLogin(w http.ResponseWriter, r *http.Request) {
	state, err := generateRandomState()
	if err != nil {
		http.Error(w, "Failed to generate state", http.StatusInternalServerError)
		return
	}

	// Store state in cookie for verification
	http.SetCookie(w, &http.Cookie{
		Name:     "oidc_state",
		Value:    state,
		Path:     "/",
		MaxAge:   300, // 5 minutes
		HttpOnly: true,
		Secure:   false, // Set to true in production with HTTPS
		SameSite: http.SameSiteLaxMode,
	})

	authURL := a.oauth2Config.AuthCodeURL(state)
	http.Redirect(w, r, authURL, http.StatusFound)
}

func (a *AuthMiddleware) HandleCallback(w http.ResponseWriter, r *http.Request) {
	if !a.enabled {
		http.Error(w, "OIDC not enabled", http.StatusBadRequest)
		return
	}

	// Verify state
	stateCookie, err := r.Cookie("oidc_state")
	if err != nil {
		http.Error(w, "State cookie not found", http.StatusBadRequest)
		return
	}

	if r.URL.Query().Get("state") != stateCookie.Value {
		http.Error(w, "State mismatch", http.StatusBadRequest)
		return
	}

	// Exchange code for token
	ctx := r.Context()
	oauth2Token, err := a.oauth2Config.Exchange(ctx, r.URL.Query().Get("code"))
	if err != nil {
		a.logger.Error("failed to exchange code", slog.Any("error", err))
		http.Error(w, "Failed to exchange code", http.StatusInternalServerError)
		return
	}

	// Extract ID token
	rawIDToken, ok := oauth2Token.Extra("id_token").(string)
	if !ok {
		http.Error(w, "No id_token in response", http.StatusInternalServerError)
		return
	}

	// Verify ID token
	idToken, err := a.verifier.Verify(ctx, rawIDToken)
	if err != nil {
		a.logger.Error("failed to verify ID token", slog.Any("error", err))
		http.Error(w, "Failed to verify ID token", http.StatusInternalServerError)
		return
	}

	// Store token in session cookie
	http.SetCookie(w, &http.Cookie{
		Name:     "aitrade_session",
		Value:    rawIDToken,
		Path:     "/",
		MaxAge:   int(time.Until(idToken.Expiry).Seconds()),
		HttpOnly: true,
		Secure:   false, // Set to true in production with HTTPS
		SameSite: http.SameSiteLaxMode,
	})

	// Clear state cookie
	http.SetCookie(w, &http.Cookie{
		Name:   "oidc_state",
		Value:  "",
		Path:   "/",
		MaxAge: -1,
	})

	a.logger.Info("user authenticated successfully")

	// Redirect to home
	http.Redirect(w, r, "/", http.StatusFound)
}

func generateRandomState() (string, error) {
	b := make([]byte, 32)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return base64.URLEncoding.EncodeToString(b), nil
}