Refactor registry code to be more correct

.dockerignore

@@ -6,4 +6,5 @@
 !go.sum
 !manifest/
 !pkg/
+!registry/
 !scripts/

cmd/bashbrew/cmd-push.go

@@ -39,6 +39,7 @@ func cmdPush(c *cli.Context) error {
 			}
 
 			// we can't use "r.Tags()" here because it will include SharedTags, which we never want to push directly (see "cmd-put-shared.go")
+		TagsLoop:
 			for i, tag := range entry.Tags {
 				if uniq && i > 0 {
 					break
@@ -47,10 +48,12 @@ func cmdPush(c *cli.Context) error {
 
 				if !force {
 					localImageId, _ := dockerInspect("{{.Id}}", tag)
-					registryImageId := fetchRegistryImageId(tag)
-					if registryImageId != "" && localImageId == registryImageId {
-						fmt.Fprintf(os.Stderr, "skipping %s (remote image matches local)\n", tag)
-						continue
+					registryImageIds := fetchRegistryImageIds(tag)
+					for _, registryImageId := range registryImageIds {
+						if localImageId == registryImageId {
+							fmt.Fprintf(os.Stderr, "skipping %s (remote image matches local)\n", tag)
+							continue TagsLoop
+						}
 					}
 				}
 				fmt.Printf("Pushing %s\n", tag)

cmd/bashbrew/registry.go

@@ -2,63 +2,48 @@ package main
 
 import (
 	"context"
-	"encoding/json"
-	"net/url"
-	"os"
 
-	"github.com/containerd/containerd/images"
-	"github.com/containerd/containerd/reference/docker"
-	"github.com/containerd/containerd/remotes"
-	dockerremote "github.com/containerd/containerd/remotes/docker"
-	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+	"github.com/docker-library/bashbrew/registry"
 )
 
-var registryImageIdCache = map[string]string{}
+var registryImageIdsCache = map[string][]string{}
 
 // assumes the provided image name is NOT a manifest list (used for testing whether we need to "bashbrew push" or whether the remote image is already up-to-date)
 // this does NOT handle authentication, and will return the empty string for repositories which require it (causing "bashbrew push" to simply shell out to "docker push" which will handle authentication appropriately)
-func fetchRegistryImageId(image string) string {
+func fetchRegistryImageIds(image string) []string {
 	ctx := context.Background()
 
-	ref, resolver, err := fetchRegistryResolveHelper(image)
+	img, err := registry.Resolve(ctx, image)
 	if err != nil {
-		return ""
-	}
-
-	name, desc, err := resolver.Resolve(ctx, ref)
-	if err != nil {
-		return ""
-	}
-
-	if desc.MediaType != images.MediaTypeDockerSchema2Manifest && desc.MediaType != ocispec.MediaTypeImageManifest {
-		return ""
+		return nil
 	}
 
-	digest := desc.Digest.String()
-	if id, ok := registryImageIdCache[digest]; ok {
-		return id
+	digest := img.Desc.Digest.String()
+	if ids, ok := registryImageIdsCache[digest]; ok {
+		return ids
 	}
 
-	fetcher, err := resolver.Fetcher(ctx, name)
+	manifests, err := img.Manifests(ctx)
 	if err != nil {
-		return ""
+		return nil
 	}
 
-	r, err := fetcher.Fetch(ctx, desc)
-	if err != nil {
-		return ""
+	ids := []string{}
+	if img.IsImageIndex() {
+		ids = append(ids, digest)
 	}
-	defer r.Close()
-
-	var manifest ocispec.Manifest
-	if err := json.NewDecoder(r).Decode(&manifest); err != nil {
-		return ""
+	for _, manifestDesc := range manifests {
+		ids = append(ids, manifestDesc.Digest.String())
+		manifest, err := img.Manifest(ctx, manifestDesc)
+		if err != nil {
+			continue
+		}
+		ids = append(ids, manifest.Config.Digest.String())
 	}
-	id := manifest.Config.Digest.String()
-	if id != "" {
-		registryImageIdCache[digest] = id
+	if len(ids) > 0 {
+		registryImageIdsCache[digest] = ids
 	}
-	return id
+	return ids
 }
 
 var registryManifestListCache = map[string][]string{}
@@ -67,46 +52,22 @@ var registryManifestListCache = map[string][]string{}
 func fetchRegistryManiestListDigests(image string) []string {
 	ctx := context.Background()
 
-	ref, resolver, err := fetchRegistryResolveHelper(image)
-	if err != nil {
-		return nil
-	}
-
-	name, desc, err := resolver.Resolve(ctx, ref)
+	img, err := registry.Resolve(ctx, image)
 	if err != nil {
 		return nil
 	}
 
-	digest := desc.Digest.String()
-	if desc.MediaType == images.MediaTypeDockerSchema2Manifest || desc.MediaType == ocispec.MediaTypeImageManifest {
-		return []string{digest}
-	}
-
-	if desc.MediaType != images.MediaTypeDockerSchema2ManifestList && desc.MediaType != ocispec.MediaTypeImageIndex {
-		return nil
-	}
-
+	digest := img.Desc.Digest.String()
 	if digests, ok := registryManifestListCache[digest]; ok {
 		return digests
 	}
 
-	fetcher, err := resolver.Fetcher(ctx, name)
-	if err != nil {
-		return nil
-	}
-
-	r, err := fetcher.Fetch(ctx, desc)
+	manifests, err := img.Manifests(ctx)
 	if err != nil {
 		return nil
 	}
-	defer r.Close()
-
-	var manifestList ocispec.Index
-	if err := json.NewDecoder(r).Decode(&manifestList); err != nil {
-		return nil
-	}
 	digests := []string{}
-	for _, manifest := range manifestList.Manifests {
+	for _, manifest := range manifests {
 		if manifest.Digest != "" {
 			digests = append(digests, manifest.Digest.String())
 		}
@@ -116,30 +77,3 @@ func fetchRegistryManiestListDigests(image string) []string {
 	}
 	return digests
 }
-
-func fetchRegistryResolveHelper(image string) (string, remotes.Resolver, error) {
-	ref, err := docker.ParseAnyReference(image)
-	if err != nil {
-		return "", nil, err
-	}
-	if namedRef, ok := ref.(docker.Named); ok {
-		// add ":latest" if necessary
-		namedRef = docker.TagNameOnly(namedRef)
-		ref = namedRef
-	}
-	return ref.String(), dockerremote.NewResolver(dockerremote.ResolverOptions{
-		Host: func(host string) (string, error) {
-			if host == "docker.io" {
-				if publicProxy := os.Getenv("DOCKERHUB_PUBLIC_PROXY"); publicProxy != "" {
-					if publicProxyURL, err := url.Parse(publicProxy); err == nil {
-						// TODO Scheme (also not sure if "host:port" will be satisfactory to containerd here, but 🤷)
-						return publicProxyURL.Host, nil
-					} else {
-						return "", err
-					}
-				}
-			}
-			return host, nil
-		},
-	}), nil
-}

registry/registry.go

@@ -0,0 +1,178 @@
+package registry
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/url"
+	"os"
+	"unicode"
+
+	// thanks, go-digest...
+	_ "crypto/sha256"
+	_ "crypto/sha512"
+
+	"github.com/containerd/containerd/images"
+	"github.com/containerd/containerd/reference/docker"
+	"github.com/containerd/containerd/remotes"
+	dockerremote "github.com/containerd/containerd/remotes/docker"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+)
+
+type ResolvedObject struct {
+	Ref  string
+	Desc ocispec.Descriptor
+
+	resolver remotes.Resolver
+	fetcher  remotes.Fetcher
+}
+
+func (obj ResolvedObject) FetchJSON(ctx context.Context, v interface{}) error {
+	// prevent go-digest panics later
+	if err := obj.Desc.Digest.Validate(); err != nil {
+		return err
+	}
+
+	r, err := obj.fetcher.Fetch(ctx, obj.Desc)
+	if err != nil {
+		return err
+	}
+	defer r.Close()
+
+	// make sure we can't possibly read (much) more than we're supposed to
+	limited := &io.LimitedReader{
+		R: r,
+		N: obj.Desc.Size + 1, // +1 to allow us to detect if we read too much (see verification below)
+	}
+
+	// copy all read data into the digest verifier so we can validate afterwards
+	verifier := obj.Desc.Digest.Verifier()
+	tee := io.TeeReader(limited, verifier)
+
+	// decode directly! (mostly avoids double memory hit for big objects)
+	// (TODO protect against malicious objects somehow?)
+	if err := json.NewDecoder(tee).Decode(v); err != nil {
+		return err
+	}
+
+	// read anything leftover ...
+	bs, err := io.ReadAll(tee)
+	if err != nil {
+		return err
+	}
+	// ... and make sure it was just whitespace, if anything
+	for _, b := range bs {
+		if !unicode.IsSpace(rune(b)) {
+			return fmt.Errorf("unexpected non-whitespace at the end of %q: %+v\n", obj.Desc.Digest.String(), rune(b))
+		}
+	}
+
+	// after reading *everything*, we should have exactly one byte left in our LimitedReader (anything else is an error)
+	if limited.N < 1 {
+		return fmt.Errorf("size of %q is bigger than it should be (%d)", obj.Desc.Digest.String(), obj.Desc.Size)
+	} else if limited.N > 1 {
+		return fmt.Errorf("size of %q is %d bytes smaller than it should be (%d)", obj.Desc.Digest.String(), limited.N-1, obj.Desc.Size)
+	}
+
+	// and finally, let's verify our checksum
+	if !verifier.Verified() {
+		return fmt.Errorf("digest of %q not correct", obj.Desc.Digest.String())
+	}
+
+	return nil
+}
+
+func (obj ResolvedObject) Manifests(ctx context.Context) ([]ocispec.Descriptor, error) {
+	if obj.IsImageManifest() {
+		return []ocispec.Descriptor{obj.Desc}, nil
+	}
+
+	if !obj.IsImageIndex() {
+		return nil, fmt.Errorf("unknown media type: %q", obj.Desc.MediaType)
+	}
+
+	// (perhaps use a containerd content store??)
+	var index ocispec.Index
+	if err := obj.FetchJSON(ctx, &index); err != nil {
+		return nil, err
+	}
+	return index.Manifests, nil
+}
+
+func (obj ResolvedObject) Manifest(ctx context.Context, desc ocispec.Descriptor) (*ocispec.Manifest, error) {
+	obj.Desc = desc // swap our object to point at this manifest
+	if !obj.IsImageManifest() {
+		return nil, fmt.Errorf("unknown media type: %q", obj.Desc.MediaType)
+	}
+
+	// (perhaps use a containerd content store??)
+	var manifest ocispec.Manifest
+	if err := obj.FetchJSON(ctx, &manifest); err != nil {
+		return nil, err
+	}
+	return &manifest, nil
+}
+
+func (obj ResolvedObject) IsImageManifest() bool {
+	return obj.Desc.MediaType == ocispec.MediaTypeImageManifest || obj.Desc.MediaType == images.MediaTypeDockerSchema2Manifest
+}
+
+func (obj ResolvedObject) IsImageIndex() bool {
+	return obj.Desc.MediaType == ocispec.MediaTypeImageIndex || obj.Desc.MediaType == images.MediaTypeDockerSchema2ManifestList
+}
+
+func Resolve(ctx context.Context, image string) (*ResolvedObject, error) {
+	var (
+		obj = ResolvedObject{
+			Ref: image,
+		}
+		err error
+	)
+
+	obj.Ref, obj.resolver, err = resolverHelper(obj.Ref)
+	if err != nil {
+		return nil, err
+	}
+
+	obj.Ref, obj.Desc, err = obj.resolver.Resolve(ctx, obj.Ref)
+	if err != nil {
+		return nil, err
+	}
+
+	obj.fetcher, err = obj.resolver.Fetcher(ctx, obj.Ref)
+	if err != nil {
+		return nil, err
+	}
+
+	return &obj, nil
+}
+
+func resolverHelper(image string) (string, remotes.Resolver, error) {
+	ref, err := docker.ParseAnyReference(image)
+	if err != nil {
+		return "", nil, err
+	}
+	if namedRef, ok := ref.(docker.Named); ok {
+		// add ":latest" if necessary
+		namedRef = docker.TagNameOnly(namedRef)
+		ref = namedRef
+	}
+	return ref.String(), dockerremote.NewResolver(dockerremote.ResolverOptions{
+		// TODO port this to "Hosts:" (especially so we can return Scheme correctly) but requires reimplementing some of https://github.com/containerd/containerd/blob/v1.6.9/remotes/docker/resolver.go#L161-L184 😞
+		Host: func(host string) (string, error) {
+			if host == "docker.io" {
+				if publicProxy := os.Getenv("DOCKERHUB_PUBLIC_PROXY"); publicProxy != "" {
+					if publicProxyURL, err := url.Parse(publicProxy); err == nil {
+						// TODO Scheme (also not sure if "host:port" will be satisfactory to containerd here, but 🤷)
+						return publicProxyURL.Host, nil
+					} else {
+						return "", err
+					}
+				}
+				return "registry-1.docker.io", nil // https://github.com/containerd/containerd/blob/1c90a442489720eec95342e1789ee8a5e1b9536f/remotes/docker/registry.go#L193
+			}
+			return host, nil
+		},
+	}), nil
+}