Set content labels based on content type

Give control of the content labeling process for children to
the client. This allows the client to control the names
associated with the labels and filter out labels.

Signed-off-by: Derek McGowan <derek@mcg.dev>

images/handlers.go

@@ -170,6 +170,19 @@ func ChildrenHandler(provider content.Provider) HandlerFunc {
 // the children returned by the handler and passes through the children.
 // Must follow a handler that returns the children to be labeled.
 func SetChildrenLabels(manager content.Manager, f HandlerFunc) HandlerFunc {
+	return SetChildrenMappedLabels(manager, f, nil)
+}
+
+// SetChildrenMappedLabels is a handler wrapper which sets labels for the content on
+// the children returned by the handler and passes through the children.
+// Must follow a handler that returns the children to be labeled.
+// The label map allows the caller to control the labels per child descriptor.
+// For returned labels, the index of the child will be appended to the end
+// except for the first index when the returned label does not end with '.'.
+func SetChildrenMappedLabels(manager content.Manager, f HandlerFunc, labelMap func(ocispec.Descriptor) []string) HandlerFunc {
+	if labelMap == nil {
+		labelMap = ChildGCLabels
+	}
 	return func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
 		children, err := f(ctx, desc)
 		if err != nil {
@@ -177,14 +190,26 @@ func SetChildrenLabels(manager content.Manager, f HandlerFunc) HandlerFunc {
 		}
 
 		if len(children) > 0 {
-			info := content.Info{
-				Digest: desc.Digest,
-				Labels: map[string]string{},
-			}
-			fields := []string{}
-			for i, ch := range children {
-				info.Labels[fmt.Sprintf("containerd.io/gc.ref.content.%d", i)] = ch.Digest.String()
-				fields = append(fields, fmt.Sprintf("labels.containerd.io/gc.ref.content.%d", i))
+			var (
+				info = content.Info{
+					Digest: desc.Digest,
+					Labels: map[string]string{},
+				}
+				fields = []string{}
+				keys   = map[string]uint{}
+			)
+			for _, ch := range children {
+				labelKeys := labelMap(ch)
+				for _, key := range labelKeys {
+					idx := keys[key]
+					keys[key] = idx + 1
+					if idx > 0 || key[len(key)-1] == '.' {
+						key = fmt.Sprintf("%s%d", key, idx)
+					}
+
+					info.Labels[key] = ch.Digest.String()
+					fields = append(fields, "labels."+key)
+				}
 			}
 
 			_, err := manager.Update(ctx, info, fields...)

images/mediatypes.go

@@ -124,3 +124,31 @@ func IsKnownConfig(mt string) bool {
 	}
 	return false
 }
+
+// ChildGCLabels returns the label for a given descriptor to reference it
+func ChildGCLabels(desc ocispec.Descriptor) []string {
+	mt := desc.MediaType
+	if IsKnownConfig(mt) {
+		return []string{"containerd.io/gc.ref.content.config"}
+	}
+
+	switch mt {
+	case MediaTypeDockerSchema2Manifest, ocispec.MediaTypeImageManifest:
+		return []string{"containerd.io/gc.ref.content.m."}
+	}
+
+	if IsLayerType(mt) {
+		return []string{"containerd.io/gc.ref.content.l."}
+	}
+
+	return []string{"containerd.io/gc.ref.content."}
+}
+
+// ChildGCLabelsFilterLayers returns the labels for a given descriptor to
+// reference it, skipping layer media types
+func ChildGCLabelsFilterLayers(desc ocispec.Descriptor) []string {
+	if IsLayerType(desc.MediaType) {
+		return nil
+	}
+	return ChildGCLabels(desc)
+}