Better trie implementations

policy/policy.go

@@ -1,6 +1,7 @@
 package policy
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"errors"
@@ -10,6 +11,7 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/go-viper/mapstructure/v2"
 	"github.com/open-policy-agent/opa/v1/ast"
@@ -168,12 +170,17 @@ func (r *Result) Response(ctx proxy.Context) (*http.Response, error) {
 	}
 }
 
-func (p *Policy) Query(input *Input) (*Result, error) {
+func (p *Policy) Query(input *Input, options ...func(*rego.Rego)) (*Result, error) {
 	log := logger.StandardLog.Value("policy", p.name)
 	log.Trace("Evaluating policy")
 
+	var regoOptions = append(p.options, rego.Input(input))
+	for _, option := range options {
+		regoOptions = append(regoOptions, option)
+	}
+
 	var (
-		rego    = rego.New(append(p.options, rego.Input(input))...)
+		rego    = rego.New(regoOptions...)
 		ctx     = context.Background()
 		rs, err = rego.Eval(ctx)
 	)
@@ -200,3 +207,27 @@ func (p *Policy) Query(input *Input) (*Result, error) {
 	}
 	return result, nil
 }
+
+// PackageFromFile reads the "package" stanza from the provided Rego policy file.
+//
+// If no stanza can be found, an error is returned.
+func PackageFromFile(name string) (string, error) {
+	f, err := os.Open(name)
+	if err != nil {
+		return "", err
+	}
+	defer func() { _ = f.Close() }()
+
+	scanner := bufio.NewScanner(f)
+	for scanner.Scan() {
+		text := strings.TrimSpace(scanner.Text())
+		part := strings.Fields(text)
+		if len(part) > 1 && part[0] == "package" {
+			return part[1], nil
+		}
+	}
+	if err := scanner.Err(); err != nil {
+		return "", err
+	}
+	return "", fmt.Errorf("policy: can't detemine package name of %s", name)
+}