From 54583b977be2fae191830d3323292eefd5082b8e Mon Sep 17 00:00:00 2001
From: Avi Deitcher <avi@deitcher.net>
Date: Fri, 12 Mar 2021 10:49:01 +0200
Subject: [PATCH] cache docker types too

Signed-off-by: Avi Deitcher <avi@deitcher.net>
---
 pkg/oras/pull.go      |  2 +-
 pkg/oras/pull_opts.go | 23 +++++++++++++++++++++--
 pkg/oras/push.go      |  2 +-
 pkg/oras/store.go     | 23 +++++++++++++----------
 4 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/pkg/oras/pull.go b/pkg/oras/pull.go
index d578e7bdf..0463db66a 100644
--- a/pkg/oras/pull.go
+++ b/pkg/oras/pull.go
@@ -61,7 +61,7 @@ func fetchContent(ctx context.Context, fetcher remotes.Fetcher, desc ocispec.Des
 
 	store := opts.contentProvideIngester
 	if store == nil {
-		store = newHybridStoreFromIngester(ingester)
+		store = newHybridStoreFromIngester(ingester, opts.cachedMediaTypes)
 	}
 	handlers := []images.Handler{
 		filterHandler(opts, opts.allowedMediaTypes...),
diff --git a/pkg/oras/pull_opts.go b/pkg/oras/pull_opts.go
index 17eab5297..7dcb104bf 100644
--- a/pkg/oras/pull_opts.go
+++ b/pkg/oras/pull_opts.go
@@ -21,6 +21,7 @@ type pullOpts struct {
 	callbackHandlers       []images.Handler
 	contentProvideIngester orascontent.ProvideIngester
 	filterName             func(ocispec.Descriptor) bool
+	cachedMediaTypes       []string
 }
 
 // PullOpt allows callers to set options on the oras pull
@@ -28,8 +29,26 @@ type PullOpt func(o *pullOpts) error
 
 func pullOptsDefaults() *pullOpts {
 	return &pullOpts{
-		dispatch:   images.Dispatch,
-		filterName: filterName,
+		dispatch:         images.Dispatch,
+		filterName:       filterName,
+		cachedMediaTypes: []string{ocispec.MediaTypeImageManifest, ocispec.MediaTypeImageIndex},
+	}
+}
+
+// WithCachedMediaTypes sets the media types normally cached in memory when pulling.
+func WithCachedMediaTypes(cachedMediaTypes ...string) PullOpt {
+	return func(o *pullOpts) error {
+		o.cachedMediaTypes = cachedMediaTypes
+		return nil
+	}
+}
+
+// WithAdditionalCachedMediaTypes adds media types normally cached in memory when pulling.
+// This does not replace the default media types, but appends to them
+func WithAdditionalCachedMediaTypes(cachedMediaTypes ...string) PullOpt {
+	return func(o *pullOpts) error {
+		o.cachedMediaTypes = append(o.cachedMediaTypes, cachedMediaTypes...)
+		return nil
 	}
 }
 
diff --git a/pkg/oras/push.go b/pkg/oras/push.go
index 42450658b..9d9271ec4 100644
--- a/pkg/oras/push.go
+++ b/pkg/oras/push.go
@@ -57,7 +57,7 @@ func Push(ctx context.Context, resolver remotes.Resolver, ref string, provider c
 
 //func pack(store *hybridStore, descriptors []ocispec.Descriptor, opts *pushOpts) (ocispec.Descriptor, error) {
 func pack(provider content.Provider, descriptors []ocispec.Descriptor, opts *pushOpts) (ocispec.Descriptor, content.Store, error) {
-	store := newHybridStoreFromProvider(provider)
+	store := newHybridStoreFromProvider(provider, nil)
 
 	// Config
 	var config ocispec.Descriptor
diff --git a/pkg/oras/store.go b/pkg/oras/store.go
index 84348d52a..95a8ee0d5 100644
--- a/pkg/oras/store.go
+++ b/pkg/oras/store.go
@@ -20,22 +20,25 @@ var (
 )
 
 type hybridStore struct {
-	cache    *orascontent.Memorystore
-	provider content.Provider
-	ingester content.Ingester
+	cache            *orascontent.Memorystore
+	cachedMediaTypes []string
+	provider         content.Provider
+	ingester         content.Ingester
 }
 
-func newHybridStoreFromProvider(provider content.Provider) *hybridStore {
+func newHybridStoreFromProvider(provider content.Provider, cachedMediaTypes []string) *hybridStore {
 	return &hybridStore{
-		cache:    orascontent.NewMemoryStore(),
-		provider: provider,
+		cache:            orascontent.NewMemoryStore(),
+		cachedMediaTypes: cachedMediaTypes,
+		provider:         provider,
 	}
 }
 
-func newHybridStoreFromIngester(ingester content.Ingester) *hybridStore {
+func newHybridStoreFromIngester(ingester content.Ingester, cachedMediaTypes []string) *hybridStore {
 	return &hybridStore{
-		cache:    orascontent.NewMemoryStore(),
-		ingester: ingester,
+		cache:            orascontent.NewMemoryStore(),
+		cachedMediaTypes: cachedMediaTypes,
+		ingester:         ingester,
 	}
 }
 
@@ -64,7 +67,7 @@ func (s *hybridStore) Writer(ctx context.Context, opts ...content.WriterOpt) (co
 		}
 	}
 
-	if isAllowedMediaType(wOpts.Desc.MediaType, ocispec.MediaTypeImageManifest, ocispec.MediaTypeImageIndex) || s.ingester == nil {
+	if isAllowedMediaType(wOpts.Desc.MediaType, s.cachedMediaTypes...) || s.ingester == nil {
 		cacheWriter, err := s.cache.Writer(ctx, opts...)
 		if err != nil {
 			return nil, err