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.

coderd/rbac/partial.go

@@ -17,8 +17,12 @@ type PartialAuthorizer struct {
 	// partialQueries is mainly used for unit testing to assert our rego policy
 	// can always be compressed into a set of queries.
 	partialQueries *rego.PartialQueries
+
 	// input is used purely for debugging and logging.
-	input map[string]interface{}
+	subjectInput        Subject
+	subjectAction       Action
+	subjectResourceType Object
+
 	// preparedQueries are the compiled set of queries after partial evaluation.
 	// Cache these prepared queries to avoid re-compiling the queries.
 	// If alwaysTrue is true, then ignore these.
@@ -54,7 +58,8 @@ func (pa *PartialAuthorizer) Authorize(ctx context.Context, object Object) error
 	// If we have no queries, then no queries can return 'true'.
 	// So the result is always 'false'.
 	if len(pa.preparedQueries) == 0 {
-		return ForbiddenWithInternal(xerrors.Errorf("policy disallows request"), pa.input, nil)
+		return ForbiddenWithInternal(xerrors.Errorf("policy disallows request"),
+			pa.subjectInput, pa.subjectAction, pa.subjectResourceType, nil)
 	}
 
 	parsed, err := ast.InterfaceToValue(map[string]interface{}{
@@ -118,7 +123,8 @@ EachQueryLoop:
 		return nil
 	}
 
-	return ForbiddenWithInternal(xerrors.Errorf("policy disallows request"), pa.input, nil)
+	return ForbiddenWithInternal(xerrors.Errorf("policy disallows request"),
+		pa.subjectInput, pa.subjectAction, pa.subjectResourceType, nil)
 }
 
 func newPartialAuthorizer(ctx context.Context, subject Subject, action Action, objectType string) (*PartialAuthorizer, error) {
@@ -129,27 +135,9 @@ func newPartialAuthorizer(ctx context.Context, subject Subject, action Action, o
 		return nil, xerrors.Errorf("subject must have a scope")
 	}
 
-	roles, err := subject.Roles.Expand()
+	input, err := regoPartialInputValue(subject, action, objectType)
 	if err != nil {
-		return nil, xerrors.Errorf("expand roles: %w", err)
-	}
-
-	scope, err := subject.Scope.Expand()
-	if err != nil {
-		return nil, xerrors.Errorf("expand scope: %w", err)
-	}
-
-	input := map[string]interface{}{
-		"subject": authSubject{
-			ID:     subject.ID,
-			Roles:  roles,
-			Scope:  scope,
-			Groups: subject.Groups,
-		},
-		"object": map[string]string{
-			"type": objectType,
-		},
-		"action": action,
+		return nil, xerrors.Errorf("prepare input: %w", err)
 	}
 
 	// Run the rego policy with a few unknown fields. This should simplify our
@@ -164,7 +152,7 @@ func newPartialAuthorizer(ctx context.Context, subject Subject, action Action, o
 			"input.object.acl_user_list",
 			"input.object.acl_group_list",
 		}),
-		rego.Input(input),
+		rego.ParsedInput(input),
 	).Partial(ctx)
 	if err != nil {
 		return nil, xerrors.Errorf("prepare: %w", err)
@@ -173,7 +161,12 @@ func newPartialAuthorizer(ctx context.Context, subject Subject, action Action, o
 	pAuth := &PartialAuthorizer{
 		partialQueries:  partialQueries,
 		preparedQueries: []rego.PreparedEvalQuery{},
-		input:           input,
+		subjectInput:    subject,
+		subjectResourceType: Object{
+			Type: objectType,
+			ID:   "prepared-object",
+		},
+		subjectAction: action,
 	}
 
 	// Prepare each query to optimize the runtime when we iterate over the objects.