Do not make rego trace unless --trace flag is passed in (#556)

* Making tracing functionality optional Signed-off-by: Sarvar Muminov <[email protected]> * Adding tests Signed-off-by: Sarvar Muminov <[email protected]> * Use Trace variable for `verify` command Signed-off-by: Sarvar Muminov <[email protected]>
open-policy-agent · May 7, 2021 · a7464de · a7464de
1 parent 8b09609
commit a7464de
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 6 deletions.
diff --git a/internal/runner/test.go b/internal/runner/test.go
@@ -62,6 +62,10 @@ func (t *TestRunner) Run(ctx context.Context, fileList []string) ([]output.Check
 		return nil, fmt.Errorf("load: %w", err)
 	}
 
+	if t.Trace {
+		engine.EnableTracing()
+	}
+
 	namespaces := t.Namespace
 	if t.AllNamespaces {
 		namespaces = engine.Namespaces()

diff --git a/internal/runner/verify.go b/internal/runner/verify.go
@@ -29,7 +29,11 @@ func (r *VerifyRunner) Run(ctx context.Context) ([]output.CheckResult, error) {
 		return nil, fmt.Errorf("load: %w", err)
 	}
 
-	runner := tester.NewRunner().SetCompiler(engine.Compiler()).SetStore(engine.Store()).SetModules(engine.Modules()).EnableTracing(true).SetRuntime(engine.Runtime())
+	if r.Trace {
+		engine.EnableTracing()
+	}
+
+	runner := tester.NewRunner().SetCompiler(engine.Compiler()).SetStore(engine.Store()).SetModules(engine.Modules()).EnableTracing(r.Trace).SetRuntime(engine.Runtime())
 	ch, err := runner.RunTests(ctx, nil)
 	if err != nil {
 		return nil, fmt.Errorf("running tests: %w", err)

diff --git a/policy/engine.go b/policy/engine.go
@@ -17,12 +17,12 @@ import (
 	"github.com/open-policy-agent/opa/loader"
 	"github.com/open-policy-agent/opa/rego"
 	"github.com/open-policy-agent/opa/storage"
-	"github.com/open-policy-agent/opa/topdown"
 	"github.com/open-policy-agent/opa/version"
 )
 
 // Engine represents the policy engine.
 type Engine struct {
+	trace    bool
 	modules  map[string]*ast.Module
 	compiler *ast.Compiler
 	store    storage.Store
@@ -107,6 +107,10 @@ func LoadWithData(ctx context.Context, policyPaths []string, dataPaths []string)
 	return engine, nil
 }
 
+func (e *Engine) EnableTracing() {
+	e.trace = true
+}
+
 // Check executes all of the loaded policies against the input and returns the results.
 func (e *Engine) Check(ctx context.Context, configs map[string]interface{}, namespace string) ([]output.CheckResult, error) {
 	var checkResults []output.CheckResult
@@ -342,24 +346,26 @@ func (e *Engine) check(ctx context.Context, path string, config interface{}, nam
 // data.main.deny to query the deny rule in the main namespace
 // data.main.warn to query the warn rule in the main namespace
 func (e *Engine) query(ctx context.Context, input interface{}, query string, namespace string) (output.QueryResult, error) {
-	stdout := topdown.NewBufferTracer()
 	options := []func(r *rego.Rego){
 		rego.Input(input),
 		rego.Query(query),
 		rego.Compiler(e.Compiler()),
 		rego.Store(e.Store()),
 		rego.Runtime(e.Runtime()),
-		rego.QueryTracer(stdout),
+		rego.Trace(e.trace),
 	}
-	resultSet, err := rego.New(options...).Eval(ctx)
+
+	regoInstance := rego.New(options...)
+	resultSet, err := regoInstance.Eval(ctx)
 	if err != nil {
 		return output.QueryResult{}, fmt.Errorf("evaluating policy: %w", err)
 	}
 
 	// After the evaluation of the policy, the results of the trace (stdout) will be populated
 	// for the query. Once populated, format the trace results into a human readable format.
 	buf := new(bytes.Buffer)
-	topdown.PrettyTrace(buf, *stdout)
+	rego.PrintTrace(buf, regoInstance)
+
 	var traces []string
 	for _, line := range strings.Split(buf.String(), "\n") {
 		if len(line) > 0 {

diff --git a/policy/engine_test.go b/policy/engine_test.go
@@ -46,6 +46,65 @@ func TestException(t *testing.T) {
 	}
 }
 
+func TestTracing(t *testing.T) {
+	t.Run("with tracing ", func(t *testing.T) {
+		ctx := context.Background()
+
+		policies := []string{"../examples/kubernetes/policy"}
+		engine, err := Load(ctx, policies)
+		if err != nil {
+			t.Fatalf("loading policies: %v", err)
+		}
+
+		engine.EnableTracing()
+
+		configFiles := []string{"../examples/kubernetes/service.yaml"}
+		configs, err := parser.ParseConfigurations(configFiles)
+		if err != nil {
+			t.Fatalf("loading configs: %v", err)
+		}
+
+		results, err := engine.Check(ctx, configs, "main")
+		if err != nil {
+			t.Fatalf("could not process policy file: %s", err)
+		}
+
+		for _, query := range results[0].Queries {
+			if len(query.Traces) == 0 {
+				t.Errorf("Tracing error: Expected trace objects, got 0 instead")
+			}
+		}
+	})
+
+	t.Run("without tracing", func(t *testing.T) {
+		ctx := context.Background()
+
+		policies := []string{"../examples/kubernetes/policy"}
+		engine, err := Load(ctx, policies)
+		if err != nil {
+			t.Fatalf("loading policies: %v", err)
+		}
+
+		configFiles := []string{"../examples/kubernetes/service.yaml"}
+		configs, err := parser.ParseConfigurations(configFiles)
+		if err != nil {
+			t.Fatalf("loading configs: %v", err)
+		}
+
+		results, err := engine.Check(ctx, configs, "main")
+		if err != nil {
+			t.Fatalf("could not process policy file: %s", err)
+		}
+
+		for _, query := range results[0].Queries {
+			if len(query.Traces) != 0 {
+				t.Errorf("Tracing error: Expected no trace objects, got %d", len(query.Traces))
+			}
+		}
+	})
+
+}
+
 func TestMultifileYaml(t *testing.T) {
 	ctx := context.Background()