diff --git a/compile/compile.go b/compile/compile.go index e878243214..fcbfcb060c 100644 --- a/compile/compile.go +++ b/compile/compile.go @@ -57,24 +57,25 @@ const resultVar = ast.Var("result") // Compiler implements bundle compilation and linking. type Compiler struct { - capabilities *ast.Capabilities // the capabilities that compiled policies may require - bundle *bundle.Bundle // the bundle that the compiler operates on - revision *string // the revision to set on the output bundle - asBundle bool // whether to assume bundle layout on file loading or not - filter loader.Filter // filter to apply to file loader - paths []string // file paths to load. TODO(tsandall): add support for supplying readers for embedded users. - entrypoints orderedStringSet // policy entrypoints required for optimization and certain targets - optimizationLevel int // how aggressive should optimization be - target string // target type (wasm, rego, etc.) - output *io.Writer // output stream to write bundle to - entrypointrefs []*ast.Term // validated entrypoints computed from default decision or manually supplied entrypoints - compiler *ast.Compiler // rego ast compiler used for semantic checks and rewriting - policy *ir.Policy // planner output when wasm or plan targets are enabled - debug debug.Debug // optionally outputs debug information produced during build - bvc *bundle.VerificationConfig // represents the key configuration used to verify a signed bundle - bsc *bundle.SigningConfig // represents the key configuration used to generate a signed bundle - keyID string // represents the name of the default key used to verify a signed bundle - metadata *map[string]interface{} // represents additional data included in .manifest file + capabilities *ast.Capabilities // the capabilities that compiled policies may require + bundle *bundle.Bundle // the bundle that the compiler operates on + revision *string // the revision to set on the output bundle + asBundle bool // whether to assume bundle layout on file loading or not + filter loader.Filter // filter to apply to file loader + paths []string // file paths to load. TODO(tsandall): add support for supplying readers for embedded users. + entrypoints orderedStringSet // policy entrypoints required for optimization and certain targets + optimizationLevel int // how aggressive should optimization be + target string // target type (wasm, rego, etc.) + output *io.Writer // output stream to write bundle to + entrypointrefs []*ast.Term // validated entrypoints computed from default decision or manually supplied entrypoints + compiler *ast.Compiler // rego ast compiler used for semantic checks and rewriting + policy *ir.Policy // planner output when wasm or plan targets are enabled + debug debug.Debug // optionally outputs debug information produced during build + enablePrintStatements bool // optionally enable rego print statements + bvc *bundle.VerificationConfig // represents the key configuration used to verify a signed bundle + bsc *bundle.SigningConfig // represents the key configuration used to generate a signed bundle + keyID string // represents the name of the default key used to verify a signed bundle + metadata *map[string]interface{} // represents additional data included in .manifest file } // New returns a new compiler instance that can be invoked. @@ -135,6 +136,14 @@ func (c *Compiler) WithDebug(sink io.Writer) *Compiler { return c } +// WithEnablePrintStatements enables print statements inside of modules compiled +// by the compiler. If print statements are not enabled, calls to print() are +// erased at compile-time. +func (c *Compiler) WithEnablePrintStatements(yes bool) *Compiler { + c.enablePrintStatements = yes + 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...) @@ -374,14 +383,15 @@ func (c *Compiler) optimize(ctx context.Context) error { if c.optimizationLevel <= 0 { var err error - c.compiler, err = compile(c.capabilities, c.bundle, c.debug) + c.compiler, err = compile(c.capabilities, c.bundle, c.debug, c.enablePrintStatements) return err } o := newOptimizer(c.capabilities, c.bundle). WithEntrypoints(c.entrypointrefs). WithDebug(c.debug.Writer()). - WithShallowInlining(c.optimizationLevel <= 1) + WithShallowInlining(c.optimizationLevel <= 1). + WithEnablePrintStatements(c.enablePrintStatements) err := o.Do(ctx) if err != nil { @@ -399,7 +409,7 @@ func (c *Compiler) compilePlan(ctx context.Context) error { // AST compiler will not be set because the default target does not require it. if c.compiler == nil { var err error - c.compiler, err = compile(c.capabilities, c.bundle, c.debug) + c.compiler, err = compile(c.capabilities, c.bundle, c.debug, c.enablePrintStatements) if err != nil { return err } @@ -615,15 +625,16 @@ func (err undefinedEntrypointErr) Error() string { } type optimizer struct { - capabilities *ast.Capabilities - bundle *bundle.Bundle - compiler *ast.Compiler - entrypoints []*ast.Term - nsprefix string - resultsymprefix string - outputprefix string - shallow bool - debug debug.Debug + capabilities *ast.Capabilities + bundle *bundle.Bundle + compiler *ast.Compiler + entrypoints []*ast.Term + nsprefix string + resultsymprefix string + outputprefix string + shallow bool + debug debug.Debug + enablePrintStatements bool } func newOptimizer(c *ast.Capabilities, b *bundle.Bundle) *optimizer { @@ -644,6 +655,11 @@ func (o *optimizer) WithDebug(sink io.Writer) *optimizer { return o } +func (o *optimizer) WithEnablePrintStatements(yes bool) *optimizer { + o.enablePrintStatements = yes + return o +} + func (o *optimizer) WithEntrypoints(es []*ast.Term) *optimizer { o.entrypoints = es return o @@ -683,7 +699,7 @@ func (o *optimizer) Do(ctx context.Context) error { for i, e := range o.entrypoints { var err error - o.compiler, err = compile(o.capabilities, o.bundle, o.debug) + o.compiler, err = compile(o.capabilities, o.bundle, o.debug, o.enablePrintStatements) if err != nil { return err } @@ -939,7 +955,7 @@ func (o *optimizer) getSupportModuleFilename(used map[string]int, module *ast.Mo var safePathPattern = regexp.MustCompile(`^[\w-_/]+$`) -func compile(c *ast.Capabilities, b *bundle.Bundle, dbg debug.Debug) (*ast.Compiler, error) { +func compile(c *ast.Capabilities, b *bundle.Bundle, dbg debug.Debug, enablePrintStatements bool) (*ast.Compiler, error) { modules := map[string]*ast.Module{} @@ -951,7 +967,7 @@ func compile(c *ast.Capabilities, b *bundle.Bundle, dbg debug.Debug) (*ast.Compi modules[mf.URL] = mf.Parsed } - compiler := ast.NewCompiler().WithCapabilities(c).WithDebug(dbg.Writer()) + compiler := ast.NewCompiler().WithCapabilities(c).WithDebug(dbg.Writer()).WithEnablePrintStatements(enablePrintStatements) compiler.Compile(modules) if compiler.Failed() { diff --git a/compile/compile_test.go b/compile/compile_test.go index af895c119c..55a30ef54b 100644 --- a/compile/compile_test.go +++ b/compile/compile_test.go @@ -215,6 +215,64 @@ func TestCompilerLoadFilesystem(t *testing.T) { }) } +func TestCompilerLoadFilesystemWithEnablePrintStatementsFalse(t *testing.T) { + files := map[string]string{ + "test.rego": ` + package test + + allow { print(1) } + `, + "data.json": ` + {"b1": {"k": "v"}}`, + } + + test.WithTempFS(files, func(root string) { + compiler := New(). + WithPaths(root). + WithTarget("plan").WithEntrypoints("test/allow"). + WithEnablePrintStatements(false) + + if err := compiler.Build(context.Background()); err != nil { + t.Fatal(err) + } + + bundle := compiler.Bundle() + + if strings.Contains(string(bundle.PlanModules[0].Raw), "internal.print") { + t.Fatalf("output different than expected:\n\ngot: %v\n\nfound: internal.print", string(bundle.PlanModules[0].Raw)) + } + }) +} + +func TestCompilerLoadFilesystemWithEnablePrintStatementsTrue(t *testing.T) { + files := map[string]string{ + "test.rego": ` + package test + + allow { print(1) } + `, + "data.json": ` + {"b1": {"k": "v"}}`, + } + + test.WithTempFS(files, func(root string) { + compiler := New(). + WithPaths(root). + WithTarget("plan").WithEntrypoints("test/allow"). + WithEnablePrintStatements(true) + + if err := compiler.Build(context.Background()); err != nil { + t.Fatal(err) + } + + bundle := compiler.Bundle() + + if !strings.Contains(string(bundle.PlanModules[0].Raw), "internal.print") { + t.Fatalf("output different than expected:\n\ngot: %v\n\nmissing: internal.print", string(bundle.PlanModules[0].Raw)) + } + }) +} + func TestCompilerLoadHonorsFilter(t *testing.T) { files := map[string]string{ "test.rego": `