chore: sql parameter to custom roles to be a (name,org) tuple

Emyrk · Emyrk · commit 8e98ceaa0328 · 2024-06-05T15:39:20.000-05:00
diff --git a/coderd/database/dbauthz/customroles_test.go b/coderd/database/dbauthz/customroles_test.go
@@ -38,7 +38,7 @@ func TestUpsertCustomRoles(t *testing.T) {
 		Name:        "can-assign",
 		DisplayName: "",
 		Site: rbac.Permissions(map[string][]policy.Action{
-			rbac.ResourceAssignRole.Type: {policy.ActionCreate},
+			rbac.ResourceAssignRole.Type: {policy.ActionRead, policy.ActionCreate},
 		}),
 	}
 
@@ -243,6 +243,20 @@ func TestUpsertCustomRoles(t *testing.T) {
 				require.ErrorContains(t, err, tc.errorContains)
 			} else {
 				require.NoError(t, err)
+
+				// Verify we can fetch the role
+				roles, err := az.CustomRoles(ctx, database.CustomRolesParams{
+					LookupRoles: []database.NameOrganizationPair{
+						{
+							Name:           "test-role",
+							OrganizationID: tc.organizationID.UUID,
+						},
+					},
+					ExcludeOrgRoles: false,
+					OrganizationID:  uuid.UUID{},
+				})
+				require.NoError(t, err)
+				require.Len(t, roles, 1)
 			}
 		})
 	}
diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -1186,12 +1186,17 @@ func (q *FakeQuerier) CustomRoles(_ context.Context, arg database.CustomRolesPar
 	for _, role := range q.data.customRoles {
 		role := role
 		if len(arg.LookupRoles) > 0 {
-			if !slices.ContainsFunc(arg.LookupRoles, func(s string) bool {
-				roleName := rbac.RoleName(role.Name, "")
-				if role.OrganizationID.UUID != uuid.Nil {
-					roleName = rbac.RoleName(role.Name, role.OrganizationID.UUID.String())
+			if !slices.ContainsFunc(arg.LookupRoles, func(pair database.NameOrganizationPair) bool {
+				if !strings.EqualFold(pair.Name, role.Name) {
+					return false
 				}
-				return strings.EqualFold(s, roleName)
+
+				if role.OrganizationID.Valid {
+					// Expect org match
+					return role.OrganizationID.UUID == pair.OrganizationID
+				}
+				// Expect no org
+				return pair.OrganizationID == uuid.Nil
 			}) {
 				continue
 			}
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/roles.sql b/coderd/database/queries/roles.sql
@@ -5,20 +5,16 @@ FROM
 	custom_roles
 WHERE
   true
-  -- Lookup roles filter expects the role names to be in the rbac package
-  -- format. Eg: name[:<organization_id>]
-  AND CASE WHEN array_length(@lookup_roles :: text[], 1) > 0  THEN
-	-- Case insensitive lookup with org_id appended (if non-null).
-    -- This will return just the name if org_id is null. It'll append
-    -- the org_id if not null
-	concat(name, NULLIF(concat(':', organization_id), ':')) ILIKE ANY(@lookup_roles :: text [])
-    ELSE true
+  -- @lookup_roles will filter for exact (role_name, org_id) pairs
+  AND CASE WHEN array_length(@lookup_roles :: name_organization_pair_list[], 1) > 0  THEN
+	(name, organization_id) ILIKE ANY (@lookup_roles::name_organization_pair_list[])
   END
-  -- Org scoping filter, to only fetch site wide roles
+  -- This allows fetching all roles, or just site wide roles
   AND CASE WHEN @exclude_org_roles :: boolean  THEN
 	organization_id IS null
 	ELSE true
   END
+  -- Allows fetching all roles to a particular organization
   AND CASE WHEN @organization_id :: uuid != '00000000-0000-0000-0000-000000000000'::uuid  THEN
       organization_id = @organization_id
     ELSE true
diff --git a/coderd/database/sqlc.yaml b/coderd/database/sqlc.yaml
@@ -28,6 +28,10 @@ sql:
         emit_enum_valid_method: true
         emit_all_enum_values: true
         overrides:
+          # Used in 'CustomRoles' query to filter by (name,organization_id)
+          - db_type: "name_organization_pair_list"
+            go_type:
+              type: "NameOrganizationPair"
           - column: "custom_roles.site_permissions"
             go_type:
               type: "CustomRolePermissions"
diff --git a/coderd/database/types.go b/coderd/database/types.go
@@ -142,3 +142,23 @@ func (a CustomRolePermission) String() string {
 	}
 	return str
 }
+
+// NameOrganizationPair is used as a lookup tuple for custom role rows.
+type NameOrganizationPair struct {
+	Name string `db:"name" json:"name"`
+	// OrganizationID if unset will assume a null column value
+	OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"`
+}
+
+func (a *NameOrganizationPair) Scan(_ interface{}) error {
+	return xerrors.Errorf("this should never happen, type 'NameOrganizationPair' should only be used as a parameter")
+}
+
+func (a NameOrganizationPair) Value() (driver.Value, error) {
+	var orgID interface{} = a.OrganizationID
+	if a.OrganizationID == uuid.Nil {
+		orgID = nil
+	}
+
+	return []interface{}{a.Name, orgID}, nil
+}
diff --git a/coderd/rbac/rolestore/rolestore.go b/coderd/rbac/rolestore/rolestore.go
@@ -69,11 +69,34 @@ func Expand(ctx context.Context, db database.Store, names []string) (rbac.Roles,
 	}
 
 	if len(lookup) > 0 {
+		// The set of roles coming in are formatted as 'rolename[:<org_id>]'.
+		// In the database, org roles are scoped with an organization column.
+		lookupArgs := make([]database.NameOrganizationPair, 0, len(lookup))
+		for _, name := range lookup {
+			roleName, orgID, err := rbac.RoleSplit(name)
+			if err != nil {
+				continue
+			}
+
+			parsedOrgID := uuid.Nil // Default to no org ID
+			if orgID != "" {
+				parsedOrgID, err = uuid.Parse(orgID)
+				if err != nil {
+					continue
+				}
+			}
+
+			lookupArgs = append(lookupArgs, database.NameOrganizationPair{
+				Name:           roleName,
+				OrganizationID: parsedOrgID,
+			})
+		}
+
 		// If some roles are missing from the database, they are omitted from
 		// the expansion. These roles are no-ops. Should we raise some kind of
 		// warning when this happens?
 		dbroles, err := db.CustomRoles(ctx, database.CustomRolesParams{
-			LookupRoles:     lookup,
+			LookupRoles:     lookupArgs,
 			ExcludeOrgRoles: false,
 			OrganizationID:  uuid.Nil,
 		})
