feat(custom-oauth): preserve non-standard IdP claims in identity_data

Author: cemalkilic

internal/api/provider/custom_oauth.go

@@ -5,11 +5,72 @@ import (
 	"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
@@ -68,10 +129,11 @@ func (p *CustomOAuthProvider) GetOAuthToken(ctx context.Context, code string, op
 
 // GetUserData fetches user data from the provider's userinfo endpoint
 func (p *CustomOAuthProvider) GetUserData(ctx context.Context, tok *oauth2.Token) (*UserProvidedData, error) {
-	var claims Claims
-	if err := makeRequest(ctx, tok, p.config, p.userinfoURL, &claims); err != nil {
+	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 {
@@ -198,6 +260,16 @@ func (p *CustomOIDCProvider) GetUserData(ctx context.Context, tok *oauth2.Token)
 			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)
@@ -208,10 +280,11 @@ func (p *CustomOIDCProvider) GetUserData(ctx context.Context, tok *oauth2.Token)
 
 	// No ID token, use userinfo endpoint
 	if p.userinfoEndpoint != "" {
-		var claims Claims
-		if err := makeRequest(ctx, tok, p.config, p.userinfoEndpoint, &claims); err != nil {
+		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 {

internal/api/provider/custom_oauth_test.go

@@ -170,6 +170,148 @@ func TestCustomOAuthProvider_GetUserDataWithAttributeMapping(t *testing.T) {
 	assert.True(t, userData.Emails[0].Verified) // Should be true from literal mapping
 }
 
+func TestCustomOAuthProvider_GetUserDataPreservesCustomClaims(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		json.NewEncoder(w).Encode(map[string]interface{}{
+			"sub":            "user-123",
+			"email":          "test@example.com",
+			"email_verified": true,
+			"name":           "Test User",
+			// Non-standard claims that previously got silently dropped.
+			"groups":    []string{"admins", "billing"},
+			"org_id":    "org_42",
+			"tenant_id": "tenant-abc",
+		})
+	}))
+	defer server.Close()
+
+	provider := NewCustomOAuthProvider(
+		"client-id",
+		"client-secret",
+		"https://example.com/authorize",
+		"https://example.com/token",
+		server.URL,
+		"https://myapp.com/callback",
+		[]string{"openid", "profile", "email"},
+		false,
+		nil,
+		nil,
+		nil,
+	)
+
+	token := &oauth2.Token{AccessToken: "test-access-token", TokenType: "Bearer"}
+	userData, err := provider.GetUserData(context.Background(), token)
+	require.NoError(t, err)
+	require.NotNil(t, userData)
+	require.NotNil(t, userData.Metadata)
+
+	// Standard fields still populated.
+	assert.Equal(t, "user-123", userData.Metadata.Subject)
+	assert.Equal(t, "test@example.com", userData.Metadata.Email)
+	assert.Equal(t, "Test User", userData.Metadata.Name)
+
+	// Non-standard claims preserved under CustomClaims.
+	require.NotNil(t, userData.Metadata.CustomClaims)
+	assert.Equal(t, "org_42", userData.Metadata.CustomClaims["org_id"])
+	assert.Equal(t, "tenant-abc", userData.Metadata.CustomClaims["tenant_id"])
+
+	groups, ok := userData.Metadata.CustomClaims["groups"].([]interface{})
+	require.True(t, ok, "groups should round-trip as []interface{}")
+	require.Len(t, groups, 2)
+	assert.Equal(t, "admins", groups[0])
+	assert.Equal(t, "billing", groups[1])
+
+	// Known fields must NOT also leak into CustomClaims.
+	_, hasEmail := userData.Metadata.CustomClaims["email"]
+	assert.False(t, hasEmail)
+	_, hasSub := userData.Metadata.CustomClaims["sub"]
+	assert.False(t, hasSub)
+}
+
+func TestCustomClaimsUnmarshalCapturesCustomClaims(t *testing.T) {
+	t.Run("standard claims fill typed fields and non-standard claims land in CustomClaims", func(t *testing.T) {
+		body := []byte(`{
+			"sub": "u-1",
+			"email": "a@b.com",
+			"email_verified": true,
+			"groups": ["x"],
+			"org_id": "o1"
+		}`)
+		var c customClaims
+		require.NoError(t, json.Unmarshal(body, &c))
+
+		assert.Equal(t, "u-1", c.Subject)
+		assert.Equal(t, "a@b.com", c.Email)
+		assert.True(t, c.EmailVerified)
+		require.NotNil(t, c.CustomClaims)
+		assert.Equal(t, "o1", c.CustomClaims["org_id"])
+		assert.Equal(t, []interface{}{"x"}, c.CustomClaims["groups"])
+
+		_, hasEmail := c.CustomClaims["email"]
+		assert.False(t, hasEmail, "standard claims must not also leak into CustomClaims")
+	})
+
+	t.Run("only standard claims means CustomClaims stays nil", func(t *testing.T) {
+		body := []byte(`{"sub":"u","email":"a@b.com"}`)
+		var c customClaims
+		require.NoError(t, json.Unmarshal(body, &c))
+		assert.Nil(t, c.CustomClaims)
+	})
+
+	t.Run("provider that literally returns custom_claims is preserved flat (not re-nested)", func(t *testing.T) {
+		body := []byte(`{"sub":"u-2","custom_claims":{"foo":"bar"}}`)
+		var c customClaims
+		require.NoError(t, json.Unmarshal(body, &c))
+		assert.Equal(t, "u-2", c.Subject)
+		require.NotNil(t, c.CustomClaims)
+		assert.Equal(t, "bar", c.CustomClaims["foo"])
+		_, nested := c.CustomClaims["custom_claims"]
+		assert.False(t, nested, "custom_claims must not be re-nested under itself")
+	})
+
+	t.Run("custom_claims object and other non-standard keys are merged at top level", func(t *testing.T) {
+		body := []byte(`{
+			"sub": "u-3",
+			"custom_claims": {"foo": "bar"},
+			"groups": ["admins"],
+			"org_id": "o1"
+		}`)
+		var c customClaims
+		require.NoError(t, json.Unmarshal(body, &c))
+		assert.Equal(t, "u-3", c.Subject)
+		require.NotNil(t, c.CustomClaims)
+		assert.Equal(t, "bar", c.CustomClaims["foo"])
+		assert.Equal(t, "o1", c.CustomClaims["org_id"])
+		assert.Equal(t, []interface{}{"admins"}, c.CustomClaims["groups"])
+		_, nested := c.CustomClaims["custom_claims"]
+		assert.False(t, nested, "custom_claims must not be re-nested under itself")
+	})
+
+	t.Run("entries inside custom_claims win over same-named top-level keys", func(t *testing.T) {
+		// IdP returns "groups" both inside custom_claims and at the top
+		// level. The typed decode places the inner value into CustomClaims
+		// first; the outer one must not silently overwrite it.
+		body := []byte(`{
+			"sub": "u-4",
+			"custom_claims": {"groups": ["from-inner"]},
+			"groups": ["from-outer"]
+		}`)
+		var c customClaims
+		require.NoError(t, json.Unmarshal(body, &c))
+		require.NotNil(t, c.CustomClaims)
+		assert.Equal(t, []interface{}{"from-inner"}, c.CustomClaims["groups"])
+	})
+
+	t.Run("plain Claims still drops non-standard claims (proves scoping)", func(t *testing.T) {
+		body := []byte(`{"sub":"u","groups":["x"]}`)
+		var c Claims
+		require.NoError(t, json.Unmarshal(body, &c))
+		assert.Equal(t, "u", c.Subject)
+		assert.Nil(t, c.CustomClaims, "non-custom providers must keep existing drop behaviour")
+	})
+}
+
 func TestApplyAttributeMapping(t *testing.T) {
 	tests := []struct {
 		name     string