custom_oauth.go

package provider

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"reflect"
	"strings"

	"github.com/coreos/go-oidc/v3/oidc"
	"golang.org/x/oauth2"
)

// standardClaimKeys is the set of JSON keys handled by the typed Claims
// struct (OIDC standard claims plus our typed extensions, and the
// custom_claims sink itself). Derived from the struct's json tags so it
// can't drift when Claims changes.
var standardClaimKeys = jsonKeysOf(reflect.TypeOf(Claims{}))

// jsonKeysOf returns the set of JSON keys declared by a struct type's tags.
func jsonKeysOf(t reflect.Type) map[string]bool {
	keys := map[string]bool{}
	for i := 0; i < t.NumField(); i++ {
		tag, _, _ := strings.Cut(t.Field(i).Tag.Get("json"), ",")
		if tag != "" && tag != "-" {
			keys[tag] = true
		}
	}
	return keys
}

// captureCustomClaims merges any top-level keys in raw that aren't part of
// the typed Claims struct into c.CustomClaims, so provider-specific claims
// (groups, roles, org_id, …) survive into identity_data downstream. Entries
// already present in c.CustomClaims (e.g. populated by the typed decode of
// a literal "custom_claims" object on the IdP response) win over any
// same-named top-level key.
func captureCustomClaims(raw map[string]interface{}, c *Claims) {
	for k, v := range raw {
		if standardClaimKeys[k] {
			continue
		}
		if _, exists := c.CustomClaims[k]; exists {
			continue
		}
		if c.CustomClaims == nil {
			c.CustomClaims = map[string]interface{}{}
		}
		c.CustomClaims[k] = v
	}
}

// customClaims is a Claims wrapper whose UnmarshalJSON also captures
// non-standard top-level keys under CustomClaims. Scoped to custom OAuth/OIDC
// providers — built-in providers (Google, Apple, …) decode into Claims
// directly and keep their existing behaviour.
type customClaims Claims

func (c *customClaims) UnmarshalJSON(data []byte) error {
	var standard Claims
	if err := json.Unmarshal(data, &standard); err != nil {
		return err
	}
	*c = customClaims(standard)

	var raw map[string]interface{}
	if err := json.Unmarshal(data, &raw); err == nil {
		captureCustomClaims(raw, (*Claims)(c))
	}
	return nil
}

// CustomOAuthProvider implements OAuthProvider for custom OAuth2 providers
type CustomOAuthProvider struct {
	config              *oauth2.Config
	userinfoURL         string
	pkceEnabled         bool
	acceptableClientIDs []string
	attributeMapping    map[string]interface{}
	authorizationParams map[string]interface{}
}

// NewCustomOAuthProvider creates a new custom OAuth provider
func NewCustomOAuthProvider(
	clientID, clientSecret, authorizationURL, tokenURL, userinfoURL, redirectURL string,
	scopes []string,
	pkceEnabled bool,
	acceptableClientIDs []string,
	attributeMapping, authorizationParams map[string]interface{},
) *CustomOAuthProvider {
	config := &oauth2.Config{
		ClientID:     clientID,
		ClientSecret: clientSecret,
		RedirectURL:  redirectURL,
		Scopes:       scopes,
		Endpoint: oauth2.Endpoint{
			AuthURL:  authorizationURL,
			TokenURL: tokenURL,
		},
	}

	return &CustomOAuthProvider{
		config:              config,
		userinfoURL:         userinfoURL,
		pkceEnabled:         pkceEnabled,
		acceptableClientIDs: acceptableClientIDs,
		attributeMapping:    attributeMapping,
		authorizationParams: authorizationParams,
	}
}

// AuthCodeURL returns the authorization URL for the OAuth flow
func (p *CustomOAuthProvider) AuthCodeURL(state string, opts ...oauth2.AuthCodeOption) string {
	// Add any additional authorization parameters (values are validated as strings at the API layer)
	for key, value := range p.authorizationParams {
		if s, ok := value.(string); ok {
			opts = append(opts, oauth2.SetAuthURLParam(key, s))
		}
	}

	return p.config.AuthCodeURL(state, opts...)
}

// GetOAuthToken exchanges the authorization code for an access token
func (p *CustomOAuthProvider) GetOAuthToken(ctx context.Context, code string, opts ...oauth2.AuthCodeOption) (*oauth2.Token, error) {
	return p.config.Exchange(ctx, code, opts...)
}

// GetUserData fetches user data from the provider's userinfo endpoint
func (p *CustomOAuthProvider) GetUserData(ctx context.Context, tok *oauth2.Token) (*UserProvidedData, error) {
	var decoded customClaims
	if err := makeRequest(ctx, tok, p.config, p.userinfoURL, &decoded); err != nil {
		return nil, err
	}
	claims := Claims(decoded)

	// Apply attribute mapping if configured
	if len(p.attributeMapping) > 0 {
		claims = applyAttributeMapping(claims, p.attributeMapping)
	}

	// Extract emails
	emails := []Email{}
	if claims.Email != "" {
		emails = append(emails, Email{
			Email:    claims.Email,
			Verified: claims.EmailVerified,
			Primary:  true,
		})
	}

	return &UserProvidedData{
		Emails:   emails,
		Metadata: &claims,
	}, nil
}

// RequiresPKCE returns whether this provider requires PKCE
func (p *CustomOAuthProvider) RequiresPKCE() bool {
	return p.pkceEnabled
}

// CustomOIDCProvider implements OAuthProvider for custom OIDC providers
type CustomOIDCProvider struct {
	config              *oauth2.Config
	oidcProvider        *oidc.Provider
	userinfoEndpoint    string
	pkceEnabled         bool
	acceptableClientIDs []string
	attributeMapping    map[string]interface{}
	authorizationParams map[string]interface{}
}

// NewCustomOIDCProvider creates a new custom OIDC provider
func NewCustomOIDCProvider(
	ctx context.Context,
	clientID, clientSecret, redirectURL string,
	scopes []string,
	issuer string,
	pkceEnabled bool,
	acceptableClientIDs []string,
	attributeMapping, authorizationParams map[string]interface{},
	cache *OIDCProviderCache,
) (*CustomOIDCProvider, error) {
	// Ensure 'openid' scope is always present for OIDC
	hasOpenID := false
	for _, scope := range scopes {
		if scope == "openid" {
			hasOpenID = true
			break
		}
	}
	if !hasOpenID {
		scopes = append([]string{"openid"}, scopes...)
	}

	// Create OIDC provider - uses cache to avoid redundant discovery fetches
	oidcProvider, err := cache.GetProvider(ctx, issuer)
	if err != nil {
		return nil, fmt.Errorf("failed to create OIDC provider: %w", err)
	}

	// Get endpoints from the OIDC provider
	endpoint := oidcProvider.Endpoint()
	userinfoEndpoint := oidcProvider.UserInfoEndpoint()

	config := &oauth2.Config{
		ClientID:     clientID,
		ClientSecret: clientSecret,
		RedirectURL:  redirectURL,
		Scopes:       scopes,
		Endpoint:     endpoint,
	}

	return &CustomOIDCProvider{
		config:              config,
		oidcProvider:        oidcProvider,
		userinfoEndpoint:    userinfoEndpoint,
		pkceEnabled:         pkceEnabled,
		acceptableClientIDs: acceptableClientIDs,
		attributeMapping:    attributeMapping,
		authorizationParams: authorizationParams,
	}, nil
}

// AuthCodeURL returns the authorization URL for the OIDC flow
func (p *CustomOIDCProvider) AuthCodeURL(state string, opts ...oauth2.AuthCodeOption) string {
	// Add any additional authorization parameters (values are validated as strings at the API layer)
	for key, value := range p.authorizationParams {
		if s, ok := value.(string); ok {
			opts = append(opts, oauth2.SetAuthURLParam(key, s))
		}
	}

	return p.config.AuthCodeURL(state, opts...)
}

// GetOAuthToken exchanges the authorization code for an access token
func (p *CustomOIDCProvider) GetOAuthToken(ctx context.Context, code string, opts ...oauth2.AuthCodeOption) (*oauth2.Token, error) {
	return p.config.Exchange(ctx, code, opts...)
}

// GetUserData fetches user data from the provider's userinfo endpoint or ID token
func (p *CustomOIDCProvider) GetUserData(ctx context.Context, tok *oauth2.Token) (*UserProvidedData, error) {
	// First, try to extract and verify claims from ID token if present
	if idToken, ok := tok.Extra("id_token").(string); ok && idToken != "" {
		// Skip client ID check in the library and validate manually to support multiple client IDs
		idTokenObj, userData, err := ParseIDToken(ctx, p.oidcProvider, &oidc.Config{
			SkipClientIDCheck: true, // We'll validate audience manually
		}, idToken, ParseIDTokenOptions{
			SkipAccessTokenCheck: true, // We don't need at_hash validation in callback flow
		})
		if err != nil {
			return nil, fmt.Errorf("failed to verify ID token: %w", err)
		}

		// Validate audience claim against acceptable client IDs
		if err := p.validateAudience(idTokenObj.Audience); err != nil {
			return nil, err
		}

		// ParseIDToken decodes into the typed Claims struct and drops anything
		// not mapped. Re-walk the raw claims to capture provider-specific
		// custom claims (groups, roles, …) for identity_data.
		if userData.Metadata != nil {
			var raw map[string]interface{}
			if err := idTokenObj.Claims(&raw); err == nil {
				captureCustomClaims(raw, userData.Metadata)
			}
		}

		// Apply attribute mapping to the metadata from ID token
		if len(p.attributeMapping) > 0 && userData.Metadata != nil {
			*userData.Metadata = applyAttributeMapping(*userData.Metadata, p.attributeMapping)
		}

		return userData, nil
	}

	// No ID token, use userinfo endpoint
	if p.userinfoEndpoint != "" {
		var decoded customClaims
		if err := makeRequest(ctx, tok, p.config, p.userinfoEndpoint, &decoded); err != nil {
			return nil, err
		}
		claims := Claims(decoded)

		// Apply attribute mapping
		if len(p.attributeMapping) > 0 {
			claims = applyAttributeMapping(claims, p.attributeMapping)
		}

		// Extract emails
		emails := []Email{}
		if claims.Email != "" {
			emails = append(emails, Email{
				Email:    claims.Email,
				Verified: claims.EmailVerified,
				Primary:  true,
			})
		}

		return &UserProvidedData{
			Emails:   emails,
			Metadata: &claims,
		}, nil
	}

	return nil, errors.New("no ID token or userinfo endpoint available")
}

// RequiresPKCE returns whether this provider requires PKCE
func (p *CustomOIDCProvider) RequiresPKCE() bool {
	return p.pkceEnabled
}

// Config returns the OAuth2 config for accessing endpoints
func (p *CustomOIDCProvider) Config() *oauth2.Config {
	return p.config
}

// validateAudience validates that the token's audience matches one of the acceptable client IDs
func (p *CustomOIDCProvider) validateAudience(audiences []string) error {
	// Build list of acceptable audiences: main client_id + acceptable_client_ids
	acceptableAudiences := append([]string{p.config.ClientID}, p.acceptableClientIDs...)

	// Check if any audience in the token matches any acceptable audience
	for _, tokenAud := range audiences {
		for _, acceptableAud := range acceptableAudiences {
			if tokenAud == acceptableAud {
				return nil // Valid audience found
			}
		}
	}

	// No valid audience found
	return fmt.Errorf("token audience %v does not match any acceptable client ID", audiences)
}

// applyAttributeMapping applies custom attribute mapping to claims
func applyAttributeMapping(claims Claims, mapping map[string]interface{}) Claims {
	// Create a map representation of claims for easier manipulation
	claimsMap := make(map[string]interface{})
	claimsBytes, _ := json.Marshal(claims)
	if err := json.Unmarshal(claimsBytes, &claimsMap); err != nil {
		// If unmarshaling fails, return original claims
		return claims
	}

	// Apply mappings
	for targetField, sourceFieldOrValue := range mapping {
		switch v := sourceFieldOrValue.(type) {
		case string:
			// If it's a string, treat it as a source field name
			if value, exists := claimsMap[v]; exists {
				claimsMap[targetField] = value
			}
		default:
			// Otherwise, use it as a literal value
			claimsMap[targetField] = v
		}
	}

	// Convert back to Claims struct
	var result Claims
	mappedBytes, _ := json.Marshal(claimsMap)
	if err := json.Unmarshal(mappedBytes, &result); err != nil {
		// If unmarshaling fails, return original claims
		return claims
	}

	return result
}