Persist downloaded bundle bytes to disk

Earlier with bundle persistence enabled, the bundle
plugin would save the bundle object to disk. In
scenarios where the downloaded bundle has multiple
data files, OPA would first read the bundle and merge
data in the bundle under the bundle.Bundle struct's
Data field. Then before persisting the bundle to disk,
the bundle plugin would use the bundle writer to write
the bundle to the provided output stream. The result
of this is that all the data files in the original
bundle are consolidated into one data.json file.

Now if signature verification is enabled, it will fail
since the files includes in the bundle's signature will not
match the ones in the persisted bundle.

This commit resolves this issue by persiting the bytes
of downloaded bundle to disk which then loaded
from disk maintain the same structure as the original.

Fixes: #3472

Signed-off-by: Ashutosh Narkar <anarkar4387@gmail.com>

download/download.go

@@ -6,8 +6,10 @@
 package download
 
 import (
+	"bytes"
 	"context"
 	"fmt"
+	"io"
 	"math/rand"
 	"net/http"
 	"path"
@@ -36,6 +38,7 @@ type Update struct {
 	Bundle  *bundle.Bundle
 	Error   error
 	Metrics metrics.Metrics
+	Raw     io.Reader
 }
 
 // Downloader implements low-level OPA bundle downloading. Downloader can be
@@ -56,6 +59,14 @@ type Downloader struct {
 	logger            logging.Logger
 	mtx               sync.Mutex
 	stopped           bool
+	persist           bool
+}
+
+type downloaderResponse struct {
+	b        *bundle.Bundle
+	raw      io.Reader
+	etag     string
+	longPoll bool
 }
 
 // New returns a new Downloader that can be started.
@@ -94,6 +105,12 @@ func (d *Downloader) WithSizeLimitBytes(n int64) *Downloader {
 	return d
 }
 
+// WithBundlePersistence specifies if the downloaded bundle will eventually be persisted to disk.
+func (d *Downloader) WithBundlePersistence(persist bool) *Downloader {
+	d.persist = persist
+	return d
+}
+
 // ClearCache resets the etag value on the downloader
 func (d *Downloader) ClearCache() {
 	d.etag = ""
@@ -183,17 +200,28 @@ func (d *Downloader) loop(ctx context.Context) {
 
 func (d *Downloader) oneShot(ctx context.Context) (bool, error) {
 	m := metrics.New()
-	b, etag, longPoll, err := d.download(ctx, m)
+	resp, err := d.download(ctx, m)
 
-	d.etag = etag
+	if err != nil {
+		d.etag = ""
+
+		if d.f != nil {
+			d.f(ctx, Update{ETag: "", Bundle: nil, Error: err, Metrics: m, Raw: nil})
+		}
+
+		return false, err
+	}
+
+	d.etag = resp.etag
 
 	if d.f != nil {
-		d.f(ctx, Update{ETag: etag, Bundle: b, Error: err, Metrics: m})
+		d.f(ctx, Update{ETag: resp.etag, Bundle: resp.b, Error: nil, Metrics: m, Raw: resp.raw})
 	}
-	return longPoll, err
+
+	return resp.longPoll, nil
 }
 
-func (d *Downloader) download(ctx context.Context, m metrics.Metrics) (*bundle.Bundle, string, bool, error) {
+func (d *Downloader) download(ctx context.Context, m metrics.Metrics) (*downloaderResponse, error) {
 	d.logger.Debug("Download starting.")
 
 	d.client = d.client.WithHeader("If-None-Match", d.etag)
@@ -213,44 +241,69 @@ func (d *Downloader) download(ctx context.Context, m metrics.Metrics) (*bundle.B
 
 	resp, err := d.client.Do(ctx, "GET", d.path)
 	if err != nil {
-		return nil, "", false, errors.Wrap(err, "request failed")
+		return nil, errors.Wrap(err, "request failed")
 	}
 
 	defer util.Close(resp)
 
 	switch resp.StatusCode {
 	case http.StatusOK:
+		var buf bytes.Buffer
 		if resp.Body != nil {
 			d.logger.Debug("Download in progress.")
 			m.Timer(metrics.RegoLoadBundles).Start()
 			defer m.Timer(metrics.RegoLoadBundles).Stop()
 			baseURL := path.Join(d.client.Config().URL, d.path)
-			loader := bundle.NewTarballLoaderWithBaseURL(resp.Body, baseURL)
+
+			var loader bundle.DirectoryLoader
+			if d.persist {
+				tee := io.TeeReader(resp.Body, &buf)
+				loader = bundle.NewTarballLoaderWithBaseURL(tee, baseURL)
+			} else {
+				loader = bundle.NewTarballLoaderWithBaseURL(resp.Body, baseURL)
+			}
+
 			reader := bundle.NewCustomReader(loader).WithMetrics(m).WithBundleVerificationConfig(d.bvc)
 			if d.sizeLimitBytes != nil {
 				reader = reader.WithSizeLimitBytes(*d.sizeLimitBytes)
 			}
 			b, err := reader.Read()
 			if err != nil {
-				return nil, "", false, err
+				return nil, err
 			}
-			return &b, resp.Header.Get("ETag"), isLongPollSupported(resp.Header), nil
+
+			return &downloaderResponse{
+				b:        &b,
+				raw:      &buf,
+				etag:     resp.Header.Get("ETag"),
+				longPoll: isLongPollSupported(resp.Header),
+			}, nil
 		}
 
 		d.logger.Debug("Server replied with empty body.")
-		return nil, "", isLongPollSupported(resp.Header), nil
+		return &downloaderResponse{
+			b:        nil,
+			raw:      nil,
+			etag:     "",
+			longPoll: isLongPollSupported(resp.Header),
+		}, nil
 	case http.StatusNotModified:
 		etag := resp.Header.Get("ETag")
 		if etag == "" {
 			etag = d.etag
 		}
-		return nil, etag, isLongPollSupported(resp.Header), nil
+		return &downloaderResponse{
+			b:        nil,
+			raw:      nil,
+			etag:     etag,
+			longPoll: isLongPollSupported(resp.Header),
+		}, nil
 	case http.StatusNotFound:
-		return nil, "", false, fmt.Errorf("server replied with not found")
+		return nil, fmt.Errorf("server replied with not found")
 	case http.StatusUnauthorized:
-		return nil, "", false, fmt.Errorf("server replied with not authorized")
+		return nil, fmt.Errorf("server replied with not authorized")
 	default:
-		return nil, "", false, fmt.Errorf("server replied with HTTP %v", resp.StatusCode)
+		return nil, fmt.Errorf("server replied with HTTP %v", resp.StatusCode)
 	}
 }
 

download/download_test.go

@@ -13,6 +13,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"reflect"
 	"strconv"
 	"strings"
 	"testing"
@@ -56,6 +57,58 @@ func TestStartStop(t *testing.T) {
 	d.Stop(ctx)
 }
 
+func TestStartStopWithBundlePersistence(t *testing.T) {
+	ctx := context.Background()
+	fixture := newTestFixture(t)
+
+	updates := make(chan *Update)
+
+	config := Config{}
+	if err := config.ValidateAndInjectDefaults(); err != nil {
+		t.Fatal(err)
+	}
+
+	d := New(config, fixture.client, "/bundles/test/bundle1").WithCallback(func(_ context.Context, u Update) {
+		updates <- &u
+	}).WithBundlePersistence(true)
+
+	d.Start(ctx)
+
+	// Give time for some download events to occur
+	time.Sleep(1 * time.Second)
+
+	u1 := <-updates
+
+	if u1.Bundle == nil || len(u1.Bundle.Modules) == 0 {
+		t.Fatal("expected bundle with at least one module but got:", u1)
+	}
+
+	if !strings.HasSuffix(u1.Bundle.Modules[0].URL, u1.Bundle.Modules[0].Path) {
+		t.Fatalf("expected URL to have path as suffix but got %v and %v", u1.Bundle.Modules[0].URL, u1.Bundle.Modules[0].Path)
+	}
+
+	if u1.Raw == nil {
+		t.Fatal("expected bundle reader to be non-nil")
+	}
+
+	r := bundle.NewReader(u1.Raw)
+
+	b, err := r.Read()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !reflect.DeepEqual(b.Data, u1.Bundle.Data) {
+		t.Fatal("expected the bundle object and reader to have the same data")
+	}
+
+	if len(b.Modules) != len(u1.Bundle.Modules) {
+		t.Fatal("expected the bundle object and reader to have the same number of bundle modules")
+	}
+
+	d.Stop(ctx)
+}
+
 func TestStopWithMultipleCalls(t *testing.T) {
 	ctx := context.Background()
 	fixture := newTestFixture(t)

plugins/bundle/plugin.go

@@ -6,11 +6,10 @@
 package bundle
 
 import (
-	"bytes"
 	"context"
 	"errors"
 	"fmt"
-	"io/ioutil"
+	"io"
 	"net/url"
 	"os"
 	"path/filepath"
@@ -328,7 +327,8 @@ func (p *Plugin) newDownloader(name string, source *Source) bundleLoader {
 	return download.New(conf, client, path).
 		WithCallback(callback).
 		WithBundleVerificationConfig(source.Signing).
-		WithSizeLimitBytes(source.SizeLimitBytes)
+		WithSizeLimitBytes(source.SizeLimitBytes).
+		WithBundlePersistence(p.persistBundle(name))
 }
 
 func (p *Plugin) oneShot(ctx context.Context, name string, u download.Update) {
@@ -394,7 +394,7 @@ func (p *Plugin) process(ctx context.Context, name string, u download.Update) {
 		if p.persistBundle(name) {
 			p.log(name).Debug("Persisting bundle to disk in progress.")
 
-			err := p.saveBundleToDisk(name, u.Bundle)
+			err := p.saveBundleToDisk(name, u.Raw)
 			if err != nil {
 				p.log(name).Error("Persisting bundle to disk failed: %v", err)
 				p.status[name].SetError(err)
@@ -524,13 +524,13 @@ func (p *Plugin) configDelta(newConfig *Config) (map[string]*Source, map[string]
 	return newBundles, updatedBundles, deletedBundles
 }
 
-func (p *Plugin) saveBundleToDisk(name string, b *bundle.Bundle) error {
+func (p *Plugin) saveBundleToDisk(name string, raw io.Reader) error {
 
 	bundleDir := filepath.Join(p.bundlePersistPath, name)
 	tmpFile := filepath.Join(bundleDir, ".bundle.tar.gz.tmp")
 	bundleFile := filepath.Join(bundleDir, "bundle.tar.gz")
 
-	saveErr := saveCurrentBundleToDisk(bundleDir, ".bundle.tar.gz.tmp", b)
+	saveErr := saveCurrentBundleToDisk(bundleDir, ".bundle.tar.gz.tmp", raw)
 	if saveErr != nil {
 		p.log(name).Error("Failed to save new bundle to disk: %v", saveErr)
 
@@ -548,13 +548,7 @@ func (p *Plugin) saveBundleToDisk(name string, b *bundle.Bundle) error {
 	return os.Rename(tmpFile, bundleFile)
 }
 
-func saveCurrentBundleToDisk(path, filename string, b *bundle.Bundle) error {
-	var buf bytes.Buffer
-
-	if err := bundle.NewWriter(&buf).UseModulePath(true).Write(*b); err != nil {
-		return err
-	}
-
+func saveCurrentBundleToDisk(path, filename string, raw io.Reader) error {
 	if _, err := os.Stat(path); os.IsNotExist(err) {
 		err = os.MkdirAll(path, os.ModePerm)
 		if err != nil {
@@ -562,7 +556,18 @@ func saveCurrentBundleToDisk(path, filename string, b *bundle.Bundle) error {
 		}
 	}
 
-	return ioutil.WriteFile(filepath.Join(path, filename), buf.Bytes(), 0644)
+	if raw == nil {
+		return fmt.Errorf("no raw bundle bytes to persist to disk")
+	}
+
+	dest, err := os.OpenFile(filepath.Join(path, filename), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
+	if err != nil {
+		return err
+	}
+	defer dest.Close()
+
+	_, err = io.Copy(dest, raw)
+	return err
 }
 
 func loadBundleFromDisk(path, name string, src *Source) (*bundle.Bundle, error) {

plugins/bundle/plugin_test.go

@@ -10,6 +10,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
@@ -22,6 +23,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/open-policy-agent/opa/internal/file/archive"
+
 	"github.com/open-policy-agent/opa/util/test"
 
 	"github.com/open-policy-agent/opa/ast"
@@ -226,7 +229,12 @@ func TestPluginOneShotBundlePersistence(t *testing.T) {
 
 	b.Manifest.Init()
 
-	plugin.oneShot(ctx, bundleName, download.Update{Bundle: &b, Metrics: metrics.New()})
+	var buf bytes.Buffer
+	if err := bundle.NewWriter(&buf).UseModulePath(true).Write(b); err != nil {
+		t.Fatal("unexpected error:", err)
+	}
+
+	plugin.oneShot(ctx, bundleName, download.Update{Bundle: &b, Metrics: metrics.New(), Raw: &buf})
 
 	ensurePluginState(t, plugin, plugins.StateOK)
 
@@ -271,6 +279,105 @@ func TestPluginOneShotBundlePersistence(t *testing.T) {
 	}
 }
 
+func TestPluginOneShotSignedBundlePersistence(t *testing.T) {
+
+	ctx := context.Background()
+	manager := getTestManager()
+
+	dir, err := ioutil.TempDir("", "")
+	if err != nil {
+		t.Fatalf("unexpected error %v", err)
+	}
+
+	defer os.RemoveAll(dir)
+
+	bundleName := "test-bundle"
+	vc := bundle.NewVerificationConfig(map[string]*bundle.KeyConfig{"foo": {Key: "secret", Algorithm: "HS256"}}, "foo", "", nil)
+	bundleSource := Source{
+		Persist: true,
+		Signing: vc,
+	}
+
+	bundles := map[string]*Source{}
+	bundles[bundleName] = &bundleSource
+
+	plugin := New(&Config{Bundles: bundles}, manager)
+
+	plugin.status[bundleName] = &Status{Name: bundleName, Metrics: metrics.New()}
+	plugin.downloaders[bundleName] = download.New(download.Config{}, plugin.manager.Client(""), bundleName)
+	plugin.bundlePersistPath = filepath.Join(dir, ".opa")
+
+	ensurePluginState(t, plugin, plugins.StateNotReady)
+
+	// simulate a bundle download error with no bundle on disk
+	plugin.oneShot(ctx, bundleName, download.Update{Error: fmt.Errorf("unknown error")})
+
+	if plugin.status[bundleName].Message == "" {
+		t.Fatal("expected error but got none")
+	}
+
+	ensurePluginState(t, plugin, plugins.StateNotReady)
+
+	// download a signed bundle and persist to disk. Then verify the bundle persisted to disk
+	signedTokenHS256 := `eyJhbGciOiJIUzI1NiJ9.eyJmaWxlcyI6W3sibmFtZSI6Ii5tYW5pZmVzdCIsImhhc2giOiI1MDdhMmMzOGExNDQxZGI1OGQyY2I4Nzk4MmM0MmFhOTFhNDM0MmVmNDIyYTZiNTQyZWRkZWJlZWY2ZjA0MTJmIiwiYWxnb3JpdGhtIjoiU0hBLTI1NiJ9LHsibmFtZSI6ImV4YW1wbGUxL2RhdGEuanNvbiIsImhhc2giOiI3YTM4YmY4MWYzODNmNjk0MzNhZDZlOTAwZDM1YjNlMjM4NTU5M2Y3NmE3YjdhYjVkNDM1NWI4YmE0MWVlMjRiIiwiYWxnb3JpdGhtIjoiU0hBLTI1NiJ9LHsibmFtZSI6ImV4YW1wbGUyL2RhdGEuanNvbiIsImhhc2giOiI5ZTRmMTg5YmY0MDc5ZDFiNmViNjQ0Njg3OTg2NmNkNWYzOWMyNjg4MGQ0ZmI1MThmNGUwMWNkMWJiZmU1MTNlIiwiYWxnb3JpdGhtIjoiU0hBLTI1NiJ9XX0.jCLRMyys5u8S2sTS2pWWY82IAeKDpLh3S641_BskCtY`
+
+	files := [][2]string{
+		{"/.manifest", `{"revision": "quickbrownfaux"}`},
+		{"/.signatures.json", fmt.Sprintf(`{"signatures": ["%v"]}`, signedTokenHS256)},
+		{"/example1/data.json", `{"foo": "bar"}`},
+		{"/example2/data.json", `{"x": true}`},
+	}
+
+	buf := archive.MustWriteTarGz(files)
+
+	var dup bytes.Buffer
+	tee := io.TeeReader(buf, &dup)
+	reader := bundle.NewReader(tee).WithBundleVerificationConfig(vc)
+	b, err := reader.Read()
+	if err != nil {
+		t.Fatal("unexpected error:", err)
+	}
+
+	plugin.oneShot(ctx, bundleName, download.Update{Bundle: &b, Metrics: metrics.New(), Raw: &dup})
+
+	ensurePluginState(t, plugin, plugins.StateOK)
+
+	// load signed bundle from disk
+	result, err := loadBundleFromDisk(plugin.bundlePersistPath, bundleName, bundles[bundleName])
+	if err != nil {
+		t.Fatal("unexpected error:", err)
+	}
+
+	if !result.Equal(b) {
+		t.Fatal("expected the downloaded bundle to be equal to the one loaded from disk")
+	}
+
+	// simulate a bundle download error and verify that the bundle on disk is activated
+	plugin.oneShot(ctx, bundleName, download.Update{Error: fmt.Errorf("unknown error")})
+
+	ensurePluginState(t, plugin, plugins.StateOK)
+
+	txn := storage.NewTransactionOrDie(ctx, manager.Store)
+	defer manager.Store.Abort(ctx, txn)
+
+	ids, err := manager.Store.ListPolicies(ctx, txn)
+	if err != nil {
+		t.Fatal(err)
+	} else if len(ids) != 0 {
+		t.Fatal("Expected no policy")
+	}
+
+	data, err := manager.Store.Read(ctx, txn, storage.Path{})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	expData := util.MustUnmarshalJSON([]byte(`{"example1": {"foo": "bar"}, "example2": {"x": true}, "system": {"bundles": {"test-bundle": {"manifest": {"revision": "quickbrownfaux", "roots": [""]}}}}}`))
+	if !reflect.DeepEqual(data, expData) {
+		t.Fatalf("Bad data content. Exp:\n%v\n\nGot:\n\n%v", expData, data)
+	}
+}
+
 func TestLoadAndActivateBundlesFromDisk(t *testing.T) {
 
 	ctx := context.Background()
@@ -321,7 +428,12 @@ func TestLoadAndActivateBundlesFromDisk(t *testing.T) {
 
 	b.Manifest.Init()
 
-	err = plugin.saveBundleToDisk(bundleName, &b)
+	var buf bytes.Buffer
+	if err := bundle.NewWriter(&buf).UseModulePath(true).Write(b); err != nil {
+		t.Fatal("unexpected error:", err)
+	}
+
+	err = plugin.saveBundleToDisk(bundleName, &buf)
 	if err != nil {
 		t.Fatalf("unexpected error %v", err)
 	}
@@ -1739,9 +1851,7 @@ func TestSaveBundleToDiskNew(t *testing.T) {
 	plugin := New(&Config{Bundles: bundles}, manager)
 	plugin.bundlePersistPath = filepath.Join(dir, ".opa")
 
-	b := getTestBundle(t)
-
-	err = plugin.saveBundleToDisk("foo", &b)
+	err = plugin.saveBundleToDisk("foo", getTestRawBundle(t))
 	if err != nil {
 		t.Fatalf("unexpected error %v", err)
 	}
@@ -1764,9 +1874,7 @@ func TestSaveBundleToDiskNewConfiguredPersistDir(t *testing.T) {
 		t.Fatalf("unexpected error %v", err)
 	}
 
-	b := getTestBundle(t)
-
-	err = plugin.saveBundleToDisk("foo", &b)
+	err = plugin.saveBundleToDisk("foo", getTestRawBundle(t))
 	if err != nil {
 		t.Fatalf("unexpected error %v", err)
 	}
@@ -1823,7 +1931,12 @@ func TestSaveBundleToDiskOverWrite(t *testing.T) {
 	}
 	newBundle.Manifest.Init()
 
-	err = plugin.saveBundleToDisk("foo", &newBundle)
+	var buf bytes.Buffer
+	if err := bundle.NewWriter(&buf).UseModulePath(true).Write(newBundle); err != nil {
+		t.Fatal("unexpected error:", err)
+	}
+
+	err = plugin.saveBundleToDisk("foo", &buf)
 	if err != nil {
 		t.Fatalf("unexpected error %v", err)
 	}
@@ -1839,8 +1952,6 @@ func TestSaveBundleToDiskOverWrite(t *testing.T) {
 }
 
 func TestSaveCurrentBundleToDisk(t *testing.T) {
-	b := getTestBundle(t)
-
 	srcDir, err := ioutil.TempDir("", "")
 	if err != nil {
 		t.Fatalf("unexpected error %v", err)
@@ -1848,7 +1959,7 @@ func TestSaveCurrentBundleToDisk(t *testing.T) {
 
 	defer os.RemoveAll(srcDir)
 
-	err = saveCurrentBundleToDisk(srcDir, "bundle.tar.gz", &b)
+	err = saveCurrentBundleToDisk(srcDir, "bundle.tar.gz", getTestRawBundle(t))
 	if err != nil {
 		t.Fatalf("unexpected error %v", err)
 	}
@@ -1856,6 +1967,16 @@ func TestSaveCurrentBundleToDisk(t *testing.T) {
 	if _, err := os.Stat(filepath.Join(srcDir, "bundle.tar.gz")); err != nil {
 		t.Fatalf("unexpected error %v", err)
 	}
+
+	err = saveCurrentBundleToDisk(srcDir, "bundle.tar.gz", nil)
+	if err == nil {
+		t.Fatal("expected error but got nil")
+	}
+
+	expErrMsg := "no raw bundle bytes to persist to disk"
+	if err.Error() != expErrMsg {
+		t.Fatalf("expected error: %v but got: %v", expErrMsg, err)
+	}
 }
 
 func TestLoadBundleFromDisk(t *testing.T) {
@@ -2082,6 +2203,19 @@ func getTestSignedBundle(t *testing.T) bundle.Bundle {
 	return b
 }
 
+func getTestRawBundle(t *testing.T) io.Reader {
+	t.Helper()
+
+	b := getTestBundle(t)
+
+	var buf bytes.Buffer
+	if err := bundle.NewWriter(&buf).UseModulePath(true).Write(b); err != nil {
+		t.Fatal("unexpected error:", err)
+	}
+
+	return &buf
+}
+
 func validateStoreState(ctx context.Context, t *testing.T, store storage.Store, root string, expData interface{}, expIds []string, expBundleName string, expBundleRev string, expMetadata map[string]interface{}) {
 	t.Helper()
 	if err := storage.Txn(ctx, store, storage.TransactionParams{}, func(txn storage.Transaction) error {