diff --git a/remotes/handlers.go b/remotes/handlers.go index 4dee1f61c..0a243eda1 100644 --- a/remotes/handlers.go +++ b/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() diff --git a/remotes/handlers_test.go b/remotes/handlers_test.go new file mode 100644 index 000000000..354c660f9 --- /dev/null +++ b/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) + } + }) +}