Skip to content

Commit 3368b8b

Browse files
EmyrkEmyrk
authored
chore: Minor rbac memory optimization (#7391)
* test: Add benchmark for static rbac roles * static roles should only be allocated once * A unit test that modifies the ast value should not mess with the globals * Cache subject AST values to avoid reallocating slices
1 parent 2e9310b commit 3368b8b

File tree

8 files changed

+177
-68
lines changed

8 files changed

+177
-68
lines changed

‎coderd/database/dbauthz/dbauthz.go

+5-3
Original file line numberDiff line numberDiff line change
@@ -144,7 +144,8 @@ var (
144144
},
145145
}),
146146
Scope: rbac.ScopeAll,
147-
}
147+
}.WithCachedASTValue()
148+
148149
subjectAutostart = rbac.Subject{
149150
ID: uuid.Nil.String(),
150151
Roles: rbac.Roles([]rbac.Role{
@@ -161,7 +162,8 @@ var (
161162
},
162163
}),
163164
Scope: rbac.ScopeAll,
164-
}
165+
}.WithCachedASTValue()
166+
165167
subjectSystemRestricted = rbac.Subject{
166168
ID: uuid.Nil.String(),
167169
Roles: rbac.Roles([]rbac.Role{
@@ -188,7 +190,7 @@ var (
188190
},
189191
}),
190192
Scope: rbac.ScopeAll,
191-
}
193+
}.WithCachedASTValue()
192194
)
193195

194196
// AsProvisionerd returns a context with an actor that has permissions required

‎coderd/httpmw/apikey.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -379,7 +379,7 @@ func ExtractAPIKey(rw http.ResponseWriter, r *http.Request, cfg ExtractAPIKeyCon
379379
Roles: rbac.RoleNames(roles.Roles),
380380
Groups: roles.Groups,
381381
Scope: rbac.ScopeName(key.Scope),
382-
},
382+
}.WithCachedASTValue(),
383383
}
384384

385385
return &key, &authz, true

‎coderd/httpmw/workspaceagent.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -110,5 +110,5 @@ func getAgentSubject(ctx context.Context, db database.Store, agent database.Work
110110
Roles: rbac.RoleNames(roles.Roles),
111111
Groups: roles.Groups,
112112
Scope: rbac.WorkspaceAgentScope(workspace.ID, user.ID),
113-
}, nil
113+
}.WithCachedASTValue(), nil
114114
}

‎coderd/rbac/astvalue.go

+17
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,10 @@ func regoPartialInputValue(subject Subject, action Action, objectType string) (a
6565

6666
// regoValue returns the ast.Object representation of the subject.
6767
func (s Subject) regoValue() (ast.Value, error) {
68+
if s.cachedASTValue != nil {
69+
return s.cachedASTValue, nil
70+
}
71+
6872
subjRoles, err := s.Roles.Expand()
6973
if err != nil {
7074
return nil, xerrors.Errorf("expand roles: %w", err)
@@ -133,7 +137,20 @@ func (z Object) regoValue() ast.Value {
133137
)
134138
}
135139

140+
// withCachedRegoValue returns a copy of the role with the cachedRegoValue.
141+
// It does not mutate the underlying role.
142+
// Avoid using this function if possible, it should only be used if the
143+
// caller can guarantee the role is static and will never change.
144+
func (role Role) withCachedRegoValue() Role {
145+
tmp := role
146+
tmp.cachedRegoValue = role.regoValue()
147+
return tmp
148+
}
149+
136150
func (role Role) regoValue() ast.Value {
151+
if role.cachedRegoValue != nil {
152+
return role.cachedRegoValue
153+
}
137154
orgMap := ast.NewObject()
138155
for k, p := range role.Org {
139156
orgMap.Insert(ast.StringTerm(k), ast.NewTerm(regoSlice(p)))

‎coderd/rbac/authz.go

+14
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,20 @@ type Subject struct {
4949
Roles ExpandableRoles
5050
Groups []string
5151
Scope ExpandableScope
52+
53+
// cachedASTValue is the cached ast value for this subject.
54+
cachedASTValue ast.Value
55+
}
56+
57+
// WithCachedASTValue can be called if the subject is static. This will compute
58+
// the ast value once and cache it for future calls.
59+
func (s Subject) WithCachedASTValue() Subject {
60+
tmp := s
61+
v, err := tmp.regoValue()
62+
if err == nil {
63+
tmp.cachedASTValue = v
64+
}
65+
return tmp
5266
}
5367

5468
func (s Subject) Equal(b Subject) bool {

‎coderd/rbac/authz_test.go

+47-1
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,21 @@ func benchmarkUserCases() (cases []benchmarkCase, users uuid.UUID, orgs []uuid.U
8686
Groups: noiseGroups,
8787
},
8888
},
89+
{
90+
Name: "ManyRolesCachedSubject",
91+
Actor: rbac.Subject{
92+
// Admin of many orgs
93+
Roles: rbac.RoleNames{
94+
rbac.RoleOrgMember(orgs[0]), rbac.RoleOrgAdmin(orgs[0]),
95+
rbac.RoleOrgMember(orgs[1]), rbac.RoleOrgAdmin(orgs[1]),
96+
rbac.RoleOrgMember(orgs[2]), rbac.RoleOrgAdmin(orgs[2]),
97+
rbac.RoleMember(),
98+
},
99+
ID: user.String(),
100+
Scope: rbac.ScopeAll,
101+
Groups: noiseGroups,
102+
}.WithCachedASTValue(),
103+
},
89104
{
90105
Name: "AdminWithScope",
91106
Actor: rbac.Subject{
@@ -96,13 +111,41 @@ func benchmarkUserCases() (cases []benchmarkCase, users uuid.UUID, orgs []uuid.U
96111
Groups: noiseGroups,
97112
},
98113
},
114+
{
115+
// This test should only use static roles. AKA no org roles.
116+
Name: "StaticRoles",
117+
Actor: rbac.Subject{
118+
// Give some extra roles that an admin might have
119+
Roles: rbac.RoleNames{
120+
"auditor", rbac.RoleOwner(), rbac.RoleMember(),
121+
rbac.RoleTemplateAdmin(), rbac.RoleUserAdmin(),
122+
},
123+
ID: user.String(),
124+
Scope: rbac.ScopeAll,
125+
Groups: noiseGroups,
126+
},
127+
},
128+
{
129+
// This test should only use static roles. AKA no org roles.
130+
Name: "StaticRolesWithCache",
131+
Actor: rbac.Subject{
132+
// Give some extra roles that an admin might have
133+
Roles: rbac.RoleNames{
134+
"auditor", rbac.RoleOwner(), rbac.RoleMember(),
135+
rbac.RoleTemplateAdmin(), rbac.RoleUserAdmin(),
136+
},
137+
ID: user.String(),
138+
Scope: rbac.ScopeAll,
139+
Groups: noiseGroups,
140+
}.WithCachedASTValue(),
141+
},
99142
}
100143
return benchCases, users, orgs
101144
}
102145

103146
// BenchmarkRBACAuthorize benchmarks the rbac.Authorize method.
104147
//
105-
// go test -bench BenchmarkRBACAuthorize -benchmem -memprofile memprofile.out -cpuprofile profile.out
148+
// go test -run=^$ -bench BenchmarkRBACAuthorize -benchmem -memprofile memprofile.out -cpuprofile profile.out
106149
func BenchmarkRBACAuthorize(b *testing.B) {
107150
benchCases, user, orgs := benchmarkUserCases()
108151
users := append([]uuid.UUID{},
@@ -111,6 +154,9 @@ func BenchmarkRBACAuthorize(b *testing.B) {
111154
uuid.MustParse("0632b012-49e0-4d70-a5b3-f4398f1dcd52"),
112155
uuid.MustParse("70dbaa7a-ea9c-4f68-a781-97b08af8461d"),
113156
)
157+
158+
// There is no caching that occurs because a fresh context is used for each
159+
// call. And the context needs 'WithCacheCtx' to work.
114160
authorizer := rbac.NewCachingAuthorizer(prometheus.NewRegistry())
115161
// This benchmarks all the simple cases using just user permissions. Groups
116162
// are added as noise, but do not do anything.

‎coderd/rbac/roles.go

+80-61
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"strings"
66

77
"github.com/google/uuid"
8+
"github.com/open-policy-agent/opa/ast"
89

910
"golang.org/x/xerrors"
1011
)
@@ -128,86 +129,100 @@ func ReloadBuiltinRoles(opts *RoleOptions) {
128129
)
129130
}
130131

132+
// Static roles that never change should be allocated in a closure.
133+
// This is to ensure these data structures are only allocated once and not
134+
// on every authorize call. 'withCachedRegoValue' can be used as well to
135+
// preallocate the rego value that is used by the rego eval engine.
136+
ownerRole := Role{
137+
Name: owner,
138+
DisplayName: "Owner",
139+
Site: allPermsExcept(ownerAndAdminExceptions...),
140+
Org: map[string][]Permission{},
141+
User: []Permission{},
142+
}.withCachedRegoValue()
143+
144+
memberRole := Role{
145+
Name: member,
146+
DisplayName: "",
147+
Site: Permissions(map[string][]Action{
148+
// All users can read all other users and know they exist.
149+
ResourceUser.Type: {ActionRead},
150+
ResourceRoleAssignment.Type: {ActionRead},
151+
// All users can see the provisioner daemons.
152+
ResourceProvisionerDaemon.Type: {ActionRead},
153+
}),
154+
Org: map[string][]Permission{},
155+
User: allPermsExcept(),
156+
}.withCachedRegoValue()
157+
158+
auditorRole := Role{
159+
Name: auditor,
160+
DisplayName: "Auditor",
161+
Site: Permissions(map[string][]Action{
162+
// Should be able to read all template details, even in orgs they
163+
// are not in.
164+
ResourceTemplate.Type: {ActionRead},
165+
ResourceAuditLog.Type: {ActionRead},
166+
}),
167+
Org: map[string][]Permission{},
168+
User: []Permission{},
169+
}.withCachedRegoValue()
170+
171+
templateAdminRole := Role{
172+
Name: templateAdmin,
173+
DisplayName: "Template Admin",
174+
Site: Permissions(map[string][]Action{
175+
ResourceTemplate.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
176+
// CRUD all files, even those they did not upload.
177+
ResourceFile.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
178+
ResourceWorkspace.Type: {ActionRead},
179+
// CRUD to provisioner daemons for now.
180+
ResourceProvisionerDaemon.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
181+
// Needs to read all organizations since
182+
ResourceOrganization.Type: {ActionRead},
183+
}),
184+
Org: map[string][]Permission{},
185+
User: []Permission{},
186+
}.withCachedRegoValue()
187+
188+
userAdminRole := Role{
189+
Name: userAdmin,
190+
DisplayName: "User Admin",
191+
Site: Permissions(map[string][]Action{
192+
ResourceRoleAssignment.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
193+
ResourceUser.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
194+
// Full perms to manage org members
195+
ResourceOrganizationMember.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
196+
ResourceGroup.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
197+
}),
198+
Org: map[string][]Permission{},
199+
User: []Permission{},
200+
}.withCachedRegoValue()
201+
131202
builtInRoles = map[string]func(orgID string) Role{
132203
// admin grants all actions to all resources.
133204
owner: func(_ string) Role {
134-
return Role{
135-
Name: owner,
136-
DisplayName: "Owner",
137-
Site: allPermsExcept(ownerAndAdminExceptions...),
138-
Org: map[string][]Permission{},
139-
User: []Permission{},
140-
}
205+
return ownerRole
141206
},
142207

143208
// member grants all actions to all resources owned by the user
144209
member: func(_ string) Role {
145-
return Role{
146-
Name: member,
147-
DisplayName: "",
148-
Site: Permissions(map[string][]Action{
149-
// All users can read all other users and know they exist.
150-
ResourceUser.Type: {ActionRead},
151-
ResourceRoleAssignment.Type: {ActionRead},
152-
// All users can see the provisioner daemons.
153-
ResourceProvisionerDaemon.Type: {ActionRead},
154-
}),
155-
Org: map[string][]Permission{},
156-
User: allPermsExcept(),
157-
}
210+
return memberRole
158211
},
159212

160213
// auditor provides all permissions required to effectively read and understand
161214
// audit log events.
162215
// TODO: Finish the auditor as we add resources.
163216
auditor: func(_ string) Role {
164-
return Role{
165-
Name: auditor,
166-
DisplayName: "Auditor",
167-
Site: Permissions(map[string][]Action{
168-
// Should be able to read all template details, even in orgs they
169-
// are not in.
170-
ResourceTemplate.Type: {ActionRead},
171-
ResourceAuditLog.Type: {ActionRead},
172-
}),
173-
Org: map[string][]Permission{},
174-
User: []Permission{},
175-
}
217+
return auditorRole
176218
},
177219

178220
templateAdmin: func(_ string) Role {
179-
return Role{
180-
Name: templateAdmin,
181-
DisplayName: "Template Admin",
182-
Site: Permissions(map[string][]Action{
183-
ResourceTemplate.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
184-
// CRUD all files, even those they did not upload.
185-
ResourceFile.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
186-
ResourceWorkspace.Type: {ActionRead},
187-
// CRUD to provisioner daemons for now.
188-
ResourceProvisionerDaemon.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
189-
// Needs to read all organizations since
190-
ResourceOrganization.Type: {ActionRead},
191-
}),
192-
Org: map[string][]Permission{},
193-
User: []Permission{},
194-
}
221+
return templateAdminRole
195222
},
196223

197224
userAdmin: func(_ string) Role {
198-
return Role{
199-
Name: userAdmin,
200-
DisplayName: "User Admin",
201-
Site: Permissions(map[string][]Action{
202-
ResourceRoleAssignment.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
203-
ResourceUser.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
204-
// Full perms to manage org members
205-
ResourceOrganizationMember.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
206-
ResourceGroup.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
207-
}),
208-
Org: map[string][]Permission{},
209-
User: []Permission{},
210-
}
225+
return userAdminRole
211226
},
212227

213228
// orgAdmin returns a role with all actions allows in a given
@@ -333,6 +348,10 @@ type Role struct {
333348
// roles.
334349
Org map[string][]Permission `json:"org"`
335350
User []Permission `json:"user"`
351+
352+
// cachedRegoValue can be used to cache the rego value for this role.
353+
// This is helpful for static roles that never change.
354+
cachedRegoValue ast.Value
336355
}
337356

338357
type Roles []Role

‎coderd/rbac/roles_internal_test.go

+12-1
Original file line numberDiff line numberDiff line change
@@ -68,8 +68,19 @@ func BenchmarkRBACValueAllocation(b *testing.B) {
6868
func TestRegoInputValue(t *testing.T) {
6969
t.Parallel()
7070

71+
// Expand all roles and make sure we have a good copy.
72+
// This is because these tests modify the roles, and we don't want to
73+
// modify the original roles.
74+
roles, err := RoleNames{RoleOrgMember(uuid.New()), RoleOrgAdmin(uuid.New()), RoleMember()}.Expand()
75+
require.NoError(t, err, "failed to expand roles")
76+
for i := range roles {
77+
// If all cached values are nil, then the role will not use
78+
// the shared cached value.
79+
roles[i].cachedRegoValue = nil
80+
}
81+
7182
actor := Subject{
72-
Roles: RoleNames{RoleOrgMember(uuid.New()), RoleOrgAdmin(uuid.New()), RoleMember()},
83+
Roles: Roles(roles),
7384
ID: uuid.NewString(),
7485
Scope: ScopeAll,
7586
Groups: []string{uuid.NewString(), uuid.NewString(), uuid.NewString()},

0 commit comments

Comments
 (0)