From 9555bb810abdcfb58c4ccfd539c4f4c489214c9e Mon Sep 17 00:00:00 2001
From: Avi Deitcher <avi@deitcher.net>
Date: Tue, 17 Nov 2020 19:38:03 +0200
Subject: [PATCH] add multireader (#199)

Signed-off-by: Avi Deitcher <avi@deitcher.net>

Co-authored-by: Josh Dolitsky <393494+jdolitsky@users.noreply.github.com>
---
 pkg/content/multireader.go      | 40 +++++++++++++++++++++
 pkg/content/multireader_test.go | 62 +++++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+)
 create mode 100644 pkg/content/multireader.go
 create mode 100644 pkg/content/multireader_test.go

diff --git a/pkg/content/multireader.go b/pkg/content/multireader.go
new file mode 100644
index 000000000..fb9cf8bfb
--- /dev/null
+++ b/pkg/content/multireader.go
@@ -0,0 +1,40 @@
+package content
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/containerd/containerd/content"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+)
+
+// MultiReader store to read content from multiple stores. It finds the content by asking each underlying
+// store to find the content, which it does based on the hash.
+//
+// Example:
+//        fileStore := NewFileStore(rootPath)
+//        memoryStore := NewMemoryStore()
+//        // load up content in fileStore and memoryStore
+//        multiStore := MultiReader([]content.Provider{fileStore, memoryStore})
+//
+// You now can use multiStore anywhere that content.Provider is accepted
+type MultiReader struct {
+	stores []content.Provider
+}
+
+// AddStore add a store to read from
+func (m *MultiReader) AddStore(store ...content.Provider) {
+	m.stores = append(m.stores, store...)
+}
+
+// ReaderAt get a reader
+func (m MultiReader) ReaderAt(ctx context.Context, desc ocispec.Descriptor) (content.ReaderAt, error) {
+	for _, store := range m.stores {
+		r, err := store.ReaderAt(ctx, desc)
+		if r != nil && err == nil {
+			return r, nil
+		}
+	}
+	// we did not find any
+	return nil, fmt.Errorf("not found")
+}
diff --git a/pkg/content/multireader_test.go b/pkg/content/multireader_test.go
new file mode 100644
index 000000000..3c74f8919
--- /dev/null
+++ b/pkg/content/multireader_test.go
@@ -0,0 +1,62 @@
+package content_test
+
+import (
+	"context"
+	"io/ioutil"
+	"testing"
+
+	ctrcontent "github.com/containerd/containerd/content"
+	"github.com/deislabs/oras/pkg/content"
+	digest "github.com/opencontainers/go-digest"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+)
+
+var (
+	testContentA     = []byte("Hello World!")
+	testContentHashA = digest.FromBytes(testContentA)
+	testContentB     = []byte("So long and thanks for all the fish!")
+	testContentHashB = digest.FromBytes(testContentB)
+	testDescriptorA  = ocispec.Descriptor{
+		MediaType: ocispec.MediaTypeImageConfig,
+		Digest:    testContentHashA,
+		Size:      int64(len(testContentA)),
+	}
+	testDescriptorB = ocispec.Descriptor{
+		MediaType: ocispec.MediaTypeImageConfig,
+		Digest:    testContentHashB,
+		Size:      int64(len(testContentB)),
+	}
+)
+
+func TestMultiReader(t *testing.T) {
+	mem1, mem2 := content.NewMemoryStore(), content.NewMemoryStore()
+	mem1.Add("a", ocispec.MediaTypeImageConfig, testContentA)
+	mem2.Add("b", ocispec.MediaTypeImageConfig, testContentB)
+	multiReader := content.MultiReader{}
+	multiReader.AddStore(mem1, mem2)
+
+	ctx := context.Background()
+	contentA, err := multiReader.ReaderAt(ctx, testDescriptorA)
+	if err != nil {
+		t.Fatalf("failed to get a reader for descriptor A: %v", err)
+	}
+	outputA, err := ioutil.ReadAll(ctrcontent.NewReader(contentA))
+	if err != nil {
+		t.Fatalf("failed to read content for descriptor A: %v", err)
+	}
+	if string(outputA) != string(testContentA) {
+		t.Errorf("mismatched content for A, actual '%s', expected '%s'", outputA, testContentA)
+	}
+
+	contentB, err := multiReader.ReaderAt(ctx, testDescriptorB)
+	if err != nil {
+		t.Fatalf("failed to get a reader for descriptor B: %v", err)
+	}
+	outputB, err := ioutil.ReadAll(ctrcontent.NewReader(contentB))
+	if err != nil {
+		t.Fatalf("failed to read content for descriptor B: %v", err)
+	}
+	if string(outputB) != string(testContentB) {
+		t.Errorf("mismatched content for B, actual '%s', expected '%s'", outputB, testContentB)
+	}
+}