loader: add WithReader to pass io.Reader directly

Signed-off-by: Stephan Renatus <stephan@styra.com>

loader/loader.go

@@ -8,6 +8,7 @@ package loader
 import (
 	"bytes"
 	"fmt"
+	"io"
 	"io/fs"
 	"os"
 	"path/filepath"
@@ -92,6 +93,7 @@ type FileLoader interface {
 	All(paths []string) (*Result, error)
 	Filtered(paths []string, filter Filter) (*Result, error)
 	AsBundle(path string) (*bundle.Bundle, error)
+	WithReader(io.Reader) FileLoader
 	WithFS(fs.FS) FileLoader
 	WithMetrics(metrics.Metrics) FileLoader
 	WithFilter(Filter) FileLoader
@@ -118,6 +120,7 @@ type fileLoader struct {
 	files      map[string]bundle.FileInfo
 	opts       ast.ParserOptions
 	fsys       fs.FS
+	reader     io.Reader
 }
 
 // WithFS provides an fs.FS to use for loading files. You can pass nil to
@@ -128,6 +131,14 @@ func (fl *fileLoader) WithFS(fsys fs.FS) FileLoader {
 	return fl
 }
 
+// WithReader provides an io.Reader to use for loading the bundle tarball.
+// An io.Reader passed via WithReader takes precedence over an fs.FS passed
+// via WithFS.
+func (fl *fileLoader) WithReader(rdr io.Reader) FileLoader {
+	fl.reader = rdr
+	return fl
+}
+
 // WithMetrics provides the metrics instance to use while loading
 func (fl *fileLoader) WithMetrics(m metrics.Metrics) FileLoader {
 	fl.metrics = m
@@ -220,7 +231,14 @@ func (fl fileLoader) AsBundle(path string) (*bundle.Bundle, error) {
 	if err != nil {
 		return nil, err
 	}
-	bundleLoader, isDir, err := GetBundleDirectoryLoaderFS(fl.fsys, path, fl.filter)
+
+	var bundleLoader bundle.DirectoryLoader
+	var isDir bool
+	if fl.reader != nil {
+		bundleLoader = bundle.NewTarballLoaderWithBaseURL(fl.reader, path).WithFilter(fl.filter)
+	} else {
+		bundleLoader, isDir, err = GetBundleDirectoryLoaderFS(fl.fsys, path, fl.filter)
+	}
 	if err != nil {
 		return nil, err
 	}

loader/loader_test.go

@@ -312,7 +312,61 @@ func TestLoadBundle(t *testing.T) {
 			t.Fatalf("Expected %v but got: %v", string(testBundle.Modules[0].Raw), loaded.Modules["/x.rego"].Raw)
 		}
 	})
+}
 
+func TestLoadBundleWithReader(t *testing.T) {
+
+	buf := bytes.Buffer{}
+	testBundle := bundle.Bundle{
+		Modules: []bundle.ModuleFile{
+			{
+				Path: "x.rego",
+				Raw: []byte(`
+				package baz
+
+				p = 1`),
+			},
+		},
+		Data: map[string]interface{}{
+			"foo": "bar",
+		},
+		Manifest: bundle.Manifest{
+			Revision: "",
+			Roots:    &[]string{"foo", "baz"},
+		},
+	}
+
+	if err := bundle.Write(&buf, testBundle); err != nil {
+		t.Fatal(err)
+	}
+
+	b, err := NewFileLoader().WithReader(&buf).AsBundle("bundle.tar.gz")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if b == nil {
+		t.Fatalf("Expected bundle to be non-nil")
+	}
+
+	if exp, act := 1, len(b.Modules); exp != act {
+		t.Fatalf("expected %d modules, got %d", exp, act)
+	}
+
+	expectedModulePaths := map[string]struct{}{
+		"/x.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)
+		}
+	}
+
+	if exp, act := map[string]any{"foo": "bar"}, b.Data; !reflect.DeepEqual(act, exp) {
+		t.Fatalf("expected data %+v, got %+v", exp, act)
+	}
+	if exp, act := []string{"foo", "baz"}, *b.Manifest.Roots; !reflect.DeepEqual(act, exp) {
+		t.Fatalf("expected roots %v, got %v", exp, act)
+	}
 }
 
 func TestLoadBundleSubDir(t *testing.T) {