From ee115dba41f24bbebde31df65a1832ec9edc5122 Mon Sep 17 00:00:00 2001
From: Stephan Renatus <stephan.renatus@gmail.com>
Date: Thu, 13 May 2021 20:25:45 +0200
Subject: [PATCH] rego+bundle: Fix issues underlying #3209 (#3444)

* bundle: cleanup path before setting it as baseDir

GetBundleDirectoryLoader will do it again, but we need the cleaned-up
`path` to avoid "file:/" parts erroneously making it into the bundle's
ModuleFile structs. GetBundleDirectoryLoader is shared with the server,
so we'll keep that as-is.

* rego: avoid (*Bundle).ParsedModules(..)

Using this will prefix a path that is already complete. Using ModuleFile's
Path instead, we get the already-resolved path.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>
---
 bundle/bundle_test.go |  2 +-
 loader/loader.go      |  4 ++++
 loader/loader_test.go | 17 +++++++++++++++--
 rego/rego.go          |  6 +++---
 rego/rego_test.go     | 40 +++++++++++++++++++++++++++++++++++++++-
 5 files changed, 62 insertions(+), 7 deletions(-)

diff --git a/bundle/bundle_test.go b/bundle/bundle_test.go
index ca6c02a272..87e65ebb80 100644
--- a/bundle/bundle_test.go
+++ b/bundle/bundle_test.go
@@ -21,7 +21,7 @@ import (
 	"github.com/open-policy-agent/opa/internal/file/archive"
 )
 
-func TestManifestAddroot(t *testing.T) {
+func TestManifestAddRoot(t *testing.T) {
 	var m Manifest
 	m.AddRoot("x/y")
 	m.AddRoot("y/z")
diff --git a/loader/loader.go b/loader/loader.go
index 605f598c56..a3b0431975 100644
--- a/loader/loader.go
+++ b/loader/loader.go
@@ -178,6 +178,10 @@ func (fl fileLoader) Filtered(paths []string, filter Filter) (*Result, error) {
 // it will be treated as a normal tarball bundle. If a directory
 // is supplied it will be loaded as an unzipped bundle tree.
 func (fl fileLoader) AsBundle(path string) (*bundle.Bundle, error) {
+	path, err := fileurl.Clean(path)
+	if err != nil {
+		return nil, err
+	}
 	bundleLoader, isDir, err := GetBundleDirectoryLoader(path)
 	if err != nil {
 		return nil, err
diff --git a/loader/loader_test.go b/loader/loader_test.go
index 7ee26dcec8..fcea7c38dc 100644
--- a/loader/loader_test.go
+++ b/loader/loader_test.go
@@ -411,8 +411,9 @@ func TestAsBundleWithDir(t *testing.T) {
 
 func TestAsBundleWithFileURLDir(t *testing.T) {
 	files := map[string]string{
-		"/foo/data.json": "[1,2,3]",
-		"/.manifest":     `{"roots": ["foo"]}`,
+		"/foo/data.json":   "[1,2,3]",
+		"/foo/policy.rego": "package foo.bar\np = 1",
+		"/.manifest":       `{"roots": ["foo"]}`,
 	}
 
 	test.WithTempFS(files, func(rootDir string) {
@@ -425,6 +426,18 @@ func TestAsBundleWithFileURLDir(t *testing.T) {
 			t.Fatalf("Expected bundle to be non-nil")
 		}
 
+		if len(b.Modules) != 1 {
+			t.Fatalf("expected 1 modules, got %d", len(b.Modules))
+		}
+		expectedModulePaths := map[string]struct{}{
+			filepath.Join(rootDir, "/foo/policy.rego"): {},
+		}
+		for _, mf := range b.Modules {
+			if _, found := expectedModulePaths[mf.Path]; !found {
+				t.Errorf("Unexpected module file with path %s in bundle modules", mf.Path)
+			}
+		}
+
 		expectedData := util.MustUnmarshalJSON([]byte(`{"foo": [1,2,3]}`))
 		if !reflect.DeepEqual(b.Data, expectedData) {
 			t.Fatalf("expected data %+v, got %+v", expectedData, b.Data)
diff --git a/rego/rego.go b/rego/rego.go
index 3d3d92b4ce..248b611eab 100644
--- a/rego/rego.go
+++ b/rego/rego.go
@@ -246,9 +246,9 @@ func (pq preparedQuery) Modules() map[string]*ast.Module {
 		mods[name] = mod
 	}
 
-	for path, b := range pq.r.bundles {
-		for name, mod := range b.ParsedModules(path) {
-			mods[name] = mod
+	for _, b := range pq.r.bundles {
+		for _, mod := range b.Modules {
+			mods[mod.Path] = mod.Parsed
 		}
 	}
 
diff --git a/rego/rego_test.go b/rego/rego_test.go
index de757c8daa..2a906c0359 100644
--- a/rego/rego_test.go
+++ b/rego/rego_test.go
@@ -10,6 +10,7 @@ import (
 	"log"
 	"net/http"
 	"net/http/httptest"
+	"path/filepath"
 	"reflect"
 	"strconv"
 	"strings"
@@ -1448,10 +1449,47 @@ func TestRegoEvalWithBundle(t *testing.T) {
 		}
 
 		assertResultSet(t, rs, `[["bar"]]`)
+
+		mods := pq.Modules()
+		if exp, act := 1, len(mods); exp != act {
+			t.Fatalf("expected %d modules, found %d", exp, act)
+		}
+		for act := range mods {
+			if exp := filepath.Join(path, "x/x.rego"); exp != act {
+				t.Errorf("expected module name %q, got %q", exp, act)
+			}
+		}
 	})
 }
 
-func TestRegoEvalPoliciesinStore(t *testing.T) {
+func TestRegoEvalWithBundleURL(t *testing.T) {
+	files := map[string]string{
+		"x/x.rego": "package x\np = data.x.b",
+	}
+
+	test.WithTempFS(files, func(path string) {
+		ctx := context.Background()
+		pq, err := New(
+			LoadBundle("file://"+path),
+			Query("data.x.p"),
+		).PrepareForEval(ctx)
+		if err != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+
+		mods := pq.Modules()
+		if exp, act := 1, len(mods); exp != act {
+			t.Fatalf("expected %d modules, found %d", exp, act)
+		}
+		for act := range mods {
+			if exp := filepath.Join(path, "x/x.rego"); exp != act {
+				t.Errorf("expected module name %q, got %q", exp, act)
+			}
+		}
+	})
+}
+
+func TestRegoEvalPoliciesInStore(t *testing.T) {
 	store := mock.New()
 	ctx := context.Background()
 	txn := storage.NewTransactionOrDie(ctx, store, storage.WriteParams)