rmi: correct removal of image by id (#43)

Also fix for tagging by image id.

Signed-off-by: Jacob Blain Christen <jacob@rancher.com>

pkg/client/image/remove.go

@@ -3,43 +3,22 @@ package image
 import (
 	"context"
 
-	"github.com/docker/distribution/reference"
 	imagesv1 "github.com/rancher/kim/pkg/apis/services/images/v1alpha1"
 	"github.com/rancher/kim/pkg/client"
 	"github.com/sirupsen/logrus"
-	criv1 "k8s.io/cri-api/pkg/apis/runtime/v1alpha2"
 )
 
 type Remove struct {
 }
 
 func (s *Remove) Do(ctx context.Context, k8s *client.Interface, image string) error {
-	if named, err := reference.ParseNormalizedNamed(image); err == nil {
-		image = reference.TagNameOnly(named).String()
-	}
 	return client.Images(ctx, k8s, func(ctx context.Context, imagesClient imagesv1.ImagesClient) error {
-		req := &imagesv1.ImageRemoveRequest{
-			Image: &criv1.ImageSpec{
-				Image: image,
-			},
-		}
-		if ref, err := reference.ParseAnyReference(image); err == nil {
-			statusRequest := imagesv1.ImageStatusRequest{
-				Image: &criv1.ImageSpec{
-					Image: ref.String(),
-				},
-			}
-			statusResponse, err := imagesClient.Status(ctx, &statusRequest)
-			logrus.Debugf("%#v", statusResponse.Image)
-			if err == nil && statusResponse.Image != nil {
-				req.Image = statusRequest.Image
-			}
-		}
-		res, err := imagesClient.Remove(ctx, req)
+		ref, err := refSpec(ctx, imagesClient, image)
 		if err != nil {
 			return err
 		}
+		res, err := imagesClient.Remove(ctx, &imagesv1.ImageRemoveRequest{Image: ref})
 		logrus.Debugf("%#v", res)
-		return nil
+		return err
 	})
 }

pkg/client/image/spec.go

@@ -0,0 +1,34 @@
+package image
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/docker/distribution/reference"
+	"github.com/sirupsen/logrus"
+
+	imagesv1 "github.com/rancher/kim/pkg/apis/services/images/v1alpha1"
+	criv1 "k8s.io/cri-api/pkg/apis/runtime/v1alpha2"
+)
+
+// refSpec attempts to normalize an arbitrary image reference by requesting the status from the cri with the
+// passed value. if it matches or is a image-id prefix then the image-id will be returned. otherwise an attempt is
+// made to normalize via reference.ParseNormalizedNamed passed through reference.TagNameOnly (handling tag-less refs)
+func refSpec(ctx context.Context, imagesClient imagesv1.ImagesClient, image string) (*criv1.ImageSpec, error) {
+	spec := &criv1.ImageSpec{
+		Image: image,
+	}
+	status, statusErr := imagesClient.Status(ctx, &imagesv1.ImageStatusRequest{Image: spec})
+	logrus.Debugf("refSpec image=%q: %#v", image, status.Image)
+	if statusErr == nil && status.Image != nil && strings.HasPrefix(status.Image.Id, fmt.Sprintf("sha256:%s", image)) {
+		spec.Image = status.Image.Id
+		return spec, nil
+	}
+	named, parseErr := reference.ParseNormalizedNamed(image)
+	if parseErr == nil {
+		spec.Image = reference.TagNameOnly(named).String()
+		return spec, nil
+	}
+	return nil, statusErr
+}

pkg/client/image/tag.go

@@ -7,30 +7,28 @@ import (
 	imagesv1 "github.com/rancher/kim/pkg/apis/services/images/v1alpha1"
 	"github.com/rancher/kim/pkg/client"
 	"github.com/sirupsen/logrus"
-	criv1 "k8s.io/cri-api/pkg/apis/runtime/v1alpha2"
 )
 
 type Tag struct {
 }
 
 func (s *Tag) Do(ctx context.Context, k8s *client.Interface, image string, tags []string) error {
-	if named, err := reference.ParseNormalizedNamed(image); err == nil {
-		image = reference.TagNameOnly(named).String()
-	}
 	normalizedTags := make([]string, len(tags))
 	for i, tag := range tags {
 		named, err := reference.ParseNormalizedNamed(tag)
 		if err != nil {
 			return err
 		}
-		normalizedTags[i] = named.String()
+		normalizedTags[i] = reference.TagNameOnly(named).String()
 	}
 	return client.Images(ctx, k8s, func(ctx context.Context, imagesClient imagesv1.ImagesClient) error {
+		ref, err := refSpec(ctx, imagesClient, image)
+		if err != nil {
+			return err
+		}
 		req := &imagesv1.ImageTagRequest{
-			Image: &criv1.ImageSpec{
-				Image: image,
-			},
-			Tags: normalizedTags,
+			Image: ref,
+			Tags:  normalizedTags,
 		}
 		res, err := imagesClient.Tag(ctx, req)
 		if err != nil {

pkg/server/images/remove.go

@@ -2,30 +2,46 @@ package images
 
 import (
 	"context"
+	"fmt"
+	"strings"
 
-	"github.com/containerd/containerd/errdefs"
 	"github.com/containerd/containerd/images"
 	"github.com/containerd/containerd/namespaces"
 	imagesv1 "github.com/rancher/kim/pkg/apis/services/images/v1alpha1"
 	"github.com/sirupsen/logrus"
-	criv1 "k8s.io/cri-api/pkg/apis/runtime/v1alpha2"
 )
 
 // Remove image server-side impl
 func (s *Server) Remove(ctx context.Context, req *imagesv1.ImageRemoveRequest) (*imagesv1.ImageRemoveResponse, error) {
-	logrus.Debugf("image-remove: %#v", req.Image)
+	logrus.Debugf("image-remove: req=%s", req)
 	ctx, done, err := s.Containerd.WithLease(namespaces.WithNamespace(ctx, "k8s.io"))
 	if err != nil {
 		return nil, err
 	}
 	defer done(ctx)
-	err = s.Containerd.ImageService().Delete(ctx, req.Image.Image, images.SynchronousDelete())
-	if errdefs.IsNotFound(err) {
-		// at this point we assume it is an image id and fallback to cri behavior, aka remove every tag/digest
-		_, err = s.ImageService().RemoveImage(ctx, &criv1.RemoveImageRequest{Image: req.Image})
+	img, err := s.Containerd.ImageService().Get(ctx, req.Image.Image)
+	if err != nil {
+		return nil, err
 	}
+	refs := []string{img.Name}
+	tags, err := s.Containerd.ImageService().List(ctx, fmt.Sprintf("target.digest==%s,name!=%s", img.Target.Digest, img.Name))
 	if err != nil {
 		return nil, err
 	}
+	switch {
+	case len(tags) == 1: // single tag
+		refs = append(refs, tags[0].Name)
+	case strings.HasPrefix(img.Name, "sha256:"): // image id
+		for _, tag := range tags {
+			refs = append(refs, tag.Name)
+		}
+	}
+	for _, ref := range refs {
+		logrus.Debugf("image-remove: ref=%s, img=%#v", ref, req.Image)
+		err = s.Containerd.ImageService().Delete(ctx, ref, images.SynchronousDelete())
+		if err != nil {
+			return nil, err
+		}
+	}
 	return &imagesv1.ImageRemoveResponse{}, nil
 }

pkg/server/images/tag.go

@@ -39,7 +39,7 @@ func (s *Server) Tag(ctx context.Context, req *imagesv1.ImageTagRequest) (*image
 				return nil, err
 			}
 		}
-		logrus.Debugf("%#v", img)
+		logrus.Debugf("image-tag: %#v", img)
 	}
 	res, err := s.ImageService().ImageStatus(ctx, &criv1.ImageStatusRequest{Image: req.Image})
 	if err != nil {