Get media type key from context

Adds a method for setting a custom media type key prefix used by the
fetch handler.
This allows both overwriting a built-in prefix (for reasons?) as well as
supplying a custom media type.

I added this because I was getting an error on `FetchHandler` when
pulling docker plugin images which have their own media type.

Signed-off-by: Brian Goff <cpuguy83@gmail.com>

remotes/handlers.go

@@ -33,13 +33,35 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
+type refKeyPrefix struct{}
+
+// WithMediaTypeKeyPrefix adds a custom key prefix for a media type which is used when storing
+// data in the content store from the FetchHandler.
+//
+// Used in `MakeRefKey` to determine what the key prefix should be.
+func WithMediaTypeKeyPrefix(ctx context.Context, mediaType, prefix string) context.Context {
+	var values map[string]string
+	if v := ctx.Value(refKeyPrefix{}); v != nil {
+		values = v.(map[string]string)
+	} else {
+		values = make(map[string]string)
+	}
+
+	values[mediaType] = prefix
+	return context.WithValue(ctx, refKeyPrefix{}, values)
+}
+
 // MakeRefKey returns a unique reference for the descriptor. This reference can be
 // used to lookup ongoing processes related to the descriptor. This function
 // may look to the context to namespace the reference appropriately.
 func MakeRefKey(ctx context.Context, desc ocispec.Descriptor) string {
-	// TODO(stevvooe): Need better remote key selection here. Should be a
-	// product of the context, which may include information about the ongoing
-	// fetch process.
+	if v := ctx.Value(refKeyPrefix{}); v != nil {
+		values := v.(map[string]string)
+		if prefix := values[desc.MediaType]; prefix != "" {
+			return prefix + "-" + desc.Digest.String()
+		}
+	}
+
 	switch desc.MediaType {
 	case images.MediaTypeDockerSchema2Manifest, ocispec.MediaTypeImageManifest:
 		return "manifest-" + desc.Digest.String()

remotes/handlers_test.go

@@ -0,0 +1,73 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package remotes
+
+import (
+	"context"
+	"testing"
+
+	"github.com/containerd/containerd/images"
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+)
+
+func TestContextCustomKeyPrefix(t *testing.T) {
+	ctx := context.Background()
+	cmt := "testing/custom.media.type"
+	ctx = WithMediaTypeKeyPrefix(ctx, images.MediaTypeDockerSchema2Layer, "bananas")
+	ctx = WithMediaTypeKeyPrefix(ctx, cmt, "apples")
+
+	// makes sure tha even though we've supplied some custom handling, the built-in still works
+	t.Run("normal supported case", func(t *testing.T) {
+		desc := ocispec.Descriptor{MediaType: ocispec.MediaTypeImageLayer}
+		expected := "layer-"
+
+		actual := MakeRefKey(ctx, desc)
+		if actual != expected {
+			t.Fatalf("unexpected ref key, expected %s, got: %s", expected, actual)
+		}
+	})
+
+	t.Run("unknown media type", func(t *testing.T) {
+		desc := ocispec.Descriptor{MediaType: "we.dont.know.what.this.is"}
+		expected := "unknown-"
+
+		actual := MakeRefKey(ctx, desc)
+		if actual != expected {
+			t.Fatalf("unexpected ref key, expected %s, got: %s", expected, actual)
+		}
+	})
+
+	t.Run("overwrite supported media type", func(t *testing.T) {
+		desc := ocispec.Descriptor{MediaType: images.MediaTypeDockerSchema2Layer}
+		expected := "bananas-"
+
+		actual := MakeRefKey(ctx, desc)
+		if actual != expected {
+			t.Fatalf("unexpected ref key, expected %s, got: %s", expected, actual)
+		}
+	})
+
+	t.Run("custom media type", func(t *testing.T) {
+		desc := ocispec.Descriptor{MediaType: cmt}
+		expected := "apples-"
+
+		actual := MakeRefKey(ctx, desc)
+		if actual != expected {
+			t.Fatalf("unexpected ref key, expected %s, got: %s", expected, actual)
+		}
+	})
+}