authorize: rewrite header evaluator to use go instead of rego (#5362)

* authorize: rewrite header evaluator to use go instead of rego

* cache signed jwt

* re-add missing trace

* address comments

authorize/evaluator/headers_evaluator.go

@@ -4,14 +4,11 @@ import (
 	"context"
 	"fmt"
 	"net/http"
-	"os"
+	"time"
 
 	envoy_config_cluster_v3 "github.com/envoyproxy/go-control-plane/envoy/config/cluster/v3"
-	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/rego"
-	"github.com/open-policy-agent/opa/types"
 
-	"github.com/pomerium/pomerium/authorize/evaluator/opa"
 	"github.com/pomerium/pomerium/authorize/internal/store"
 	"github.com/pomerium/pomerium/config"
 	"github.com/pomerium/pomerium/internal/telemetry/trace"
@@ -69,115 +66,30 @@ type HeadersResponse struct {
 	Headers http.Header
 }
 
-var variableSubstitutionFunctionRegoOption = rego.Function2(&rego.Function{
-	Name: "pomerium.variable_substitution",
-	Decl: types.NewFunction(
-		types.Args(
-			types.Named("input_string", types.S),
-			types.Named("replacements",
-				types.NewObject(nil, types.NewDynamicProperty(types.S, types.S))),
-		),
-		types.Named("output", types.S),
-	),
-}, func(_ rego.BuiltinContext, op1 *ast.Term, op2 *ast.Term) (*ast.Term, error) {
-	inputString, ok := op1.Value.(ast.String)
-	if !ok {
-		return nil, fmt.Errorf("invalid input_string type: %T", op1.Value)
-	}
-
-	replacements, ok := op2.Value.(ast.Object)
-	if !ok {
-		return nil, fmt.Errorf("invalid replacements type: %T", op2.Value)
-	}
-
-	var err error
-	output := os.Expand(string(inputString), func(key string) string {
-		if key == "$" {
-			return "$" // allow a dollar sign to be escaped using $$
-		}
-		r := replacements.Get(ast.StringTerm(key))
-		if r == nil {
-			return ""
-		}
-		s, ok := r.Value.(ast.String)
-		if !ok {
-			err = fmt.Errorf("invalid replacement value type for key %q: %T", key, r.Value)
-		}
-		return string(s)
-	})
-	if err != nil {
-		return nil, err
-	}
-	return ast.StringTerm(output), nil
-})
-
 // A HeadersEvaluator evaluates the headers.rego script.
 type HeadersEvaluator struct {
-	q rego.PreparedEvalQuery
+	store *store.Store
 }
 
 // NewHeadersEvaluator creates a new HeadersEvaluator.
-func NewHeadersEvaluator(ctx context.Context, store *store.Store, options ...func(rego *rego.Rego)) (*HeadersEvaluator, error) {
-	r := rego.New(append([]func(*rego.Rego){
-		rego.Store(store),
-		rego.Module("pomerium.headers", opa.HeadersRego),
-		rego.Query("result := data.pomerium.headers"),
-		rego.EnablePrintStatements(true),
-		getGoogleCloudServerlessHeadersRegoOption,
-		variableSubstitutionFunctionRegoOption,
-		store.GetDataBrokerRecordOption(),
-		rego.SetRegoVersion(ast.RegoV1),
-	}, options...)...)
-
-	q, err := r.PrepareForEval(ctx)
-	if err != nil {
-		return nil, err
-	}
-
+func NewHeadersEvaluator(store *store.Store) *HeadersEvaluator {
 	return &HeadersEvaluator{
-		q: q,
-	}, nil
+		store: store,
+	}
 }
 
 // Evaluate evaluates the headers.rego script.
 func (e *HeadersEvaluator) Evaluate(ctx context.Context, req *HeadersRequest, options ...rego.EvalOption) (*HeadersResponse, error) {
 	ctx, span := trace.StartSpan(ctx, "authorize.HeadersEvaluator.Evaluate")
 	defer span.End()
-	rs, err := safeEval(ctx, e.q, append([]rego.EvalOption{rego.EvalInput(req)}, options...)...)
-	if err != nil {
-		return nil, fmt.Errorf("authorize: error evaluating headers.rego: %w", err)
-	}
 
-	if len(rs) == 0 {
-		return nil, fmt.Errorf("authorize: unexpected empty result from evaluating headers.rego")
+	ectx := new(rego.EvalContext)
+	for _, option := range options {
+		option(ectx)
 	}
-
-	return &HeadersResponse{
-		Headers: e.getHeader(rs[0].Bindings),
-	}, nil
-}
-
-func (e *HeadersEvaluator) getHeader(vars rego.Vars) http.Header {
-	h := make(http.Header)
-
-	m, ok := vars["result"].(map[string]any)
-	if !ok {
-		return h
-	}
-
-	m, ok = m["identity_headers"].(map[string]any)
-	if !ok {
-		return h
-	}
-
-	for k := range m {
-		vs, ok := m[k].([]any)
-		if !ok {
-			continue
-		}
-		for _, v := range vs {
-			h.Add(k, fmt.Sprintf("%v", v))
-		}
-	}
-	return h
+	now := ectx.Time()
+	if now.IsZero() {
+		now = time.Now()
+	}
+	return newHeadersEvaluatorEvaluation(e, req, now).execute(ctx)
 }