From ad7a605584609fda997779d69c405323af060b66 Mon Sep 17 00:00:00 2001
From: Sergey-Kizimov <96204102+Sergey-Kizimov@users.noreply.github.com>
Date: Wed, 21 Aug 2024 14:49:34 -0700
Subject: [PATCH] download: Fix saving OCI bundles on disk

This commit fixes an issue related to zero-sized bundles being saved to disk,
which can cause OPA to fail to start if a remote OCI repository is unavailable.

Fixes: #6939

Signed-off-by: Sergey-Kizimov <serget.kizimov@hiya.com>
---
 download/oci_download.go      | 14 +++++++---
 download/oci_download_test.go | 50 +++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/download/oci_download.go b/download/oci_download.go
index b5c744d30e..b514a41b58 100644
--- a/download/oci_download.go
+++ b/download/oci_download.go
@@ -3,6 +3,7 @@
 package download
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -210,13 +211,14 @@ func (d *OCIDownloader) oneShot(ctx context.Context) error {
 	d.SetCache(resp.etag) // set the current etag sha to the cache
 
 	if d.f != nil {
-		d.f(ctx, Update{ETag: resp.etag, Bundle: resp.b, Error: nil, Metrics: m, Raw: resp.raw})
+		d.f(ctx, Update{ETag: resp.etag, Bundle: resp.b, Error: nil, Metrics: m, Raw: resp.raw, Size: resp.size})
 	}
 	return nil
 }
 
 func (d *OCIDownloader) download(ctx context.Context, m metrics.Metrics) (*downloaderResponse, error) {
 	d.logger.Debug("OCI - Download starting.")
+	var buf bytes.Buffer
 
 	preferences := []string{fmt.Sprintf("modes=%v,%v", defaultBundleMode, deltaBundleMode)}
 
@@ -256,10 +258,15 @@ func (d *OCIDownloader) download(ctx context.Context, m metrics.Metrics) (*downl
 		}, nil
 	}
 	fileReader, err := os.Open(bundleFilePath)
+
+	cnt := &count{}
+	r := io.TeeReader(fileReader, cnt)
+	tee := io.TeeReader(r, &buf)
+
 	if err != nil {
 		return nil, err
 	}
-	loader := bundle.NewTarballLoaderWithBaseURL(fileReader, d.localStorePath)
+	loader := bundle.NewTarballLoaderWithBaseURL(tee, d.localStorePath)
 	reader := bundle.NewCustomReader(loader).
 		WithMetrics(m).
 		WithBundleVerificationConfig(d.bvc).
@@ -274,9 +281,10 @@ func (d *OCIDownloader) download(ctx context.Context, m metrics.Metrics) (*downl
 
 	return &downloaderResponse{
 		b:        &bundleInfo,
-		raw:      fileReader,
+		raw:      &buf,
 		etag:     etag,
 		longPoll: false,
+		size:     cnt.Bytes(),
 	}, nil
 }
 
diff --git a/download/oci_download_test.go b/download/oci_download_test.go
index 6b5e2ff3cd..86b68c1cc2 100644
--- a/download/oci_download_test.go
+++ b/download/oci_download_test.go
@@ -8,6 +8,7 @@ import (
 	"encoding/base64"
 	"fmt"
 	"net/http"
+	"reflect"
 	"strings"
 	"testing"
 	"time"
@@ -382,6 +383,55 @@ func TestOCICustomAuthPlugin(t *testing.T) {
 	}
 }
 
+func TestOCIValidateAndInjectDefaults(t *testing.T) {
+	ctx := context.Background()
+	fixture := newTestFixture(t)
+	fixture.server.expEtag = "sha256:c5834dbce332cabe6ae68a364de171a50bf5b08024c27d7c08cc72878b4df7ff"
+
+	updates := make(chan *Update)
+
+	config := Config{}
+	if err := config.ValidateAndInjectDefaults(); err != nil {
+		t.Fatal(err)
+	}
+
+	d := NewOCI(config, fixture.client, "ghcr.io/org/repo:latest", t.TempDir()).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.Size == 0 {
+		t.Fatal("expected non-0 size")
+	}
+
+	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 mockAuthPluginLookup(string) rest.HTTPAuthPlugin {
 	return &mockAuthPlugin{}
 }