Add target flag to OPA subcommands

This commit adds a target flag to the
bench, eval, test and run (repl) commands
which allows users to exercise the wasm
rumtime.

Fixes #2878

Signed-off-by: Ashutosh Narkar <anarkar4387@gmail.com>

CHANGELOG.md

@@ -56,6 +56,8 @@ This release contains a number of enhancements and bug fixes.
 
 ## Unreleased
 
+Previously, the `opa test` command used `t` as shorthand for the `timeout` flag. This change adds a new `target` flag to `opa eval`, `opa bench` and `opa test`([#2878](https://github.com/open-policy-agent/opa/issues/2878)). `t` is now a shorthand for the `target` flag and is no longer used for the `timeout` flag in `opa test`.
+
 ## 0.25.2
 
 This release extends the HTTP server authorizer (`--authorization=basic`) to supply the HTTP message body in the `input` document. See the [Authentication and Authorization](https://www.openpolicyagent.org/docs/edge/security/#authentication-and-authorization) section in the security documentation for details.

cmd/bench.go

@@ -13,6 +13,8 @@ import (
 	"sort"
 	"testing"
 
+	"github.com/open-policy-agent/opa/compile"
+
 	"github.com/olekukonko/tablewriter"
 	"github.com/spf13/cobra"
 
@@ -43,6 +45,7 @@ func newBenchmarkEvalParams() benchmarkCommandParams {
 				evalPrettyOutput,
 				benchmarkGoBenchOutput,
 			}),
+			target: util.NewEnumFlag(compile.TargetRego, []string{compile.TargetRego, compile.TargetWasm}),
 		},
 	}
 }
@@ -90,6 +93,7 @@ The optional "gobench" output format conforms to the Go Benchmark Data Format.
 	addOutputFormat(benchCommand.Flags(), params.outputFormat)
 	addIgnoreFlag(benchCommand.Flags(), &params.ignore)
 	addSchemaFlag(benchCommand.Flags(), &params.schemaPath)
+	addTargetFlag(benchCommand.Flags(), params.target)
 
 	// Shared benchmark flags
 	addCountFlag(benchCommand.Flags(), &params.count, "benchmark")

cmd/eval.go

@@ -14,6 +14,8 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/open-policy-agent/opa/compile"
+
 	"github.com/spf13/cobra"
 
 	"github.com/open-policy-agent/opa/ast"
@@ -58,6 +60,7 @@ type evalCommandParams struct {
 	failDefined         bool
 	bundlePaths         repeatedStringFlag
 	schemaPath          string
+	target              *util.EnumFlag
 }
 
 func newEvalCommandParams() evalCommandParams {
@@ -70,6 +73,7 @@ func newEvalCommandParams() evalCommandParams {
 			evalSourceOutput,
 		}),
 		explain: newExplainFlag([]string{explainModeOff, explainModeFull, explainModeNotes, explainModeFails}),
+		target:  util.NewEnumFlag(compile.TargetRego, []string{compile.TargetRego, compile.TargetWasm}),
 	}
 }
 
@@ -254,6 +258,7 @@ The -s/--schema flag provides a single JSON Schema used to validate references t
 	addIgnoreFlag(evalCommand.Flags(), &params.ignore)
 	setExplainFlag(evalCommand.Flags(), params.explain)
 	addSchemaFlag(evalCommand.Flags(), &params.schemaPath)
+	addTargetFlag(evalCommand.Flags(), params.target)
 
 	RootCommand.AddCommand(evalCommand)
 }
@@ -403,6 +408,8 @@ func setupEval(args []string, params evalCommandParams) (*evalContext, error) {
 	// skip bundle verification
 	regoArgs = append(regoArgs, rego.SkipBundleVerification(true))
 
+	regoArgs = append(regoArgs, rego.Target(params.target.String()))
+
 	inputBytes, err := readInputBytes(params)
 	if err != nil {
 		return nil, err
@@ -439,6 +446,10 @@ func setupEval(args []string, params evalCommandParams) (*evalContext, error) {
 	if params.explain != nil && params.explain.String() != explainModeOff {
 		tracer = topdown.NewBufferTracer()
 		evalArgs = append(evalArgs, rego.EvalQueryTracer(tracer))
+
+		if params.target.String() == compile.TargetWasm {
+			fmt.Fprintf(os.Stdout, "warning: explain mode \"%v\" is not supported with wasm target\n", params.explain.String())
+		}
 	}
 
 	if params.disableIndexing {

cmd/flags.go

@@ -134,6 +134,10 @@ func addSchemaFlag(fs *pflag.FlagSet, schemaPath *string) {
 	fs.StringVarP(schemaPath, "schema", "s", "", "set schema file path")
 }
 
+func addTargetFlag(fs *pflag.FlagSet, target *util.EnumFlag) {
+	fs.VarP(target, "target", "t", "set the runtime to exercise")
+}
+
 const (
 	explainModeOff   = "off"
 	explainModeFull  = "full"

cmd/test.go

@@ -10,6 +10,8 @@ import (
 	"os"
 	"time"
 
+	"github.com/open-policy-agent/opa/compile"
+
 	"github.com/open-policy-agent/opa/storage/inmem"
 	"github.com/open-policy-agent/opa/topdown/lineage"
 
@@ -47,12 +49,14 @@ type testCommandParams struct {
 	benchMem     bool
 	runRegex     string
 	count        int
+	target       *util.EnumFlag
 }
 
 func newTestCommandParams() *testCommandParams {
 	return &testCommandParams{
 		outputFormat: util.NewEnumFlag(testPrettyOutput, []string{testPrettyOutput, testJSONOutput, benchmarkGoBenchOutput}),
 		explain:      newExplainFlag([]string{explainModeFails, explainModeFull, explainModeNotes}),
+		target:       util.NewEnumFlag(compile.TargetRego, []string{compile.TargetRego, compile.TargetWasm}),
 	}
 }
 
@@ -217,7 +221,8 @@ func opaTest(args []string) int {
 		SetModules(modules).
 		SetBundles(bundles).
 		SetTimeout(testParams.timeout).
-		Filter(testParams.runRegex)
+		Filter(testParams.runRegex).
+		Target(testParams.target.String())
 
 	var reporter tester.Reporter
 
@@ -337,7 +342,7 @@ func init() {
 	testCommand.Flags().BoolVarP(&testParams.verbose, "verbose", "v", false, "set verbose reporting mode")
 	testCommand.Flags().BoolVarP(&testParams.failureLine, "show-failure-line", "l", false, "show test failure line")
 	testCommand.Flags().MarkDeprecated("show-failure-line", "use -v instead")
-	testCommand.Flags().DurationVarP(&testParams.timeout, "timeout", "t", time.Second*5, "set test timeout")
+	testCommand.Flags().DurationVarP(&testParams.timeout, "timeout", "", time.Second*5, "set test timeout")
 	testCommand.Flags().VarP(testParams.outputFormat, "format", "f", "set output format")
 	testCommand.Flags().BoolVarP(&testParams.coverage, "coverage", "c", false, "report coverage (overrides debug tracing)")
 	testCommand.Flags().Float64VarP(&testParams.threshold, "threshold", "", 0, "set coverage threshold and exit with non-zero status if coverage is less than threshold %")
@@ -349,5 +354,6 @@ func init() {
 	addMaxErrorsFlag(testCommand.Flags(), &testParams.errLimit)
 	addIgnoreFlag(testCommand.Flags(), &testParams.ignore)
 	setExplainFlag(testCommand.Flags(), testParams.explain)
+	addTargetFlag(testCommand.Flags(), testParams.target)
 	RootCommand.AddCommand(testCommand)
 }

internal/rego/opa/nop.go

@@ -0,0 +1,53 @@
+// Copyright 2021 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.
+
+// +build !opa_wasm
+
+package opa
+
+import (
+	"context"
+
+	"github.com/open-policy-agent/opa/metrics"
+)
+
+// OPA is a stub implementation of a opa.OPA.
+type OPA struct {
+}
+
+// Result is a stub implementation of a opa.Result.
+type Result struct {
+	Result []byte
+}
+
+// EvalOpts is a stub implementation of a opa.EvalOpts.
+type EvalOpts struct {
+	Input   *interface{}
+	Metrics metrics.Metrics
+}
+
+// New unimplemented.
+func New() *OPA {
+	panic("WebAssembly runtime not supported in this build")
+}
+
+// WithPolicyBytes unimplemented.
+func (o *OPA) WithPolicyBytes(policy []byte) *OPA {
+	panic("unreachable")
+}
+
+// WithDataJSON unimplemented.
+func (o *OPA) WithDataJSON(data interface{}) *OPA {
+	panic("unreachable")
+}
+
+// Init unimplemented.
+func (o *OPA) Init() (*OPA, error) {
+	panic("unreachable")
+}
+
+// Eval unimplemented.
+func (o *OPA) Eval(ctx context.Context, opts EvalOpts) (*Result, error) {
+	panic("unreachable")
+}

internal/rego/opa/opa.go

@@ -0,0 +1,71 @@
+// Copyright 2021 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.
+
+// +build opa_wasm
+
+package opa
+
+import (
+	"context"
+	wopa "github.com/open-policy-agent/opa/internal/wasm/sdk/opa"
+	"github.com/open-policy-agent/opa/metrics"
+)
+
+// OPA is an implementation of the OPA SDK.
+type OPA struct {
+	opa *wopa.OPA
+}
+
+// Result holds the evaluation result.
+type Result struct {
+	Result []byte
+}
+
+// EvalOpts define options for performing an evaluation.
+type EvalOpts struct {
+	Input   *interface{}
+	Metrics metrics.Metrics
+}
+
+// New constructs a new OPA instance.
+func New() *OPA {
+	return &OPA{opa: wopa.New()}
+}
+
+// WithPolicyBytes configures the compiled policy to load.
+func (o *OPA) WithPolicyBytes(policy []byte) *OPA {
+	o.opa = o.opa.WithPolicyBytes(policy)
+	return o
+}
+
+// WithDataJSON configures the JSON data to load.
+func (o *OPA) WithDataJSON(data interface{}) *OPA {
+	o.opa = o.opa.WithDataJSON(data)
+	return o
+}
+
+// Init initializes the OPA instance.
+func (o *OPA) Init() (*OPA, error) {
+	i, err := o.opa.Init()
+	if err != nil {
+		return nil, err
+	}
+	o.opa = i
+	return o, nil
+}
+
+// Eval evaluates the policy.
+func (o *OPA) Eval(ctx context.Context, opts EvalOpts) (*Result, error) {
+	evalOptions := wopa.EvalOpts{
+		Input:   opts.Input,
+		Metrics: opts.Metrics,
+	}
+
+	res, err := o.opa.Eval(ctx, evalOptions)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Result{Result: res.Result}, nil
+}

rego/rego.go

@@ -20,6 +20,7 @@ import (
 	"github.com/open-policy-agent/opa/internal/compiler/wasm"
 	"github.com/open-policy-agent/opa/internal/ir"
 	"github.com/open-policy-agent/opa/internal/planner"
+	"github.com/open-policy-agent/opa/internal/rego/opa"
 	"github.com/open-policy-agent/opa/internal/wasm/encoding"
 	"github.com/open-policy-agent/opa/loader"
 	"github.com/open-policy-agent/opa/metrics"
@@ -32,7 +33,12 @@ import (
 	"github.com/open-policy-agent/opa/util"
 )
 
-const defaultPartialNamespace = "partial"
+const (
+	defaultPartialNamespace = "partial"
+	targetWasm              = "wasm"
+)
+
+var wasmVarPrefix = "^"
 
 // CompileResult represents the result of compiling a Rego query, zero or more
 // Rego modules, and arbitrary contextual data into an executable.
@@ -519,6 +525,8 @@ type Rego struct {
 	strictBuiltinErrors    bool
 	resolvers              []refResolver
 	schemaSet              *ast.SchemaSet
+	target                 string // target type (wasm, rego, etc.)
+	opa                    *opa.OPA
 }
 
 // Function represents a built-in function that is callable in Rego.
@@ -1030,6 +1038,13 @@ func Schemas(x *ast.SchemaSet) func(r *Rego) {
 	}
 }
 
+// Target sets the runtime to exercise.
+func Target(t string) func(r *Rego) {
+	return func(r *Rego) {
+		r.target = t
+	}
+}
+
 // New returns a new Rego object.
 func New(options ...func(r *Rego)) *Rego {
 
@@ -1284,6 +1299,10 @@ func (r *Rego) Compile(ctx context.Context, opts ...CompileOption) (*CompileResu
 		queries = []ast.Body{r.compiledQueries[compileQueryType].query}
 	}
 
+	return r.compileWasm(modules, queries, compileQueryType)
+}
+
+func (r *Rego) compileWasm(modules []*ast.Module, queries []ast.Body, qType queryType) (*CompileResult, error) {
 	decls := make(map[string]*ast.Builtin, len(r.builtinDecls)+len(ast.BuiltinMap))
 
 	for k, v := range ast.BuiltinMap {
@@ -1301,7 +1320,7 @@ func (r *Rego) Compile(ctx context.Context, opts ...CompileOption) (*CompileResu
 			{
 				Name:          queryName,
 				Queries:       queries,
-				RewrittenVars: r.compiledQueries[compileQueryType].compiler.RewrittenVars(),
+				RewrittenVars: r.compiledQueries[qType].compiler.RewrittenVars(),
 			},
 		}).
 		WithModules(modules).
@@ -1416,6 +1435,37 @@ func (r *Rego) PrepareForEval(ctx context.Context, opts ...PrepareOption) (Prepa
 			},
 		},
 	})
+
+	if r.target == targetWasm {
+
+		if r.hasWasmModule() {
+			return PreparedEvalQuery{}, fmt.Errorf("wasm target not supported")
+		}
+
+		var modules []*ast.Module
+		for _, module := range r.compiler.Modules {
+			modules = append(modules, module)
+		}
+
+		queries := []ast.Body{r.compiledQueries[evalQueryType].query}
+
+		cr, err := r.compileWasm(modules, queries, evalQueryType)
+		if err != nil {
+			return PreparedEvalQuery{}, err
+		}
+
+		data, err := r.store.Read(ctx, r.txn, storage.Path{})
+		if err != nil {
+			return PreparedEvalQuery{}, err
+		}
+
+		o, err := opa.New().WithPolicyBytes(cr.Bytes).WithDataJSON(data).Init()
+		if err != nil {
+			return PreparedEvalQuery{}, err
+		}
+		r.opa = o
+	}
+
 	txnErr := txnClose(ctx, err) // Always call closer
 	if err != nil {
 		return PreparedEvalQuery{}, err
@@ -1463,6 +1513,7 @@ func (r *Rego) PrepareForPartial(ctx context.Context, opts ...PrepareOption) (Pr
 	if txnErr != nil {
 		return PreparedPartialQuery{}, txnErr
 	}
+
 	return PreparedPartialQuery{preparedQuery{r, pCfg}}, err
 }
 
@@ -1688,6 +1739,7 @@ func (r *Rego) compileModules(ctx context.Context, txn storage.Transaction, m me
 		if err != nil {
 			return err
 		}
+
 		for _, rslvr := range resolvers {
 			for _, ep := range rslvr.Entrypoints() {
 				r.resolvers = append(r.resolvers, refResolver{ep, rslvr})
@@ -1766,6 +1818,9 @@ func (r *Rego) compileQuery(query ast.Body, m metrics.Metrics, extras []extraSta
 }
 
 func (r *Rego) eval(ctx context.Context, ectx *EvalContext) (ResultSet, error) {
+	if r.opa != nil {
+		return r.evalWasm(ctx, ectx)
+	}
 
 	q := topdown.NewQuery(ectx.compiledQuery.query).
 		WithQueryCompiler(ectx.compiledQuery.compiler).
@@ -1805,36 +1860,11 @@ func (r *Rego) eval(ctx context.Context, ectx *EvalContext) (ResultSet, error) {
 		c.Cancel()
 	})
 
-	rewritten := ectx.compiledQuery.compiler.RewrittenVars()
 	var rs ResultSet
 	err := q.Iter(ctx, func(qr topdown.QueryResult) error {
-		result := newResult()
-		for k := range qr {
-			v, err := ast.JSON(qr[k].Value)
-			if err != nil {
-				return err
-			}
-			if rw, ok := rewritten[k]; ok {
-				k = rw
-			}
-			if isTermVar(k) || k.IsGenerated() || k.IsWildcard() {
-				continue
-			}
-			result.Bindings[string(k)] = v
-		}
-		for _, expr := range ectx.compiledQuery.query {
-			if expr.Generated {
-				continue
-			}
-			if k, ok := r.capture[expr]; ok {
-				v, err := ast.JSON(qr[k].Value)
-				if err != nil {
-					return err
-				}
-				result.Expressions = append(result.Expressions, newExpressionValue(expr, v))
-			} else {
-				result.Expressions = append(result.Expressions, newExpressionValue(expr, true))
-			}
+		result, err := r.generateResult(qr, ectx)
+		if err != nil {
+			return err
 		}
 		rs = append(rs, result)
 		return nil
@@ -1851,6 +1881,107 @@ func (r *Rego) eval(ctx context.Context, ectx *EvalContext) (ResultSet, error) {
 	return rs, nil
 }
 
+func (r *Rego) evalWasm(ctx context.Context, ectx *EvalContext) (ResultSet, error) {
+
+	var input *interface{}
+	if ectx.parsedInput != nil {
+		i, err := ast.JSON(ectx.parsedInput)
+		if err != nil {
+			return nil, err
+		}
+		input = &i
+	}
+
+	result, err := r.opa.Eval(ctx, opa.EvalOpts{Metrics: r.metrics, Input: input})
+	if err != nil {
+		return nil, err
+	}
+
+	parsed, err := ast.ParseTerm(string(result.Result))
+	if err != nil {
+		return nil, err
+	}
+
+	resultSet, ok := parsed.Value.(ast.Set)
+	if !ok {
+		return nil, fmt.Errorf("illegal result type")
+	}
+
+	if resultSet.Len() == 0 {
+		return nil, nil
+	}
+
+	qr := topdown.QueryResult{}
+	err = resultSet.Iter(func(term *ast.Term) error {
+		obj, ok := term.Value.(ast.Object)
+		if !ok {
+			return fmt.Errorf("illegal result type")
+		}
+
+		obj.Foreach(func(k, v *ast.Term) {
+			kvt := ast.VarTerm(string(k.Value.(ast.String)))
+			qr[kvt.Value.(ast.Var)] = v
+		})
+
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	res, err := r.generateResult(qr, ectx)
+	if err != nil {
+		return nil, err
+	}
+
+	rs := ResultSet{res}
+
+	if len(rs) == 0 {
+		return nil, nil
+	}
+
+	return rs, nil
+}
+
+func (r *Rego) generateResult(qr topdown.QueryResult, ectx *EvalContext) (Result, error) {
+
+	rewritten := ectx.compiledQuery.compiler.RewrittenVars()
+
+	result := newResult()
+	for k := range qr {
+		v, err := ast.JSON(qr[k].Value)
+		if err != nil {
+			return result, err
+		}
+
+		if rw, ok := rewritten[k]; ok {
+			k = rw
+		}
+		if isTermVar(k) || isTermWasmVar(k) || k.IsGenerated() || k.IsWildcard() {
+			continue
+		}
+		result.Bindings[string(k)] = v
+	}
+
+	for _, expr := range ectx.compiledQuery.query {
+		if expr.Generated {
+			continue
+		}
+
+		if k, ok := r.capture[expr]; ok {
+			v, err := ast.JSON(qr[k].Value)
+			if err != nil {
+				return result, err
+			}
+			result.Expressions = append(result.Expressions, newExpressionValue(expr, v))
+		} else {
+			result.Expressions = append(result.Expressions, newExpressionValue(expr, true))
+		}
+
+	}
+	return result, nil
+}
+
 func (r *Rego) partialResult(ctx context.Context, pCfg *PrepareConfig) (PartialResult, error) {
 
 	err := r.prepare(ctx, partialResultQueryType, []extraStage{
@@ -2104,6 +2235,10 @@ func (r *Rego) rewriteEqualsForPartialQueryCompile(_ ast.QueryCompiler, query as
 
 func (r *Rego) generateTermVar() *ast.Term {
 	r.termVarID++
+
+	if r.target == targetWasm {
+		return ast.VarTerm(wasmVarPrefix + fmt.Sprintf("term%v", r.termVarID))
+	}
 	return ast.VarTerm(ast.WildcardPrefix + fmt.Sprintf("term%v", r.termVarID))
 }
 
@@ -2111,6 +2246,15 @@ func (r Rego) hasQuery() bool {
 	return len(r.query) != 0 || len(r.parsedQuery) != 0
 }
 
+func (r Rego) hasWasmModule() bool {
+	for _, b := range r.bundles {
+		if len(b.WasmModules) > 0 {
+			return true
+		}
+	}
+	return false
+}
+
 type transactionCloser func(ctx context.Context, err error) error
 
 // getTxn will conditionally create a read or write transaction suitable for
@@ -2189,6 +2333,10 @@ func isTermVar(v ast.Var) bool {
 	return strings.HasPrefix(string(v), ast.WildcardPrefix+"term")
 }
 
+func isTermWasmVar(v ast.Var) bool {
+	return strings.HasPrefix(string(v), wasmVarPrefix+"term")
+}
+
 func waitForDone(ctx context.Context, exit chan struct{}, f func()) {
 	select {
 	case <-exit:

rego/rego_test.go

@@ -1712,6 +1712,105 @@ func TestPrepareWithEmptyModule(t *testing.T) {
 	}
 }
 
+func TestPrepareWithWasmTargetNotSupported(t *testing.T) {
+	files := map[string]string{
+		"x/x.rego":     "package x\np = data.x.b",
+		"x/data.json":  `{"b": "bar"}`,
+		"/policy.wasm": `modules-compiled-as-wasm-binary`,
+	}
+
+	test.WithTempFS(files, func(path string) {
+		ctx := context.Background()
+
+		_, err := New(
+			LoadBundle(path),
+			Query("data.x.p"),
+			Target("wasm"),
+		).PrepareForEval(ctx)
+
+		expected := "wasm target not supported"
+		if err == nil || err.Error() != expected {
+			t.Fatalf("Expected error %s, got %s", expected, err)
+		}
+	})
+}
+
+func TestPrepareAndEvalWithWasmTarget(t *testing.T) {
+	mod := `
+	package test
+	default p = false
+	p {
+		input.x == 1
+	}
+	`
+
+	ctx := context.Background()
+
+	pq, err := New(
+		Query("data.test.p = x"),
+		Target("wasm"),
+		Module("a.rego", mod),
+	).PrepareForEval(ctx)
+
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+
+	assertPreparedEvalQueryEval(t, pq, []EvalOption{
+		EvalInput(map[string]int{"x": 1}),
+	}, "[[true]]")
+
+	pq, err = New(
+		Query("a = [1,2]; x = a[i]"),
+		Target("wasm"),
+	).PrepareForEval(ctx)
+
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+
+	assertPreparedEvalQueryEval(t, pq, []EvalOption{}, "[[true, true]]")
+}
+
+func TestPrepareAndEvalWithWasmTargetModulesOnCompiler(t *testing.T) {
+	mod := `
+	package test
+	default p = false
+	p {
+		input.x == data.x.p
+	}
+	`
+
+	compiler := ast.NewCompiler()
+
+	compiler.Compile(map[string]*ast.Module{
+		"a.rego": ast.MustParseModule(mod),
+	})
+
+	if len(compiler.Errors) > 0 {
+		t.Fatalf("Unexpected compile errors: %s", compiler.Errors)
+	}
+
+	ctx := context.Background()
+
+	pq, err := New(
+		Compiler(compiler),
+		Query("data.test.p"),
+		Target("wasm"),
+		Store(inmem.NewFromObject(map[string]interface{}{
+			"x": map[string]interface{}{"p": 1},
+		})),
+	).PrepareForEval(ctx)
+
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+
+	assertPreparedEvalQueryEval(t, pq, []EvalOption{
+		EvalInput(map[string]int{"x": 1}),
+	}, "[[true]]")
+}
+
 func TestEvalWithInterQueryCache(t *testing.T) {
 	query := `http.send({"method": "get", "url": "%URL%", "force_json_decode": true, "cache": true})`
 	newHeaders := map[string][]string{"Cache-Control": {"max-age=290304000, public"}}

repl/repl.go

@@ -17,6 +17,8 @@ import (
 	"strings"
 	"sync"
 
+	"github.com/open-policy-agent/opa/compile"
+
 	"github.com/open-policy-agent/opa/version"
 
 	"github.com/peterh/liner"
@@ -62,6 +64,7 @@ type REPL struct {
 	errLimit          int
 	prettyLimit       int
 	report            [][2]string
+	target            string // target type (wasm, rego, etc.)
 	mtx               sync.Mutex
 }
 
@@ -76,6 +79,8 @@ const (
 
 const defaultPrettyLimit = 80
 
+var allowedTargets = map[string]bool{compile.TargetRego: true, compile.TargetWasm: true}
+
 const exitPromptMessage = "Do you want to exit ([y]/n)? "
 
 // New returns a new instance of the REPL.
@@ -93,6 +98,7 @@ func New(store storage.Store, historyPath string, output io.Writer, outputFormat
 		banner:       banner,
 		errLimit:     errLimit,
 		prettyLimit:  defaultPrettyLimit,
+		target:       compile.TargetRego,
 	}
 }
 
@@ -252,6 +258,8 @@ func (r *REPL) OneShot(ctx context.Context, line string) error {
 				return r.cmdUnknown(cmd.args)
 			case "strict-builtin-errors":
 				return r.cmdStrictBuiltinErrors()
+			case "target":
+				return r.cmdTarget(cmd.args)
 			case "help":
 				return r.cmdHelp(cmd.args)
 			case "exit":
@@ -388,6 +396,21 @@ func (r *REPL) cmdFormat(s string) error {
 	return nil
 }
 
+func (r *REPL) cmdTarget(t []string) error {
+	if len(t) != 1 {
+		return newBadArgsErr("target <mode>: expects exactly one argument")
+	}
+
+	if _, ok := allowedTargets[t[0]]; !ok {
+		return fmt.Errorf("invalid target \"%v\":must be one of {rego,wasm}", t[0])
+	}
+
+	r.target = t[0]
+
+	r.checkTraceSupported()
+	return nil
+}
+
 func (r *REPL) cmdPrettyLimit(s []string) error {
 	if len(s) != 1 {
 		return fmt.Errorf("usage: pretty-limit <n>")
@@ -459,9 +482,17 @@ func (r *REPL) cmdTrace(mode explainMode) error {
 	} else {
 		r.explain = mode
 	}
+
+	r.checkTraceSupported()
 	return nil
 }
 
+func (r *REPL) checkTraceSupported() {
+	if r.explain != explainOff && r.target == compile.TargetWasm {
+		fmt.Fprintf(r.output, "warning: trace mode \"%v\" is not supported with wasm target\n", r.explain)
+	}
+}
+
 func (r *REPL) metricsEnabled() bool {
 	if r.metrics != nil {
 		return true
@@ -916,6 +947,7 @@ func (r *REPL) evalBody(ctx context.Context, compiler *ast.Compiler, input ast.V
 		rego.Instrument(r.instrument),
 		rego.Runtime(r.runtime),
 		rego.StrictBuiltinErrors(r.strictBuiltinErrors),
+		rego.Target(r.target),
 	}
 
 	if r.explain != explainOff {
@@ -1214,6 +1246,7 @@ var builtin = [...]commandDesc{
 	{"strict-builtin-errors", []string{}, "toggle strict built-in error mode"},
 	{"dump", []string{"[path]"}, "dump raw data in storage"},
 	{"help", []string{"[topic]"}, "print this message"},
+	{"target", []string{"[mode]"}, "set the runtime to exercise {rego,wasm} (default rego)"},
 	{"exit", []string{}, "exit out of shell (or ctrl+d)"},
 	{"ctrl+l", []string{}, "clear the screen"},
 }

repl/repl_test.go

@@ -1625,6 +1625,70 @@ func TestEvalBodyWith(t *testing.T) {
 	}
 }
 
+func TestReplWasmTarget(t *testing.T) {
+	ctx := context.Background()
+	store := newTestStore()
+	var buffer bytes.Buffer
+	repl := newRepl(store, &buffer)
+
+	err := repl.OneShot(ctx, "target foo bar")
+
+	expected := "code bad arguments: target <mode>: expects exactly one argument"
+	if err == nil || err.Error() != expected {
+		t.Fatalf("Expected error %s, got %s", expected, err)
+	}
+
+	err = repl.OneShot(ctx, "target foo")
+
+	expected = "invalid target \"foo\":must be one of {rego,wasm}"
+	if err == nil || err.Error() != expected {
+		t.Fatalf("Expected error %s, got %s", expected, err)
+	}
+
+	buffer.Reset()
+	err = repl.OneShot(ctx, "target wasm")
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+
+	repl.OneShot(ctx, `p = true { input.foo = "bar" }`)
+	buffer.Reset()
+	repl.OneShot(ctx, "p")
+
+	if buffer.String() != "undefined\n" {
+		t.Fatalf("Expected undefined but got: %v", buffer.String())
+	}
+
+	buffer.Reset()
+	repl.OneShot(ctx, `p with input as {"foo": "bar"}`)
+
+	result := buffer.String()
+	expected = "true\n"
+
+	if result != expected {
+		t.Fatalf("Expected true but got: %v", result)
+	}
+
+	buffer.Reset()
+	repl.OneShot(ctx, `p with input.foo as "bar"`)
+
+	result = buffer.String()
+
+	if result != expected {
+		t.Fatalf("Expected true but got: %v", result)
+	}
+
+	buffer.Reset()
+	repl.OneShot(ctx, `trace`)
+
+	result = buffer.String()
+
+	expected = "warning: trace mode \"full\" is not supported with wasm target\n"
+	if result != expected {
+		t.Fatalf("Expected true but got: %v", result)
+	}
+}
+
 func TestEvalBodyRewrittenBuiltin(t *testing.T) {
 	ctx := context.Background()
 	store := newTestStore()

runtime/runtime.go

@@ -436,7 +436,8 @@ func (rt *Runtime) StartREPL(ctx context.Context) {
 	defer rt.Manager.Stop(ctx)
 
 	banner := rt.getBanner()
-	repl := repl.New(rt.Store, rt.Params.HistoryPath, rt.Params.Output, rt.Params.OutputFormat, rt.Params.ErrorLimit, banner).WithRuntime(rt.Manager.Info)
+	repl := repl.New(rt.Store, rt.Params.HistoryPath, rt.Params.Output, rt.Params.OutputFormat, rt.Params.ErrorLimit, banner).
+		WithRuntime(rt.Manager.Info)
 
 	if rt.Params.Watch {
 		if err := rt.startWatcher(ctx, rt.Params.Paths, onReloadPrinter(rt.Params.Output)); err != nil {

tester/runner.go

@@ -121,6 +121,7 @@ type Runner struct {
 	modules     map[string]*ast.Module
 	bundles     map[string]*bundle.Bundle
 	filter      string
+	target      string // target type (wasm, rego, etc.)
 }
 
 // NewRunner returns a new runner.
@@ -216,6 +217,12 @@ func (r *Runner) Filter(regex string) *Runner {
 	return r
 }
 
+// Target sets the output target type to use.
+func (r *Runner) Target(target string) *Runner {
+	r.target = target
+	return r
+}
+
 func getFailedAtFromTrace(bufFailureLineTracer *topdown.BufferTracer) *ast.Expr {
 	events := *bufFailureLineTracer
 	const SecondToLast = 2
@@ -415,6 +422,7 @@ func (r *Runner) runTest(ctx context.Context, txn storage.Transaction, mod *ast.
 		rego.Query(rule.Path().String()),
 		rego.QueryTracer(tracer),
 		rego.Runtime(r.runtime),
+		rego.Target(r.target),
 	)
 
 	t0 := time.Now()
@@ -466,6 +474,7 @@ func (r *Runner) runBenchmark(ctx context.Context, txn storage.Transaction, mod
 			rego.Compiler(r.compiler),
 			rego.Query(rule.Path().String()),
 			rego.Runtime(r.runtime),
+			rego.Target(r.target),
 		).PrepareForEval(ctx)
 
 		if err != nil {