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

package oidc

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"net/http"
	"strings"
	"time"

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

// TokenHandler implements POST /token.
type TokenHandler struct {
	ClientConfig  map[string]*domain.Client
	Sessions      *SessionStore
	Users         domain.UserRepository
	SigningKey     *rsa.PrivateKey
	Issuer        string
	TokenLifetime time.Duration
	Emitter       telemetry.Emitter
}

// tokenResponse is the JSON body returned on a successful token exchange.
type tokenResponse struct {
	AccessToken string `json:"access_token"`
	TokenType   string `json:"token_type"`
	ExpiresIn   int    `json:"expires_in"`
	IDToken     string `json:"id_token"`
}

// ServeHTTP handles POST /token.
func (h *TokenHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	if err := r.ParseForm(); err != nil {
		http.Error(w, "invalid form body", http.StatusBadRequest)
		return
	}

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

	// 1. Validate grant_type.
	if grantType != "authorization_code" {
		profileerrors.FeatureNotSupported(
			"only grant_type=authorization_code is supported",
			"grant_type="+grantType,
		).Write(w, http.StatusBadRequest)
		return
	}

	// 2. Validate client exists (basic check; secret auth delegated to future work).
	if _, ok := h.ClientConfig[clientID]; !ok {
		profileerrors.InvalidProfileUsage("unknown client_id", "client_id").
			Write(w, http.StatusBadRequest)
		return
	}

	// 3. Look up PKCE session.
	sess, ok := h.Sessions.Get(code)
	if !ok {
		profileerrors.InvalidProfileUsage(
			"authorization code not found or expired",
			"code",
		).Write(w, http.StatusBadRequest)
		return
	}

	// Verify client_id matches the session.
	if sess.ClientID != clientID {
		profileerrors.InvalidProfileUsage(
			"client_id does not match the authorization code",
			"client_id",
		).Write(w, http.StatusBadRequest)
		return
	}

	// 4. Verify PKCE code_verifier.
	if !verifyPKCE(codeVerifier, sess.PKCEChallenge) {
		profileerrors.InvalidProfileUsage(
			"code_verifier does not match code_challenge",
			"code_verifier",
		).Write(w, http.StatusBadRequest)
		return
	}

	// 5. Look up user.
	user, err := h.Users.LookupUser(ctx, sess.Username)
	if err != nil {
		http.Error(w, "user not found", http.StatusInternalServerError)
		return
	}

	// 6. Build JWT claims.
	now := time.Now()
	exp := now.Add(h.TokenLifetime)

	claims := map[string]interface{}{
		"iss": h.Issuer,
		"sub": user.ID,
		"aud": clientID,
		"exp": exp.Unix(),
		"iat": now.Unix(),
	}
	if sess.Nonce != "" {
		claims["nonce"] = sess.Nonce
	}

	scopeSet := make(map[string]bool)
	for _, s := range sess.Scopes {
		scopeSet[s] = true
	}

	if scopeSet["profile"] {
		claims["preferred_username"] = user.Username
	}
	if scopeSet["email"] {
		claims["email"] = user.Email
	}
	if scopeSet["groups"] {
		claims["groups"] = user.Groups
	}

	// 7. Sign JWT with RSA-SHA256.
	kid := "key-1" // static kid for v0.1
	jwtToken, err := buildJWT(claims, kid, h.SigningKey)
	if err != nil {
		http.Error(w, "failed to build JWT", http.StatusInternalServerError)
		return
	}

	// 8. Delete used PKCE session (prevent replay).
	h.Sessions.Delete(code)

	// 9. Build response.
	resp := tokenResponse{
		AccessToken: jwtToken,
		TokenType:   "Bearer",
		ExpiresIn:   int(h.TokenLifetime.Seconds()),
		IDToken:     jwtToken,
	}

	// 10. Emit token_issued telemetry.
	h.Emitter.Emit(ctx, telemetry.Event{
		Timestamp: time.Now(),
		EventType: telemetry.EventTokenIssued,
		ClientID:  clientID,
		Endpoint:  "/token",
		Result:    "success",
		Scopes:    sess.Scopes,
		GrantType: grantType,
	})

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	_ = json.NewEncoder(w).Encode(resp)
}

// ---------------------------------------------------------------------------
// PKCE verification
// ---------------------------------------------------------------------------

// verifyPKCE checks BASE64URL(SHA256(verifier)) == challenge (S256 method).
func verifyPKCE(verifier, challenge string) bool {
	h := sha256.New()
	h.Write([]byte(verifier))
	computed := base64.RawURLEncoding.EncodeToString(h.Sum(nil))
	return computed == challenge
}

// ---------------------------------------------------------------------------
// JWT construction (stdlib only — no external JWT library)
// ---------------------------------------------------------------------------

type jwtHeader struct {
	Alg string `json:"alg"`
	Typ string `json:"typ"`
	Kid string `json:"kid"`
}

// buildJWT constructs and signs a JWT using RSA-SHA256 with the standard library.
// Format: base64url(header) + "." + base64url(payload) + "." + base64url(signature)
func buildJWT(claims map[string]interface{}, kid string, key *rsa.PrivateKey) (string, error) {
	// Header.
	hdr := jwtHeader{Alg: "RS256", Typ: "JWT", Kid: kid}
	hdrJSON, err := json.Marshal(hdr)
	if err != nil {
		return "", err
	}
	hdrB64 := base64.RawURLEncoding.EncodeToString(hdrJSON)

	// Payload.
	payloadJSON, err := json.Marshal(claims)
	if err != nil {
		return "", err
	}
	payloadB64 := base64.RawURLEncoding.EncodeToString(payloadJSON)

	// Signing input.
	signingInput := hdrB64 + "." + payloadB64

	// Digest.
	digest := sha256.Sum256([]byte(signingInput))

	// Sign with PKCS1v15 / SHA256.
	sig, err := rsa.SignPKCS1v15(rand.Reader, key, crypto.SHA256, digest[:])
	if err != nil {
		return "", err
	}
	sigB64 := base64.RawURLEncoding.EncodeToString(sig)

	return strings.Join([]string{hdrB64, payloadB64, sigB64}, "."), nil
}