Skip to content

Commit 22c6f50

Browse files
EmyrkEmyrk
committed
chore: Optimize rego policy evaluation allocations
Manually convert to ast.Value instead of using generic json.Marshal conversion.
1 parent 049984c commit 22c6f50

File tree

4 files changed

+299
-16
lines changed

4 files changed

+299
-16
lines changed

‎coderd/rbac/astvalue.go

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
1+
package rbac
2+
3+
import (
4+
"github.com/open-policy-agent/opa/ast"
5+
"golang.org/x/xerrors"
6+
)
7+
8+
func regoInputValue(subject Subject, action Action, object Object) (ast.Value, error) {
9+
regoSubj, err := subject.regoValue()
10+
if err != nil {
11+
return nil, xerrors.Errorf("subject: %w", err)
12+
}
13+
14+
s := [2]*ast.Term{
15+
ast.StringTerm("subject"),
16+
ast.NewTerm(regoSubj),
17+
}
18+
a := [2]*ast.Term{
19+
ast.StringTerm("action"),
20+
ast.StringTerm(string(action)),
21+
}
22+
o := [2]*ast.Term{
23+
ast.StringTerm("object"),
24+
ast.NewTerm(object.regoValue()),
25+
}
26+
27+
input := ast.NewObject(s, a, o)
28+
29+
return input, nil
30+
}
31+
32+
func (s Subject) regoValue() (ast.Value, error) {
33+
subjRoles, err := s.Roles.Expand()
34+
if err != nil {
35+
return nil, xerrors.Errorf("expand roles: %w", err)
36+
}
37+
38+
subjScope, err := s.Scope.Expand()
39+
if err != nil {
40+
return nil, xerrors.Errorf("expand scope: %w", err)
41+
}
42+
subj := ast.NewObject(
43+
[2]*ast.Term{
44+
ast.StringTerm("id"),
45+
ast.StringTerm(s.ID),
46+
},
47+
[2]*ast.Term{
48+
ast.StringTerm("roles"),
49+
ast.NewTerm(regoSlice(subjRoles)),
50+
},
51+
[2]*ast.Term{
52+
ast.StringTerm("scope"),
53+
ast.NewTerm(subjScope.regoValue()),
54+
},
55+
[2]*ast.Term{
56+
ast.StringTerm("groups"),
57+
ast.NewTerm(regoSliceString(s.Groups...)),
58+
},
59+
)
60+
61+
return subj, nil
62+
}
63+
64+
func (z Object) regoValue() ast.Value {
65+
userACL := ast.NewObject()
66+
for k, v := range z.ACLUserList {
67+
userACL.Insert(ast.StringTerm(k), ast.NewTerm(regoSlice(v)))
68+
}
69+
grpACL := ast.NewObject()
70+
for k, v := range z.ACLGroupList {
71+
grpACL.Insert(ast.StringTerm(k), ast.NewTerm(regoSlice(v)))
72+
}
73+
return ast.NewObject(
74+
[2]*ast.Term{
75+
ast.StringTerm("id"),
76+
ast.StringTerm(z.ID),
77+
},
78+
[2]*ast.Term{
79+
ast.StringTerm("owner"),
80+
ast.StringTerm(z.Owner),
81+
},
82+
[2]*ast.Term{
83+
ast.StringTerm("org_owner"),
84+
ast.StringTerm(z.OrgID),
85+
},
86+
[2]*ast.Term{
87+
ast.StringTerm("type"),
88+
ast.StringTerm(z.Type),
89+
},
90+
[2]*ast.Term{
91+
ast.StringTerm("acl_user_list"),
92+
ast.NewTerm(userACL),
93+
},
94+
[2]*ast.Term{
95+
ast.StringTerm("acl_group_list"),
96+
ast.NewTerm(grpACL),
97+
},
98+
)
99+
}
100+
101+
func (role Role) regoValue() ast.Value {
102+
orgMap := ast.NewObject()
103+
for k, p := range role.Org {
104+
orgMap.Insert(ast.StringTerm(k), ast.NewTerm(regoSlice(p)))
105+
}
106+
return ast.NewObject(
107+
[2]*ast.Term{
108+
ast.StringTerm("site"),
109+
ast.NewTerm(regoSlice(role.Site)),
110+
},
111+
[2]*ast.Term{
112+
ast.StringTerm("org"),
113+
ast.NewTerm(orgMap),
114+
},
115+
[2]*ast.Term{
116+
ast.StringTerm("user"),
117+
ast.NewTerm(regoSlice(role.User)),
118+
},
119+
)
120+
}
121+
122+
func (s Scope) regoValue() ast.Value {
123+
r := s.Role.regoValue().(ast.Object)
124+
r.Insert(
125+
ast.StringTerm("allow_list"),
126+
ast.NewTerm(regoSliceString(s.AllowIDList...)),
127+
)
128+
return r
129+
}
130+
131+
func (perm Permission) regoValue() ast.Value {
132+
return ast.NewObject(
133+
[2]*ast.Term{
134+
ast.StringTerm("negate"),
135+
ast.BooleanTerm(perm.Negate),
136+
},
137+
[2]*ast.Term{
138+
ast.StringTerm("resource_type"),
139+
ast.StringTerm(perm.ResourceType),
140+
},
141+
[2]*ast.Term{
142+
ast.StringTerm("action"),
143+
ast.StringTerm(string(perm.Action)),
144+
},
145+
)
146+
}
147+
148+
func (act Action) regoValue() ast.Value {
149+
return ast.StringTerm(string(act)).Value
150+
}
151+
152+
type regoValue interface {
153+
regoValue() ast.Value
154+
}
155+
156+
func regoSlice[T regoValue](slice []T) *ast.Array {
157+
terms := make([]*ast.Term, len(slice))
158+
for i, v := range slice {
159+
terms[i] = ast.NewTerm(v.regoValue())
160+
}
161+
return ast.NewArray(terms...)
162+
}
163+
164+
func regoSliceString(slice ...string) *ast.Array {
165+
terms := make([]*ast.Term, len(slice))
166+
for i, v := range slice {
167+
terms[i] = ast.StringTerm(v)
168+
}
169+
return ast.NewArray(terms...)
170+
}

‎coderd/rbac/authz.go

Lines changed: 70 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"sync"
77
"time"
88

9+
"github.com/open-policy-agent/opa/ast"
910
"github.com/open-policy-agent/opa/rego"
1011
"github.com/prometheus/client_golang/prometheus"
1112
"github.com/prometheus/client_golang/prometheus/promauto"
@@ -251,6 +252,43 @@ func (a RegoAuthorizer) Authorize(ctx context.Context, subject Subject, action A
251252
return nil
252253
}
253254

255+
type inputType struct {
256+
Subject authSubject `json:"subject"`
257+
Action Action `json:"action"`
258+
Object Object `json:"object"`
259+
}
260+
261+
func (i inputType) Value() (ast.Value, error) {
262+
s := [2]*ast.Term{
263+
ast.StringTerm("subject"),
264+
ast.NewTerm(ast.NewObject([2]*ast.Term{
265+
// ID
266+
ast.StringTerm("id"),
267+
ast.NewTerm(ast.NewObject([2]*ast.Term{
268+
ast.StringTerm("id"),
269+
ast.StringTerm(i.Subject.ID),
270+
})),
271+
////
272+
//ast.NewTerm(ast.NewObject([2]*ast.Term{
273+
// ast.StringTerm("groups"),
274+
// ast.NewTerm(ast.NewArray(i.Subject.Groups)),
275+
//})),
276+
})),
277+
}
278+
o := [2]*ast.Term{
279+
ast.StringTerm("subject"),
280+
ast.StringTerm(i.Subject.ID),
281+
}
282+
a := [2]*ast.Term{
283+
ast.StringTerm("action"),
284+
ast.StringTerm(string(i.Action)),
285+
}
286+
287+
input := ast.NewObject(s, a, o)
288+
289+
return input, nil
290+
}
291+
254292
// authorize is the internal function that does the actual authorization.
255293
// It is a different function so the exported one can add tracing + metrics.
256294
// That code tends to clutter up the actual logic, so it's separated out.
@@ -272,25 +310,44 @@ func (a RegoAuthorizer) authorize(ctx context.Context, subject Subject, action A
272310
if err != nil {
273311
return xerrors.Errorf("expand scope: %w", err)
274312
}
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,
313+
var _, _ = subjRoles, subjScope
314+
315+
//input := inputType{
316+
// Subject: authSubject{
317+
// ID: subject.ID,
318+
// Roles: subjRoles,
319+
// Groups: subject.Groups,
320+
// Scope: subjScope,
321+
// },
322+
// Action: action,
323+
// Object: object,
324+
//}
325+
326+
//jinput := map[string]interface{}{
327+
// "subject": authSubject{
328+
// ID: subject.ID,
329+
// Roles: subjRoles,
330+
// Groups: subject.Groups,
331+
// Scope: subjScope,
332+
// },
333+
// "object": object,
334+
// "action": action,
335+
//}
336+
//var _ = jinput
337+
338+
astV, err := regoInputValue(subject, action, object)
339+
if err != nil {
340+
return xerrors.Errorf("convert input to value: %w", err)
285341
}
342+
var _ = astV
286343

287-
results, err := a.query.Eval(ctx, rego.EvalInput(input))
344+
results, err := a.query.Eval(ctx, rego.EvalParsedInput(astV))
288345
if err != nil {
289-
return ForbiddenWithInternal(xerrors.Errorf("eval rego: %w", err), input, results)
346+
return ForbiddenWithInternal(xerrors.Errorf("eval rego: %w", err), nil, results)
290347
}
291348

292349
if !results.Allowed() {
293-
return ForbiddenWithInternal(xerrors.Errorf("policy disallows request"), input, results)
350+
return ForbiddenWithInternal(xerrors.Errorf("policy disallows request"), nil, results)
294351
}
295352
return nil
296353
}

‎coderd/rbac/builtin_internal_test.go

Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,53 @@ package rbac
33
import (
44
"testing"
55

6-
"github.com/stretchr/testify/require"
7-
86
"github.com/google/uuid"
7+
"github.com/open-policy-agent/opa/ast"
8+
"github.com/stretchr/testify/require"
99
)
1010

11+
func BenchmarkRBACValueAllocation(b *testing.B) {
12+
actor := Subject{
13+
Roles: RoleNames{RoleOrgMember(uuid.New()), RoleOrgAdmin(uuid.New()), RoleMember()},
14+
ID: uuid.NewString(),
15+
Scope: ScopeAll,
16+
Groups: []string{uuid.NewString(), uuid.NewString(), uuid.NewString()},
17+
}
18+
obj := ResourceTemplate.
19+
WithID(uuid.New()).
20+
InOrg(uuid.New()).
21+
WithOwner(uuid.NewString()).
22+
WithGroupACL(map[string][]Action{
23+
uuid.NewString(): {ActionRead, ActionCreate},
24+
uuid.NewString(): {ActionRead, ActionCreate},
25+
uuid.NewString(): {ActionRead, ActionCreate},
26+
}).WithACLUserList(map[string][]Action{
27+
uuid.NewString(): {ActionRead, ActionCreate},
28+
uuid.NewString(): {ActionRead, ActionCreate},
29+
})
30+
31+
jsonSubject := authSubject{
32+
ID: actor.ID,
33+
Roles: must(actor.Roles.Expand()),
34+
Groups: actor.Groups,
35+
Scope: must(actor.Scope.Expand()),
36+
}
37+
38+
b.Run("ManualRegoValue", func(b *testing.B) {
39+
for i := 0; i < b.N; i++ {
40+
_, err := regoInputValue(actor, ActionRead, obj)
41+
require.NoError(b, err)
42+
}
43+
})
44+
b.Run("JSONRegoValue", func(b *testing.B) {
45+
for i := 0; i < b.N; i++ {
46+
_, err := ast.InterfaceToValue(jsonSubject)
47+
require.NoError(b, err)
48+
}
49+
})
50+
51+
}
52+
1153
func TestRoleByName(t *testing.T) {
1254
t.Parallel()
1355

‎coderd/rbac/input.json

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,21 @@
1212
},
1313
"subject": {
1414
"id": "10d03e62-7703-4df5-a358-4f76577d4e2f",
15-
"roles": [],
15+
"roles": [
16+
{
17+
"name": "owner",
18+
"display_name": "Owner",
19+
"site": [
20+
{
21+
"negate": false,
22+
"resource_type": "*",
23+
"action": "*"
24+
}
25+
],
26+
"org": {},
27+
"user": [],
28+
}
29+
],
1630
"groups": ["b617a647-b5d0-4cbe-9e40-26f89710bf18"],
1731
"scope": {
1832
"name": "Scope_all",

0 commit comments

Comments
 (0)