auth0mock/handlers/oauth.go

package handlers

import (
	"embed"
	"encoding/json"
	"fmt"
	"html/template"
	"log/slog"
	"net/http"
	"net/url"

	"gitlab.com/unboundsoftware/auth0mock/auth"
	"gitlab.com/unboundsoftware/auth0mock/store"
)

//go:embed templates/login.html
var templateFS embed.FS

// OAuthHandler handles OAuth/OIDC endpoints
type OAuthHandler struct {
	jwtService    *auth.JWTService
	sessionStore  *store.SessionStore
	loginTemplate *template.Template
	logger        *slog.Logger
}

// NewOAuthHandler creates a new OAuth handler
func NewOAuthHandler(jwtService *auth.JWTService, sessionStore *store.SessionStore, logger *slog.Logger) (*OAuthHandler, error) {
	tmpl, err := template.ParseFS(templateFS, "templates/login.html")
	if err != nil {
		return nil, fmt.Errorf("parse login template: %w", err)
	}

	return &OAuthHandler{
		jwtService:    jwtService,
		sessionStore:  sessionStore,
		loginTemplate: tmpl,
		logger:        logger,
	}, nil
}

// TokenRequest represents the token endpoint request body
type TokenRequest struct {
	GrantType    string `json:"grant_type"`
	ClientID     string `json:"client_id"`
	ClientSecret string `json:"client_secret"`
	Code         string `json:"code"`
	CodeVerifier string `json:"code_verifier"`
	RedirectURI  string `json:"redirect_uri"`
}

// TokenResponse represents the token endpoint response
type TokenResponse struct {
	AccessToken string `json:"access_token"`
	IDToken     string `json:"id_token"`
	Scope       string `json:"scope"`
	ExpiresIn   int    `json:"expires_in"`
	TokenType   string `json:"token_type"`
}

// Token handles the POST /oauth/token endpoint
func (h *OAuthHandler) Token(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		http.Error(w, "Invalid request body", http.StatusBadRequest)
		return
	}

	grantType := r.FormValue("grant_type")
	clientID := r.FormValue("client_id")
	code := r.FormValue("code")
	codeVerifier := r.FormValue("code_verifier")

	w.Header().Set("Content-Type", "application/json")

	switch grantType {
	case "client_credentials":
		h.handleClientCredentials(w, clientID)
	case "authorization_code", "":
		if code != "" {
			h.handleAuthorizationCode(w, code, codeVerifier, clientID)
		} else {
			w.WriteHeader(http.StatusUnauthorized)
			json.NewEncoder(w).Encode(map[string]string{"error": "Missing code"})
		}
	default:
		w.WriteHeader(http.StatusBadRequest)
		json.NewEncoder(w).Encode(map[string]string{"error": "Unsupported grant type"})
	}
}

func (h *OAuthHandler) handleClientCredentials(w http.ResponseWriter, clientID string) {
	if clientID == "" {
		w.WriteHeader(http.StatusUnauthorized)
		json.NewEncoder(w).Encode(map[string]string{"error": "Missing client_id"})
		return
	}

	adminClaim := map[string]interface{}{
		h.jwtService.AdminClaim(): true,
	}

	accessToken, err := h.jwtService.SignAccessToken(
		"auth0|management",
		clientID,
		"management@example.org",
		[]map[string]interface{}{adminClaim},
	)
	if err != nil {
		h.logger.Error("failed to sign access token", "error", err)
		http.Error(w, "Failed to generate token", http.StatusInternalServerError)
		return
	}

	idToken, err := h.jwtService.SignIDToken(
		"auth0|management",
		clientID,
		"",
		"management@example.org",
		"Management API",
		"Management",
		"API",
		"",
		[]map[string]interface{}{adminClaim},
	)
	if err != nil {
		h.logger.Error("failed to sign ID token", "error", err)
		http.Error(w, "Failed to generate token", http.StatusInternalServerError)
		return
	}

	h.logger.Info("signed token for management API")

	json.NewEncoder(w).Encode(TokenResponse{
		AccessToken: accessToken,
		IDToken:     idToken,
		Scope:       "openid%20profile%20email",
		ExpiresIn:   7200,
		TokenType:   "Bearer",
	})
}

func (h *OAuthHandler) handleAuthorizationCode(w http.ResponseWriter, code, codeVerifier, clientID string) {
	session, ok := h.sessionStore.Get(code)
	if !ok {
		w.WriteHeader(http.StatusUnauthorized)
		json.NewEncoder(w).Encode(map[string]string{"error": "Invalid code"})
		return
	}

	// Verify PKCE if code_verifier is provided
	if codeVerifier != "" && session.CodeChallenge != "" {
		// Determine method - default to S256 if challenge looks like a hash
		method := auth.PKCEMethodS256
		if len(session.CodeChallenge) < 43 {
			method = auth.PKCEMethodPlain
		}

		if !auth.VerifyPKCE(codeVerifier, session.CodeChallenge, method) {
			w.WriteHeader(http.StatusUnauthorized)
			json.NewEncoder(w).Encode(map[string]string{"error": "Invalid code_verifier"})
			return
		}
	}

	accessToken, err := h.jwtService.SignAccessToken(
		"auth0|"+session.Email,
		session.ClientID,
		session.Email,
		session.CustomClaims,
	)
	if err != nil {
		h.logger.Error("failed to sign access token", "error", err)
		http.Error(w, "Failed to generate token", http.StatusInternalServerError)
		return
	}

	idToken, err := h.jwtService.SignIDToken(
		"auth0|"+session.Email,
		session.ClientID,
		session.Nonce,
		session.Email,
		"Example Person",
		"Example",
		"Person",
		"https://cdn.playbuzz.com/cdn/5458360f-32ea-460e-a707-1a2d26760558/70bda687-cb84-4756-8a44-8cf735ed87b3.jpg",
		session.CustomClaims,
	)
	if err != nil {
		h.logger.Error("failed to sign ID token", "error", err)
		http.Error(w, "Failed to generate token", http.StatusInternalServerError)
		return
	}

	h.logger.Info("signed token", "email", session.Email)

	// Clean up session after successful token exchange
	h.sessionStore.Delete(code)

	json.NewEncoder(w).Encode(TokenResponse{
		AccessToken: accessToken,
		IDToken:     idToken,
		Scope:       "openid%20profile%20email",
		ExpiresIn:   7200,
		TokenType:   "Bearer",
	})
}

// Code handles the POST /code endpoint (form submission from login page)
func (h *OAuthHandler) Code(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		http.Error(w, "Invalid request body", http.StatusBadRequest)
		return
	}

	email := r.FormValue("email")
	password := r.FormValue("password")
	codeChallenge := r.FormValue("codeChallenge")
	redirect := r.FormValue("redirect")
	state := r.FormValue("state")
	nonce := r.FormValue("nonce")
	clientID := r.FormValue("clientId")
	admin := r.FormValue("admin") == "true"

	if email == "" || password == "" || codeChallenge == "" {
		h.logger.Debug("invalid code request", "email", email, "hasPassword", password != "", "hasChallenge", codeChallenge != "")
		http.Error(w, "Email, password, or code challenge is missing", http.StatusBadRequest)
		return
	}

	adminClaim := map[string]interface{}{
		h.jwtService.AdminClaim(): admin,
	}

	session := &store.Session{
		Email:         email,
		Password:      password,
		State:         state,
		Nonce:         nonce,
		ClientID:      clientID,
		CodeChallenge: codeChallenge,
		CustomClaims:  []map[string]interface{}{adminClaim},
	}

	h.sessionStore.Create(codeChallenge, session)

	redirectURL := fmt.Sprintf("%s?code=%s&state=%s", redirect, codeChallenge, url.QueryEscape(state))
	http.Redirect(w, r, redirectURL, http.StatusFound)
}

// Authorize handles the GET /authorize endpoint
func (h *OAuthHandler) Authorize(w http.ResponseWriter, r *http.Request) {
	redirect := r.URL.Query().Get("redirect_uri")
	state := r.URL.Query().Get("state")
	nonce := r.URL.Query().Get("nonce")
	clientID := r.URL.Query().Get("client_id")
	codeChallenge := r.URL.Query().Get("code_challenge")
	prompt := r.URL.Query().Get("prompt")
	responseMode := r.URL.Query().Get("response_mode")

	// Try to get existing session from cookie
	cookie, err := r.Cookie("auth0")
	var existingCode string
	if err == nil {
		existingCode = cookie.Value
	}

	// Handle response_mode=query with existing session
	if responseMode == "query" && existingCode != "" {
		if h.sessionStore.Update(existingCode, codeChallenge, func(s *store.Session) {
			s.Nonce = nonce
			s.State = state
			s.CodeChallenge = codeChallenge
		}) {
			redirectURL := fmt.Sprintf("%s?code=%s&state=%s", redirect, codeChallenge, state)
			http.Redirect(w, r, redirectURL, http.StatusFound)
			return
		}
	}

	// Handle prompt=none with response_mode=web_message (silent auth)
	if prompt == "none" && responseMode == "web_message" && existingCode != "" {
		session, ok := h.sessionStore.Get(existingCode)
		if ok {
			h.sessionStore.Update(existingCode, existingCode, func(s *store.Session) {
				s.Nonce = nonce
				s.State = state
				s.CodeChallenge = codeChallenge
			})

			// Send postMessage response
			w.Header().Set("Content-Type", "text/html")
			fmt.Fprintf(w, `<!DOCTYPE html>
<html>
  <body>
  <script type="text/javascript">
(() => {
    const msg = {
      type: 'authorization_response',
      response: {
        code: '%s',
        state: '%s'
      }
    }
    parent.postMessage(msg, "*")
})()
  </script>
  </body>
</html>`, existingCode, session.State)
			return
		}
	}

	// Set cookie for session tracking
	http.SetCookie(w, &http.Cookie{
		Name:     "auth0",
		Value:    codeChallenge,
		Path:     "/",
		SameSite: http.SameSiteNoneMode,
		Secure:   true,
		HttpOnly: true,
	})

	// Render login form
	w.Header().Set("Content-Type", "text/html")
	data := map[string]string{
		"Redirect":      redirect,
		"State":         state,
		"Nonce":         nonce,
		"ClientID":      clientID,
		"CodeChallenge": codeChallenge,
	}

	if err := h.loginTemplate.Execute(w, data); err != nil {
		h.logger.Error("failed to render login template", "error", err)
		http.Error(w, "Internal server error", http.StatusInternalServerError)
	}
}