zitadel/internal/api/oidc/integration_test/userinfo_test.go

//go:build integration

package oidc_test

import (
	"fmt"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/zitadel/oidc/v3/pkg/client/rp"
	"github.com/zitadel/oidc/v3/pkg/oidc"
	"golang.org/x/oauth2"
	"google.golang.org/grpc/metadata"

	oidc_api "github.com/zitadel/zitadel/internal/api/oidc"
	"github.com/zitadel/zitadel/internal/domain"
	"github.com/zitadel/zitadel/internal/integration"
	"github.com/zitadel/zitadel/pkg/grpc/management"
	oidc_pb "github.com/zitadel/zitadel/pkg/grpc/oidc/v2"
)

func TestServer_UserInfo(t *testing.T) {
	const (
		roleFoo = "foo"
		roleBar = "bar"
	)

	clientID, projectID := createClient(t, Instance)
	addProjectRolesGrants(t, User.GetUserId(), projectID, roleFoo, roleBar)

	tests := []struct {
		name       string
		prepare    func(t *testing.T, clientID string, scope []string) *oidc.Tokens[*oidc.IDTokenClaims]
		scope      []string
		assertions []func(*testing.T, *oidc.UserInfo)
		wantErr    bool
	}{
		{
			name: "invalid token",
			prepare: func(*testing.T, string, []string) *oidc.Tokens[*oidc.IDTokenClaims] {
				return &oidc.Tokens[*oidc.IDTokenClaims]{
					Token: &oauth2.Token{
						AccessToken: "DEAFBEEFDEADBEEF",
						TokenType:   oidc.BearerToken,
					},
					IDTokenClaims: &oidc.IDTokenClaims{
						TokenClaims: oidc.TokenClaims{
							Subject: User.GetUserId(),
						},
					},
				}
			},
			scope: []string{oidc.ScopeProfile, oidc.ScopeOpenID, oidc.ScopeEmail, oidc.ScopeOfflineAccess},
			assertions: []func(*testing.T, *oidc.UserInfo){
				func(t *testing.T, ui *oidc.UserInfo) {
					assert.Nil(t, ui)
				},
			},
			wantErr: true,
		},
		{
			name:    "standard scopes",
			prepare: getTokens,
			scope:   []string{oidc.ScopeProfile, oidc.ScopeOpenID, oidc.ScopeEmail, oidc.ScopeOfflineAccess},
			assertions: []func(*testing.T, *oidc.UserInfo){
				assertUserinfo,
				func(t *testing.T, ui *oidc.UserInfo) {
					assertNoReservedScopes(t, ui.Claims)
				},
			},
		},
		{
			name: "project role assertion",
			prepare: func(t *testing.T, clientID string, scope []string) *oidc.Tokens[*oidc.IDTokenClaims] {
				_, err := Instance.Client.Mgmt.UpdateProject(CTX, &management.UpdateProjectRequest{
					Id:                   projectID,
					Name:                 integration.ProjectName(),
					ProjectRoleAssertion: true,
				})
				require.NoError(t, err)
				t.Cleanup(func() {
					_, err := Instance.Client.Mgmt.UpdateProject(CTX, &management.UpdateProjectRequest{
						Id:                   projectID,
						Name:                 integration.ProjectName(),
						ProjectRoleAssertion: false,
					})
					require.NoError(t, err)
				})
				resp, err := Instance.Client.Mgmt.GetProjectByID(CTX, &management.GetProjectByIDRequest{Id: projectID})
				require.NoError(t, err)
				require.True(t, resp.GetProject().GetProjectRoleAssertion(), "project role assertion")

				return getTokens(t, clientID, scope)
			},
			scope: []string{oidc.ScopeProfile, oidc.ScopeOpenID, oidc.ScopeEmail, oidc.ScopeOfflineAccess},
			assertions: []func(*testing.T, *oidc.UserInfo){
				assertUserinfo,
				func(t *testing.T, ui *oidc.UserInfo) {
					assertProjectRoleClaims(t, projectID, ui.Claims, true, []string{roleFoo, roleBar}, []string{Instance.DefaultOrg.Id})
				},
			},
		},
		{
			name:    "project role scope",
			prepare: getTokens,
			scope: []string{oidc.ScopeProfile, oidc.ScopeOpenID, oidc.ScopeEmail, oidc.ScopeOfflineAccess,
				oidc_api.ScopeProjectRolePrefix + roleFoo,
			},
			assertions: []func(*testing.T, *oidc.UserInfo){
				assertUserinfo,
				func(t *testing.T, ui *oidc.UserInfo) {
					assertProjectRoleClaims(t, projectID, ui.Claims, true, []string{roleFoo, roleBar}, []string{Instance.DefaultOrg.Id})
				},
			},
		},
		{
			name:    "project role and audience scope",
			prepare: getTokens,
			scope: []string{oidc.ScopeProfile, oidc.ScopeOpenID, oidc.ScopeEmail, oidc.ScopeOfflineAccess,
				oidc_api.ScopeProjectRolePrefix + roleFoo,
				domain.ProjectIDScope + projectID + domain.AudSuffix,
			},
			assertions: []func(*testing.T, *oidc.UserInfo){
				assertUserinfo,
				func(t *testing.T, ui *oidc.UserInfo) {
					assertProjectRoleClaims(t, projectID, ui.Claims, true, []string{roleFoo, roleBar}, []string{Instance.DefaultOrg.Id})
				},
			},
		},
		{
			name: "PAT",
			prepare: func(t *testing.T, clientID string, scope []string) *oidc.Tokens[*oidc.IDTokenClaims] {
				user := Instance.Users.Get(integration.UserTypeOrgOwner)
				return &oidc.Tokens[*oidc.IDTokenClaims]{
					Token: &oauth2.Token{
						AccessToken: user.Token,
						TokenType:   oidc.BearerToken,
					},
					IDTokenClaims: &oidc.IDTokenClaims{
						TokenClaims: oidc.TokenClaims{
							Subject: user.ID,
						},
					},
				}
			},
			assertions: []func(*testing.T, *oidc.UserInfo){
				func(t *testing.T, ui *oidc.UserInfo) {
					user := Instance.Users.Get(integration.UserTypeOrgOwner)
					assert.Equal(t, user.ID, ui.Subject)
					assert.NotEmpty(t, ui.Claims[oidc_api.ClaimResourceOwnerName])
					assert.NotEmpty(t, ui.Claims[oidc_api.ClaimResourceOwnerPrimaryDomain])
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tokens := tt.prepare(t, clientID, tt.scope)
			provider, err := Instance.CreateRelyingParty(CTX, clientID, redirectURI)
			require.NoError(t, err)
			userinfo, err := rp.Userinfo[*oidc.UserInfo](CTX, tokens.AccessToken, tokens.TokenType, tokens.IDTokenClaims.Subject, provider)
			if tt.wantErr {
				assert.Error(t, err)
				return
			}
			require.NoError(t, err)
			for _, assertion := range tt.assertions {
				assertion(t, userinfo)
			}
		})
	}
}

// TestServer_UserInfo_OrgIDRoles tests the [domain.OrgRoleIDScope] functionality
// it is a separate test because it is not supported in legacy mode.
func TestServer_UserInfo_OrgIDRoles(t *testing.T) {
	const (
		roleFoo = "foo"
		roleBar = "bar"
	)
	clientID, projectID := createClient(t, Instance)
	addProjectRolesGrants(t, User.GetUserId(), projectID, roleFoo, roleBar)
	grantedOrgID := addProjectOrgGrant(t, User.GetUserId(), projectID, roleFoo, roleBar)

	_, err := Instance.Client.Mgmt.UpdateProject(CTX, &management.UpdateProjectRequest{
		Id:                   projectID,
		Name:                 integration.ProjectName(),
		ProjectRoleAssertion: true,
	})
	require.NoError(t, err)
	resp, err := Instance.Client.Mgmt.GetProjectByID(CTX, &management.GetProjectByIDRequest{Id: projectID})
	require.NoError(t, err)
	require.True(t, resp.GetProject().GetProjectRoleAssertion(), "project role assertion")

	tests := []struct {
		name           string
		scope          []string
		wantRoleOrgIDs []string
	}{
		{
			name: "default returns all role orgs",
			scope: []string{
				oidc.ScopeOpenID, oidc.ScopeOfflineAccess,
			},
			wantRoleOrgIDs: []string{Instance.DefaultOrg.Id, grantedOrgID},
		},
		{
			name: "only granted org",
			scope: []string{
				oidc.ScopeOpenID, oidc.ScopeOfflineAccess,
				domain.OrgRoleIDScope + grantedOrgID},
			wantRoleOrgIDs: []string{grantedOrgID},
		},
		{
			name: "only own org",
			scope: []string{
				oidc.ScopeOpenID, oidc.ScopeOfflineAccess,
				domain.OrgRoleIDScope + Instance.DefaultOrg.Id,
			},
			wantRoleOrgIDs: []string{Instance.DefaultOrg.Id},
		},
		{
			name: "request both orgs",
			scope: []string{
				oidc.ScopeOpenID, oidc.ScopeOfflineAccess,
				domain.OrgRoleIDScope + Instance.DefaultOrg.Id,
				domain.OrgRoleIDScope + grantedOrgID,
			},
			wantRoleOrgIDs: []string{Instance.DefaultOrg.Id, grantedOrgID},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tokens := getTokens(t, clientID, tt.scope)
			provider, err := Instance.CreateRelyingParty(CTX, clientID, redirectURI)
			require.NoError(t, err)
			userinfo, err := rp.Userinfo[*oidc.UserInfo](CTX, tokens.AccessToken, tokens.TokenType, tokens.IDTokenClaims.Subject, provider)
			require.NoError(t, err)
			assertProjectRoleClaims(t, projectID, userinfo.Claims, true, []string{roleFoo, roleBar}, tt.wantRoleOrgIDs)
		})
	}
}

// https://github.com/zitadel/zitadel/issues/6662
func TestServer_UserInfo_Issue6662(t *testing.T) {
	const (
		roleFoo = "foo"
		roleBar = "bar"
	)

	projectID := Instance.CreateProject(CTX, t, "", integration.ProjectName(), false, false).GetId()
	user, _, clientID, clientSecret, err := Instance.CreateOIDCCredentialsClient(CTX)
	require.NoError(t, err)
	addProjectRolesGrants(t, user.GetUserId(), projectID, roleFoo, roleBar)

	scope := []string{oidc.ScopeProfile, oidc.ScopeOpenID, oidc.ScopeEmail, oidc.ScopeOfflineAccess,
		oidc_api.ScopeProjectRolePrefix + roleFoo,
		domain.ProjectIDScope + projectID + domain.AudSuffix,
	}

	provider, err := rp.NewRelyingPartyOIDC(CTX, Instance.OIDCIssuer(), clientID, clientSecret, redirectURI, scope)
	require.NoError(t, err)
	tokens, err := rp.ClientCredentials(CTX, provider, nil)
	require.NoError(t, err)

	userinfo, err := rp.Userinfo[*oidc.UserInfo](CTX, tokens.AccessToken, tokens.TokenType, user.GetUserId(), provider)
	require.NoError(t, err)
	assertProjectRoleClaims(t, projectID, userinfo.Claims, false, []string{roleFoo, roleBar}, []string{Instance.DefaultOrg.Id})
}

func addProjectRolesGrants(t *testing.T, userID, projectID string, roles ...string) {
	t.Helper()
	bulkRoles := make([]*management.BulkAddProjectRolesRequest_Role, len(roles))
	for i, role := range roles {
		bulkRoles[i] = &management.BulkAddProjectRolesRequest_Role{
			Key:         role,
			DisplayName: role,
		}
	}
	_, err := Instance.Client.Mgmt.BulkAddProjectRoles(CTX, &management.BulkAddProjectRolesRequest{
		ProjectId: projectID,
		Roles:     bulkRoles,
	})
	require.NoError(t, err)
	_, err = Instance.Client.Mgmt.AddUserGrant(CTX, &management.AddUserGrantRequest{
		UserId:    userID,
		ProjectId: projectID,
		RoleKeys:  roles,
	})
	require.NoError(t, err)
}

// addProjectOrgGrant adds a new organization which will be granted on the projectID with the specified roles.
// The userID will be granted in the new organization to the project with the same roles.
func addProjectOrgGrant(t *testing.T, userID, projectID string, roles ...string) (grantedOrgID string) {
	grantedOrg := Instance.CreateOrganization(CTXIAM, integration.OrganizationName(), integration.Email())
	projectGrant, err := Instance.Client.Mgmt.AddProjectGrant(CTX, &management.AddProjectGrantRequest{
		ProjectId:    projectID,
		GrantedOrgId: grantedOrg.GetOrganizationId(),
		RoleKeys:     roles,
	})
	require.NoError(t, err)

	ctxOrg := metadata.AppendToOutgoingContext(CTXIAM, "x-zitadel-orgid", grantedOrg.GetOrganizationId())
	_, err = Instance.Client.Mgmt.AddUserGrant(ctxOrg, &management.AddUserGrantRequest{
		UserId:         userID,
		ProjectId:      projectID,
		ProjectGrantId: projectGrant.GetGrantId(),
		RoleKeys:       roles,
	})
	require.NoError(t, err)
	return grantedOrg.GetOrganizationId()
}

func getTokens(t *testing.T, clientID string, scope []string) *oidc.Tokens[*oidc.IDTokenClaims] {
	authRequestID := createAuthRequest(t, Instance, clientID, redirectURI, scope...)
	sessionID, sessionToken, startTime, changeTime := Instance.CreateVerifiedWebAuthNSession(t, CTXLOGIN, User.GetUserId())
	linkResp, err := Instance.Client.OIDCv2.CreateCallback(CTXLOGIN, &oidc_pb.CreateCallbackRequest{
		AuthRequestId: authRequestID,
		CallbackKind: &oidc_pb.CreateCallbackRequest_Session{
			Session: &oidc_pb.Session{
				SessionId:    sessionID,
				SessionToken: sessionToken,
			},
		},
	})
	require.NoError(t, err)

	// code exchange
	code := assertCodeResponse(t, linkResp.GetCallbackUrl())
	tokens, err := exchangeTokens(t, Instance, clientID, code, redirectURI)
	require.NoError(t, err)
	assertTokens(t, tokens, true)
	assertIDTokenClaims(t, tokens.IDTokenClaims, User.GetUserId(), armPasskey, startTime, changeTime, sessionID)

	return tokens
}

func assertUserinfo(t *testing.T, userinfo *oidc.UserInfo) {
	t.Helper()
	assert.Equal(t, User.GetUserId(), userinfo.Subject)
	assert.Equal(t, "Mickey", userinfo.GivenName)
	assert.Equal(t, "Mouse", userinfo.FamilyName)
	assert.Equal(t, "Mickey Mouse", userinfo.Name)
	assert.NotEmpty(t, userinfo.PreferredUsername)
	assert.Equal(t, userinfo.PreferredUsername, userinfo.Email)
	assert.False(t, bool(userinfo.EmailVerified))
	assertOIDCTime(t, userinfo.UpdatedAt, User.GetDetails().GetChangeDate().AsTime())
}

func assertNoReservedScopes(t *testing.T, claims map[string]any) {
	t.Helper()
	t.Log(claims)
	for claim := range claims {
		assert.Falsef(t, strings.HasPrefix(claim, oidc_api.ClaimPrefix), "claim %s has prefix %s", claim, oidc_api.ClaimPrefix)
	}
}

// assertProjectRoleClaims asserts the projectRoles in the claims.
// By default it searches for the [oidc_api.ClaimProjectRolesFormat] claim with a project ID,
// and optionally for the [oidc_api.ClaimProjectRoles] claim if claimProjectRole is true.
// Each claim should contain the roles expected by wantRoles and
// each role should contain the org IDs expected by wantRoleOrgIDs.
//
// In the claim map, each project role claim is expected to be a map of multiple roles and
// each role is expected to be a map of multiple Org IDs to Org Domains.
func assertProjectRoleClaims(t *testing.T, projectID string, claims map[string]any, claimProjectRole bool, wantRoles, wantRoleOrgIDs []string) {
	t.Helper()
	projectRoleClaims := []string{fmt.Sprintf(oidc_api.ClaimProjectRolesFormat, projectID)}
	if claimProjectRole {
		projectRoleClaims = append(projectRoleClaims, oidc_api.ClaimProjectRoles)
	}
	for _, claim := range projectRoleClaims {
		roleMap, ok := claims[claim].(map[string]any) // map of multiple roles
		require.Truef(t, ok, "claim %s not found or wrong type %T", claim, claims[claim])

		gotRoles := make([]string, 0, len(roleMap))
		for roleKey := range roleMap {
			role, ok := roleMap[roleKey].(map[string]any) // map of multiple org IDs to org domains
			require.Truef(t, ok, "role %s not found or wrong type %T", roleKey, roleMap[roleKey])
			gotRoles = append(gotRoles, roleKey)

			gotRoleOrgIDs := make([]string, 0, len(role))
			for orgID := range role {
				gotRoleOrgIDs = append(gotRoleOrgIDs, orgID)
			}
			assert.ElementsMatch(t, wantRoleOrgIDs, gotRoleOrgIDs)
		}
		assert.ElementsMatch(t, wantRoles, gotRoles)
	}
}