key-cape/src/internal/server/oidc/authorize.go

package oidc

import (
	"context"
	"html/template"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"keycape/internal/domain"
	profileerrors "keycape/internal/errors"
	"keycape/internal/server/telemetry"
)

// PendingState holds the authorization request parameters while the user is
// being authenticated by the upstream provider (e.g. Authelia). It is keyed
// by the opaque state value that is round-tripped through the upstream.
type PendingState struct {
	ClientID            string
	RedirectURI         string
	PKCEChallenge       string
	PKCEChallengeMethod string
	State               string
	Nonce               string
	Scopes              []string
	ExpiresAt           time.Time
	AuthenticatedUser   string
}

// pendingStateStore is a thread-safe map of state → PendingState.
type pendingStateStore struct {
	mu    sync.Mutex
	store map[string]*PendingState
}

func newPendingStateStore() *pendingStateStore {
	return &pendingStateStore{store: make(map[string]*PendingState)}
}

func (p *pendingStateStore) Store(state string, ps *PendingState) {
	p.mu.Lock()
	p.store[state] = ps
	p.mu.Unlock()
}

func (p *pendingStateStore) Load(state string) (*PendingState, bool) {
	p.mu.Lock()
	ps, ok := p.store[state]
	p.mu.Unlock()
	return ps, ok
}

func (p *pendingStateStore) Delete(state string) {
	p.mu.Lock()
	delete(p.store, state)
	p.mu.Unlock()
}

// AuthorizeHandler implements GET /authorize and GET /authorize/callback.
type AuthorizeHandler struct {
	ClientConfig map[string]*domain.Client
	Auth         domain.AuthProvider
	MFA          domain.MFAProvider
	Sessions     *SessionStore
	Emitter      telemetry.Emitter

	pending *pendingStateStore
	once    sync.Once
}

// PendingStates returns the underlying pending-state store so tests can seed it.
func (h *AuthorizeHandler) PendingStates() *pendingStateStore {
	h.init()
	return h.pending
}

func (h *AuthorizeHandler) init() {
	h.once.Do(func() {
		if h.pending == nil {
			h.pending = newPendingStateStore()
		}
	})
}

// ServeHTTP dispatches to the authorize or callback handler based on path.
func (h *AuthorizeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h.init()
	if strings.HasSuffix(r.URL.Path, "/callback") {
		h.ServeHTTPCallback(w, r)
		return
	}
	h.serveAuthorize(w, r)
}

// serveAuthorize handles the initial GET /authorize request.
func (h *AuthorizeHandler) serveAuthorize(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	q := r.URL.Query()

	clientID := q.Get("client_id")
	redirectURI := q.Get("redirect_uri")
	responseType := q.Get("response_type")
	scope := q.Get("scope")
	state := q.Get("state")
	nonce := q.Get("nonce")
	codeChallenge := q.Get("code_challenge")
	codeChallengeMethod := q.Get("code_challenge_method")

	// Emit auth_start telemetry immediately.
	h.Emitter.Emit(ctx, telemetry.Event{
		Timestamp: time.Now(),
		EventType: telemetry.EventAuthStart,
		ClientID:  clientID,
		Endpoint:  "/authorize",
		Result:    "pending",
	})

	// 1. Validate client_id.
	client, ok := h.ClientConfig[clientID]
	if !ok {
		profileerrors.InvalidProfileUsage("unknown client_id", "client_id").
			Write(w, http.StatusBadRequest)
		return
	}

	// 2. Validate redirect_uri — check for wildcards first, then exact match.
	for _, registered := range client.RedirectURIs {
		if strings.ContainsAny(registered, "*?") {
			profileerrors.RejectedForSafety(
				"wildcard redirect URIs are not permitted",
				"redirect_uri",
			).Write(w, http.StatusBadRequest)
			return
		}
	}
	if !uriRegistered(client.RedirectURIs, redirectURI) {
		profileerrors.InvalidProfileUsage(
			"redirect_uri does not match any registered URI",
			"redirect_uri",
		).Write(w, http.StatusBadRequest)
		return
	}

	// 3. Validate response_type.
	if responseType != "code" {
		profileerrors.FeatureNotSupported(
			"only response_type=code is supported",
			"response_type="+responseType,
		).Write(w, http.StatusBadRequest)
		return
	}

	// 4. Validate scope contains openid.
	if !scopeContains(scope, "openid") {
		profileerrors.InvalidProfileUsage(
			"scope must include openid",
			"scope",
		).Write(w, http.StatusBadRequest)
		return
	}

	// 5. Validate code_challenge is present.
	if codeChallenge == "" {
		profileerrors.InvalidProfileUsage(
			"code_challenge is required (PKCE S256)",
			"code_challenge",
		).Write(w, http.StatusBadRequest)
		return
	}

	// 6. Validate code_challenge_method.
	if codeChallengeMethod == "plain" {
		profileerrors.RejectedForSafety(
			"code_challenge_method=plain is rejected for security; use S256",
			"code_challenge_method",
		).Write(w, http.StatusBadRequest)
		return
	}
	if codeChallengeMethod != "S256" {
		profileerrors.InvalidProfileUsage(
			"code_challenge_method must be S256",
			"code_challenge_method",
		).Write(w, http.StatusBadRequest)
		return
	}

	// Store pending state so the callback can reconstruct the session.
	h.pending.Store(state, &PendingState{
		ClientID:            clientID,
		RedirectURI:         redirectURI,
		PKCEChallenge:       codeChallenge,
		PKCEChallengeMethod: codeChallengeMethod,
		State:               state,
		Nonce:               nonce,
		Scopes:              strings.Fields(scope),
		ExpiresAt:           time.Now().Add(10 * time.Minute),
	})

	// Delegate to Auth provider.
	authURL, err := h.Auth.AuthorizeURL(ctx, domain.AuthRequest{
		ClientID:            clientID,
		RedirectURI:         redirectURI,
		State:               state,
		Scopes:              strings.Fields(scope),
		PKCEChallenge:       codeChallenge,
		PKCEChallengeMethod: codeChallengeMethod,
	})
	if err != nil {
		http.Error(w, "upstream auth provider error", http.StatusBadGateway)
		return
	}

	http.Redirect(w, r, authURL, http.StatusFound)
}

// ServeHTTPCallback handles GET /authorize/callback.
func (h *AuthorizeHandler) ServeHTTPCallback(w http.ResponseWriter, r *http.Request) {
	h.init()
	ctx := r.Context()

	if r.Method == http.MethodPost {
		h.serveMFASubmission(w, r)
		return
	}
	if r.Method != http.MethodGet {
		w.Header().Set("Allow", "GET, POST")
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	q := r.URL.Query()

	state := q.Get("state")
	code := q.Get("code")
	mfaToken := q.Get("mfa_token")

	// Recover pending state keyed by state param.
	ps, ok := h.pending.Load(state)
	if !ok {
		http.Error(w, "unknown or expired state", http.StatusBadRequest)
		return
	}
	if time.Now().After(ps.ExpiresAt) {
		h.pending.Delete(state)
		http.Error(w, "authorization request expired", http.StatusBadRequest)
		return
	}

	// Handle upstream callback.
	result, err := h.Auth.HandleCallback(ctx, domain.CallbackParams{
		Code:  code,
		State: state,
	})
	if err != nil {
		h.Emitter.Emit(ctx, telemetry.Event{
			Timestamp: time.Now(),
			EventType: telemetry.EventAuthFailure,
			ClientID:  ps.ClientID,
			Endpoint:  "/authorize/callback",
			Result:    "failure",
			ErrorType: "auth_failed",
		})
		http.Error(w, "authentication failed", http.StatusUnauthorized)
		return
	}
	if result == nil || result.Username == "" {
		h.pending.Delete(state)
		h.Emitter.Emit(ctx, telemetry.Event{
			Timestamp: time.Now(),
			EventType: telemetry.EventAuthFailure,
			ClientID:  ps.ClientID,
			Endpoint:  "/authorize/callback",
			Result:    "failure",
			ErrorType: "auth_failed",
		})
		http.Error(w, "authentication failed", http.StatusUnauthorized)
		return
	}

	// Check MFA requirement.
	mfaRequired, err := h.MFA.CheckMFARequired(ctx, result.Username)
	if err != nil {
		h.Emitter.Emit(ctx, telemetry.Event{
			Timestamp: time.Now(),
			EventType: telemetry.EventAuthFailure,
			ClientID:  ps.ClientID,
			Endpoint:  "/authorize/callback",
			Result:    "failure",
			ErrorType: "mfa_check_error",
		})
		http.Error(w, "mfa check error", http.StatusInternalServerError)
		return
	}
	if mfaRequired {
		if mfaToken == "" {
			ps.AuthenticatedUser = result.Username
			h.pending.Store(state, ps)
			h.renderMFAChallenge(w, ps, "")
			return
		}
		if err := h.MFA.ValidateMFAToken(ctx, result.Username, mfaToken); err != nil {
			h.pending.Delete(state)
			h.emitMFAFailure(ctx, ps.ClientID)
			http.Error(w, "MFA validation failed", http.StatusUnauthorized)
			return
		}
	}

	h.pending.Delete(state)
	h.completeAuthorization(w, r, ps, result.Username)
}

func (h *AuthorizeHandler) serveMFASubmission(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	if err := r.ParseForm(); err != nil {
		http.Error(w, "invalid form", http.StatusBadRequest)
		return
	}

	state := r.Form.Get("state")
	mfaToken := r.Form.Get("mfa_token")

	ps, ok := h.pending.Load(state)
	if !ok {
		http.Error(w, "unknown or expired state", http.StatusBadRequest)
		return
	}
	if time.Now().After(ps.ExpiresAt) {
		h.pending.Delete(state)
		http.Error(w, "authorization request expired", http.StatusBadRequest)
		return
	}
	if ps.AuthenticatedUser == "" {
		h.pending.Delete(state)
		http.Error(w, "mfa challenge not active", http.StatusBadRequest)
		return
	}
	if strings.TrimSpace(mfaToken) == "" {
		h.renderMFAChallenge(w, ps, "Enter the one-time code.")
		return
	}

	if err := h.MFA.ValidateMFAToken(ctx, ps.AuthenticatedUser, mfaToken); err != nil {
		h.pending.Delete(state)
		h.emitMFAFailure(ctx, ps.ClientID)
		http.Error(w, "MFA validation failed", http.StatusUnauthorized)
		return
	}

	h.pending.Delete(state)
	h.completeAuthorization(w, r, ps, ps.AuthenticatedUser)
}

func (h *AuthorizeHandler) completeAuthorization(w http.ResponseWriter, r *http.Request, ps *PendingState, username string) {
	// Generate authorization code and store PKCE session.
	sess := &PKCESession{
		ClientID:            ps.ClientID,
		RedirectURI:         ps.RedirectURI,
		PKCEChallenge:       ps.PKCEChallenge,
		PKCEChallengeMethod: ps.PKCEChallengeMethod,
		State:               ps.State,
		Nonce:               ps.Nonce,
		Username:            username,
		Scopes:              ps.Scopes,
		ExpiresAt:           time.Now().Add(10 * time.Minute),
	}
	authCode := h.Sessions.Create(sess)

	h.Emitter.Emit(r.Context(), telemetry.Event{
		Timestamp: time.Now(),
		EventType: telemetry.EventAuthSuccess,
		ClientID:  ps.ClientID,
		Endpoint:  "/authorize/callback",
		Result:    "success",
		Scopes:    ps.Scopes,
	})

	// Redirect to client with code and state.
	redirectTo, err := url.Parse(ps.RedirectURI)
	if err != nil {
		http.Error(w, "invalid redirect_uri", http.StatusInternalServerError)
		return
	}
	q := redirectTo.Query()
	q.Set("code", authCode)
	q.Set("state", ps.State)
	redirectTo.RawQuery = q.Encode()
	http.Redirect(w, r, redirectTo.String(), http.StatusFound)
}

func (h *AuthorizeHandler) emitMFAFailure(ctx context.Context, clientID string) {
	h.Emitter.Emit(ctx, telemetry.Event{
		Timestamp: time.Now(),
		EventType: telemetry.EventAuthFailure,
		ClientID:  clientID,
		Endpoint:  "/authorize/callback",
		Result:    "failure",
		ErrorType: "mfa_failed",
	})
}

func (h *AuthorizeHandler) renderMFAChallenge(w http.ResponseWriter, ps *PendingState, errorMessage string) {
	clientName := ps.ClientID
	if client, ok := h.ClientConfig[ps.ClientID]; ok && client.DisplayName != "" {
		clientName = client.DisplayName
	}
	status := http.StatusOK
	if errorMessage != "" {
		status = http.StatusBadRequest
	}
	w.Header().Set("Cache-Control", "no-store")
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.WriteHeader(status)
	_ = mfaChallengeTemplate.Execute(w, struct {
		State        string
		Username     string
		ClientName   string
		ErrorMessage string
	}{
		State:        ps.State,
		Username:     ps.AuthenticatedUser,
		ClientName:   clientName,
		ErrorMessage: errorMessage,
	})
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

var mfaChallengeTemplate = template.Must(template.New("mfa-challenge").Parse(`<!doctype html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width, initial-scale=1">
  <title>KeyCape MFA</title>
  <style>
    :root { color-scheme: light; font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; }
    body { margin: 0; min-height: 100vh; display: grid; place-items: center; background: #f6f7f9; color: #17202a; }
    main { width: min(420px, calc(100vw - 32px)); background: #fff; border: 1px solid #dfe4ea; border-radius: 8px; padding: 28px; box-shadow: 0 18px 45px rgba(23, 32, 42, .08); }
    h1 { margin: 0 0 6px; font-size: 22px; font-weight: 650; letter-spacing: 0; }
    p { margin: 0 0 20px; color: #52606d; line-height: 1.45; }
    label { display: block; margin: 0 0 8px; font-size: 13px; font-weight: 650; color: #344054; }
    input[type="text"] { width: 100%; box-sizing: border-box; height: 44px; border: 1px solid #c9d3df; border-radius: 6px; padding: 0 12px; font: inherit; background: #fff; }
    input[type="text"]:focus { outline: 2px solid #2f80ed; outline-offset: 2px; border-color: #2f80ed; }
    button { width: 100%; height: 44px; border: 0; border-radius: 6px; margin-top: 16px; background: #17324d; color: #fff; font: inherit; font-weight: 650; cursor: pointer; }
    button:focus { outline: 2px solid #2f80ed; outline-offset: 2px; }
    .meta { font-size: 13px; color: #667085; }
    .error { margin: 0 0 12px; color: #b42318; font-size: 13px; font-weight: 650; }
  </style>
</head>
<body>
  <main>
    <h1>Verify sign-in</h1>
    <p class="meta">{{.Username}} for {{.ClientName}}</p>
    {{if .ErrorMessage}}<p class="error">{{.ErrorMessage}}</p>{{end}}
    <form method="post" action="/authorize/callback" autocomplete="off">
      <input type="hidden" name="state" value="{{.State}}">
      <label for="mfa_token">One-time code</label>
      <input id="mfa_token" name="mfa_token" type="text" inputmode="numeric" autocomplete="one-time-code" required autofocus>
      <button type="submit">Verify</button>
    </form>
  </main>
</body>
</html>`))

func uriRegistered(registered []string, target string) bool {
	for _, u := range registered {
		if u == target {
			return true
		}
	}
	return false
}

func scopeContains(scope, want string) bool {
	for _, s := range strings.Fields(scope) {
		if s == want {
			return true
		}
	}
	return false
}