coder · Emyrk · May 29, 2024 · May 23, 2024 · May 23, 2024 · May 24, 2024
diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go
diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -424,6 +424,7 @@ func New(options *Options) *API {
 		TemplateScheduleStore:       options.TemplateScheduleStore,
 		UserQuietHoursScheduleStore: options.UserQuietHoursScheduleStore,
 		AccessControlStore:          options.AccessControlStore,
+		CustomRoleHandler:           atomic.Pointer[CustomRoleHandler]{},
 		Experiments:                 experiments,
 		healthCheckGroup:            &singleflight.Group[string, *healthsdk.HealthcheckReport]{},
 		Acquirer: provisionerdserver.NewAcquirer(
@@ -436,6 +437,8 @@ func New(options *Options) *API {
 		workspaceUsageTracker: options.WorkspaceUsageTracker,
 	}
 
+	var customRoleHandler CustomRoleHandler = &agplCustomRoleHandler{}
+	api.CustomRoleHandler.Store(&customRoleHandler)
 	api.AppearanceFetcher.Store(&appearance.DefaultFetcher)
 	api.PortSharer.Store(&portsharing.DefaultPortSharer)
 	buildInfo := codersdk.BuildInfoResponse{
@@ -828,7 +831,12 @@ func New(options *Options) *API {
 					})
 				})
 				r.Route("/members", func(r chi.Router) {
-					r.Get("/roles", api.assignableOrgRoles)
+					r.Route("/roles", func(r chi.Router) {
+						r.Get("/", api.assignableOrgRoles)
+						r.With(httpmw.RequireExperiment(api.Experiments, codersdk.ExperimentCustomRoles)).
+							Patch("/", api.patchOrgRoles)
+					})
+
 					r.Route("/{user}", func(r chi.Router) {
 						r.Use(
 							httpmw.ExtractOrganizationMemberParam(options.Database),
@@ -1249,6 +1257,8 @@ type API struct {
 	// passed to dbauthz.
 	AccessControlStore *atomic.Pointer[dbauthz.AccessControlStore]
 	PortSharer         atomic.Pointer[portsharing.PortSharer]
+	// CustomRoleHandler is the AGPL/Enterprise implementation for custom roles.
+	CustomRoleHandler atomic.Pointer[CustomRoleHandler]
 
 	HTTPAuth *HTTPAuthorizer
 

diff --git a/coderd/database/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go
@@ -531,12 +531,16 @@ func Role(role rbac.Role) codersdk.Role {
 	if err != nil {
 		roleName = role.Name
 	}
+
 	return codersdk.Role{
-		Name:                    roleName,
-		OrganizationID:          orgIDStr,
-		DisplayName:             role.DisplayName,
-		SitePermissions:         List(role.Site, Permission),
-		OrganizationPermissions: Map(role.Org, ListLazy(Permission)),
+		Name:            roleName,
+		OrganizationID:  orgIDStr,
+		DisplayName:     role.DisplayName,
+		SitePermissions: List(role.Site, Permission),
+		// This is not perfect. If there are organization permissions in another
+		// organization, they will be omitted. This should not be allowed, so
+		// should never happen.
+		OrganizationPermissions: List(role.Org[orgIDStr], Permission),
 		UserPermissions:         List(role.User, Permission),
 	}
 }
@@ -550,11 +554,18 @@ func Permission(permission rbac.Permission) codersdk.Permission {
 }
 
 func RoleToRBAC(role codersdk.Role) rbac.Role {
+	orgPerms := map[string][]rbac.Permission{}
+	if role.OrganizationID != "" {
+		orgPerms = map[string][]rbac.Permission{
+			role.OrganizationID: List(role.OrganizationPermissions, PermissionToRBAC),
+		}
+	}
+
 	return rbac.Role{
 		Name:        rbac.RoleName(role.Name, role.OrganizationID),
 		DisplayName: role.DisplayName,
 		Site:        List(role.SitePermissions, PermissionToRBAC),
-		Org:         Map(role.OrganizationPermissions, ListLazy(PermissionToRBAC)),
+		Org:         orgPerms,
 		User:        List(role.UserPermissions, PermissionToRBAC),
 	}
 }

diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -600,14 +600,29 @@ func (q *querier) canAssignRoles(ctx context.Context, orgID *uuid.UUID, added, r
 	customRoles := make([]string, 0)
 	// Validate that the roles being assigned are valid.
 	for _, r := range grantedRoles {
-		_, isOrgRole := rbac.IsOrgRole(r)
+		roleOrgIDStr, isOrgRole := rbac.IsOrgRole(r)
 		if shouldBeOrgRoles && !isOrgRole {
 			return xerrors.Errorf("Must only update org roles")
 		}
 		if !shouldBeOrgRoles && isOrgRole {
 			return xerrors.Errorf("Must only update site wide roles")
 		}
 
+		if shouldBeOrgRoles {
+			roleOrgID, err := uuid.Parse(roleOrgIDStr)
+			if err != nil {
+				return xerrors.Errorf("role %q has invalid uuid for org: %w", r, err)
+			}
+
+			if orgID == nil {
+				return xerrors.Errorf("should never happen, orgID is nil, but trying to assign an organization role")
+			}
+
+			if roleOrgID != *orgID {
+				return xerrors.Errorf("attempted to assign role from a different org, role %q to %q", r, orgID.String())
+			}
+		}
+
 		// All roles should be valid roles
 		if _, err := rbac.RoleByName(r); err != nil {
 			customRoles = append(customRoles, r)

diff --git a/coderd/members.go b/coderd/members.go
@@ -1,13 +1,8 @@
 package coderd
 
 import (
-	"context"
 	"net/http"
 
-	"github.com/google/uuid"
-
-	"golang.org/x/xerrors"
-
 	"github.com/coder/coder/v2/coderd/database/db2sdk"
 	"github.com/coder/coder/v2/coderd/rbac"
 
@@ -48,7 +43,7 @@ func (api *API) putMemberRoles(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	updatedUser, err := api.updateOrganizationMemberRoles(ctx, database.UpdateMemberRolesParams{
+	updatedUser, err := api.Database.UpdateMemberRoles(ctx, database.UpdateMemberRolesParams{
 		GrantedRoles: params.Roles,
 		UserID:       member.UserID,
 		OrgID:        organization.ID,
@@ -63,36 +58,6 @@ func (api *API) putMemberRoles(rw http.ResponseWriter, r *http.Request) {
 	httpapi.Write(ctx, rw, http.StatusOK, convertOrganizationMember(updatedUser))
 }
 
-func (api *API) updateOrganizationMemberRoles(ctx context.Context, args database.UpdateMemberRolesParams) (database.OrganizationMember, error) {
-	// Enforce only site wide roles
-	for _, r := range args.GrantedRoles {
-		// Must be an org role for the org in the args
-		orgID, ok := rbac.IsOrgRole(r)
-		if !ok {
-			return database.OrganizationMember{}, xerrors.Errorf("must only update organization roles")
-		}
-
-		roleOrg, err := uuid.Parse(orgID)
-		if err != nil {
-			return database.OrganizationMember{}, xerrors.Errorf("Role must have proper UUIDs for organization, %q does not", r)
-		}
-
-		if roleOrg != args.OrgID {
-			return database.OrganizationMember{}, xerrors.Errorf("Must only pass roles for org %q", args.OrgID.String())
-		}
-
-		if _, err := rbac.RoleByName(r); err != nil {
-			return database.OrganizationMember{}, xerrors.Errorf("%q is not a supported organization role", r)
-		}
-	}
-
-	updatedUser, err := api.Database.UpdateMemberRoles(ctx, args)
-	if err != nil {
-		return database.OrganizationMember{}, xerrors.Errorf("Update site roles: %w", err)
-	}
-	return updatedUser, nil
-}
-
 func convertOrganizationMember(mem database.OrganizationMember) codersdk.OrganizationMember {
 	convertedMember := codersdk.OrganizationMember{
 		UserID:         mem.UserID,