diff --git a/internal/runtime/init/init.go b/internal/runtime/init/init.go index 5b4bb21b8b..7c9f6b096a 100644 --- a/internal/runtime/init/init.go +++ b/internal/runtime/init/init.go @@ -31,6 +31,7 @@ type InsertAndCompileOptions struct { EnablePrintStatements bool ParserOptions ast.ParserOptions BundleActivatorPlugin string + ExternalSources *util.HasherMap[ast.Ref, ast.ExternalRuleSource] } // InsertAndCompileResult contains the output of the operation. @@ -59,18 +60,31 @@ func InsertAndCompile(ctx context.Context, opts InsertAndCompileOptions) (*Inser SetErrorLimit(opts.MaxErrors). WithPathConflictsCheck(storage.NonEmpty(ctx, opts.Store, opts.Txn)). WithEnablePrintStatements(opts.EnablePrintStatements) + + // Apply external sources to the compiler before bundle activation. + // Bundle activation applies them again via compileModules, but we need them + // here too: there may be no bundles, or a custom activator plugin may not + // call compileModules. + if opts.ExternalSources != nil { + opts.ExternalSources.Iter(func(ref ast.Ref, source ast.ExternalRuleSource) bool { + compiler = compiler.WithExternalSource(ref, source) + return false + }) + } + m := metrics.New() activation := &bundle.ActivateOpts{ - Ctx: ctx, - Store: opts.Store, - Txn: opts.Txn, - Compiler: compiler, - Metrics: m, - Bundles: opts.Bundles, - ExtraModules: policies, - ParserOptions: opts.ParserOptions, - Plugin: opts.BundleActivatorPlugin, + Ctx: ctx, + Store: opts.Store, + Txn: opts.Txn, + Compiler: compiler, + Metrics: m, + Bundles: opts.Bundles, + ExtraModules: policies, + ExternalSources: opts.ExternalSources, + ParserOptions: opts.ParserOptions, + Plugin: opts.BundleActivatorPlugin, } err := bundle.Activate(activation) diff --git a/v1/ast/compile.go b/v1/ast/compile.go index 1b2ab09489..4069b2445b 100644 --- a/v1/ast/compile.go +++ b/v1/ast/compile.go @@ -5,6 +5,7 @@ package ast import ( + "context" "errors" "fmt" "io" @@ -128,6 +129,7 @@ type Compiler struct { localvargen *localVarGenerator moduleLoader ModuleLoader + externalSources *util.HasherMap[Ref, ExternalRuleSource] stages []stage maxErrs int errCount uint32 @@ -135,6 +137,7 @@ type Compiler struct { sorted []string // list of sorted module names pathExists func([]string) (bool, error) pathConflictCheckRoots []string + injectedVirtual func(Ref) bool // optional custom virtual document checker after map[string][]CompilerStageDefinition metrics metrics.Metrics capabilities *Capabilities // user-supplied capabilities @@ -419,6 +422,7 @@ func NewCompiler() *Compiler { Modules: map[string]*Module{}, RewrittenVars: map[Var]Var{}, Required: &Capabilities{}, + externalSources: util.NewHasherMap[Ref, ExternalRuleSource](RefEqual), maxErrs: CompileErrorLimitDefault, mu: &sync.Mutex{}, after: map[string][]CompilerStageDefinition{}, @@ -651,6 +655,14 @@ func (c *Compiler) QueryCompiler() QueryCompiler { return newQueryCompiler(&c0) } +// WithVirtual sets a custom virtual document checker on the compiler. +// The provided function will be called during rule index building to determine +// if additional refs should be considered virtual documents. +func (c *Compiler) WithVirtual(fn func(Ref) bool) *Compiler { + c.injectedVirtual = fn + return c +} + // Compile runs the compilation process on the input modules. The compiled // version of the modules and associated data structures are stored on the // compiler. If the compilation process fails for any reason, the compiler will @@ -894,9 +906,8 @@ func (c *Compiler) GetRulesDynamic(ref Ref) []*Rule { // Without the options, it would be excluded. func (c *Compiler) GetRulesDynamicWithOpts(ref Ref, opts RulesOptions) []*Rule { node := c.RuleTree - set := map[*Rule]struct{}{} - var walk func(node *TreeNode, i int) + var walk func(*TreeNode, int) walk = func(node *TreeNode, i int) { switch { case i >= len(ref): @@ -1041,6 +1052,19 @@ func (c *Compiler) WithModuleLoader(f ModuleLoader) *Compiler { return c } +// WithExternalSource registers an external rule source for the given package +// reference. When rules under this package are queried via RuleIndex, the +// external source will be invoked to fetch all rules for the package. The +// fetched rules are cached so the external source is only called once per +// package. +// +// The package reference should be a fully qualified path (e.g., data.foo.bar). +// All rule queries under this package will be handled by the external source. +func (c *Compiler) WithExternalSource(packageRef Ref, source ExternalRuleSource) *Compiler { + c.externalSources.Put(packageRef, source) + return c +} + // WithDefaultRegoVersion sets the default Rego version to use when a module doesn't specify one; // such as when it's hand-crafted instead of parsed. func (c *Compiler) WithDefaultRegoVersion(regoVersion RegoVersion) *Compiler { @@ -1113,10 +1137,14 @@ func (c *Compiler) counterAdd(name string, n uint64) { func (c *Compiler) buildRuleIndices() { c.RuleTree.DepthFirst(func(node *TreeNode) bool { - if len(node.Values) == 0 { + if len(node.Values) == 0 && node.External == nil { return false } - rules := node.Values + if node.External != nil { + // Skip external sources - they build indices dynamically + return true + } + rules := node.Values // must be len > 0 here hasNonGroundRef := false for _, r := range rules { hasNonGroundRef = !r.Head.Ref().IsGround() @@ -1141,9 +1169,7 @@ func (c *Compiler) buildRuleIndices() { } } - index := newBaseDocEqIndex(func(ref Ref) bool { - return isVirtual(c.RuleTree, ref.GroundPrefix()) - }) + index := newBaseDocEqIndex(c.isVirtual) if index.Build(rules) { node.Index = index } @@ -1302,15 +1328,19 @@ func (c *Compiler) checkRuleConflicts() { return false // go deeper } - kinds := make(map[RuleKind]struct{}, len(node.Values)) + rules := node.Values + if len(rules) == 0 { + return true // ?? right + } + kinds := make(map[RuleKind]struct{}, len(rules)) completeRules := 0 partialRules := 0 - arities := make(map[int]struct{}, len(node.Values)) + arities := make(map[int]struct{}, len(rules)) name := "" var conflicts []Ref defaultRules := make([]*Rule, 0) - for _, rule := range node.Values { + for _, rule := range rules { r := rule ref := r.Ref() name = rw(ref.CopyNonGround()).String() // varRewriter operates in-place @@ -1383,10 +1413,10 @@ func (c *Compiler) checkRuleConflicts() { switch { case conflicts != nil: - return !c.err(NewError(TypeErr, node.Values[0].Loc(), "rule %v conflicts with %v", name, conflicts)) + return !c.err(NewError(TypeErr, rules[0].Loc(), "rule %v conflicts with %v", name, conflicts)) case len(kinds) > 1 || len(arities) > 1 || (completeRules >= 1 && partialRules >= 1): - return !c.err(NewError(TypeErr, node.Values[0].Loc(), "conflicting rules %v found", name)) + return !c.err(NewError(TypeErr, rules[0].Loc(), "conflicting rules %v found", name)) case len(defaultRules) > 1: buf := append(append(append(make([]byte, 0, 64), "multiple default rules "...), name...), " found at "...) @@ -3455,6 +3485,17 @@ func (c *Compiler) setModuleTree() { func (c *Compiler) setRuleTree() { c.RuleTree = NewRuleTree(c.ModuleTree) + + // Add tree nodes for external source paths so evaluation knows to look there + c.externalSources.Iter(func(pkgRef Ref, source ExternalRuleSource) bool { + ri, err := source.Init(context.TODO(), pkgRef) + if err != nil { + c.err(NewError(CompileErr, nil, "failed to initialize external rule source for ref %v: %v", pkgRef, err)) + return true + } + c.RuleTree.add(pkgRef, ri) + return false + }) } func (c *Compiler) setGraph() { @@ -4123,6 +4164,7 @@ func (n *ModuleTreeNode) DepthFirst(f func(*ModuleTreeNode) bool) { // rule path. type TreeNode struct { Key Value + External *ExternalIndex Values []*Rule Children map[Value]*TreeNode Sorted []Value @@ -4167,20 +4209,104 @@ func NewRuleTree(mtree *ModuleTreeNode) *TreeNode { return &root } -func (n *TreeNode) add(path Ref, rule *Rule) { +func (n *TreeNode) add(path Ref, val any) { node, tail := n.find(path) if len(tail) > 0 { - sub := treeNodeFromRef(tail, rule) + sub := treeNodeFromRef(path, tail, val) if node.Children == nil { node.Children = make(map[Value]*TreeNode, 1) } node.Children[sub.Key] = sub node.Sorted = append(node.Sorted, sub.Key) - } else if rule != nil { - node.Values = append(node.Values, rule) + } else if val != nil { + switch val := val.(type) { + case *Rule: + node.Values = append(node.Values, val) + case ExternalRuleIndex: + node.External = &ExternalIndex{ + Index: val, + Ref: path, + } + } } } +type ExternalIndex struct { + Index ExternalRuleIndex + Ref Ref +} + +func (ei *ExternalIndex) Tree(ctx context.Context, rt *TreeNode, prefix Ref, input *Term, m metrics.Metrics, reqMD map[string]any, respMD map[string]any) (*TreeNode, ExternalRuleIndex, error) { + resolver := &termResolver{input: input} + + rules, updatedIndex, err := ei.Index.Lookup(ctx, + LookupResolver(resolver), + LookupMetrics(m), + LookupRequestMetadata(reqMD), + LookupResponseMetadata(respMD), + ) + if err != nil { + return nil, nil, err + } + c0 := NewCompiler() + + if o := ei.Index.Opts(); o != nil { + if len(o.SkippedStages) > 0 { + c0.WithSkipStages(o.SkippedStages...) + } + + if len(o.VisibleRefs) > 0 { + visible := o.VisibleRefs + c0.WithVirtual(func(ref Ref) bool { + return slices.ContainsFunc(visible, ref.HasPrefix) && rt.isVirtual(ref) + }) + } + } + + modules := make(map[string]*Module) + for _, rule := range rules { + pkgPathStr := rule.Module.Package.Path.String() + if mod, exists := modules[pkgPathStr]; exists { + mod.Rules = append(mod.Rules, rule) + } else { + modules[pkgPathStr] = &Module{ + Package: &Package{Path: rule.Module.Package.Path}, + Rules: []*Rule{rule}, + } + } + } + if m != nil { + t := m.Timer("external_lookup_compile_module") + t.Start() + defer t.Stop() + } + c0.Compile(modules) + if c0.Failed() { + return nil, nil, c0.Errors + } + + node := c0.RuleTree.Find(prefix) + return node, updatedIndex, nil +} + +type termResolver struct { + input *Term +} + +func (r *termResolver) Resolve(ref Ref) (Value, error) { + if ref.HasPrefix(InputRootRef) { + if r.input == nil { + return nil, UnknownValueErr{} + } + v, err := r.input.Value.Find(ref[1:]) + if err != nil { + return nil, UnknownValueErr{} + } + return v, nil + } + return nil, UnknownValueErr{} +} + // Size returns the number of rules in the tree. func (n *TreeNode) Size() (s int) { for _, c := range n.Children { @@ -4240,28 +4366,67 @@ func (n *TreeNode) DepthFirst(f func(*TreeNode) bool) { } } -func treeNodeFromRef(ref Ref, rule *Rule) *TreeNode { - depth := len(ref) - 1 - key := ref[depth].Value - node := &TreeNode{ - Key: key, - Children: nil, +func (c *Compiler) isVirtual(ref Ref) bool { + return (c.injectedVirtual != nil && c.injectedVirtual(ref)) || + c.RuleTree.isVirtual(ref.GroundPrefix()) +} + +// isVirtual returns true if the ref is virtual (has rules). +func (n *TreeNode) isVirtual(ref Ref) bool { + node := n + for i := range ref { + child := node.Child(ref[i].Value) + if child == nil { + return false + } else if len(child.Values) > 0 || child.External != nil { + return true + } + node = child } - if rule != nil { - node.Values = []*Rule{rule} + return true +} + +func treeNodeFromRef(ref, tail Ref, val any) *TreeNode { + if len(tail) == 0 { + node := &TreeNode{ + Children: make(map[Value]*TreeNode), + } + attachValueToNode(node, ref, val) + return node } - for i := len(ref) - 2; i >= 0; i-- { - key := ref[i].Value + depth := len(tail) - 1 + node := &TreeNode{ + Key: tail[depth].Value, + } + attachValueToNode(node, ref, val) + + for i := depth - 1; i >= 0; i-- { + childKey := tail[i+1].Value node = &TreeNode{ - Key: key, - Children: map[Value]*TreeNode{ref[i+1].Value: node}, - Sorted: []Value{ref[i+1].Value}, + Key: tail[i].Value, + Children: map[Value]*TreeNode{childKey: node}, + Sorted: []Value{childKey}, } } return node } +func attachValueToNode(node *TreeNode, ref Ref, val any) { + if val == nil { + return + } + switch val := val.(type) { + case *Rule: + node.Values = append(node.Values, val) + case ExternalRuleIndex: + node.External = &ExternalIndex{ + Index: val, + Ref: ref, + } + } +} + // flattenChildren flattens all children's rule refs into a sorted array. func (n *TreeNode) flattenChildren() []Ref { return n.flattenMatchingChildren(func(_ *Rule) bool { return true }) @@ -4288,6 +4453,60 @@ func (n *TreeNode) flattenMatchingChildren(f func(*Rule) bool) []Ref { return util.SortedFunc(ret.s, RefCompare) } +// Copy creates a shallow copy of the TreeNode suitable for augmentation. +// Children map is copied recursively. Values slices are initially shared but +// reallocated on modification (e.g., by MergeChild's append operation). +func (n *TreeNode) Copy() *TreeNode { + if n == nil { + return nil + } + + result := &TreeNode{ + Key: n.Key, + External: n.External, + Values: n.Values, + Hide: n.Hide, + Index: n.Index, + } + + if n.Children != nil { + result.Children = make(map[Value]*TreeNode, len(n.Children)) + for k, v := range n.Children { + result.Children[k] = v.Copy() + } + } + + if n.Sorted != nil { + result.Sorted = make([]Value, len(n.Sorted)) + copy(result.Sorted, n.Sorted) + } + + return result +} + +// MergeChild merges another TreeNode into this node's children. +func (n *TreeNode) MergeChild(key Value, other *TreeNode) { + if other == nil { + return + } + + existing := n.Child(key) + if existing == nil { + if n.Children == nil { + n.Children = make(map[Value]*TreeNode) + } + n.Children[key] = other + n.Sorted = append(n.Sorted, key) + return + } + + existing.Values = append(existing.Values, other.Values...) + + for childKey, childNode := range other.Children { + existing.MergeChild(childKey, childNode) + } +} + // Graph represents the graph of dependencies between rules. type Graph struct { adj map[util.T]map[util.T]struct{} @@ -4299,7 +4518,6 @@ type Graph struct { // NewGraph returns a new Graph based on modules. The list function must return // the rules referred to directly by the ref. func NewGraph(modules map[string]*Module, list func(Ref) []*Rule) *Graph { - graph := &Graph{ adj: map[util.T]map[util.T]struct{}{}, radj: map[util.T]map[util.T]struct{}{}, @@ -6676,8 +6894,8 @@ func validateWith(c *Compiler, unsafeBuiltinsMap map[string]struct{}, expr *Expr // target is a function. It's probably wrong for arity-0 functions, but those are // and edge case anyways. if child := targetNode.Child(ref[len(ref)-1].Value); child != nil { - for _, v := range child.Values { - if len(v.Head.Args) > 0 { + for _, r := range child.Values { + if len(r.Head.Args) > 0 { if ok, err := validateWithFunctionValue(c.builtins, unsafeBuiltinsMap, c.RuleTree, value); err != nil || ok { return false, err // err may be nil } @@ -6690,8 +6908,8 @@ func validateWith(c *Compiler, unsafeBuiltinsMap map[string]struct{}, expr *Expr if r, ok := value.Value.(Ref); ok { // TODO: check that target ref doesn't exist? if valueNode := c.RuleTree.Find(r); valueNode != nil { - for _, v := range valueNode.Values { - if len(v.Head.Args) > 0 { + for _, r := range valueNode.Values { + if len(r.Head.Args) > 0 { return false, nil } } @@ -6776,19 +6994,6 @@ func isBuiltinRefOrVar(bs map[string]*Builtin, unsafeBuiltinsMap map[string]stru return false, nil } -func isVirtual(node *TreeNode, ref Ref) bool { - for i := range ref { - child := node.Child(ref[i].Value) - if child == nil { - return false - } else if len(child.Values) > 0 { - return true - } - node = child - } - return true -} - func safetyErrorSlice(unsafe unsafeVars, rewritten map[Var]Var) (result Errors) { if len(unsafe) == 0 { return diff --git a/v1/ast/compile_test.go b/v1/ast/compile_test.go index c5707ae155..e22b098ef3 100644 --- a/v1/ast/compile_test.go +++ b/v1/ast/compile_test.go @@ -819,14 +819,15 @@ func TestRuleIndices(t *testing.T) { c.sorted = append(c.sorted, strconv.Itoa(i)) } compileStages(c, StageBuildRuleIndices) + t.Log(c.RuleTree.Dump()) for k, expIndex := range tc.exp { kref := MustParseRef(k) - node := c.RuleTree.Find(kref) - if node == nil || node.Index == nil { + i := c.RuleIndex(kref) + if i == nil { t.Fatalf("expected rule indices for %v", k) } - index := node.Index.(*baseDocEqIndex) + index := i.(*baseDocEqIndex) for _, expRef := range expIndex { found := false for _, r := range index.root.rules { @@ -1199,7 +1200,7 @@ func TestRuleTree(t *testing.T) { t.Fatalf("Expected user.system node to be visible") } - if !isVirtual(tree, MustParseRef("data.a.b.empty")) { + if !tree.isVirtual(MustParseRef("data.a.b.empty")) { t.Fatal("Expected data.a.b.empty to be virtual") } diff --git a/v1/ast/external_source.go b/v1/ast/external_source.go new file mode 100644 index 0000000000..ce3433cce0 --- /dev/null +++ b/v1/ast/external_source.go @@ -0,0 +1,128 @@ +// Copyright 2026 The OPA Authors. All rights reserved. +// Use of this source code is governed by an Apache2 +// license that can be found in the LICENSE file. + +package ast + +import ( + "context" + + "github.com/open-policy-agent/opa/v1/metrics" +) + +type ExternalRuleSource interface { + // Refs returns the package refs that this source provides rules for. + // A source can provide rules for multiple packages. + Refs() []Ref + + // Init returns an initialized [ExternalRuleIndex]. A `Ref` is provided + // so we know which package we're preparing if multiple Refs are external. + Init(context.Context, Ref) (ExternalRuleIndex, error) +} + +// ExternalRuleIndex mirrors RuleIndex.Lookup(), but add a [context.Context] parameter. +type ExternalRuleIndex interface { + // Opts returns the options for the ExternalRuleIndex. Returns nil if no + // options are configured. + Opts() *ExternalSourceOptions + + // Lookup returns rules and optionally an updated ExternalRuleIndex instance. + // The returned ExternalRuleIndex (if non-nil) will be used for subsequent + // Lookup calls within the same evaluation context, allowing plugins to + // maintain per-evaluation state. + // + // Plugins can use two strategies: + // 1. Immutable: Return a new ExternalRuleIndex instance with updated state + // 2. Mutable: Update internal state and return self + // + // If the plugin does not need per-evaluation state, it can return nil for + // the ExternalRuleIndex, and the original instance will continue to be used. + Lookup(context.Context, ...LookupOption) ([]*Rule, ExternalRuleIndex, error) +} + +// ExternalRuleIndexCloser is an optional interface for resource cleanup. +type ExternalRuleIndexCloser interface { + ExternalRuleIndex + Close() error +} + +// ExternalSourceOptions contains options for registering an external rule source. +type ExternalSourceOptions struct { + // VisibleRefs controls which parts of the surrounding rule tree the external + // source can reference during compilation. By default (nil), the source is + // fully isolated and cannot access any surrounding policy. An empty slice + // is equivalent to nil (fully isolated). + // + // To allow access to the entire rule tree, use []Ref{MustParseRef("data")}. + // To allow access to specific subtrees only, list them explicitly, e.g. + // []Ref{MustParseRef("data.helpers")}. The external source can then + // reference rules under those prefixes but nothing else. + VisibleRefs []Ref + + // SkippedStages allows external sources to skip stages in the dynamic compiler + // used with the externally-provided Rego. If, for example, the `[]*Rule` returned + // has already been compiled, we can skip all stages. + // + // For pre-compiled rules, prefer starting from AllStages() and removing only + // the stages you need (e.g. SetModuleTree, SetRuleTree, BuildRuleIndices). + // This is forward-compatible: new compiler stages added in future releases + // will be skipped automatically rather than running unexpectedly. + SkippedStages []StageID +} + +// LookupOption is a functional option for ExternalRuleIndex.Lookup calls. +type LookupOption func(*LookupOptions) + +// LookupOptions contains options for ExternalRuleIndex.Lookup calls. +type LookupOptions struct { + metrics metrics.Metrics + resolver ValueResolver + requestMetadata map[string]any + responseMetadata map[string]any +} + +// Metrics returns the metrics instance from the options, or nil if not set. +func (o *LookupOptions) Metrics() metrics.Metrics { + if o == nil { + return nil + } + return o.metrics +} + +func (o *LookupOptions) Resolver() ValueResolver { + return o.resolver +} + +func (o *LookupOptions) RequestMetadata() map[string]any { + return o.requestMetadata +} + +func (o *LookupOptions) ResponseMetadata() map[string]any { + return o.responseMetadata +} + +// LookupMetrics returns a LookupOption that sets the metrics instance +// for the Lookup call. +func LookupMetrics(m metrics.Metrics) LookupOption { + return func(opts *LookupOptions) { + opts.metrics = m + } +} + +func LookupResolver(r ValueResolver) LookupOption { + return func(opts *LookupOptions) { + opts.resolver = r + } +} + +func LookupRequestMetadata(m map[string]any) LookupOption { + return func(opts *LookupOptions) { + opts.requestMetadata = m + } +} + +func LookupResponseMetadata(m map[string]any) LookupOption { + return func(opts *LookupOptions) { + opts.responseMetadata = m + } +} diff --git a/v1/ast/external_source_test.go b/v1/ast/external_source_test.go new file mode 100644 index 0000000000..45adb822ae --- /dev/null +++ b/v1/ast/external_source_test.go @@ -0,0 +1,72 @@ +package ast + +import ( + "context" + "sync/atomic" + "testing" +) + +type mockExternalSource struct { + refs []Ref + rules []*Rule + callCount int32 +} + +func newMockExternalSource(refs []Ref, rules []*Rule) *mockExternalSource { + return &mockExternalSource{ + refs: refs, + rules: rules, + } +} + +func (m *mockExternalSource) Refs() []Ref { + return m.refs +} + +func (m *mockExternalSource) Init(context.Context, Ref) (ExternalRuleIndex, error) { + return &mockExternalIndex{rules: m.rules, callCount: &m.callCount}, nil +} + +type mockExternalIndex struct { + rules []*Rule + callCount *int32 +} + +func (*mockExternalIndex) Opts() *ExternalSourceOptions { + return nil +} + +func (m *mockExternalIndex) Lookup(context.Context, ...LookupOption) ([]*Rule, ExternalRuleIndex, error) { + atomic.AddInt32(m.callCount, 1) + return m.rules, nil, nil +} + +func (m *mockExternalSource) getCallCount() int { + return int(atomic.LoadInt32(&m.callCount)) +} + +func TestCompilerRuleIndexReturnsNilForExternalSources(t *testing.T) { + rule := &Rule{ + Head: &Head{ + Reference: MustParseRef("data.external.test.foo"), + Value: BooleanTerm(true), + }, + Body: NewBody( + Equality.Expr(VarTerm("x"), IntNumberTerm(1)), + ), + } + + packageRef := MustParseRef("data.external.test") + source := newMockExternalSource([]Ref{packageRef}, []*Rule{rule}) + compiler := NewCompiler() + compiler.WithExternalSource(packageRef, source) + + index := compiler.RuleIndex(packageRef) + if index != nil { + t.Error("Expected RuleIndex to return nil for external source path (delegation to evaluation-time)") + } + + if source.getCallCount() != 0 { + t.Errorf("Expected GetRules NOT to be called at compile-time, got %d calls", source.getCallCount()) + } +} diff --git a/v1/ast/treenode_dump.go b/v1/ast/treenode_dump.go new file mode 100644 index 0000000000..f22367bc81 --- /dev/null +++ b/v1/ast/treenode_dump.go @@ -0,0 +1,56 @@ +package ast + +import ( + "fmt" + "sort" + "strings" +) + +// Dump returns a string representation of the tree structure rooted at this node. +func (n *TreeNode) Dump() string { + var sb strings.Builder + n.dumpRecursive(&sb, "", "") + return sb.String() +} + +func (n *TreeNode) dumpRecursive(sb *strings.Builder, prefix, childPrefix string) { + sb.WriteString(prefix) + fmt.Fprintf(sb, "%v", n.Key) + + if n.Hide { + sb.WriteString(" [hidden]") + } + if n.External != nil { + fmt.Fprintf(sb, " ext:%v", n.External.Ref) + } + if len(n.Values) > 0 { + fmt.Fprintf(sb, " rules:%d", len(n.Values)) + } + sb.WriteString("\n") + + if len(n.Children) == 0 { + return + } + + keys := make([]Value, 0, len(n.Children)) + for k := range n.Children { + keys = append(keys, k) + } + sort.Slice(keys, func(i, j int) bool { + return Compare(keys[i], keys[j]) < 0 + }) + + for i, key := range keys { + child := n.Children[key] + isLast := i == len(keys)-1 + var newPrefix, newChildPrefix string + if isLast { + newPrefix = childPrefix + "└── " + newChildPrefix = childPrefix + " " + } else { + newPrefix = childPrefix + "├── " + newChildPrefix = childPrefix + "│ " + } + child.dumpRecursive(sb, newPrefix, newChildPrefix) + } +} diff --git a/v1/bundle/store.go b/v1/bundle/store.go index f6c13c75fc..e549340408 100644 --- a/v1/bundle/store.go +++ b/v1/bundle/store.go @@ -359,8 +359,9 @@ type ActivateOpts struct { TxnCtx *storage.Context Compiler *ast.Compiler Metrics metrics.Metrics - Bundles map[string]*Bundle // Optional - ExtraModules map[string]*ast.Module // Optional + Bundles map[string]*Bundle // Optional + ExtraModules map[string]*ast.Module // Optional + ExternalSources *util.HasherMap[ast.Ref, ast.ExternalRuleSource] // Optional AuthorizationDecisionRef ast.Ref ParserOptions ast.ParserOptions Plugin string @@ -535,7 +536,7 @@ func activateBundles(opts *ActivateOpts) error { maps.Copy(remainingAndExtra, remaining) maps.Copy(remainingAndExtra, opts.ExtraModules) - err = compileModules(opts.Compiler, opts.Metrics, snapshotBundles, remainingAndExtra, opts.legacy, opts.AuthorizationDecisionRef) + err = compileModules(opts.Compiler, opts.Metrics, snapshotBundles, remainingAndExtra, opts.legacy, opts.AuthorizationDecisionRef, opts.ExternalSources) if err != nil { return err } @@ -965,11 +966,19 @@ func writeData(ctx context.Context, store storage.Store, txn storage.Transaction return nil } -func compileModules(compiler *ast.Compiler, m metrics.Metrics, bundles map[string]*Bundle, extraModules map[string]*ast.Module, legacy bool, authorizationDecisionRef ast.Ref) error { +func compileModules(compiler *ast.Compiler, m metrics.Metrics, bundles map[string]*Bundle, extraModules map[string]*ast.Module, legacy bool, authorizationDecisionRef ast.Ref, externalSources *util.HasherMap[ast.Ref, ast.ExternalRuleSource]) error { m.Timer(metrics.RegoModuleCompile).Start() defer m.Timer(metrics.RegoModuleCompile).Stop() + // Apply external sources before compilation + if externalSources != nil { + externalSources.Iter(func(ref ast.Ref, source ast.ExternalRuleSource) bool { + compiler = compiler.WithExternalSource(ref, source) + return false + }) + } + modules := make(map[string]*ast.Module, len(compiler.Modules)+len(extraModules)+len(bundles)) // preserve any modules already on the compiler @@ -1000,11 +1009,19 @@ func compileModules(compiler *ast.Compiler, m metrics.Metrics, bundles map[strin return iCompiler.VerifyAuthorizationPolicySchema(compiler, authorizationDecisionRef) } -func writeModules(ctx context.Context, store storage.Store, txn storage.Transaction, compiler *ast.Compiler, m metrics.Metrics, bundles map[string]*Bundle, extraModules map[string]*ast.Module, legacy bool) error { +func writeModules(ctx context.Context, store storage.Store, txn storage.Transaction, compiler *ast.Compiler, m metrics.Metrics, bundles map[string]*Bundle, extraModules map[string]*ast.Module, legacy bool, externalSources *util.HasherMap[ast.Ref, ast.ExternalRuleSource]) error { m.Timer(metrics.RegoModuleCompile).Start() defer m.Timer(metrics.RegoModuleCompile).Stop() + // Apply external sources before compilation + if externalSources != nil { + externalSources.Iter(func(ref ast.Ref, source ast.ExternalRuleSource) bool { + compiler = compiler.WithExternalSource(ref, source) + return false + }) + } + modules := map[string]*ast.Module{} // preserve any modules already on the compiler diff --git a/v1/bundle/store_test.go b/v1/bundle/store_test.go index 51c101e0c8..6d76d13294 100644 --- a/v1/bundle/store_test.go +++ b/v1/bundle/store_test.go @@ -3714,7 +3714,7 @@ func testWriteData(t *testing.T, tc testWriteModuleCase, legacy bool) { } } - err := writeModules(t.Context(), mockStore, txn, compiler, metrics.NoOp(), tc.bundles, tc.extraMods, legacy) + err := writeModules(t.Context(), mockStore, txn, compiler, metrics.NoOp(), tc.bundles, tc.extraMods, legacy, nil) if !tc.expectErr && err != nil { t.Fatalf("unepected error: %s", err) } else if tc.expectErr && err == nil { diff --git a/v1/compile/compile.go b/v1/compile/compile.go index 7779f2435a..ec90426b71 100644 --- a/v1/compile/compile.go +++ b/v1/compile/compile.go @@ -87,7 +87,8 @@ type Compiler struct { fsys fs.FS // file system to use when loading paths ns string regoVersion ast.RegoVersion - followSymlinks bool // optionally follow symlinks in the bundle directory when building the bundle + followSymlinks bool // optionally follow symlinks in the bundle directory when building the bundle + externalRefs []ast.Ref // external entrypoints provided dynamically } // New returns a new compiler instance that can be invoked. @@ -179,6 +180,12 @@ func (c *Compiler) WithEnablePrintStatements(yes bool) *Compiler { return c } +// WithExternalRefs sets the external entrypoints that are provided dynamically. +func (c *Compiler) WithExternalRefs(refs []ast.Ref) *Compiler { + c.externalRefs = refs + return c +} + // WithPaths adds input filepaths to read policy and data from. func (c *Compiler) WithPaths(p ...string) *Compiler { c.paths = append(c.paths, p...) diff --git a/v1/hooks/hooks.go b/v1/hooks/hooks.go index cb756e5020..329959e0ad 100644 --- a/v1/hooks/hooks.go +++ b/v1/hooks/hooks.go @@ -8,6 +8,7 @@ import ( "context" "fmt" + "github.com/open-policy-agent/opa/v1/bundle" "github.com/open-policy-agent/opa/v1/config" topdown_cache "github.com/open-policy-agent/opa/v1/topdown/cache" ) @@ -44,6 +45,13 @@ func New(hs ...Hook) Hooks { return h } +func (hs *Hooks) Append(h Hook) { + if hs.m == nil { + hs.m = make(map[Hook]struct{}) + } + hs.m[h] = struct{}{} +} + func (hs Hooks) Each(fn func(Hook)) { for h := range hs.m { fn(h) @@ -82,13 +90,21 @@ type InterQueryValueCacheHook interface { OnInterQueryValueCache(context.Context, topdown_cache.InterQueryValueCache) error } +// BundlePreActivateHook is called before a bundle is activated and its policies +// are compiled. This allows hooks to inspect the bundle manifest (e.g. metadata) +// and register external rule sources that will be available during compilation. +type BundlePreActivateHook interface { + OnBundlePreActivate(ctx context.Context, bundleName string, manifest bundle.Manifest) error +} + func (hs Hooks) Validate() error { for h := range hs.m { switch h.(type) { case InterQueryCacheHook, InterQueryValueCacheHook, ConfigHook, - ConfigDiscoveryHook: // OK + ConfigDiscoveryHook, + BundlePreActivateHook: // OK default: return fmt.Errorf("unknown hook type %T", h) } diff --git a/v1/plugins/bundle/plugin.go b/v1/plugins/bundle/plugin.go index 88da1b1eb2..717da0b9d7 100644 --- a/v1/plugins/bundle/plugin.go +++ b/v1/plugins/bundle/plugin.go @@ -24,6 +24,7 @@ import ( "github.com/open-policy-agent/opa/v1/ast" "github.com/open-policy-agent/opa/v1/bundle" "github.com/open-policy-agent/opa/v1/download" + "github.com/open-policy-agent/opa/v1/hooks" "github.com/open-policy-agent/opa/v1/logging" "github.com/open-policy-agent/opa/v1/metrics" "github.com/open-policy-agent/opa/v1/plugins" @@ -635,15 +636,26 @@ func (p *Plugin) activate(ctx context.Context, name string, b *bundle.Bundle, is var activateErr error + // Call pre-activation hooks so plugins can inspect the bundle manifest + // and register external sources before compilation. + p.manager.Hooks().Each(func(h hooks.Hook) { + if f, ok := h.(hooks.BundlePreActivateHook); ok { + if err := f.OnBundlePreActivate(ctx, name, b.Manifest); err != nil { + p.log(name).Warn("Pre-activation hook failed: %v", err) + } + } + }) + opts := &bundle.ActivateOpts{ - Ctx: ctx, - Store: p.manager.Store, - Txn: txn, - TxnCtx: params.Context, - Compiler: compiler, - Metrics: p.status[name].Metrics, - Bundles: map[string]*bundle.Bundle{name: b}, - ParserOptions: p.manager.ParserOptions(), + Ctx: ctx, + Store: p.manager.Store, + Txn: txn, + TxnCtx: params.Context, + Compiler: compiler, + Metrics: p.status[name].Metrics, + Bundles: map[string]*bundle.Bundle{name: b}, + ExternalSources: p.manager.GetExternalSources(), + ParserOptions: p.manager.ParserOptions(), } if p.manager.Info != nil { diff --git a/v1/plugins/plugins.go b/v1/plugins/plugins.go index e4c8a1af48..90bd26cb9c 100644 --- a/v1/plugins/plugins.go +++ b/v1/plugins/plugins.go @@ -36,6 +36,7 @@ import ( "github.com/open-policy-agent/opa/v1/topdown/cache" "github.com/open-policy-agent/opa/v1/topdown/print" "github.com/open-policy-agent/opa/v1/tracing" + "github.com/open-policy-agent/opa/v1/util" ) // Factory defines the interface OPA uses to instantiate your plugin. @@ -239,6 +240,8 @@ type Manager struct { extraMiddlewares []func(http.Handler) http.Handler extraAuthorizerRoutes []func(string, []any) bool bundleActivatorPlugin string + externalSources *util.HasherMap[ast.Ref, ast.ExternalRuleSource] + externalSourcesMux sync.RWMutex } type pluginStatusMsg interface { @@ -459,6 +462,16 @@ func WithHooks(hs hooks.Hooks) func(*Manager) { } } +// Hooks returns the hooks configured on the Manager. +func (m *Manager) Hooks() hooks.Hooks { + return m.hooks +} + +// AppendHook allows adding to the hooks configured on the Manager. +func (m *Manager) AppendHook(h hooks.Hook) { + m.hooks.Append(h) +} + // WithParserOptions sets the parser options to be used by the plugin manager. func WithParserOptions(opts ast.ParserOptions) func(*Manager) { return func(m *Manager) { @@ -629,6 +642,7 @@ func (m *Manager) Init(ctx context.Context) error { EnablePrintStatements: m.enablePrintStatements, ParserOptions: m.parserOptions, BundleActivatorPlugin: m.bundleActivatorPlugin, + ExternalSources: m.GetExternalSources(), }) if err != nil { return err @@ -828,6 +842,30 @@ func (m *Manager) setWasmResolvers(rs []*wasm.Resolver) { m.wasmResolvers = rs } +// RegisterExternalSource registers an external rule source with the manager. +// The source will be applied to all compilers created by the manager. +// This should be called from a plugin's constructor or Start() method. +func (m *Manager) RegisterExternalSource(pkgRef ast.Ref, source ast.ExternalRuleSource) { + m.externalSourcesMux.Lock() + defer m.externalSourcesMux.Unlock() + + if m.externalSources == nil { + m.externalSources = util.NewHasherMap[ast.Ref, ast.ExternalRuleSource](ast.RefEqual) + } + + m.externalSources.Put(pkgRef, source) + + m.logger.Debug("Registered external source for package: %s", pkgRef) +} + +// GetExternalSources returns the registered external sources +func (m *Manager) GetExternalSources() *util.HasherMap[ast.Ref, ast.ExternalRuleSource] { + m.externalSourcesMux.RLock() + defer m.externalSourcesMux.RUnlock() + + return m.externalSources +} + // Start starts the manager. Init() should be called once before Start(). func (m *Manager) Start(ctx context.Context) error { if m == nil { @@ -857,6 +895,25 @@ func (m *Manager) Start(ctx context.Context) error { } } + // After starting plugins, check if any external sources were registered + // and recompile if necessary to include them in the rule tree + externalSources := m.GetExternalSources() + if externalSources != nil && externalSources.Len() > 0 { + err := storage.Txn(ctx, m.Store, storage.TransactionParams{}, func(txn storage.Transaction) error { + compiler, err := loadCompilerFromStore(ctx, m.Store, txn, m.enablePrintStatements, m.ParserOptions(), externalSources) + if err != nil { + return err + } + m.setCompiler(compiler) + m.logger.Debug("Recompiled policies with %d external source(s) after plugin startup", externalSources.Len()) + return nil + }) + + if err != nil { + return fmt.Errorf("failed to recompile with external sources: %w", err) + } + } + return nil } @@ -1038,7 +1095,7 @@ func (m *Manager) onCommit(ctx context.Context, txn storage.Transaction, event s // compiler on the context but the server does not (nor would users // implementing their own policy loading.) if compiler == nil && event.PolicyChanged() { - compiler, _ = loadCompilerFromStore(ctx, m.Store, txn, m.enablePrintStatements, m.ParserOptions()) + compiler, _ = loadCompilerFromStore(ctx, m.Store, txn, m.enablePrintStatements, m.ParserOptions(), m.GetExternalSources()) } if compiler != nil { @@ -1076,7 +1133,7 @@ func (m *Manager) onCommit(ctx context.Context, txn storage.Transaction, event s } } -func loadCompilerFromStore(ctx context.Context, store storage.Store, txn storage.Transaction, enablePrintStatements bool, popts ast.ParserOptions) (*ast.Compiler, error) { +func loadCompilerFromStore(ctx context.Context, store storage.Store, txn storage.Transaction, enablePrintStatements bool, popts ast.ParserOptions, externalSources *util.HasherMap[ast.Ref, ast.ExternalRuleSource]) (*ast.Compiler, error) { policies, err := store.ListPolicies(ctx, txn) if err != nil { return nil, err @@ -1102,6 +1159,14 @@ func loadCompilerFromStore(ctx context.Context, store storage.Store, txn storage compiler = compiler.WithDefaultRegoVersion(popts.RegoVersion) } + // Apply external sources BEFORE compilation + if externalSources != nil { + externalSources.Iter(func(ref ast.Ref, source ast.ExternalRuleSource) bool { + compiler = compiler.WithExternalSource(ref, source) + return false + }) + } + compiler.Compile(modules) return compiler, nil } diff --git a/v1/plugins/plugins_test.go b/v1/plugins/plugins_test.go index c27fb32f68..3d8b2a0f3f 100644 --- a/v1/plugins/plugins_test.go +++ b/v1/plugins/plugins_test.go @@ -14,6 +14,7 @@ import ( internal_tracing "github.com/open-policy-agent/opa/internal/distributedtracing" "github.com/open-policy-agent/opa/internal/storage/mock" + "github.com/open-policy-agent/opa/v1/ast" "github.com/open-policy-agent/opa/v1/logging" "github.com/open-policy-agent/opa/v1/logging/test" "github.com/open-policy-agent/opa/v1/plugins/rest" @@ -614,3 +615,156 @@ func (p prometheusRegisterMock) Unregister(collector prom.Collector) bool { delete(p.Collectors, collector) return true } + +// mockExternalSource is a simple implementation for testing +type mockExternalSource struct { + refs []ast.Ref + rules []*ast.Rule +} + +func (m *mockExternalSource) Refs() []ast.Ref { + return m.refs +} + +func (m *mockExternalSource) Init(context.Context, ast.Ref) (ast.ExternalRuleIndex, error) { + return &mockExternalIndex{rules: m.rules}, nil +} + +type mockExternalIndex struct { + rules []*ast.Rule +} + +func (*mockExternalIndex) Opts() *ast.ExternalSourceOptions { + return nil +} + +func (m *mockExternalIndex) Lookup(context.Context, ...ast.LookupOption) ([]*ast.Rule, ast.ExternalRuleIndex, error) { + return m.rules, nil, nil +} + +// testExternalSourcePlugin registers an external source during construction +type testExternalSourcePlugin struct { + manager *Manager + started bool +} + +func (p *testExternalSourcePlugin) Start(context.Context) error { + p.started = true + return nil +} + +func (*testExternalSourcePlugin) Stop(context.Context) {} + +func (*testExternalSourcePlugin) Reconfigure(context.Context, any) {} + +// TestExternalSourceIntegration verifies external source behavior during plugin lifecycle +func TestExternalSourceIntegration(t *testing.T) { + t.Run("sources wired after plugin start", func(t *testing.T) { + ctx := context.Background() + m, err := New([]byte(`{}`), "test", inmem.New()) + if err != nil { + t.Fatalf("Failed to create manager: %v", err) + } + + if err := m.Init(ctx); err != nil { + t.Fatalf("Failed to initialize manager: %v", err) + } + + module := ast.MustParseModule(`package external.test +test_rule := true`) + pkgRef := ast.MustParseRef("data.external.test") + + source := &mockExternalSource{ + refs: []ast.Ref{pkgRef}, + rules: module.Rules, + } + + plugin := &testExternalSourcePlugin{manager: m} + m.Register("test_external_source", plugin) + m.RegisterExternalSource(pkgRef, source) + + if m.GetExternalSources() == nil || m.GetExternalSources().Len() != 1 { + t.Fatalf("Expected 1 external source, got %d", m.GetExternalSources().Len()) + } + + if err := m.Start(ctx); err != nil { + t.Fatalf("Failed to start manager: %v", err) + } + + if !plugin.started { + t.Fatal("Expected plugin to be started") + } + + compiler := m.GetCompiler() + if compiler == nil || compiler.RuleTree == nil { + t.Fatal("Expected compiler with rule tree after Start()") + } + }) + + t.Run("stop cleans up external source plugins", func(t *testing.T) { + ctx := t.Context() + m, err := New([]byte(`{}`), "test", inmem.New()) + if err != nil { + t.Fatalf("Failed to create manager: %v", err) + } + + if err := m.Init(ctx); err != nil { + t.Fatalf("Failed to initialize manager: %v", err) + } + + module := ast.MustParseModule(`package external.test +test_rule := true`) + pkgRef := ast.MustParseRef("data.external.test") + + source := &mockExternalSource{ + refs: []ast.Ref{pkgRef}, + rules: module.Rules, + } + + plugin := &testExternalSourcePlugin{manager: m} + m.Register("test_external_source", plugin) + m.RegisterExternalSource(pkgRef, source) + + if err := m.Start(ctx); err != nil { + t.Fatalf("Failed to start manager: %v", err) + } + + if !plugin.started { + t.Fatal("Expected plugin to be started") + } + + m.Stop(ctx) + + if m.GetExternalSources() == nil || m.GetExternalSources().Len() != 1 { + t.Fatalf("Expected external sources to still be registered after stop, got %d", m.GetExternalSources().Len()) + } + }) + + t.Run("no recompilation when no sources registered", func(t *testing.T) { + ctx := context.Background() + m, err := New([]byte(`{}`), "test", inmem.New()) + if err != nil { + t.Fatalf("Failed to create manager: %v", err) + } + + if err := m.Init(ctx); err != nil { + t.Fatalf("Failed to initialize manager: %v", err) + } + + compilerBeforeStart := m.GetCompiler() + if compilerBeforeStart == nil { + t.Fatal("Expected compiler to be initialized after Init()") + } + + plugin := &testPlugin{m: m} + m.Register("test_plugin", plugin) + + if err := m.Start(ctx); err != nil { + t.Fatalf("Failed to start manager: %v", err) + } + + if m.GetCompiler() != compilerBeforeStart { + t.Fatal("Expected compiler to remain the same when no external sources registered") + } + }) +} diff --git a/v1/rego/rego.go b/v1/rego/rego.go index 6318ec70aa..309cba490c 100644 --- a/v1/rego/rego.go +++ b/v1/rego/rego.go @@ -673,6 +673,7 @@ type Rego struct { strictBuiltinErrors bool builtinErrorList *[]topdown.Error resolvers []refResolver + externalSources []ast.ExternalRuleSource schemaSet *ast.SchemaSet target string // target type (wasm, rego, etc.) opa opa.EvalEngine @@ -1287,6 +1288,15 @@ func Resolver(ref ast.Ref, r resolver.Resolver) func(r *Rego) { } } +// ExternalSource adds an external rule source that provides rules dynamically. +// The source declares which package refs it handles via its Refs() method. +// A single source can provide rules for multiple packages. +func ExternalSource(source ast.ExternalRuleSource) func(r *Rego) { + return func(rego *Rego) { + rego.externalSources = append(rego.externalSources, source) + } +} + // Schemas sets the schemaSet func Schemas(x *ast.SchemaSet) func(r *Rego) { return func(r *Rego) { @@ -2142,6 +2152,18 @@ func parserOptionsFromRegoVersionImport(imports []*ast.Import, popts ast.ParserO } func (r *Rego) compileModules(ctx context.Context, txn storage.Transaction, m metrics.Metrics) error { + if len(r.externalSources) > 0 && r.target != "" && r.target != targetRego { + return fmt.Errorf("external rule sources are not supported with target %q: only the default (rego) target is supported", r.target) + } + + // Apply external sources to the compiler before compilation + for i := range r.externalSources { + source := r.externalSources[i] + for _, ref := range source.Refs() { + r.compiler.WithExternalSource(ref, source) + } + } + // Only compile again if there are new modules. if len(r.bundles) > 0 || len(r.parsedModules) > 0 { diff --git a/v1/rego/rego_external_source_test.go b/v1/rego/rego_external_source_test.go new file mode 100644 index 0000000000..8e87e200de --- /dev/null +++ b/v1/rego/rego_external_source_test.go @@ -0,0 +1,456 @@ +package rego + +import ( + "context" + "errors" + "slices" + "strings" + "testing" + "time" + + "github.com/open-policy-agent/opa/v1/ast" +) + +type mockExternalSource struct { + refRules map[string][]*ast.Rule + visibleRefs []ast.Ref +} + +func newMockExternalSource(refs []ast.Ref, rules []*ast.Rule) *mockExternalSource { + refRules := make(map[string][]*ast.Rule) + for _, ref := range refs { + refRules[ref.String()] = rules + } + return &mockExternalSource{ + refRules: refRules, + } +} + +func (m *mockExternalSource) Refs() []ast.Ref { + refs := make([]ast.Ref, 0, len(m.refRules)) + for refStr := range m.refRules { + refs = append(refs, ast.MustParseRef(refStr)) + } + return refs +} + +func (m *mockExternalSource) Init(_ context.Context, ref ast.Ref) (ast.ExternalRuleIndex, error) { + rules, ok := m.refRules[ref.String()] + if !ok { + return nil, nil + } + return &mockExternalIndex{rules: rules, visibleRefs: m.visibleRefs}, nil +} + +type mockExternalIndex struct { + rules []*ast.Rule + visibleRefs []ast.Ref +} + +func (m *mockExternalIndex) Opts() *ast.ExternalSourceOptions { + if m.visibleRefs == nil { + return nil + } + return &ast.ExternalSourceOptions{VisibleRefs: m.visibleRefs} +} + +func (m *mockExternalIndex) Lookup(ctx context.Context, _ ...ast.LookupOption) ([]*ast.Rule, ast.ExternalRuleIndex, error) { + return m.rules, nil, nil +} + +func evalWithExternalSource(t *testing.T, ctx context.Context, query, module string, source *mockExternalSource, input map[string]any) ResultSet { + t.Helper() + opts := []func(*Rego){ + Query(query), + Module("test.rego", module), + ExternalSource(source), + } + if input != nil { + opts = append(opts, Input(input)) + } + r := New(opts...) + rs, err := r.Eval(ctx) + if err != nil { + t.Fatalf("Eval failed: %v", err) + } + return rs +} + +func partialEvalWithExternalSource(t *testing.T, ctx context.Context, query, module string, source *mockExternalSource) *PartialQueries { + t.Helper() + r := New( + Query(query), + Module("test.rego", module), + ExternalSource(source), + ) + pq, err := r.Partial(ctx) + if err != nil { + t.Fatalf("Partial failed: %v", err) + } + return pq +} + +func assertBoolResult(t *testing.T, rs ResultSet, expected bool, msg string) { + t.Helper() + if len(rs) != 1 { + t.Fatalf("Expected 1 result, got %d", len(rs)) + } + result, ok := rs[0].Expressions[0].Value.(bool) + if !ok { + t.Fatalf("Expected boolean result, got %T", rs[0].Expressions[0].Value) + } + if result != expected { + t.Errorf("%s: expected %v, got %v", msg, expected, result) + } +} + +func assertNoResults(t *testing.T, rs ResultSet, msg string) { + t.Helper() + if len(rs) != 0 { + t.Errorf("%s: expected 0 results, got %d", msg, len(rs)) + } +} + +func assertPartialQuery(t *testing.T, pq *PartialQueries, expectedQueries string) { + t.Helper() + expected := ast.MustParseBody(expectedQueries) + if len(pq.Queries) != 1 { + t.Fatalf("Expected 1 query, got %d", len(pq.Queries)) + } + if !pq.Queries[0].Equal(expected) { + t.Errorf("Expected PE result:\n%v\n\nGot:\n%v", expected, pq.Queries[0]) + } +} + +func TestExternalSourceDecisionMaking(t *testing.T) { + ctx := t.Context() + + externalModule := ast.MustParseModule(`package external.authz +allow if input.role == "admin"`) + + packageRef := ast.MustParseRef("data.external.authz") + source := newMockExternalSource([]ast.Ref{packageRef}, externalModule.Rules) + + staticModule := `package authz +default allow := false +allow if data.external.authz.allow` + + t.Run("admin role allowed", func(t *testing.T) { + rs := evalWithExternalSource(t, ctx, "data.authz.allow", staticModule, source, map[string]any{"role": "admin"}) + assertBoolResult(t, rs, true, "Expected allow=true for admin role") + }) + + t.Run("non-admin role not allowed", func(t *testing.T) { + rs := evalWithExternalSource(t, ctx, "data.authz.allow", staticModule, source, map[string]any{"role": "user"}) + assertBoolResult(t, rs, false, "Expected allow=false for user role") + }) +} + +func TestExternalSourcePartialEval(t *testing.T) { + ctx := t.Context() + + externalModule := ast.MustParseModule(`package external.authz +allow if input.role == "admin"`) + + packageRef := ast.MustParseRef("data.external.authz") + source := newMockExternalSource([]ast.Ref{packageRef}, externalModule.Rules) + + staticModule := `package authz +default allow := false +allow if data.external.authz.allow` + + t.Run("partial eval into external rule", func(t *testing.T) { + pq := partialEvalWithExternalSource(t, ctx, "data.authz.allow", staticModule, source) + assertPartialQuery(t, pq, `input.role = "admin"`) + }) +} + +func TestExternalSourceCallBackIntoStaticRego(t *testing.T) { + ctx := t.Context() + + externalModule := ast.MustParseModule(`package external.authz +allow if data.static.authz.foo == "bar"`) + + packageRef := ast.MustParseRef("data.external.authz") + source := &mockExternalSource{ + refRules: map[string][]*ast.Rule{ + packageRef.String(): externalModule.Rules, + }, + visibleRefs: []ast.Ref{ast.MustParseRef("data")}, + } + + staticModule := `package static.authz +default allow := false +allow if data.external.authz.allow + +foo := "bar"` + + rs := evalWithExternalSource(t, ctx, "data.static.authz.allow", staticModule, source, nil) + assertBoolResult(t, rs, true, "Expected allow=true") +} + +func TestExternalSourceCallBackIntoStaticRegoWithRecursion(t *testing.T) { + externalModule := ast.MustParseModule(`package external.authz +allow if data.static.authz.allow`) + + packageRef := ast.MustParseRef("data.external.authz") + + staticModule := `package static.authz +default allow := false +allow if data.external.authz.allow` + + t.Run("isolated by default prevents recursion", func(t *testing.T) { + ctx := t.Context() + source := newMockExternalSource([]ast.Ref{packageRef}, externalModule.Rules) + rs := evalWithExternalSource(t, ctx, "data.static.authz.allow", staticModule, source, nil) + assertBoolResult(t, rs, false, "Expected allow=false (isolated external source cannot access static policy)") + }) + + t.Run("visible refs allows recursion and hits deadline", func(t *testing.T) { + ctx := t.Context() + source := &mockExternalSource{ + refRules: map[string][]*ast.Rule{ + packageRef.String(): externalModule.Rules, + }, + visibleRefs: []ast.Ref{ast.MustParseRef("data")}, + } + + r := New( + Query("data.static.authz.allow"), + Module("authz.rego", staticModule), + ExternalSource(source), + ) + + ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond) + t.Cleanup(cancel) + _, err := r.Eval(ctx) + if !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("expected deadline-exceeded, got err: %v", err) + } + }) +} + +func TestExternalSourceCrossRefCalls(t *testing.T) { + ctx := t.Context() + + fooModule := ast.MustParseModule(`package external.foo +result if data.external.bar.value == 42`) + + barModule := ast.MustParseModule(`package external.bar +value := 42`) + + staticModule := `package main +result := data.external.foo.result` + + t.Run("isolated prevents cross-ref calls", func(t *testing.T) { + source := &mockExternalSource{ + refRules: map[string][]*ast.Rule{ + "data.external.foo": fooModule.Rules, + "data.external.bar": barModule.Rules, + }, + visibleRefs: []ast.Ref{}, // explicitly isolated + } + + rs := evalWithExternalSource(t, ctx, "data.main.result", staticModule, source, nil) + assertNoResults(t, rs, "Expected 0 results when isolated (external.foo cannot access external.bar)") + }) + + t.Run("visible refs allows cross-ref calls", func(t *testing.T) { + source := &mockExternalSource{ + refRules: map[string][]*ast.Rule{ + "data.external.foo": fooModule.Rules, + "data.external.bar": barModule.Rules, + }, + visibleRefs: []ast.Ref{ast.MustParseRef("data")}, + } + + rs := evalWithExternalSource(t, ctx, "data.main.result", staticModule, source, nil) + assertBoolResult(t, rs, true, "Expected result=true when all refs visible") + }) +} + +func TestExternalSourceNestedPackage(t *testing.T) { + ctx := t.Context() + + externalModule := ast.MustParseModule(`package external.project.authz + +allowed if other_rule +other_rule if input.foo == "bar"`) + + parentRef := ast.MustParseRef("data.external") + source := newMockExternalSource([]ast.Ref{parentRef}, externalModule.Rules) + + staticModule := `package main +result := data.external.project.authz.allowed` + + t.Run("eval allowed when foo is bar", func(t *testing.T) { + rs := evalWithExternalSource(t, ctx, "data.main.result", staticModule, source, map[string]any{"foo": "bar"}) + assertBoolResult(t, rs, true, "Expected allowed=true when foo is bar") + }) + + t.Run("eval denied when foo is not bar", func(t *testing.T) { + rs := evalWithExternalSource(t, ctx, "data.main.result", staticModule, source, map[string]any{"foo": "baz"}) + assertNoResults(t, rs, "Expected no results when foo is not bar") + }) + + t.Run("partial eval", func(t *testing.T) { + pq := partialEvalWithExternalSource(t, ctx, "data.main.result", staticModule, source) + assertPartialQuery(t, pq, `input.foo = "bar"`) + }) + + t.Run("call external function with argument fails", func(t *testing.T) { + externalModuleWithFunc := ast.MustParseModule(`package external.authz + +foo(x) if x == "bar"`) + + funcSource := newMockExternalSource([]ast.Ref{ast.MustParseRef("data.external.authz")}, externalModuleWithFunc.Rules) + + staticModuleWithFuncCall := `package main +allow if data.external.authz.foo("bar")` + + r := New( + Query("data.main.allow"), + Module("test.rego", staticModuleWithFuncCall), + ExternalSource(funcSource), + ) + _, err := r.Eval(ctx) + if err == nil { + t.Fatal("Expected error when calling external function with argument, but got none") + } + var errs ast.Errors + if !errors.As(err, &errs) { + t.Fatalf("Expected ast.Errors, got: %T: %v", err, err) + } + if !slices.ContainsFunc(errs, func(e *ast.Error) bool { return e.Code == ast.TypeErr }) { + t.Errorf("Expected type error for undefined function, got: %v", err) + } + }) + + t.Run("call into external rule that uses function internally", func(t *testing.T) { + externalModuleWithInternalFunc := ast.MustParseModule(`package external.authz + +allowed if foo("bar") +foo(x) if x == "bar"`) + + funcSource := newMockExternalSource([]ast.Ref{ast.MustParseRef("data.external.authz")}, externalModuleWithInternalFunc.Rules) + + staticModuleWithRuleCall := `package main +allow if data.external.authz.allowed` + + rs := evalWithExternalSource(t, ctx, "data.main.allow", staticModuleWithRuleCall, funcSource, nil) + assertBoolResult(t, rs, true, "Expected allow=true when external rule internally uses function") + }) +} + +func TestExternalSourcePartialVisibility(t *testing.T) { + ctx := t.Context() + + externalModule := ast.MustParseModule(`package external.authz +allow if { + data.pkg_a.check + data.pkg_b.check +}`) + + packageRef := ast.MustParseRef("data.external.authz") + + modA := `package pkg_a +check if input.role == "admin"` + + modB := `package pkg_b +check := true` + + t.Run("only pkg_a visible", func(t *testing.T) { + source := &mockExternalSource{ + refRules: map[string][]*ast.Rule{ + packageRef.String(): externalModule.Rules, + }, + visibleRefs: []ast.Ref{ast.MustParseRef("data.pkg_a")}, + } + + r := New( + Query("data.external.authz.allow"), + Module("a.rego", modA), + Module("b.rego", modB), + ExternalSource(source), + Input(map[string]any{"role": "admin"}), + ) + rs, err := r.Eval(ctx) + if err != nil { + t.Fatal(err) + } + assertNoResults(t, rs, "Expected 0 results (pkg_b not visible)") + }) + + t.Run("both visible", func(t *testing.T) { + source := &mockExternalSource{ + refRules: map[string][]*ast.Rule{ + packageRef.String(): externalModule.Rules, + }, + visibleRefs: []ast.Ref{ + ast.MustParseRef("data.pkg_a"), + ast.MustParseRef("data.pkg_b"), + }, + } + + r := New( + Query("data.external.authz.allow"), + Module("a.rego", modA), + Module("b.rego", modB), + ExternalSource(source), + Input(map[string]any{"role": "admin"}), + ) + rs, err := r.Eval(ctx) + if err != nil { + t.Fatal(err) + } + assertBoolResult(t, rs, true, "Expected allow=true when both packages visible") + }) + + t.Run("all of data visible", func(t *testing.T) { + source := &mockExternalSource{ + refRules: map[string][]*ast.Rule{ + packageRef.String(): externalModule.Rules, + }, + visibleRefs: []ast.Ref{ast.MustParseRef("data")}, + } + + r := New( + Query("data.external.authz.allow"), + Module("a.rego", modA), + Module("b.rego", modB), + ExternalSource(source), + Input(map[string]any{"role": "admin"}), + ) + rs, err := r.Eval(ctx) + if err != nil { + t.Fatal(err) + } + assertBoolResult(t, rs, true, "Expected allow=true when all of data is visible") + }) +} + +func TestExternalSourceUnsupportedTarget(t *testing.T) { + externalModule := ast.MustParseModule(`package external.authz +allow if input.role == "admin"`) + + packageRef := ast.MustParseRef("data.external.authz") + source := newMockExternalSource([]ast.Ref{packageRef}, externalModule.Rules) + + for _, target := range []string{"wasm", "plan", "custom-plugin"} { + t.Run(target, func(t *testing.T) { + r := New( + Query("data.external.authz.allow"), + Module("test.rego", `package test`), + ExternalSource(source), + Target(target), + ) + _, err := r.Eval(t.Context()) + if err == nil { + t.Fatal("Expected error for non-rego target with external sources") + } + if !strings.Contains(err.Error(), "external rule sources are not supported") { + t.Fatalf("Unexpected error: %v", err) + } + }) + } +} diff --git a/v1/topdown/eval.go b/v1/topdown/eval.go index 61e11ec32f..8447bba081 100644 --- a/v1/topdown/eval.go +++ b/v1/topdown/eval.go @@ -92,8 +92,9 @@ type eval struct { input *ast.Term data *ast.Term external *resolverTrie + externalTreeStack *externalTreeStack targetStack *refStack - traceLastLocation *ast.Location // Last location of a trace event. + traceLastLocation *ast.Location instr *Instrumentation builtins map[string]*Builtin builtinCache builtins.Cache @@ -267,12 +268,15 @@ func (e *eval) unknown(x any, b *bindings) bool { x = ast.NewTerm(v) } - return saveRequired(e.compiler, e.inliningControl, true, e.saveSet, b, x, false) + return saveRequired(e.compiler.RuleTree, e.externalTreeStack, e.inliningControl, true, e.saveSet, b, x, false) } // exactly like `unknown` above` but without the cost of `any` boxing when arg is known to be a ref func (e *eval) unknownRef(ref ast.Ref, b *bindings) bool { - return e.partial() && saveRequired(e.compiler, e.inliningControl, true, e.saveSet, b, ast.NewTerm(ref), false) + if !e.partial() { + return false + } + return saveRequired(e.compiler.RuleTree, e.externalTreeStack, e.inliningControl, true, e.saveSet, b, ast.NewTerm(ref), false) } func (e *eval) traceEnter(x ast.Node) { @@ -755,26 +759,33 @@ func (e *eval) evalWith(iter evalIterator) error { } } - oldInput, oldData := e.evalWithPush(input, data, functionMocks, targets, disable) + oldInput, oldData, pushedFrame := e.evalWithPush(input, data, functionMocks, targets, disable) err = e.evalStep(func(e *eval) error { - e.evalWithPop(oldInput, oldData) + e.evalWithPop(oldInput, oldData, pushedFrame) err := e.next(iter) - oldInput, oldData = e.evalWithPush(input, data, functionMocks, targets, disable) + oldInput, oldData, pushedFrame = e.evalWithPush(input, data, functionMocks, targets, disable) return err }) - e.evalWithPop(oldInput, oldData) + e.evalWithPop(oldInput, oldData, pushedFrame) return err } -func (e *eval) evalWithPush(input, data *ast.Term, functionMocks [][2]*ast.Term, targets, disable []ast.Ref) (*ast.Term, *ast.Term) { +func (e *eval) evalWithPush(input, data *ast.Term, functionMocks [][2]*ast.Term, targets, disable []ast.Ref) (*ast.Term, *ast.Term, bool) { var oldInput *ast.Term + var pushedFrame bool if input != nil { oldInput = e.input e.input = input + + // When input changes, push a new frame for external tree caching + if e.externalTreeStack != nil { + e.externalTreeStack.PushFrame() + pushedFrame = true + } } var oldData *ast.Term @@ -804,16 +815,22 @@ func (e *eval) evalWithPush(input, data *ast.Term, functionMocks [][2]*ast.Term, e.functionMocks.PutPairs(functionMocks) - return oldInput, oldData + return oldInput, oldData, pushedFrame } -func (e *eval) evalWithPop(input, data *ast.Term) { +func (e *eval) evalWithPop(input, data *ast.Term, popFrame bool) { // NOTE(ae) no nil checks here as we assume evalWithPush always called first e.inliningControl.PopDisable() e.targetStack.Pop() e.virtualCache.Pop() e.comprehensionCache.Pop() e.functionMocks.PopPairs() + + // When input is restored, pop the external tree frame + if popFrame { + e.externalTreeStack.PopFrame() + } + e.data = data e.input = input } @@ -980,15 +997,19 @@ func (e *eval) evalCall(terms []*ast.Term, iter unifyIterator) error { var ir *ast.IndexResult var err error + index := e.ruleIndex(ref) if e.partial() { - ir, err = e.getRules(ref, nil) + ir, err = e.getRules(ref, nil, index) } else { - ir, err = e.getRules(ref, terms[1:]) + ir, err = e.getRules(ref, terms[1:], index) } defer ast.IndexResultPool.Put(ir) if err != nil { return err } + if ir == nil { + return nil + } eval := evalFuncPool.Get() defer evalFuncPool.Put(eval) @@ -1724,11 +1745,10 @@ func (e *eval) saveInlinedNegatedExprs(exprs []*ast.Expr, iter unifyIterator) er return err } -func (e *eval) getRules(ref ast.Ref, args []*ast.Term) (*ast.IndexResult, error) { +func (e *eval) getRules(ref ast.Ref, args []*ast.Term, index ast.RuleIndex) (*ast.IndexResult, error) { e.instr.startTimer(evalOpRuleIndex) defer e.instr.stopTimer(evalOpRuleIndex) - index := e.ruleIndex(ref) if index == nil { return nil, nil } @@ -1742,6 +1762,7 @@ func (e *eval) getRules(ref ast.Ref, args []*ast.Term) (*ast.IndexResult, error) var result *ast.IndexResult var err error + resolver.e = e if e.indexing { resolver.args = args @@ -1779,11 +1800,6 @@ func (e *eval) getRules(ref ast.Ref, args []*ast.Term) (*ast.IndexResult, error) return result, err } -// ruleIndex performs a lookup for a RuleIndex in the compiler's RuleTree. -func (e *eval) ruleIndex(ref ast.Ref) ast.RuleIndex { - return e.compiler.RuleIndex(ref) -} - func (e *eval) Resolve(ref ast.Ref) (ast.Value, error) { return (&evalResolver{e: e}).Resolve(ref) } @@ -1996,7 +2012,7 @@ func (e *eval) getDeclArgsLen(x *ast.Expr) (int, error) { return bi.Decl.Arity(), nil } - ir, err := e.getRules(operator, nil) + ir, err := e.getRules(operator, nil, e.ruleIndex(operator)) defer ast.IndexResultPool.Put(ir) if err != nil { return -1, err @@ -2536,10 +2552,48 @@ func (e evalTree) next(iter unifyIterator, plugged *ast.Term) error { cpy.plugged[e.pos] = plugged cpy.pos++ + // Track whether we pushed an external tree that needs cleanup + pushedExternalTree := false if !e.e.targetStack.Prefixed(cpy.plugged[:cpy.pos]) { if e.node != nil { node = e.node.Child(plugged.Value) - if node != nil && len(node.Values) > 0 { + + // Handle external sources transparently + if node != nil && node.External != nil { + externalRef := node.External.Ref + externalIndex := node.External.Index + + // Initialize externalTreeStack if needed + if e.e.externalTreeStack == nil { + e.e.externalTreeStack = newExternalTreeStack(e.e) + } + + // Check cache first + cachedNode, _, found := e.e.externalTreeStack.findCached(externalRef) + if found { + node = cachedNode + } else { + // Call Tree() and cache the result + e.e.instr.startTimer(evalOpExternalRuleSource) + tree, updatedIndex, err := node.External.Tree(e.e.ctx, e.e.compiler.RuleTree, externalRef, e.e.input, e.e.metrics, e.e.requestMetadata, e.e.responseMetadata) + e.e.instr.stopTimer(evalOpExternalRuleSource) + if err != nil { + return err + } + if tree != nil { + if updatedIndex != nil { + externalIndex = updatedIndex + } + e.e.externalTreeStack.Push(externalRef, tree, externalIndex, e.e.input) + node = tree + pushedExternalTree = true + } + } + } + + hasRules := node != nil && len(node.Values) > 0 + + if hasRules { r := evalVirtual{ e: e.e, ref: e.ref, @@ -2550,13 +2604,21 @@ func (e evalTree) next(iter unifyIterator, plugged *ast.Term) error { rbindings: e.rbindings, } r.plugged[e.pos] = plugged - return r.eval(iter) + err := r.eval(iter) + if pushedExternalTree { + e.e.externalTreeStack.Pop() + } + return err } } } cpy.node = node - return cpy.eval(iter) + err := cpy.eval(iter) + if pushedExternalTree { + e.e.externalTreeStack.Pop() + } + return err } // enumerateNext is a helper to avoid closure allocation in enumerate loops. @@ -2732,12 +2794,16 @@ type evalVirtual struct { func (e evalVirtual) eval(iter unifyIterator) error { - ir, err := e.e.getRules(e.plugged[:e.pos+1], nil) + ir, err := e.e.getRules(e.plugged[:e.pos+1], nil, e.e.ruleIndex(e.plugged[:e.pos+1])) defer ast.IndexResultPool.Put(ir) if err != nil { return err } + if ir == nil { + return nil + } + // Partial evaluation of ordered rules is not supported currently. Save the // expression and continue. This could be revisited in the future. if len(ir.Else) > 0 && e.e.unknownRef(e.ref, e.bindings) { @@ -4490,3 +4556,180 @@ func (e *eval) updateSavedMocks(withs []*ast.With) []*ast.With { } return ret } + +// simpleTreeNode provides minimal tree structure for navigation +type simpleTreeNode struct { + tree *ast.TreeNode + children map[ast.Value]*simpleTreeNode +} + +func newSimpleTreeNode() *simpleTreeNode { + return &simpleTreeNode{ + children: make(map[ast.Value]*simpleTreeNode), + } +} + +// externalTreeStack caches external rule trees and tracks frames for input changes. +// It maintains both a flat cache for lookups and a tree structure for navigation. +type externalTreeStack struct { + eval *eval + entries []externalTreeEntry // flat list of cached entries + frames []int // frame markers for input changes (indices into entries) + root *simpleTreeNode // tree structure for navigation +} + +type externalTreeEntry struct { + ref ast.Ref + input *ast.Term + tree *ast.TreeNode + index ast.ExternalRuleIndex +} + +func newExternalTreeStack(e *eval) *externalTreeStack { + return &externalTreeStack{ + eval: e, + entries: make([]externalTreeEntry, 0, 4), + frames: make([]int, 0, 4), + } +} + +// findCached checks if we already have a cached tree for this ref. +// Frame tracking ensures any cached entry has the correct input. +func (s *externalTreeStack) findCached(ref ast.Ref) (*ast.TreeNode, ast.ExternalRuleIndex, bool) { + // Determine search boundary: only search within current frame if one exists + startIdx := 0 + if len(s.frames) > 0 { + startIdx = s.frames[len(s.frames)-1] + } + + // Search from most recent to the frame boundary + for i := len(s.entries) - 1; i >= startIdx; i-- { + entry := &s.entries[i] + if entry.ref.Equal(ref) { + return entry.tree, entry.index, true + } + } + return nil, nil, false +} + +func (s *externalTreeStack) Push(ref ast.Ref, tree *ast.TreeNode, index ast.ExternalRuleIndex, input *ast.Term) { + // Add entry to cache (we never have duplicates in the same frame) + s.entries = append(s.entries, externalTreeEntry{ + ref: ref, + input: input, + tree: tree, + index: index, + }) + + // Update root tree structure + if s.root == nil { + s.root = newSimpleTreeNode() + } + node := s.root + for _, term := range ref { + key := term.Value + if node.children[key] == nil { + node.children[key] = newSimpleTreeNode() + } + node = node.children[key] + } + node.tree = tree +} + +// Pop removes the most recent entry from the stack and closes its index. +// If a frame is active and entries are at or below the frame boundary, +// Pop is a no-op — PopFrame already cleaned up (or will clean up) those entries. +func (s *externalTreeStack) Pop() { + if len(s.entries) == 0 { + return + } + + // Don't pop at or below the current frame boundary. + if len(s.frames) > 0 && len(s.entries) <= s.frames[len(s.frames)-1] { + return + } + + // Close the external index if it supports closing + lastEntry := &s.entries[len(s.entries)-1] + if closer, ok := lastEntry.index.(ast.ExternalRuleIndexCloser); ok { + _ = closer.Close() + } + + // Remove the most recent entry + s.entries = s.entries[:len(s.entries)-1] + + // Rebuild tree from remaining entries + s.root = newSimpleTreeNode() + for i := range s.entries { + entry := &s.entries[i] + node := s.root + for _, term := range entry.ref { + key := term.Value + if node.children[key] == nil { + node.children[key] = newSimpleTreeNode() + } + node = node.children[key] + } + node.tree = entry.tree + } +} + +// PushFrame marks the current stack position for input changes +func (s *externalTreeStack) PushFrame() { + s.frames = append(s.frames, len(s.entries)) +} + +// PopFrame restores the cache to the marked position +func (s *externalTreeStack) PopFrame() { + if len(s.frames) == 0 { + return + } + + // Get the frame marker and truncate entries + targetSize := s.frames[len(s.frames)-1] + s.frames = s.frames[:len(s.frames)-1] + + // Close indices of entries being removed + for i := len(s.entries) - 1; i >= targetSize; i-- { + if closer, ok := s.entries[i].index.(ast.ExternalRuleIndexCloser); ok { + _ = closer.Close() + } + } + + // Rebuild tree from remaining entries + s.root = newSimpleTreeNode() + for i := range targetSize { + entry := &s.entries[i] + node := s.root + for _, term := range entry.ref { + key := term.Value + if node.children[key] == nil { + node.children[key] = newSimpleTreeNode() + } + node = node.children[key] + } + node.tree = entry.tree + } + + s.entries = s.entries[:targetSize] +} + +// ruleIndex performs a shadowed lookup for a RuleIndex, checking external trees first. +// It searches through the pushStack (most recent to oldest), navigating the tree structure +// to find matching rules, then falls back to the compiler's static RuleTree. +func (e *eval) ruleIndex(ref ast.Ref) ast.RuleIndex { + if e.externalTreeStack != nil && len(e.externalTreeStack.entries) > 0 { + // Search from most recent to oldest + for i := len(e.externalTreeStack.entries) - 1; i >= 0; i-- { + entry := &e.externalTreeStack.entries[i] + if ref.HasPrefix(entry.ref) { + // Look for the relative ref in the cached tree + relativeRef := ref[len(entry.ref):] + if found := entry.tree.Find(relativeRef); found != nil { + return found.Index + } + } + } + } + return e.compiler.RuleIndex(ref) +} diff --git a/v1/topdown/external_source_test.go b/v1/topdown/external_source_test.go new file mode 100644 index 0000000000..3e4c8ccd07 --- /dev/null +++ b/v1/topdown/external_source_test.go @@ -0,0 +1,429 @@ +package topdown + +import ( + "context" + "errors" + "slices" + "sync/atomic" + "testing" + + "github.com/open-policy-agent/opa/v1/ast" + "github.com/open-policy-agent/opa/v1/metrics" + "github.com/open-policy-agent/opa/v1/storage/inmem" +) + +type countingExternalSource struct { + refs []ast.Ref + rules []*ast.Rule + callCount int32 +} + +func (m *countingExternalSource) Init(context.Context, ast.Ref) (ast.ExternalRuleIndex, error) { + return &countingExternalIndex{rules: m.rules, callCount: &m.callCount}, nil +} + +func (m *countingExternalSource) Refs() []ast.Ref { + return m.refs +} + +type countingExternalIndex struct { + rules []*ast.Rule + callCount *int32 +} + +func (*countingExternalIndex) Opts() *ast.ExternalSourceOptions { + return nil +} + +func (m *countingExternalIndex) Lookup(context.Context, ...ast.LookupOption) ([]*ast.Rule, ast.ExternalRuleIndex, error) { + atomic.AddInt32(m.callCount, 1) + return m.rules, nil, nil +} + +func (m *countingExternalSource) getCallCount() int { + return int(atomic.LoadInt32(&m.callCount)) +} + +func setupCompiler(t *testing.T, packageRef ast.Ref, source ast.ExternalRuleSource, staticModule *ast.Module) *ast.Compiler { + t.Helper() + compiler := ast.NewCompiler() + compiler.WithExternalSource(packageRef, source) + modules := map[string]*ast.Module{} + if staticModule != nil { + modules["main.rego"] = staticModule + } + compiler.Compile(modules) + if compiler.Failed() { + t.Fatalf("Compiler failed: %v", compiler.Errors) + } + return compiler +} + +func runQuery(t *testing.T, compiler *ast.Compiler, queryStr string, input *ast.Term) QueryResultSet { + t.Helper() + store := inmem.New() + ctx := t.Context() + txn, err := store.NewTransaction(ctx) + if err != nil { + t.Fatal(err) + } + defer store.Abort(ctx, txn) + m := metrics.New() + instr := NewInstrumentation(m) + + query := ast.MustParseBody(queryStr) + q := NewQuery(query). + WithCompiler(compiler). + WithStore(store). + WithTransaction(txn). + WithInput(input). + WithInstrumentation(instr) + + qrs, err := q.Run(ctx) + if err != nil { + t.Fatalf("Query failed: %v", err) + } + t.Logf("metrics: %v", m.All()) + return qrs +} + +func TestExternalSourceE2EWithInputOverride(t *testing.T) { + t.Parallel() + + externalModule := ast.MustParseModule(`package authz +allowed if input.user == "alice"`) + + packageRef := ast.MustParseRef("data.authz") + source := &countingExternalSource{refs: []ast.Ref{packageRef}, rules: externalModule.Rules} + + staticModule := ast.MustParseModule(`package main +check if { + data.authz.allowed + data.authz.allowed with input as {"user": "bob"} +}`) + + compiler := setupCompiler(t, packageRef, source, staticModule) + + input := ast.MustParseTerm(`{"user": "alice"}`) + qrs := runQuery(t, compiler, "data.main.check", input) + + if len(qrs) != 0 { + t.Errorf("Expected 0 results (second check with bob should fail), got %d", len(qrs)) + } + + if callCount := source.getCallCount(); callCount != 2 { + t.Errorf("Expected external source to be called twice (once per input), got %d calls", callCount) + } +} + +func TestExternalSourceE2EWithMultipleRulesFromSamePackage(t *testing.T) { + t.Parallel() + + externalModule := ast.MustParseModule(`package authz +allow if input.user == "alice" +deny if input.action == "delete" +allowed if { + allow + not deny +}`) + + packageRef := ast.MustParseRef("data.authz") + source := &countingExternalSource{refs: []ast.Ref{packageRef}, rules: externalModule.Rules} + + staticModule := ast.MustParseModule(`package main +check if data.authz.allowed`) + + compiler := setupCompiler(t, packageRef, source, staticModule) + + input := ast.MustParseTerm(`{"user": "alice", "action": "read"}`) + qrs := runQuery(t, compiler, "data.main.check", input) + + if len(qrs) != 1 { + t.Errorf("Expected 1 result, got %d", len(qrs)) + } + + if callCount := source.getCallCount(); callCount != 1 { + t.Errorf("Expected external source to be called once (cached for same ref and input), got %d calls", callCount) + } +} + +type closableExternalSource struct { + refs []ast.Ref + rules []*ast.Rule + closeCalls int32 +} + +func (m *closableExternalSource) Init(context.Context, ast.Ref) (ast.ExternalRuleIndex, error) { + return &closableExternalIndex{rules: m.rules, closeCalls: &m.closeCalls}, nil +} + +func (m *closableExternalSource) Refs() []ast.Ref { + return m.refs +} + +func (m *closableExternalSource) getCloseCalls() int { + return int(atomic.LoadInt32(&m.closeCalls)) +} + +type closableExternalIndex struct { + rules []*ast.Rule + closeCalls *int32 +} + +func (*closableExternalIndex) Opts() *ast.ExternalSourceOptions { + return nil +} + +func (m *closableExternalIndex) Lookup(context.Context, ...ast.LookupOption) ([]*ast.Rule, ast.ExternalRuleIndex, error) { + return m.rules, nil, nil +} + +func (m *closableExternalIndex) Close() error { + atomic.AddInt32(m.closeCalls, 1) + return nil +} + +func TestExternalSourceCloseCalled(t *testing.T) { + t.Parallel() + + externalModule := ast.MustParseModule(`package authz +allowed if input.user == "alice"`) + + packageRef := ast.MustParseRef("data.authz") + source := &closableExternalSource{refs: []ast.Ref{packageRef}, rules: externalModule.Rules} + + staticModule := ast.MustParseModule(`package main +check if data.authz.allowed`) + + compiler := setupCompiler(t, packageRef, source, staticModule) + + input := ast.MustParseTerm(`{"user": "alice"}`) + qrs := runQuery(t, compiler, "data.main.check", input) + + if len(qrs) != 1 { + t.Errorf("Expected 1 result, got %d", len(qrs)) + } + + if closeCalls := source.getCloseCalls(); closeCalls != 1 { + t.Errorf("Expected Close() to be called once, got %d calls", closeCalls) + } +} + +type preCompiledRulesSource struct { + refs []ast.Ref + compiledRules []*ast.Rule +} + +func (s *preCompiledRulesSource) Init(context.Context, ast.Ref) (ast.ExternalRuleIndex, error) { + return &preCompiledRulesIndex{compiledRules: s.compiledRules}, nil +} + +func (s *preCompiledRulesSource) Refs() []ast.Ref { + return s.refs +} + +type preCompiledRulesIndex struct { + compiledRules []*ast.Rule +} + +func (*preCompiledRulesIndex) Opts() *ast.ExternalSourceOptions { + // For pre-compiled rules, skip all stages except those essential for + // integrating the rules into the compiler + var skippedStages []ast.StageID + essentialStages := []ast.StageID{ + ast.StageSetModuleTree, + ast.StageSetRuleTree, ast.StageBuildRuleIndices, + } + + for _, stage := range ast.AllStages() { + if !slices.Contains(essentialStages, stage) { + skippedStages = append(skippedStages, stage) + } + } + + return &ast.ExternalSourceOptions{ + SkippedStages: skippedStages, + } +} + +func (idx *preCompiledRulesIndex) Lookup(_ context.Context, _ ...ast.LookupOption) ([]*ast.Rule, ast.ExternalRuleIndex, error) { + return idx.compiledRules, nil, nil +} + +func TestExternalSourceWithPreCompiledRules(t *testing.T) { + t.Parallel() + + // Create and pre-compile an external module + externalModule := ast.MustParseModule(`package authz + +allow if input.user == "admin" +deny if input.action == "delete" + +permitted if { + allow + not deny +}`) + + // Pre-compile the rules using a separate compiler + preCompiler := ast.NewCompiler() + preCompiler.Compile(map[string]*ast.Module{"authz.rego": externalModule}) + + if preCompiler.Failed() { + t.Fatalf("Pre-compilation failed: %v", preCompiler.Errors) + } + + // Extract the pre-compiled rules + compiledRules := make([]*ast.Rule, 0, len(preCompiler.Modules)) + for _, mod := range preCompiler.Modules { + compiledRules = append(compiledRules, mod.Rules...) + } + + if len(compiledRules) == 0 { + t.Fatal("No compiled rules found") + } + + // Create an external source that returns pre-compiled rules + packageRef := ast.MustParseRef("data.authz") + source := &preCompiledRulesSource{ + refs: []ast.Ref{packageRef}, + compiledRules: compiledRules, + } + + // Create a static module that uses the externally-provided rules + staticModule := ast.MustParseModule(`package main + +check if data.authz.permitted`) + + // Set up compiler with the external source + compiler := setupCompiler(t, packageRef, source, staticModule) + + t.Run("admin with read action should be allowed", func(t *testing.T) { + input := ast.MustParseTerm(`{"user": "admin", "action": "read"}`) + qrs := runQuery(t, compiler, "data.main.check", input) + + if len(qrs) != 1 { + t.Errorf("Expected 1 result (allowed), got %d", len(qrs)) + } + }) + + t.Run("admin with delete action should be denied", func(t *testing.T) { + input := ast.MustParseTerm(`{"user": "admin", "action": "delete"}`) + qrs := runQuery(t, compiler, "data.main.check", input) + + if len(qrs) != 0 { + t.Errorf("Expected 0 results (denied), got %d", len(qrs)) + } + }) + + t.Run("non-admin user should be denied", func(t *testing.T) { + input := ast.MustParseTerm(`{"user": "bob", "action": "read"}`) + qrs := runQuery(t, compiler, "data.main.check", input) + + if len(qrs) != 0 { + t.Errorf("Expected 0 results (denied), got %d", len(qrs)) + } + }) +} + +func TestExternalSourceDirectQuery(t *testing.T) { + t.Parallel() + + externalModule := ast.MustParseModule(`package authz +allow if input.user == "alice" +deny if input.user == "bob"`) + + packageRef := ast.MustParseRef("data.authz") + source := &countingExternalSource{refs: []ast.Ref{packageRef}, rules: externalModule.Rules} + + compiler := setupCompiler(t, packageRef, source, nil) + + t.Run("direct query allowed", func(t *testing.T) { + input := ast.MustParseTerm(`{"user": "alice"}`) + qrs := runQuery(t, compiler, "data.authz.allow", input) + if len(qrs) != 1 { + t.Errorf("Expected 1 result for direct external query, got %d", len(qrs)) + } + }) + + t.Run("direct query denied", func(t *testing.T) { + input := ast.MustParseTerm(`{"user": "bob"}`) + qrs := runQuery(t, compiler, "data.authz.deny", input) + if len(qrs) != 1 { + t.Errorf("Expected 1 result for direct external query, got %d", len(qrs)) + } + }) + + t.Run("direct query no match", func(t *testing.T) { + input := ast.MustParseTerm(`{"user": "charlie"}`) + qrs := runQuery(t, compiler, "data.authz.allow", input) + if len(qrs) != 0 { + t.Errorf("Expected 0 results, got %d", len(qrs)) + } + }) +} + +func TestExternalSourceCompilationFailure(t *testing.T) { + t.Parallel() + + // Rules with an unsafe variable should fail compilation + externalModule := ast.MustParseModule(`package authz +allow if { x }`) + + packageRef := ast.MustParseRef("data.authz") + source := &countingExternalSource{refs: []ast.Ref{packageRef}, rules: externalModule.Rules} + + staticModule := ast.MustParseModule(`package main +check if data.authz.allow`) + + compiler := setupCompiler(t, packageRef, source, staticModule) + + store := inmem.New() + ctx := t.Context() + txn, err := store.NewTransaction(ctx) + if err != nil { + t.Fatal(err) + } + defer store.Abort(ctx, txn) + + query := ast.MustParseBody("data.main.check") + q := NewQuery(query). + WithCompiler(compiler). + WithStore(store). + WithTransaction(txn) + + _, err = q.Run(ctx) + if err == nil { + t.Fatal("Expected compilation error from external source, got nil") + } + var errs ast.Errors + if !errors.As(err, &errs) { + t.Fatalf("Expected ast.Errors, got: %T: %v", err, err) + } + if !slices.ContainsFunc(errs, func(e *ast.Error) bool { return e.Code == ast.UnsafeVarErr }) { + t.Errorf("Expected unsafe var error, got: %v", err) + } +} + +func TestExternalSourceE2EWithInputOverrideNilInput(t *testing.T) { + t.Parallel() + + externalModule := ast.MustParseModule(`package authz +allowed if input.user == "alice"`) + + packageRef := ast.MustParseRef("data.authz") + source := &countingExternalSource{refs: []ast.Ref{packageRef}, rules: externalModule.Rules} + + staticModule := ast.MustParseModule(`package main +check if data.authz.allowed with input as {"user": "alice"}`) + + compiler := setupCompiler(t, packageRef, source, staticModule) + + // Query with nil input — the with clause provides it. + // This previously panicked because evalWithPop skipped PopFrame + // when oldInput was nil, leaking a frame on the externalTreeStack. + qrs := runQuery(t, compiler, "data.main.check", nil) + + if len(qrs) != 1 { + t.Errorf("Expected 1 result, got %d", len(qrs)) + } +} diff --git a/v1/topdown/instrumentation.go b/v1/topdown/instrumentation.go index 93da1d0022..3fc7e0138e 100644 --- a/v1/topdown/instrumentation.go +++ b/v1/topdown/instrumentation.go @@ -19,6 +19,7 @@ const ( evalOpComprehensionCacheBuild = "eval_op_comprehension_cache_build" evalOpComprehensionCacheHit = "eval_op_comprehension_cache_hit" evalOpComprehensionCacheMiss = "eval_op_comprehension_cache_miss" + evalOpExternalRuleSource = "eval_op_external_rule_source" partialOpSaveUnify = "partial_op_save_unify" partialOpSaveSetContains = "partial_op_save_set_contains" partialOpSaveSetContainsRec = "partial_op_save_set_contains_rec" diff --git a/v1/topdown/save.go b/v1/topdown/save.go index 47bf7521b4..7a922b5f7b 100644 --- a/v1/topdown/save.go +++ b/v1/topdown/save.go @@ -357,7 +357,7 @@ func splitPackageAndRule(path ast.Ref) (ast.Ref, ast.Ref) { // being saved. This check allows the evaluator to evaluate statements // completely during partial evaluation as long as they do not depend on any // kind of unknown value or statements that would generate saves. -func saveRequired(c *ast.Compiler, ic *inliningControl, icIgnoreInternal bool, ss *saveSet, b *bindings, x any, rec bool) bool { +func saveRequired(compilerTree *ast.TreeNode, extStack *externalTreeStack, ic *inliningControl, icIgnoreInternal bool, ss *saveSet, b *bindings, x any, rec bool) bool { var found bool @@ -389,8 +389,9 @@ func saveRequired(c *ast.Compiler, ic *inliningControl, icIgnoreInternal bool, s } else if ic.Disabled(v.ConstantPrefix(), icIgnoreInternal) { found = true } else { - for _, rule := range c.GetRulesDynamicWithOpts(v, ast.RulesOptions{IncludeHiddenModules: false}) { - if saveRequired(c, ic, icIgnoreInternal, ss, b, rule, true) { + rules := getRulesDynamic(compilerTree, extStack, v, ast.RulesOptions{IncludeHiddenModules: false}) + for _, rule := range rules { + if saveRequired(compilerTree, extStack, ic, icIgnoreInternal, ss, b, rule, true) { found = true break } @@ -406,6 +407,76 @@ func saveRequired(c *ast.Compiler, ic *inliningControl, icIgnoreInternal bool, s return found } +// getRulesDynamic looks up rules in both the compiler tree and external sources. +func getRulesDynamic(compilerTree *ast.TreeNode, extStack *externalTreeStack, ref ast.Ref, opts ast.RulesOptions) []*ast.Rule { + var rules []*ast.Rule + + // Check external trees + if extStack != nil { + for i := range extStack.entries { + entry := &extStack.entries[i] + if entry.tree != nil && ref.HasPrefix(entry.ref) { + // Navigate into the external tree using the remaining path + remaining := ref[len(entry.ref):] + rules = append(rules, getRulesFromTree(entry.tree, remaining, opts)...) + } + } + } + + // Then check compiler tree + rules = append(rules, getRulesFromTree(compilerTree, ref, opts)...) + + return rules +} + +// getRulesFromTree walks a tree to find all rules matching the given ref. +func getRulesFromTree(node *ast.TreeNode, ref ast.Ref, opts ast.RulesOptions) []*ast.Rule { + set := map[*ast.Rule]struct{}{} + var walk func(*ast.TreeNode, int) + walk = func(nav *ast.TreeNode, i int) { + switch { + case i >= len(ref): + nav.DepthFirst(func(descendant *ast.TreeNode) bool { + for _, rule := range descendant.Values { + set[rule] = struct{}{} + } + if opts.IncludeHiddenModules { + return false + } + return descendant.Hide + }) + + case i == 0 || ast.IsConstant(ref[i].Value): + if child := nav.Child(ref[i].Value); child != nil { + for _, rule := range child.Values { + set[rule] = struct{}{} + } + walk(child, i+1) + } else { + return + } + + default: + for _, child := range nav.Children { + if child.Hide && !opts.IncludeHiddenModules { + continue + } + for _, rule := range child.Values { + set[rule] = struct{}{} + } + walk(child, i+1) + } + } + } + + walk(node, 0) + rules := make([]*ast.Rule, 0, len(set)) + for rule := range set { + rules = append(rules, rule) + } + return rules +} + func ignoreExprDuringPartial(expr *ast.Expr) bool { if !expr.IsCall() { return false