Skip to content

chore: Optimize parial rego execution byte allocations #6144

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Feb 10, 2023
Merged

chore: Optimize parial rego execution byte allocations #6144

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
647651e
chore: Implement benchmark for authorizer.Prepare
Emyrk Feb 9, 2023
816a889
Merge remote-tracking branch 'origin/main' into stevenmasley/rego_pre…
Emyrk Feb 9, 2023
891ecc5
Pre-parse all rego inputs to partial execution
Emyrk Feb 9, 2023
46f9a84
Prepare in the same place as the normal rego policy
Emyrk Feb 9, 2023
8420094
Fix benchmark to match objects
Emyrk Feb 9, 2023
0ef4b5c
Rule index org allow expansion
Emyrk Feb 9, 2023
402e31f
Revert "Rule index org allow expansion"
Emyrk Feb 9, 2023
fd29978
Correct unit test json logs
Emyrk Feb 9, 2023
62875c6
Remove unused file
Emyrk Feb 9, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 24 additions & 6 deletions coderd/rbac/authz.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,8 @@ func Filter[O Objecter](ctx context.Context, auth Authorizer, subject Subject, a

// RegoAuthorizer will use a prepared rego query for performing authorize()
type RegoAuthorizer struct {
query rego.PreparedEvalQuery
query rego.PreparedEvalQuery
partialQuery rego.PreparedPartialQuery

authorizeHist *prometheus.HistogramVec
prepareHist prometheus.Histogram
Expand All @@ -151,9 +152,10 @@ var (
// Load the policy from policy.rego in this directory.
//
//go:embed policy.rego
policy string
queryOnce sync.Once
query rego.PreparedEvalQuery
policy string
queryOnce sync.Once
query rego.PreparedEvalQuery
partialQuery rego.PreparedPartialQuery
)

func NewAuthorizer(registry prometheus.Registerer) *RegoAuthorizer {
Expand All @@ -166,6 +168,21 @@ func NewAuthorizer(registry prometheus.Registerer) *RegoAuthorizer {
if err != nil {
panic(xerrors.Errorf("compile rego: %w", err))
}

partialQuery, err = rego.New(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

review: this is inside queryOnce.Do()

rego.Unknowns([]string{
"input.object.id",
"input.object.owner",
"input.object.org_owner",
"input.object.acl_user_list",
"input.object.acl_group_list",
}),
rego.Query("data.authz.allow = true"),
rego.Module("policy.rego", policy),
).PrepareForPartial(context.Background())
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

review: I wondered whether it would make sense to pass through a higher-level context here (e.g. NewAuthorizer(ctx, registry)) but the place we call NewAuthorizer doesn't appear to have any application-level context already available.

if err != nil {
panic(xerrors.Errorf("compile partial rego: %w", err))
}
})

// Register metrics to prometheus.
Expand Down Expand Up @@ -207,7 +224,8 @@ func NewAuthorizer(registry prometheus.Registerer) *RegoAuthorizer {
})

return &RegoAuthorizer{
query: query,
query: query,
partialQuery: partialQuery,

authorizeHist: authorizeHistogram,
prepareHist: prepareHistogram,
Expand Down Expand Up @@ -289,7 +307,7 @@ func (a RegoAuthorizer) Prepare(ctx context.Context, subject Subject, action Act
)
defer span.End()

prepared, err := newPartialAuthorizer(ctx, subject, action, objectType)
prepared, err := a.newPartialAuthorizer(ctx, subject, action, objectType)
if err != nil {
return nil, xerrors.Errorf("new partial authorizer: %w", err)
}
Expand Down
11 changes: 8 additions & 3 deletions coderd/rbac/authz_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -977,9 +977,14 @@ func testAuthorize(t *testing.T, name string, subject Subject, sets ...[]authTes

d, _ := json.Marshal(map[string]interface{}{
// This is not perfect marshal, but it is good enough for debugging this test.
"subject": subject,
"object": c.resource,
"action": a,
"subject": authSubject{
ID: subject.ID,
Roles: must(subject.Roles.Expand()),
Groups: subject.Groups,
Scope: must(subject.Scope.Expand()),
},
"object": c.resource,
"action": a,
})

// Logging only
Expand Down
11 changes: 11 additions & 0 deletions coderd/rbac/authz_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -189,6 +189,17 @@ func BenchmarkRBACFilter(b *testing.B) {
)

authorizer := rbac.NewAuthorizer(prometheus.NewRegistry())

for _, c := range benchCases {
b.Run("PrepareOnly-"+c.Name, func(b *testing.B) {
obType := rbac.ResourceWorkspace.Type
for i := 0; i < b.N; i++ {
_, err := authorizer.Prepare(context.Background(), c.Actor, rbac.ActionRead, obType)
require.NoError(b, err)
}
})
}

for _, c := range benchCases {
b.Run(c.Name, func(b *testing.B) {
objects := benchmarkSetup(orgs, users, b.N)
Expand Down
6 changes: 5 additions & 1 deletion coderd/rbac/builtin_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,11 @@ func BenchmarkRBACValueAllocation(b *testing.B) {
})
b.Run("JSONRegoValue", func(b *testing.B) {
for i := 0; i < b.N; i++ {
_, err := ast.InterfaceToValue(jsonSubject)
_, err := ast.InterfaceToValue(map[string]interface{}{
"subject": jsonSubject,
"action": ActionRead,
"object": obj,
})
require.NoError(b, err)
}
})
Expand Down
18 changes: 3 additions & 15 deletions coderd/rbac/partial.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,7 @@ EachQueryLoop:
pa.subjectInput, pa.subjectAction, pa.subjectResourceType, nil)
}

func newPartialAuthorizer(ctx context.Context, subject Subject, action Action, objectType string) (*PartialAuthorizer, error) {
func (a RegoAuthorizer) newPartialAuthorizer(ctx context.Context, subject Subject, action Action, objectType string) (*PartialAuthorizer, error) {
if subject.Roles == nil {
return nil, xerrors.Errorf("subject must have roles")
}
Expand All @@ -140,20 +140,8 @@ func newPartialAuthorizer(ctx context.Context, subject Subject, action Action, o
return nil, xerrors.Errorf("prepare input: %w", err)
}

// Run the rego policy with a few unknown fields. This should simplify our
// policy to a set of queries.
partialQueries, err := rego.New(
rego.Query("data.authz.allow = true"),
rego.Module("policy.rego", policy),
rego.Unknowns([]string{
"input.object.id",
"input.object.owner",
"input.object.org_owner",
"input.object.acl_user_list",
"input.object.acl_group_list",
}),
rego.ParsedInput(input),
).Partial(ctx)
partialQueries, err := a.partialQuery.Partial(ctx, rego.EvalParsedInput(input))

if err != nil {
return nil, xerrors.Errorf("prepare: %w", err)
}
Expand Down