Prepare in the same place as the normal rego policy

Emyrk · Emyrk · commit 46f9a84168ca · 2023-02-09T14:21:37.000-06:00
diff --git a/coderd/rbac/authz.go b/coderd/rbac/authz.go
@@ -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
@@ -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 {
@@ -166,6 +168,21 @@ func NewAuthorizer(registry prometheus.Registerer) *RegoAuthorizer {
 		if err != nil {
 			panic(xerrors.Errorf("compile rego: %w", err))
 		}
+
+		partialQuery, err = rego.New(
+			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())
+		if err != nil {
+			panic(xerrors.Errorf("compile partial rego: %w", err))
+		}
 	})
 
 	// Register metrics to prometheus.
@@ -207,7 +224,8 @@ func NewAuthorizer(registry prometheus.Registerer) *RegoAuthorizer {
 	})
 
 	return &RegoAuthorizer{
-		query: query,
+		query:        query,
+		partialQuery: partialQuery,
 
 		authorizeHist: authorizeHistogram,
 		prepareHist:   prepareHistogram,
@@ -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)
 	}
diff --git a/coderd/rbac/partial.go b/coderd/rbac/partial.go
@@ -127,24 +127,7 @@ EachQueryLoop:
 		pa.subjectInput, pa.subjectAction, pa.subjectResourceType, nil)
 }
 
-// Precompiled values to be reused for each Prepare call.
-// These values are static and do not change.
-var (
-	// unknownTerms are the unknown values in the rego input.
-	// These values are pre-parsed to prevent reparsing on every Prepare call.
-	unknownTerms = []*ast.Term{
-		ast.MustParseTerm("input.object.id"),
-		ast.MustParseTerm("input.object.owner"),
-		ast.MustParseTerm("input.object.org_owner"),
-		ast.MustParseTerm("input.object.acl_user_list"),
-		ast.MustParseTerm("input.object.acl_group_list"),
-	}
-
-	partialQuery = ast.MustParseBody("data.authz.allow = true")
-	policyModule = ast.MustParseModule(policy)
-)
-
-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")
 	}
@@ -157,14 +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.ParsedQuery(partialQuery),
-		rego.ParsedModule(policyModule),
-		rego.ParsedUnknowns(unknownTerms),
-		rego.ParsedInput(input),
-	).Partial(ctx)
+	partialQueries, err := a.partialQuery.Partial(ctx, rego.EvalParsedInput(input))
+
 	if err != nil {
 		return nil, xerrors.Errorf("prepare: %w", err)
 	}
