Skip to content

Commit af59e2b

Browse files
EmyrkEmyrk
authored
chore: Optimize rego policy input allocations (#6135)
* chore: Optimize rego policy evaluation allocations Manually convert to ast.Value instead of using generic json.Marshal conversion. * Add a unit test that prevents regressions of rego input The optimized input is always compared to the normal json marshal parser.
1 parent 22f6400 commit af59e2b

File tree

8 files changed

+466
-58
lines changed

8 files changed

+466
-58
lines changed

‎coderd/coderdtest/authorize.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -445,7 +445,7 @@ func NewAuthTester(ctx context.Context, t *testing.T, client *codersdk.Client, a
445445
func (a *AuthTester) Test(ctx context.Context, assertRoute map[string]RouteCheck, skipRoutes map[string]string) {
446446
// Always fail auth from this point forward
447447
a.authorizer.Wrapped = &FakeAuthorizer{
448-
AlwaysReturn: rbac.ForbiddenWithInternal(xerrors.New("fake implementation"), nil, nil),
448+
AlwaysReturn: rbac.ForbiddenWithInternal(xerrors.New("fake implementation"), rbac.Subject{}, "", rbac.Object{}, nil),
449449
}
450450

451451
routeMissing := make(map[string]bool)

‎coderd/rbac/astvalue.go

Lines changed: 210 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,210 @@
1+
package rbac
2+
3+
import (
4+
"github.com/open-policy-agent/opa/ast"
5+
"golang.org/x/xerrors"
6+
)
7+
8+
// regoInputValue returns a rego input value for the given subject, action, and
9+
// object. This rego input is already parsed and can be used directly in a
10+
// rego query.
11+
func regoInputValue(subject Subject, action Action, object Object) (ast.Value, error) {
12+
regoSubj, err := subject.regoValue()
13+
if err != nil {
14+
return nil, xerrors.Errorf("subject: %w", err)
15+
}
16+
17+
s := [2]*ast.Term{
18+
ast.StringTerm("subject"),
19+
ast.NewTerm(regoSubj),
20+
}
21+
a := [2]*ast.Term{
22+
ast.StringTerm("action"),
23+
ast.StringTerm(string(action)),
24+
}
25+
o := [2]*ast.Term{
26+
ast.StringTerm("object"),
27+
ast.NewTerm(object.regoValue()),
28+
}
29+
30+
input := ast.NewObject(s, a, o)
31+
32+
return input, nil
33+
}
34+
35+
// regoPartialInputValue is the same as regoInputValue but only includes the
36+
// object type. This is for partial evaluations.
37+
func regoPartialInputValue(subject Subject, action Action, objectType string) (ast.Value, error) {
38+
regoSubj, err := subject.regoValue()
39+
if err != nil {
40+
return nil, xerrors.Errorf("subject: %w", err)
41+
}
42+
43+
s := [2]*ast.Term{
44+
ast.StringTerm("subject"),
45+
ast.NewTerm(regoSubj),
46+
}
47+
a := [2]*ast.Term{
48+
ast.StringTerm("action"),
49+
ast.StringTerm(string(action)),
50+
}
51+
o := [2]*ast.Term{
52+
ast.StringTerm("object"),
53+
ast.NewTerm(ast.NewObject(
54+
[2]*ast.Term{
55+
ast.StringTerm("type"),
56+
ast.StringTerm(objectType),
57+
}),
58+
),
59+
}
60+
61+
input := ast.NewObject(s, a, o)
62+
63+
return input, nil
64+
}
65+
66+
// regoValue returns the ast.Object representation of the subject.
67+
func (s Subject) regoValue() (ast.Value, error) {
68+
subjRoles, err := s.Roles.Expand()
69+
if err != nil {
70+
return nil, xerrors.Errorf("expand roles: %w", err)
71+
}
72+
73+
subjScope, err := s.Scope.Expand()
74+
if err != nil {
75+
return nil, xerrors.Errorf("expand scope: %w", err)
76+
}
77+
subj := ast.NewObject(
78+
[2]*ast.Term{
79+
ast.StringTerm("id"),
80+
ast.StringTerm(s.ID),
81+
},
82+
[2]*ast.Term{
83+
ast.StringTerm("roles"),
84+
ast.NewTerm(regoSlice(subjRoles)),
85+
},
86+
[2]*ast.Term{
87+
ast.StringTerm("scope"),
88+
ast.NewTerm(subjScope.regoValue()),
89+
},
90+
[2]*ast.Term{
91+
ast.StringTerm("groups"),
92+
ast.NewTerm(regoSliceString(s.Groups...)),
93+
},
94+
)
95+
96+
return subj, nil
97+
}
98+
99+
func (z Object) regoValue() ast.Value {
100+
userACL := ast.NewObject()
101+
for k, v := range z.ACLUserList {
102+
userACL.Insert(ast.StringTerm(k), ast.NewTerm(regoSlice(v)))
103+
}
104+
grpACL := ast.NewObject()
105+
for k, v := range z.ACLGroupList {
106+
grpACL.Insert(ast.StringTerm(k), ast.NewTerm(regoSlice(v)))
107+
}
108+
return ast.NewObject(
109+
[2]*ast.Term{
110+
ast.StringTerm("id"),
111+
ast.StringTerm(z.ID),
112+
},
113+
[2]*ast.Term{
114+
ast.StringTerm("owner"),
115+
ast.StringTerm(z.Owner),
116+
},
117+
[2]*ast.Term{
118+
ast.StringTerm("org_owner"),
119+
ast.StringTerm(z.OrgID),
120+
},
121+
[2]*ast.Term{
122+
ast.StringTerm("type"),
123+
ast.StringTerm(z.Type),
124+
},
125+
[2]*ast.Term{
126+
ast.StringTerm("acl_user_list"),
127+
ast.NewTerm(userACL),
128+
},
129+
[2]*ast.Term{
130+
ast.StringTerm("acl_group_list"),
131+
ast.NewTerm(grpACL),
132+
},
133+
)
134+
}
135+
136+
func (role Role) regoValue() ast.Value {
137+
orgMap := ast.NewObject()
138+
for k, p := range role.Org {
139+
orgMap.Insert(ast.StringTerm(k), ast.NewTerm(regoSlice(p)))
140+
}
141+
return ast.NewObject(
142+
[2]*ast.Term{
143+
ast.StringTerm("site"),
144+
ast.NewTerm(regoSlice(role.Site)),
145+
},
146+
[2]*ast.Term{
147+
ast.StringTerm("org"),
148+
ast.NewTerm(orgMap),
149+
},
150+
[2]*ast.Term{
151+
ast.StringTerm("user"),
152+
ast.NewTerm(regoSlice(role.User)),
153+
},
154+
)
155+
}
156+
157+
func (s Scope) regoValue() ast.Value {
158+
r, ok := s.Role.regoValue().(ast.Object)
159+
if !ok {
160+
panic("developer error: role is not an object")
161+
}
162+
r.Insert(
163+
ast.StringTerm("allow_list"),
164+
ast.NewTerm(regoSliceString(s.AllowIDList...)),
165+
)
166+
return r
167+
}
168+
169+
func (perm Permission) regoValue() ast.Value {
170+
return ast.NewObject(
171+
[2]*ast.Term{
172+
ast.StringTerm("negate"),
173+
ast.BooleanTerm(perm.Negate),
174+
},
175+
[2]*ast.Term{
176+
ast.StringTerm("resource_type"),
177+
ast.StringTerm(perm.ResourceType),
178+
},
179+
[2]*ast.Term{
180+
ast.StringTerm("action"),
181+
ast.StringTerm(string(perm.Action)),
182+
},
183+
)
184+
}
185+
186+
func (act Action) regoValue() ast.Value {
187+
return ast.StringTerm(string(act)).Value
188+
}
189+
190+
type regoValue interface {
191+
regoValue() ast.Value
192+
}
193+
194+
// regoSlice returns the ast.Array representation of the slice.
195+
// The slice must contain only types that implement the regoValue interface.
196+
func regoSlice[T regoValue](slice []T) *ast.Array {
197+
terms := make([]*ast.Term, len(slice))
198+
for i, v := range slice {
199+
terms[i] = ast.NewTerm(v.regoValue())
200+
}
201+
return ast.NewArray(terms...)
202+
}
203+
204+
func regoSliceString(slice ...string) *ast.Array {
205+
terms := make([]*ast.Term, len(slice))
206+
for i, v := range slice {
207+
terms[i] = ast.StringTerm(v)
208+
}
209+
return ast.NewArray(terms...)
210+
}

‎coderd/rbac/authz.go

Lines changed: 5 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -263,34 +263,18 @@ func (a RegoAuthorizer) authorize(ctx context.Context, subject Subject, action A
263263
return xerrors.Errorf("subject must have a scope")
264264
}
265265

266-
subjRoles, err := subject.Roles.Expand()
266+
astV, err := regoInputValue(subject, action, object)
267267
if err != nil {
268-
return xerrors.Errorf("expand roles: %w", err)
268+
return xerrors.Errorf("convert input to value: %w", err)
269269
}
270270

271-
subjScope, err := subject.Scope.Expand()
271+
results, err := a.query.Eval(ctx, rego.EvalParsedInput(astV))
272272
if err != nil {
273-
return xerrors.Errorf("expand scope: %w", err)
274-
}
275-
276-
input := map[string]interface{}{
277-
"subject": authSubject{
278-
ID: subject.ID,
279-
Roles: subjRoles,
280-
Groups: subject.Groups,
281-
Scope: subjScope,
282-
},
283-
"object": object,
284-
"action": action,
285-
}
286-
287-
results, err := a.query.Eval(ctx, rego.EvalInput(input))
288-
if err != nil {
289-
return ForbiddenWithInternal(xerrors.Errorf("eval rego: %w", err), input, results)
273+
return ForbiddenWithInternal(xerrors.Errorf("eval rego: %w", err), subject, action, object, results)
290274
}
291275

292276
if !results.Allowed() {
293-
return ForbiddenWithInternal(xerrors.Errorf("policy disallows request"), input, results)
277+
return ForbiddenWithInternal(xerrors.Errorf("policy disallows request"), subject, action, object, results)
294278
}
295279
return nil
296280
}

‎coderd/rbac/builtin.go

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package rbac
22

33
import (
4+
"sort"
45
"strings"
56

67
"github.com/google/uuid"
@@ -79,6 +80,8 @@ var (
7980
Site: permissions(map[string][]Action{
8081
ResourceWildcard.Type: {WildcardSymbol},
8182
}),
83+
Org: map[string][]Permission{},
84+
User: []Permission{},
8285
}
8386
},
8487

@@ -94,6 +97,7 @@ var (
9497
// All users can see the provisioner daemons.
9598
ResourceProvisionerDaemon.Type: {ActionRead},
9699
}),
100+
Org: map[string][]Permission{},
97101
User: permissions(map[string][]Action{
98102
ResourceWildcard.Type: {WildcardSymbol},
99103
}),
@@ -113,6 +117,8 @@ var (
113117
ResourceTemplate.Type: {ActionRead},
114118
ResourceAuditLog.Type: {ActionRead},
115119
}),
120+
Org: map[string][]Permission{},
121+
User: []Permission{},
116122
}
117123
},
118124

@@ -128,6 +134,8 @@ var (
128134
// CRUD to provisioner daemons for now.
129135
ResourceProvisionerDaemon.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
130136
}),
137+
Org: map[string][]Permission{},
138+
User: []Permission{},
131139
}
132140
},
133141

@@ -142,6 +150,8 @@ var (
142150
ResourceOrganizationMember.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
143151
ResourceGroup.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
144152
}),
153+
Org: map[string][]Permission{},
154+
User: []Permission{},
145155
}
146156
},
147157

@@ -151,6 +161,7 @@ var (
151161
return Role{
152162
Name: roleName(orgAdmin, organizationID),
153163
DisplayName: "Organization Admin",
164+
Site: []Permission{},
154165
Org: map[string][]Permission{
155166
organizationID: {
156167
{
@@ -160,6 +171,7 @@ var (
160171
},
161172
},
162173
},
174+
User: []Permission{},
163175
}
164176
},
165177

@@ -169,6 +181,7 @@ var (
169181
return Role{
170182
Name: roleName(orgMember, organizationID),
171183
DisplayName: "",
184+
Site: []Permission{},
172185
Org: map[string][]Permission{
173186
organizationID: {
174187
{
@@ -192,6 +205,7 @@ var (
192205
},
193206
},
194207
},
208+
User: []Permission{},
195209
}
196210
},
197211
}
@@ -422,5 +436,9 @@ func permissions(perms map[string][]Action) []Permission {
422436
})
423437
}
424438
}
439+
// Deterministic ordering of permissions
440+
sort.Slice(list, func(i, j int) bool {
441+
return list[i].ResourceType < list[j].ResourceType
442+
})
425443
return list
426444
}

0 commit comments

Comments
 (0)