Update registry cache to (ab)use Data field of Descriptor objects

This allows us to store full descriptors (necessary to implement `Resolve*`), but with a net decrease in the number of fields we have to juggle / keep in sync.

This does mean consumers need to be careful about how they use the `Descriptor` objects we return (esp. WRT `Data`), but it makes it easier for them to then have `Data` available if they want it (which is something I'd like to use in the future).  This is a net win anyhow because the upstream objects might've contained `Data` fields so this forces us to deal with them in a sane way we're comfortable with instead of potentially just including them verbatim unintentionally. 🚀

Author: tianon

registry/cache.go

@@ -21,8 +21,7 @@ func RegistryCache(r ociregistry.Interface) ociregistry.Interface {
 		registry: r, // TODO support "nil" here so this can be a poor-man's ocimem implementation? 👀  see also https://github.com/cue-labs/oci/issues/24
 		has:      map[string]bool{},
 		tags:     map[string]ociregistry.Digest{},
-		types:    map[ociregistry.Digest]string{},
-		data:     map[ociregistry.Digest][]byte{},
+		data:     map[ociregistry.Digest]ociregistry.Descriptor{},
 	}
 }
 
@@ -32,11 +31,10 @@ type registryCache struct {
 	registry ociregistry.Interface
 
 	// https://github.com/cue-labs/oci/issues/24
-	mu    sync.Mutex                    // TODO some kind of per-object/name/digest mutex so we don't request the same object from the upstream registry concurrently (on *top* of our maps mutex)?
-	has   map[string]bool               // "repo/name@digest" => true (whether a given repo has the given digest)
-	tags  map[string]ociregistry.Digest // "repo/name:tag" => digest
-	types map[ociregistry.Digest]string // digest => "mediaType" (most recent *storing* / "cache-miss" lookup wins, in the case of upstream/cross-repo ambiguity)
-	data  map[ociregistry.Digest][]byte // digest => data
+	mu   sync.Mutex                                    // TODO some kind of per-object/name/digest mutex so we don't request the same object from the upstream registry concurrently (on *top* of our maps mutex)?
+	has  map[string]bool                               // "repo/name@digest" => true (whether a given repo has the given digest)
+	tags map[string]ociregistry.Digest                 // "repo/name:tag" => digest
+	data map[ociregistry.Digest]ociregistry.Descriptor // digest => mediaType+size(+data) (most recent *storing* / "cache-miss" lookup wins, in the case of upstream/cross-repo ambiguity)
 }
 
 func cacheKeyDigest(repo string, digest ociregistry.Digest) string {
@@ -52,41 +50,38 @@ func (rc *registryCache) getBlob(ctx context.Context, repo string, digest ocireg
 	rc.mu.Lock()
 	defer rc.mu.Unlock()
 
-	if b, ok := rc.data[digest]; ok && rc.has[cacheKeyDigest(repo, digest)] {
-		return ocimem.NewBytesReader(b, ociregistry.Descriptor{
-			MediaType: rc.types[digest],
-			Digest:    digest,
-			Size:      int64(len(b)),
-		}), nil
+	if desc, ok := rc.data[digest]; ok && desc.Data != nil && rc.has[cacheKeyDigest(repo, digest)] {
+		return ocimem.NewBytesReader(desc.Data, desc), nil
 	}
 
 	r, err := f(ctx, repo, digest)
 	if err != nil {
 		return nil, err
 	}
-	//defer r.Close()
+	// defer r.Close() happens later when we know we aren't making Close the caller's responsibility
 
 	desc := r.Descriptor()
+	digest = desc.Digest // if this isn't a no-op, we've got a naughty registry
 
-	rc.has[cacheKeyDigest(repo, desc.Digest)] = true
-	rc.types[desc.Digest] = desc.MediaType
+	rc.has[cacheKeyDigest(repo, digest)] = true
+
+	if desc.Size > manifestSizeLimit {
+		rc.data[digest] = desc
+		return r, nil
+	}
+	defer r.Close()
 
-	b, err := io.ReadAll(r)
+	desc.Data, err = io.ReadAll(r)
 	if err != nil {
-		r.Close()
 		return nil, err
 	}
 	if err := r.Close(); err != nil {
 		return nil, err
 	}
 
-	if len(b) <= manifestSizeLimit {
-		rc.data[desc.Digest] = b
-	} else {
-		delete(rc.data, desc.Digest)
-	}
+	rc.data[digest] = desc
 
-	return ocimem.NewBytesReader(b, desc), nil
+	return ocimem.NewBytesReader(desc.Data, desc), nil
 }
 
 func (rc *registryCache) GetBlob(ctx context.Context, repo string, digest ociregistry.Digest) (ociregistry.BlobReader, error) {
@@ -104,43 +99,39 @@ func (rc *registryCache) GetTag(ctx context.Context, repo string, tag string) (o
 	tagKey := cacheKeyTag(repo, tag)
 
 	if digest, ok := rc.tags[tagKey]; ok {
-		if b, ok := rc.data[digest]; ok {
-			return ocimem.NewBytesReader(b, ociregistry.Descriptor{
-				MediaType: rc.types[digest],
-				Digest:    digest,
-				Size:      int64(len(b)),
-			}), nil
+		if desc, ok := rc.data[digest]; ok && desc.Data != nil {
+			return ocimem.NewBytesReader(desc.Data, desc), nil
 		}
 	}
 
 	r, err := rc.registry.GetTag(ctx, repo, tag)
 	if err != nil {
 		return nil, err
 	}
-	//defer r.Close()
+	// defer r.Close() happens later when we know we aren't making Close the caller's responsibility
 
 	desc := r.Descriptor()
 
 	rc.has[cacheKeyDigest(repo, desc.Digest)] = true
 	rc.tags[tagKey] = desc.Digest
-	rc.types[desc.Digest] = desc.MediaType
 
-	b, err := io.ReadAll(r)
+	if desc.Size > manifestSizeLimit {
+		rc.data[desc.Digest] = desc
+		return r, nil
+	}
+	defer r.Close()
+
+	desc.Data, err = io.ReadAll(r)
 	if err != nil {
-		r.Close()
 		return nil, err
 	}
 	if err := r.Close(); err != nil {
 		return nil, err
 	}
 
-	if len(b) <= manifestSizeLimit {
-		rc.data[desc.Digest] = b
-	} else {
-		delete(rc.data, desc.Digest)
-	}
+	rc.data[desc.Digest] = desc
 
-	return ocimem.NewBytesReader(b, desc), nil
+	return ocimem.NewBytesReader(desc.Data, desc), nil
 }
 
 // TODO more methods (currently only implements what's actually necessary for SynthesizeIndex)

registry/read-helpers.go

@@ -20,6 +20,8 @@ func readJSONHelper(r ociregistry.BlobReader, v interface{}) error {
 		return err
 	}
 
+	// TODO if desc.Data != nil and len() == desc.Size, we should probably check/use that? 👀
+
 	// make sure we can't possibly read (much) more than we're supposed to
 	limited := &io.LimitedReader{
 		R: r,

registry/synthesize-index.go

@@ -9,6 +9,7 @@ import (
 	"github.com/docker-library/bashbrew/architecture"
 
 	"cuelabs.dev/go/oci/ociregistry"
+	"cuelabs.dev/go/oci/ociregistry/ocimem"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 )
 
@@ -127,6 +128,11 @@ func SynthesizeIndex(ctx context.Context, ref Reference) (*ocispec.Index, error)
 			}
 		}
 
+		// TODO if m.Size > 2048 {
+		// make sure we don't return any (big) data fields, now that we know we don't need them for sure (they might exist in the index we queried, but they're also used as an implementation detail in our registry cache code to store the original upstream data)
+		m.Data = nil
+		// }
+
 		index.Manifests[i] = m
 		seen[string(m.Digest)] = &index.Manifests[i]
 		i++
@@ -158,9 +164,13 @@ func normalizeManifestPlatform(ctx context.Context, m *ocispec.Descriptor, r oci
 		case ocispec.MediaTypeImageManifest, mediaTypeDockerImageManifest:
 			var err error
 			if r == nil {
-				r, err = client.GetManifest(ctx, ref.Repository, m.Digest)
-				if err != nil {
-					return err
+				if m.Data != nil && int64(len(m.Data)) == m.Size {
+					r = ocimem.NewBytesReader(m.Data, *m)
+				} else {
+					r, err = client.GetManifest(ctx, ref.Repository, m.Digest)
+					if err != nil {
+						return err
+					}
 				}
 				defer r.Close()
 			}
@@ -172,9 +182,14 @@ func normalizeManifestPlatform(ctx context.Context, m *ocispec.Descriptor, r oci
 
 			switch manifest.Config.MediaType {
 			case ocispec.MediaTypeImageConfig, mediaTypeDockerImageConfig:
-				r, err := client.GetBlob(ctx, ref.Repository, manifest.Config.Digest)
-				if err != nil {
-					return err
+				var r ociregistry.BlobReader
+				if manifest.Config.Data != nil && int64(len(manifest.Config.Data)) == manifest.Config.Size {
+					r = ocimem.NewBytesReader(manifest.Config.Data, manifest.Config)
+				} else {
+					r, err = client.GetBlob(ctx, ref.Repository, manifest.Config.Digest)
+					if err != nil {
+						return err
+					}
 				}
 				defer r.Close()