chore: add query to fetch top level idp claim fields

Used for idp sync settings. Tests WIP
coder · Emyrk · Nov 18, 2024 · Nov 14, 2024 · Nov 14, 2024 · Nov 14, 2024
commit 1ae65f324e7133fc6d63160f916da7e9ee6a477d
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -3283,6 +3283,18 @@ func (q *querier) ListWorkspaceAgentPortShares(ctx context.Context, workspaceID
 	return q.db.ListWorkspaceAgentPortShares(ctx, workspaceID)
 }
 
+func (q *querier) OIDCClaimFields(ctx context.Context, organizationID uuid.UUID) ([]string, error) {
+	resource := rbac.ResourceIdpsyncSettings
+	if organizationID != uuid.Nil {
+		resource = resource.InOrg(organizationID)
+	}
+
+	if err := q.authorizeContext(ctx, policy.ActionRead, resource); err != nil {
+		return nil, err
+	}
+	return q.db.OIDCClaimFields(ctx, organizationID)
+}
+
 func (q *querier) OrganizationMembers(ctx context.Context, arg database.OrganizationMembersParams) ([]database.OrganizationMembersRow, error) {
 	return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.OrganizationMembers)(ctx, arg)
 }

diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -8409,6 +8409,35 @@ func (q *FakeQuerier) ListWorkspaceAgentPortShares(_ context.Context, workspaceI
 	return shares, nil
 }
 
+func (q *FakeQuerier) OIDCClaimFields(ctx context.Context, organizationID uuid.UUID) ([]string, error) {
+	orgMembers := q.getOrganizationMemberNoLock(organizationID)
+
+	var fields []string
+	for _, link := range q.userLinks {
+		if organizationID != uuid.Nil {
+			inOrg := slices.ContainsFunc(orgMembers, func(organizationMember database.OrganizationMember) bool {
+				return organizationMember.UserID == link.UserID
+			})
+			if !inOrg {
+				continue
+			}
+		}
+
+		if link.LoginType != database.LoginTypeOIDC {
+			continue
+		}
+
+		for k := range link.Claims.IDTokenClaims {
+			fields = append(fields, k)
+		}
+		for k := range link.Claims.UserInfoClaims {
+			fields = append(fields, k)
+		}
+	}
+
+	return slice.Unique(fields), nil
+}
+
 func (q *FakeQuerier) OrganizationMembers(_ context.Context, arg database.OrganizationMembersParams) ([]database.OrganizationMembersRow, error) {
 	if err := validateDatabaseType(arg); err != nil {
 		return []database.OrganizationMembersRow{}, err

diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/coderd/database/modelqueries.go b/coderd/database/modelqueries.go
@@ -3,6 +3,7 @@ package database
 import (
 	"context"
 	"database/sql"
+	"encoding/json"
 	"fmt"
 	"strings"
 
@@ -527,3 +528,9 @@ func insertAuthorizedFilter(query string, replaceWith string) (string, error) {
 	filtered := strings.Replace(query, authorizedQueryPlaceholder, replaceWith, 1)
 	return filtered, nil
 }
+
+// UpdateUserLinkRawJSON is a custom query for unit testing. Do not ever expose this
+func (q *sqlQuerier) UpdateUserLinkRawJSON(ctx context.Context, userID uuid.UUID, data json.RawMessage) error {
+	_, err := q.sdb.Exec("INSERT INTO user_links (user_id, claims) VALUES ($1, $2)", userID, data)
+	return err
+}
diff --git a/coderd/database/oidcclaims_test.go b/coderd/database/oidcclaims_test.go
@@ -0,0 +1,186 @@
+package database_test
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbfake"
+	"github.com/coder/coder/v2/coderd/database/dbtestutil"
+	"github.com/coder/coder/v2/coderd/database/dbtime"
+	"github.com/coder/coder/v2/testutil"
+)
+
+type extraKeys struct {
+	database.UserLinkClaims
+	Foo string `json:"foo"`
+}
+
+func TestOIDCClaims(t *testing.T) {
+	t.Parallel()
+
+	toJSON := func(a any) json.RawMessage {
+		b, _ := json.Marshal(a)
+		return b
+	}
+
+	db, _ := dbtestutil.NewDB(t)
+	g := userGenerator{t: t, db: db}
+
+	// https://en.wikipedia.org/wiki/Alice_and_Bob#Cast_of_characters
+	alice := g.withLink(database.LoginTypeOIDC, toJSON(extraKeys{
+		UserLinkClaims: database.UserLinkClaims{
+			IDTokenClaims: map[string]interface{}{
+				"sub": "alice",
+			},
+			UserInfoClaims: map[string]interface{}{
+				"sub": "alice",
+			},
+		},
+		// Always should be a no-op
+		Foo: "bar",
+	}))
+	bob := g.withLink(database.LoginTypeOIDC, toJSON(database.UserLinkClaims{
+		IDTokenClaims: map[string]interface{}{
+			"sub": "bob",
+		},
+		UserInfoClaims: map[string]interface{}{
+			"sub": "bob",
+		},
+	}))
+	charlie := g.withLink(database.LoginTypeOIDC, toJSON(database.UserLinkClaims{
+		IDTokenClaims: map[string]interface{}{
+			"sub": "charlie",
+		},
+		UserInfoClaims: map[string]interface{}{
+			"sub": "charlie",
+		},
+	}))
+
+	// users that just try to cause problems, but should not affect the output of
+	// queries.
+	problematics := []database.User{
+		g.withLink(database.LoginTypeOIDC, toJSON(database.UserLinkClaims{})), // null claims
+		g.withLink(database.LoginTypeOIDC, []byte(`{}`)),                      // empty claims
+		g.withLink(database.LoginTypeOIDC, []byte(`{"foo": "bar"}`)),          // random keys
+		g.noLink(database.LoginTypeOIDC),                                      // no link
+
+		g.withLink(database.LoginTypeGithub, toJSON(database.UserLinkClaims{
+			IDTokenClaims: map[string]interface{}{
+				"not": "allowed",
+			},
+			UserInfoClaims: map[string]interface{}{
+				"do-not": "look",
+			},
+		})), // github should be omitted
+
+		// extra random users
+		g.noLink(database.LoginTypeGithub),
+		g.noLink(database.LoginTypePassword),
+	}
+
+	// Insert some orgs, users, and links
+	orgA := dbfake.Organization(t, db).Members(
+		append(problematics,
+			alice,
+			bob)...,
+	).Do()
+	orgB := dbfake.Organization(t, db).Members(
+		append(problematics,
+			charlie,
+		)...,
+	).Do()
+
+	// Verify the OIDC claim fields
+	requireClaims(t, db, orgA.Org.ID, []string{"sub"})
+	requireClaims(t, db, orgB.Org.ID, []string{"sub"})
+}
+
+func requireClaims(t *testing.T, db database.Store, orgID uuid.UUID, want []string) {
+	t.Helper()
+
+	ctx := testutil.Context(t, testutil.WaitMedium)
+	got, err := db.OIDCClaimFields(ctx, orgID)
+	require.NoError(t, err)
+
+	require.ElementsMatch(t, want, got)
+}
+
+type userGenerator struct {
+	t  *testing.T
+	db database.Store
+}
+
+func (g userGenerator) noLink(lt database.LoginType) database.User {
+	return g.user(lt, false, nil)
+}
+
+func (g userGenerator) withLink(lt database.LoginType, rawJSON json.RawMessage) database.User {
+	return g.user(lt, true, rawJSON)
+}
+
+func (g userGenerator) user(lt database.LoginType, createLink bool, rawJSON json.RawMessage) database.User {
+	t := g.t
+	db := g.db
+
+	t.Helper()
+
+	u, err := db.InsertUser(context.Background(), database.InsertUserParams{
+		ID:        uuid.New(),
+		Email:     testutil.GetRandomName(t),
+		Username:  testutil.GetRandomName(t),
+		Name:      testutil.GetRandomName(t),
+		CreatedAt: dbtime.Now(),
+		UpdatedAt: dbtime.Now(),
+		RBACRoles: []string{},
+		LoginType: lt,
+		Status:    string(database.UserStatusActive),
+	})
+	require.NoError(t, err)
+
+	if !createLink {
+		return u
+	}
+
+	link, err := db.InsertUserLink(context.Background(), database.InsertUserLinkParams{
+		UserID:    u.ID,
+		LoginType: lt,
+		Claims:    database.UserLinkClaims{},
+	})
+	require.NoError(t, err)
+
+	if sql, ok := db.(rawUpdater); ok {
+		// The only way to put arbitrary json into the db for testing edge cases.
+		// Making this a public API would be a mistake.
+		err = sql.UpdateUserLinkRawJSON(context.Background(), u.ID, rawJSON)
+		require.NoError(t, err)
+	} else {
+		// no need to test the json key logic in dbmem. Everything is type safe.
+		var claims database.UserLinkClaims
+		err := json.Unmarshal(rawJSON, &claims)
+		require.NoError(t, err)
+
+		_, err = db.UpdateUserLink(context.Background(), database.UpdateUserLinkParams{
+			OAuthAccessToken:       link.OAuthAccessToken,
+			OAuthAccessTokenKeyID:  link.OAuthAccessTokenKeyID,
+			OAuthRefreshToken:      link.OAuthRefreshToken,
+			OAuthRefreshTokenKeyID: link.OAuthRefreshTokenKeyID,
+			OAuthExpiry:            link.OAuthExpiry,
+			UserID:                 link.UserID,
+			LoginType:              link.LoginType,
+			// The new claims
+			Claims: claims,
+		})
+		require.NoError(t, err)
+	}
+
+	return u
+}
+
+type rawUpdater interface {
+	UpdateUserLinkRawJSON(ctx context.Context, userID uuid.UUID, data json.RawMessage) error
+}
diff --git a/coderd/database/querier.go b/coderd/database/querier.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go