Cache subject AST values to avoid reallocating slices

coder · Emyrk · May 3, 2023 · May 3, 2023 · May 3, 2023 · May 3, 2023
commit 2bc9c7e680e4d286ed7636a89eac2a34bff8b70b
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -144,7 +144,8 @@ var (
 			},
 		}),
 		Scope: rbac.ScopeAll,
-	}
+	}.WithCachedASTValue()
+
 	subjectAutostart = rbac.Subject{
 		ID: uuid.Nil.String(),
 		Roles: rbac.Roles([]rbac.Role{
@@ -161,7 +162,8 @@ var (
 			},
 		}),
 		Scope: rbac.ScopeAll,
-	}
+	}.WithCachedASTValue()
+
 	subjectSystemRestricted = rbac.Subject{
 		ID: uuid.Nil.String(),
 		Roles: rbac.Roles([]rbac.Role{
@@ -188,7 +190,7 @@ var (
 			},
 		}),
 		Scope: rbac.ScopeAll,
-	}
+	}.WithCachedASTValue()
 )
 
 // AsProvisionerd returns a context with an actor that has permissions required

diff --git a/coderd/httpmw/apikey.go b/coderd/httpmw/apikey.go
@@ -379,7 +379,7 @@ func ExtractAPIKey(rw http.ResponseWriter, r *http.Request, cfg ExtractAPIKeyCon
 			Roles:  rbac.RoleNames(roles.Roles),
 			Groups: roles.Groups,
 			Scope:  rbac.ScopeName(key.Scope),
-		},
+		}.WithCachedASTValue(),
 	}
 
 	return &key, &authz, true

diff --git a/coderd/httpmw/workspaceagent.go b/coderd/httpmw/workspaceagent.go
@@ -110,5 +110,5 @@ func getAgentSubject(ctx context.Context, db database.Store, agent database.Work
 		Roles:  rbac.RoleNames(roles.Roles),
 		Groups: roles.Groups,
 		Scope:  rbac.WorkspaceAgentScope(workspace.ID, user.ID),
-	}, nil
+	}.WithCachedASTValue(), nil
 }
diff --git a/coderd/rbac/astvalue.go b/coderd/rbac/astvalue.go
@@ -65,6 +65,10 @@ func regoPartialInputValue(subject Subject, action Action, objectType string) (a
 
 // regoValue returns the ast.Object representation of the subject.
 func (s Subject) regoValue() (ast.Value, error) {
+	if s.cachedASTValue != nil {
+		return s.cachedASTValue, nil
+	}
+
 	subjRoles, err := s.Roles.Expand()
 	if err != nil {
 		return nil, xerrors.Errorf("expand roles: %w", err)

diff --git a/coderd/rbac/authz.go b/coderd/rbac/authz.go
@@ -49,6 +49,20 @@ type Subject struct {
 	Roles  ExpandableRoles
 	Groups []string
 	Scope  ExpandableScope
+
+	// cachedASTValue is the cached ast value for this subject.
+	cachedASTValue ast.Value
+}
+
+// WithCachedASTValue can be called if the subject is static. This will compute
+// the ast value once and cache it for future calls.
+func (s Subject) WithCachedASTValue() Subject {
+	tmp := s
+	v, err := tmp.regoValue()
+	if err == nil {
+		tmp.cachedASTValue = v
+	}
+	return tmp
 }
 
 func (s Subject) Equal(b Subject) bool {

diff --git a/coderd/rbac/authz_test.go b/coderd/rbac/authz_test.go
@@ -86,6 +86,21 @@ func benchmarkUserCases() (cases []benchmarkCase, users uuid.UUID, orgs []uuid.U
 				Groups: noiseGroups,
 			},
 		},
+		{
+			Name: "ManyRolesCachedSubject",
+			Actor: rbac.Subject{
+				// Admin of many orgs
+				Roles: rbac.RoleNames{
+					rbac.RoleOrgMember(orgs[0]), rbac.RoleOrgAdmin(orgs[0]),
+					rbac.RoleOrgMember(orgs[1]), rbac.RoleOrgAdmin(orgs[1]),
+					rbac.RoleOrgMember(orgs[2]), rbac.RoleOrgAdmin(orgs[2]),
+					rbac.RoleMember(),
+				},
+				ID:     user.String(),
+				Scope:  rbac.ScopeAll,
+				Groups: noiseGroups,
+			}.WithCachedASTValue(),
+		},
 		{
 			Name: "AdminWithScope",
 			Actor: rbac.Subject{