coder · Emyrk · Jun 21, 2024 · Jun 21, 2024 · Jun 21, 2024 · Jun 21, 2024
diff --git a/coderd/httpmw/apikey.go b/coderd/httpmw/apikey.go
@@ -406,16 +406,15 @@ func ExtractAPIKey(rw http.ResponseWriter, r *http.Request, cfg ExtractAPIKeyCon
 	// If the key is valid, we also fetch the user roles and status.
 	// The roles are used for RBAC authorize checks, and the status
 	// is to block 'suspended' users from accessing the platform.
-	//nolint:gocritic // system needs to update user roles
-	roles, err := cfg.DB.GetAuthorizationUserRoles(dbauthz.AsSystemRestricted(ctx), key.UserID)
+	actor, userStatus, err := UserRBACSubject(ctx, cfg.DB, key.UserID, rbac.ScopeName(key.Scope))
 	if err != nil {
 		return write(http.StatusUnauthorized, codersdk.Response{
 			Message: internalErrorMessage,
 			Detail:  fmt.Sprintf("Internal error fetching user's roles. %s", err.Error()),
 		})
 	}
 
-	if roles.Status == database.UserStatusDormant {
+	if userStatus == database.UserStatusDormant {
 		// If coder confirms that the dormant user is valid, it can switch their account to active.
 		// nolint:gocritic
 		u, err := cfg.DB.UpdateUserStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateUserStatusParams{
@@ -429,47 +428,50 @@ func ExtractAPIKey(rw http.ResponseWriter, r *http.Request, cfg ExtractAPIKeyCon
 				Detail:  fmt.Sprintf("can't activate a dormant user: %s", err.Error()),
 			})
 		}
-		roles.Status = u.Status
+		userStatus = u.Status
 	}
 
-	if roles.Status != database.UserStatusActive {
+	if userStatus != database.UserStatusActive {
 		return write(http.StatusUnauthorized, codersdk.Response{
-			Message: fmt.Sprintf("User is not active (status = %q). Contact an admin to reactivate your account.", roles.Status),
+			Message: fmt.Sprintf("User is not active (status = %q). Contact an admin to reactivate your account.", userStatus),
 		})
 	}
 
+	if cfg.PostAuthAdditionalHeadersFunc != nil {
+		cfg.PostAuthAdditionalHeadersFunc(actor, rw.Header())
+	}
+
+	return key, &actor, true
+}
+
+// UserRBACSubject fetches a user's rbac.Subject from the database. It pulls all roles from both
+// site and organization scopes. It also pulls the groups, and the user's status.
+func UserRBACSubject(ctx context.Context, db database.Store, userID uuid.UUID, scope rbac.ExpandableScope) (rbac.Subject, database.UserStatus, error) {
+	//nolint:gocritic // system needs to update user roles
+	roles, err := db.GetAuthorizationUserRoles(dbauthz.AsSystemRestricted(ctx), userID)
+	if err != nil {
+		return rbac.Subject{}, "", xerrors.Errorf("get authorization user roles: %w", err)
+	}
+
 	roleNames, err := roles.RoleNames()
 	if err != nil {
-		return write(http.StatusInternalServerError, codersdk.Response{
-			Message: "Internal Server Error",
-			Detail:  err.Error(),
-		})
+		return rbac.Subject{}, "", xerrors.Errorf("expand role names: %w", err)
 	}
 
 	//nolint:gocritic // Permission to lookup custom roles the user has assigned.
-	rbacRoles, err := rolestore.Expand(dbauthz.AsSystemRestricted(ctx), cfg.DB, roleNames)
+	rbacRoles, err := rolestore.Expand(dbauthz.AsSystemRestricted(ctx), db, roleNames)
 	if err != nil {
-		return write(http.StatusInternalServerError, codersdk.Response{
-			Message:     "Failed to expand authenticated user roles",
-			Detail:      err.Error(),
-			Validations: nil,
-		})
+		return rbac.Subject{}, "", xerrors.Errorf("expand role names: %w", err)
 	}
 
-	// Actor is the user's authorization context.
 	actor := rbac.Subject{
 		FriendlyName: roles.Username,
-		ID:           key.UserID.String(),
+		ID:           userID.String(),
 		Roles:        rbacRoles,
 		Groups:       roles.Groups,
-		Scope:        rbac.ScopeName(key.Scope),
+		Scope:        scope,
 	}.WithCachedASTValue()
-
-	if cfg.PostAuthAdditionalHeadersFunc != nil {
-		cfg.PostAuthAdditionalHeadersFunc(actor, rw.Header())
-	}
-
-	return key, &actor, true
+	return actor, roles.Status, nil
 }
 
 // APITokenFromRequest returns the api token from the request.

diff --git a/coderd/httpmw/apikey_test.go b/coderd/httpmw/apikey_test.go
@@ -14,16 +14,20 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/exp/slices"
 	"golang.org/x/oauth2"
 
 	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbauthz"
 	"github.com/coder/coder/v2/coderd/database/dbgen"
 	"github.com/coder/coder/v2/coderd/database/dbmem"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/httpapi"
 	"github.com/coder/coder/v2/coderd/httpmw"
+	"github.com/coder/coder/v2/coderd/rbac"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/cryptorand"
 	"github.com/coder/coder/v2/testutil"
@@ -38,6 +42,37 @@ func randomAPIKeyParts() (id string, secret string) {
 func TestAPIKey(t *testing.T) {
 	t.Parallel()
 
+	// assertActorOk asserts all the properties of the user auth are ok.
+	assertActorOk := func(t *testing.T, r *http.Request) {
+		t.Helper()
+
+		actor, ok := dbauthz.ActorFromContext(r.Context())
+		assert.True(t, ok, "dbauthz actor ok")
+		if ok {
+			_, err := actor.Roles.Expand()
+			assert.NoError(t, err, "actor roles ok")
+
+			_, err = actor.Scope.Expand()
+			assert.NoError(t, err, "actor scope ok")
+
+			err = actor.RegoValueOk()
+			assert.NoError(t, err, "actor rego ok")
+		}
+
+		auth, ok := httpmw.UserAuthorizationOptional(r)
+		assert.True(t, ok, "httpmw auth ok")
+		if ok {
+			_, err := auth.Roles.Expand()
+			assert.NoError(t, err, "auth roles ok")
+
+			_, err = auth.Scope.Expand()
+			assert.NoError(t, err, "auth scope ok")
+
+			err = auth.RegoValueOk()
+			assert.NoError(t, err, "auth rego ok")
+		}
+	}
+
 	successHandler := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
 		// Only called if the API key passes through the handler.
 		httpapi.Write(context.Background(), rw, http.StatusOK, codersdk.Response{
@@ -256,6 +291,7 @@ func TestAPIKey(t *testing.T) {
 		})(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
 			// Checks that it exists on the context!
 			_ = httpmw.APIKey(r)
+			assertActorOk(t, r)
 			httpapi.Write(r.Context(), rw, http.StatusOK, codersdk.Response{
 				Message: "It worked!",
 			})
@@ -296,6 +332,7 @@ func TestAPIKey(t *testing.T) {
 			// Checks that it exists on the context!
 			apiKey := httpmw.APIKey(r)
 			assert.Equal(t, database.APIKeyScopeApplicationConnect, apiKey.Scope)
+			assertActorOk(t, r)
 
 			httpapi.Write(r.Context(), rw, http.StatusOK, codersdk.Response{
 				Message: "it worked!",
@@ -330,6 +367,8 @@ func TestAPIKey(t *testing.T) {
 		})(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
 			// Checks that it exists on the context!
 			_ = httpmw.APIKey(r)
+			assertActorOk(t, r)
+
 			httpapi.Write(r.Context(), rw, http.StatusOK, codersdk.Response{
 				Message: "It worked!",
 			})
@@ -633,7 +672,7 @@ func TestAPIKey(t *testing.T) {
 		require.Equal(t, sentAPIKey.LoginType, gotAPIKey.LoginType)
 	})
 
-	t.Run("MissongConfig", func(t *testing.T) {
+	t.Run("MissingConfig", func(t *testing.T) {
 		t.Parallel()
 		var (
 			db       = dbmem.New()
@@ -667,4 +706,133 @@ func TestAPIKey(t *testing.T) {
 		out, _ := io.ReadAll(res.Body)
 		require.Contains(t, string(out), "Unable to refresh")
 	})
+
+	t.Run("CustomRoles", func(t *testing.T) {
+		t.Parallel()
+		var (
+			db         = dbmem.New()
+			org        = dbgen.Organization(t, db, database.Organization{})
+			customRole = dbgen.CustomRole(t, db, database.CustomRole{
+				Name:           "custom-role",
+				OrgPermissions: []database.CustomRolePermission{},
+				OrganizationID: uuid.NullUUID{
+					UUID:  org.ID,
+					Valid: true,
+				},
+			})
+			user = dbgen.User(t, db, database.User{
+				RBACRoles: []string{},
+			})
+			_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
+				UserID:         user.ID,
+				OrganizationID: org.ID,
+				CreatedAt:      time.Time{},
+				UpdatedAt:      time.Time{},
+				Roles: []string{
+					rbac.RoleOrgAdmin(),
+					customRole.Name,
+				},
+			})
+			_, token = dbgen.APIKey(t, db, database.APIKey{
+				UserID:    user.ID,
+				ExpiresAt: dbtime.Now().AddDate(0, 0, 1),
+			})
+
+			r  = httptest.NewRequest("GET", "/", nil)
+			rw = httptest.NewRecorder()
+		)
+		r.Header.Set(codersdk.SessionTokenHeader, token)
+
+		httpmw.ExtractAPIKeyMW(httpmw.ExtractAPIKeyConfig{
+			DB:              db,
+			RedirectToLogin: false,
+		})(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+			assertActorOk(t, r)
+
+			auth := httpmw.UserAuthorization(r)
+
+			roles, err := auth.Roles.Expand()
+			assert.NoError(t, err, "expand user roles")
+			// Assert built in org role
+			assert.True(t, slices.ContainsFunc(roles, func(role rbac.Role) bool {
+				return role.Identifier.Name == rbac.RoleOrgAdmin() && role.Identifier.OrganizationID == org.ID
+			}), "org admin role")
+			// Assert custom role
+			assert.True(t, slices.ContainsFunc(roles, func(role rbac.Role) bool {
+				return role.Identifier.Name == customRole.Name && role.Identifier.OrganizationID == org.ID
+			}), "custom org role")
+
+			httpapi.Write(r.Context(), rw, http.StatusOK, codersdk.Response{
+				Message: "It worked!",
+			})
+		})).ServeHTTP(rw, r)
+		res := rw.Result()
+		defer res.Body.Close()
+		require.Equal(t, http.StatusOK, res.StatusCode)
+	})
+
+	// There is no sql foreign key constraint to require all assigned roles
+	// still exist in the database. We need to handle deleted roles.
+	t.Run("RoleNotExists", func(t *testing.T) {
+		t.Parallel()
+		var (
+			roleNotExistsName = "role-not-exists"
+			db                = dbmem.New()
+			org               = dbgen.Organization(t, db, database.Organization{})
+			user              = dbgen.User(t, db, database.User{
+				RBACRoles: []string{
+					// Also provide an org not exists. In practice this makes no sense
+					// to store org roles in the user table, but there is no org to
+					// store it in. So just throw this here for even more unexpected
+					// behavior handling!
+					rbac.RoleIdentifier{Name: roleNotExistsName, OrganizationID: uuid.New()}.String(),
+				},
+			})
+			_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
+				UserID:         user.ID,
+				OrganizationID: org.ID,
+				CreatedAt:      time.Time{},
+				UpdatedAt:      time.Time{},
+				Roles: []string{
+					rbac.RoleOrgAdmin(),
+					roleNotExistsName,
+				},
+			})
+			_, token = dbgen.APIKey(t, db, database.APIKey{
+				UserID:    user.ID,
+				ExpiresAt: dbtime.Now().AddDate(0, 0, 1),
+			})
+
+			r  = httptest.NewRequest("GET", "/", nil)
+			rw = httptest.NewRecorder()
+		)
+		r.Header.Set(codersdk.SessionTokenHeader, token)
+
+		httpmw.ExtractAPIKeyMW(httpmw.ExtractAPIKeyConfig{
+			DB:              db,
+			RedirectToLogin: false,
+		})(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+			assertActorOk(t, r)
+			auth := httpmw.UserAuthorization(r)
+
+			roles, err := auth.Roles.Expand()
+			assert.NoError(t, err, "expand user roles")
+			// Assert built in org role
+			assert.True(t, slices.ContainsFunc(roles, func(role rbac.Role) bool {
+				return role.Identifier.Name == rbac.RoleOrgAdmin() && role.Identifier.OrganizationID == org.ID
+			}), "org admin role")
+
+			// Assert the role-not-exists is not returned
+			assert.False(t, slices.ContainsFunc(roles, func(role rbac.Role) bool {
+				return role.Identifier.Name == roleNotExistsName
+			}), "role should not exist")
+
+			httpapi.Write(r.Context(), rw, http.StatusOK, codersdk.Response{
+				Message: "It worked!",
+			})
+		})).ServeHTTP(rw, r)
+		res := rw.Result()
+		defer res.Body.Close()
+		require.Equal(t, http.StatusOK, res.StatusCode)
+	})
 }
diff --git a/coderd/identityprovider/tokens.go b/coderd/identityprovider/tokens.go
@@ -209,27 +209,14 @@ func authorizationCodeGrant(ctx context.Context, db database.Store, app database
 	}
 
 	// Grab the user roles so we can perform the exchange as the user.
-	//nolint:gocritic // In the token exchange, there is no user actor.
-	roles, err := db.GetAuthorizationUserRoles(dbauthz.AsSystemRestricted(ctx), dbCode.UserID)
+	actor, _, err := httpmw.UserRBACSubject(ctx, db, dbCode.UserID, rbac.ScopeAll)
 	if err != nil {
-		return oauth2.Token{}, err
-	}
-
-	roleNames, err := roles.RoleNames()
-	if err != nil {
-		return oauth2.Token{}, xerrors.Errorf("role names: %w", err)
-	}
-
-	userSubj := rbac.Subject{
-		ID:     dbCode.UserID.String(),
-		Roles:  rbac.RoleIdentifiers(roleNames),
-		Groups: roles.Groups,
-		Scope:  rbac.ScopeAll,
+		return oauth2.Token{}, xerrors.Errorf("fetch user actor: %w", err)
 	}
 
 	// Do the actual token exchange in the database.
 	err = db.InTx(func(tx database.Store) error {
-		ctx := dbauthz.As(ctx, userSubj)
+		ctx := dbauthz.As(ctx, actor)
 		err = tx.DeleteOAuth2ProviderAppCodeByID(ctx, dbCode.ID)
 		if err != nil {
 			return xerrors.Errorf("delete oauth2 app code: %w", err)
@@ -311,22 +298,10 @@ func refreshTokenGrant(ctx context.Context, db database.Store, app database.OAut
 	if err != nil {
 		return oauth2.Token{}, err
 	}
-	//nolint:gocritic // There is no user yet so we must use the system.
-	roles, err := db.GetAuthorizationUserRoles(dbauthz.AsSystemRestricted(ctx), prevKey.UserID)
-	if err != nil {
-		return oauth2.Token{}, err
-	}
 
-	roleNames, err := roles.RoleNames()
+	actor, _, err := httpmw.UserRBACSubject(ctx, db, prevKey.UserID, rbac.ScopeAll)
 	if err != nil {
-		return oauth2.Token{}, xerrors.Errorf("role names: %w", err)
-	}
-
-	userSubj := rbac.Subject{
-		ID:     prevKey.UserID.String(),
-		Roles:  rbac.RoleIdentifiers(roleNames),
-		Groups: roles.Groups,
-		Scope:  rbac.ScopeAll,
+		return oauth2.Token{}, xerrors.Errorf("fetch user actor: %w", err)
 	}
 
 	// Generate a new refresh token.
@@ -351,7 +326,7 @@ func refreshTokenGrant(ctx context.Context, db database.Store, app database.OAut
 
 	// Replace the token.
 	err = db.InTx(func(tx database.Store) error {
-		ctx := dbauthz.As(ctx, userSubj)
+		ctx := dbauthz.As(ctx, actor)
 		err = tx.DeleteAPIKeyByID(ctx, prevKey.ID) // This cascades to the token.
 		if err != nil {
 			return xerrors.Errorf("delete oauth2 app token: %w", err)