Skip to content

Commit

Permalink
Add ability to specify namespace for optimized files
Browse files Browse the repository at this point in the history 
Currently the namespace for partially evaluated files
in an optimized bundle cannot be modified. As a result
if more than one optimized bundle is loaded in OPA, a root
conflict error would occur as the optimized bundles have a
root called "partial" automatially added to their
manifest. This change allows the namespace to be configured via
the build command.

Fixes: open-policy-agent#5933

Signed-off-by: Ashutosh Narkar <[email protected]>
  • Loading branch information
@ashutosh-narkar
ashutosh-narkar committed Jun 25, 2023
1 parent 39c3ab4 commit 7ea3ee0
Show file tree
Hide file tree
Showing 3 changed files with 129 additions and 7 deletions.
8 changes: 6 additions & 2 deletions cmd/build.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ type buildParams struct {
claimsFile string
excludeVerifyFiles []string
plugin string
ns string
}

func newBuildParams() buildParams {
Expand Down Expand Up @@ -102,7 +103,8 @@ The -O flag controls the optimization level. By default, optimization is disable
When optimization is enabled the 'build' command generates a bundle that is semantically
equivalent to the input files however the structure of the files in the bundle may have
been changed by rewriting, inlining, pruning, etc. Higher optimization levels may result
in longer build times.
in longer build times. The --partial-namespace flag can used in conjunction with the -O flag
to specify the namespace for the partially evaluated files in the optimized bundle.
The 'build' command supports targets (specified by -t):
Expand Down Expand Up @@ -233,6 +235,7 @@ against OPA v0.22.0:
buildCommand.Flags().VarP(&buildParams.entrypoints, "entrypoint", "e", "set slash separated entrypoint path")
buildCommand.Flags().VarP(&buildParams.revision, "revision", "r", "set output bundle revision")
buildCommand.Flags().StringVarP(&buildParams.outputFile, "output", "o", "bundle.tar.gz", "set the output filename")
buildCommand.Flags().StringVar(&buildParams.ns, "partial-namespace", "partial", "set the namespace to use for partially evaluated files in an optimized bundle")

addBundleModeFlag(buildCommand.Flags(), &buildParams.bundleMode, false)
addIgnoreFlag(buildCommand.Flags(), &buildParams.ignore)
Expand Down Expand Up @@ -294,7 +297,8 @@ func dobuild(params buildParams, args []string) error {
WithPaths(args...).
WithFilter(buildCommandLoaderFilter(params.bundleMode, params.ignore)).
WithBundleVerificationConfig(bvc).
WithBundleSigningConfig(bsc)
WithBundleSigningConfig(bsc).
WithPartialNamespace(params.ns)

if params.revision.isSet {
compiler = compiler.WithRevision(*params.revision.v)
Expand Down
16 changes: 16 additions & 0 deletions compile/compile.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@ type Compiler struct {
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
fsys fs.FS // file system to use when loading paths
ns string
}

// New returns a new compiler instance that can be invoked.
Expand Down Expand Up @@ -228,6 +229,12 @@ func (c *Compiler) WithFS(fsys fs.FS) *Compiler {
return c
}

// WithPartialNamespace sets the namespace to use for partial evaluation results
func (c *Compiler) WithPartialNamespace(ns string) *Compiler {
c.ns = ns
return c
}

func addEntrypointsFromAnnotations(c *Compiler, ar []*ast.AnnotationsRef) error {
for _, ref := range ar {
var entrypoint ast.Ref
Expand Down Expand Up @@ -490,6 +497,10 @@ func (c *Compiler) optimize(ctx context.Context) error {
WithShallowInlining(c.optimizationLevel <= 1).
WithEnablePrintStatements(c.enablePrintStatements)

if c.ns != "" {
o = o.WithPartialNamespace(c.ns)
}

err := o.Do(ctx)
if err != nil {
return err
Expand Down Expand Up @@ -844,6 +855,11 @@ func (o *optimizer) WithShallowInlining(yes bool) *optimizer {
return o
}

func (o *optimizer) WithPartialNamespace(ns string) *optimizer {
o.nsprefix = ns
return o
}

func (o *optimizer) Do(ctx context.Context) error {

// NOTE(tsandall): if there are multiple entrypoints, copy the bundle because
Expand Down
112 changes: 107 additions & 5 deletions compile/compile_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -538,6 +538,57 @@ func TestCompilerOptimizationL2(t *testing.T) {
}
}

func TestCompilerOptimizationWithConfiguredNamespace(t *testing.T) {

files := map[string]string{
"test.rego": `
package test
p { not q }
q { k[input.a]; k[input.b] } # generate a product that is not inlined
k = {1,2,3}
`,
}

for _, useMemoryFS := range []bool{false, true} {
test.WithTestFS(files, useMemoryFS, func(root string, fsys fs.FS) {

compiler := New().
WithFS(fsys).
WithPaths(root).
WithOptimizationLevel(1).
WithEntrypoints("test/p").
WithPartialNamespace("custom")

err := compiler.Build(context.Background())
if err != nil {
t.Fatal(err)
}

if len(compiler.bundle.Modules) != 2 {
t.Fatalf("expected two modules but got: %v", len(compiler.bundle.Modules))
}

optimizedExp := ast.MustParseModule(`package custom
__not1_0_2__ = true { data.test.q = _; _ }`)

if optimizedExp.String() != compiler.bundle.Modules[0].Parsed.String() {
t.Fatalf("expected optimized module to be:\n\n%v\n\ngot:\n\n%v", optimizedExp, compiler.bundle.Modules[0])
}

expected := ast.MustParseModule(`package test
k = {1, 2, 3} { true }
p = true { not data.custom.__not1_0_2__ }
q = true { __local0__3 = input.a; data.test.k[__local0__3] = _; _; __local1__3 = input.b; data.test.k[__local1__3] = _; _ }`)

if expected.String() != compiler.bundle.Modules[1].Parsed.String() {
t.Fatalf("expected module to be:\n\n%v\n\ngot:\n\n%v", expected, compiler.bundle.Modules[1])
}
})
}
}

// NOTE(sr): we override this to not depend on build tags in tests
func wasmABIVersions(vs ...ast.WasmABIVersion) *ast.Capabilities {
caps := ast.CapabilitiesForThisVersion()
Expand Down Expand Up @@ -1428,7 +1479,7 @@ func TestOptimizerNoops(t *testing.T) {

for _, tc := range tests {
t.Run(tc.note, func(t *testing.T) {
o := getOptimizer(tc.modules, "", tc.entrypoints, nil)
o := getOptimizer(tc.modules, "", tc.entrypoints, nil, "")
cpy := o.bundle.Copy()
err := o.Do(context.Background())
if err != nil {
Expand Down Expand Up @@ -1479,7 +1530,7 @@ func TestOptimizerErrors(t *testing.T) {

for _, tc := range tests {
t.Run(tc.note, func(t *testing.T) {
o := getOptimizer(tc.modules, "", tc.entrypoints, nil)
o := getOptimizer(tc.modules, "", tc.entrypoints, nil, "")
cpy := o.bundle.Copy()
got := o.Do(context.Background())
if got == nil || got.Error() != tc.wantErr.Error() {
Expand All @@ -1499,6 +1550,7 @@ func TestOptimizerOutput(t *testing.T) {
modules map[string]string
data string
roots []string
namespace string
wantModules map[string]string
}{
{
Expand Down Expand Up @@ -1767,6 +1819,46 @@ func TestOptimizerOutput(t *testing.T) {
`,
},
},
{
note: "configured package namespace",
entrypoints: []string{"data.test.p"},
namespace: "custom",
modules: map[string]string{
"test.rego": `
package test
p { not q }
q { k[input.a]; k[input.b] } # generate a product that is not inlined
k = {1,2,3}
`,
},
wantModules: map[string]string{
"optimized/custom.rego": `
package custom
__not1_0_2__ = true { 1 = input.a; 1 = input.b }
__not1_0_2__ = true { 1 = input.a; 2 = input.b }
__not1_0_2__ = true { 1 = input.a; 3 = input.b }
__not1_0_2__ = true { 2 = input.a; 1 = input.b }
__not1_0_2__ = true { 2 = input.a; 2 = input.b }
__not1_0_2__ = true { 2 = input.a; 3 = input.b }
__not1_0_2__ = true { 3 = input.a; 1 = input.b }
__not1_0_2__ = true { 3 = input.a; 2 = input.b }
__not1_0_2__ = true { 3 = input.a; 3 = input.b }
`,
"optimized/test.rego": `
package test
p = __result__ { not data.custom.__not1_0_2__; __result__ = true }
`,
"test.rego": `
package test
q = true { k[input.a]; k[input.b] }
k = {1, 2, 3} { true }
`,
},
},
{
note: "infer unknowns from roots",
entrypoints: []string{"data.test.p"},
Expand Down Expand Up @@ -1935,7 +2027,7 @@ func TestOptimizerOutput(t *testing.T) {
for _, tc := range tests {
t.Run(tc.note, func(t *testing.T) {

o := getOptimizer(tc.modules, tc.data, tc.entrypoints, tc.roots)
o := getOptimizer(tc.modules, tc.data, tc.entrypoints, tc.roots, tc.namespace)
original := o.bundle.Copy()
err := o.Do(context.Background())
if err != nil {
Expand All @@ -1949,7 +2041,13 @@ func TestOptimizerOutput(t *testing.T) {

if len(tc.roots) > 0 {
exp.Manifest.Roots = &tc.roots
exp.Manifest.AddRoot("partial") // optimizer will add this automatically

// optimizer will add the manifest root in the optimized bundle automatically
if tc.namespace != "" {
exp.Manifest.AddRoot(tc.namespace)
} else {
exp.Manifest.AddRoot("partial")
}
}

exp.Manifest.Revision = "" // optimizations must reset the revision.
Expand Down Expand Up @@ -2017,7 +2115,7 @@ func TestRefSet(t *testing.T) {

}

func getOptimizer(modules map[string]string, data string, entries []string, roots []string) *optimizer {
func getOptimizer(modules map[string]string, data string, entries []string, roots []string, ns string) *optimizer {

b := &bundle.Bundle{
Modules: getModuleFiles(modules, true),
Expand All @@ -2042,6 +2140,10 @@ func getOptimizer(modules map[string]string, data string, entries []string, root
o := newOptimizer(ast.CapabilitiesForThisVersion(), b).
WithEntrypoints(entrypoints)

if ns != "" {
o = o.WithPartialNamespace(ns)
}

o.resultsymprefix = ""

return o
Expand Down

0 comments on commit 7ea3ee0

Please sign in to comment.