loader: Support for loading bundle dirs

This adds a new API to load a bundle from a path which can be either
a tarball file or a directory to load as a bundle.

Signed-off-by: Patrick East <east.patrick@gmail.com>

loader/loader.go

@@ -15,6 +15,8 @@ import (
 	"runtime"
 	"strings"
 
+	"github.com/open-policy-agent/opa/internal/file"
+
 	"github.com/ghodss/yaml"
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/bundle"
@@ -131,6 +133,37 @@ func Rego(path string) (*RegoFile, error) {
 	return loadRego(path, bs)
 }
 
+// AsBundle loads a path as a bundle. If it is a single file
+// it will be treated as a normal tarball bundle. If a directory
+// is supplied it will be loaded as an unzipped bundle tree.
+func AsBundle(path string) (*bundle.Bundle, error) {
+	path, err := cleanFileURL(path)
+	if err != nil {
+		return nil, err
+	}
+
+	fi, err := os.Stat(path)
+	if err != nil {
+		return nil, fmt.Errorf("error reading %q: %s", path, err)
+	}
+
+	var bundleLoader file.DirectoryLoader
+
+	if fi.IsDir() {
+		bundleLoader = file.NewDirectoryLoader(path)
+	} else {
+		fh, err := os.Open(path)
+		if err != nil {
+			return nil, err
+		}
+		bundleLoader = file.NewTarballLoader(fh)
+	}
+
+	br := bundle.NewCustomReader(bundleLoader)
+	b, err := br.Read()
+	return &b, err
+}
+
 // CleanPath returns the normalized version of a path that can be used as an identifier.
 func CleanPath(path string) string {
 	return strings.Trim(path, "/")
@@ -328,7 +361,7 @@ func loadKnownTypes(path string, bs []byte) (interface{}, error) {
 		return loadYAML(path, bs)
 	default:
 		if strings.HasSuffix(path, ".tar.gz") {
-			return loadBundle(bs)
+			return loadBundleFile(bs)
 		}
 	}
 	return nil, unrecognizedFile(path)
@@ -350,11 +383,16 @@ func loadFileForAnyType(path string, bs []byte) (interface{}, error) {
 	return nil, unrecognizedFile(path)
 }
 
-func loadBundle(bs []byte) (bundle.Bundle, error) {
-	br := bundle.NewReader(bytes.NewBuffer(bs)).IncludeManifestInData(true)
+func loadBundleFile(bs []byte) (bundle.Bundle, error) {
+	tl := file.NewTarballLoader(bytes.NewBuffer(bs))
+	br := bundle.NewCustomReader(tl).IncludeManifestInData(true)
 	return br.Read()
 }
 
+func loadBundleDir(path string) (bundle.Bundle, error) {
+	return bundle.Bundle{}, nil
+}
+
 func loadRego(path string, bs []byte) (*RegoFile, error) {
 	module, err := ast.ParseModule(path, string(bs))
 	if err != nil {

loader/loader_test.go

@@ -259,6 +259,126 @@ func TestLoadBundleSubDir(t *testing.T) {
 	})
 }
 
+func TestAsBundleWithDir(t *testing.T) {
+	files := map[string]string{
+		"/foo/data.json":    "[1,2,3]",
+		"/bar/bar.yaml":     "abc",  // Should be ignored
+		"/baz/qux/qux.json": "null", // Should be ignored
+		"/foo/policy.rego":  "package foo\np = 1",
+		"base.rego":         "package bar\nx = 1",
+		"/.manifest":        `{"roots": ["foo", "bar", "baz"]}`,
+	}
+
+	test.WithTempFS(files, func(rootDir string) {
+		b, err := AsBundle(rootDir)
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		if b == nil {
+			t.Fatalf("Expected bundle to be non-nil")
+		}
+
+		if len(b.Modules) != 2 {
+			t.Fatalf("expected 2 modules, got %d", len(b.Modules))
+		}
+
+		expectedData := util.MustUnmarshalJSON([]byte(`{"foo": [1,2,3]}`))
+		if !reflect.DeepEqual(b.Data, expectedData) {
+			t.Fatalf("expected data %+v, got %+v", expectedData, b.Data)
+		}
+
+		expectedRoots := []string{"foo", "bar", "baz"}
+		if !reflect.DeepEqual(*b.Manifest.Roots, expectedRoots) {
+			t.Fatalf("expected roots %s, got: %s", expectedRoots, *b.Manifest.Roots)
+		}
+	})
+}
+
+func TestAsBundleWithFileURLDir(t *testing.T) {
+	files := map[string]string{
+		"/foo/data.json": "[1,2,3]",
+		"/.manifest":     `{"roots": ["foo"]}`,
+	}
+
+	test.WithTempFS(files, func(rootDir string) {
+		b, err := AsBundle("file://" + rootDir)
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		if b == nil {
+			t.Fatalf("Expected bundle to be non-nil")
+		}
+
+		expectedData := util.MustUnmarshalJSON([]byte(`{"foo": [1,2,3]}`))
+		if !reflect.DeepEqual(b.Data, expectedData) {
+			t.Fatalf("expected data %+v, got %+v", expectedData, b.Data)
+		}
+
+		expectedRoots := []string{"foo"}
+		if !reflect.DeepEqual(*b.Manifest.Roots, expectedRoots) {
+			t.Fatalf("expected roots %s, got: %s", expectedRoots, *b.Manifest.Roots)
+		}
+	})
+}
+
+func TestAsBundleWithFile(t *testing.T) {
+	files := map[string]string{
+		"bundle.tar.gz": "",
+	}
+
+	mod := "package b.c\np=1"
+
+	b := &bundle.Bundle{
+		Manifest: bundle.Manifest{
+			Roots:    &[]string{"a", "b/c"},
+			Revision: "123",
+		},
+		Data: map[string]interface{}{
+			"a": map[string]interface{}{
+				"b": []int{4, 5, 6},
+			},
+		},
+		Modules: []bundle.ModuleFile{
+			{
+				Path:   "/policy.rego",
+				Raw:    []byte(mod),
+				Parsed: ast.MustParseModule(mod),
+			},
+		},
+	}
+
+	test.WithTempFS(files, func(rootDir string) {
+		path := filepath.Join(rootDir, "bundle.tar.gz")
+		f, err := os.Create(path)
+		if err != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+
+		err = bundle.Write(f, *b)
+		f.Close()
+		if err != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+
+		actual, err := AsBundle(path)
+		if err != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+
+		var tmp interface{} = b
+		err = util.RoundTrip(&tmp)
+		if err != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+
+		if !actual.Equal(*b) {
+			t.Fatalf("Loaded bundle doesn't match expected.\n\nExpected: %+v\n\nActual: %+v\n\n", b, actual)
+		}
+	})
+}
+
 func TestLoadRooted(t *testing.T) {
 	files := map[string]string{
 		"/foo.json":         "[1,2,3]",