pull the image with the <image>@<digest> format

aws · Jul 26, 2024 · 454e9b0 · 454e9b0
1 parent 522561b
commit 454e9b0
Show file tree

Hide file tree

Showing 3 changed files with 231 additions and 2 deletions.
diff --git a/agent/engine/docker_task_engine.go b/agent/engine/docker_task_engine.go
@@ -1581,6 +1581,29 @@ func (engine *DockerTaskEngine) pullAndUpdateContainerReference(task *apitask.Ta
 			field.ImageDigest: imageDigest,
 			field.ImageRef:    imageRef,
 		})
+	} else if referenceutil.TagExists(imageRef) {
+		digest := referenceutil.GetDigestFromImageRef(imageRef)
+		if digest != "" {
+			// if image reference contains both tag and digest, i.e. busybox:latest@sha256:abcd, remove the tag
+			canonicalRefWithoutTag, err := referenceutil.GetCanonicalRefWithoutTag(imageRef, digest)
+			if err != nil {
+				logger.Error("Failed to prepare a canonical reference without a tag. Cannot pull image.", logger.Fields{
+					field.TaskID:      task.GetID(),
+					field.Container:   container.Name,
+					field.Image:       imageRef,
+					field.ImageDigest: digest.String(),
+					field.Error:       err,
+				})
+				return dockerapi.DockerContainerMetadata{
+					Error: dockerapi.CannotPullContainerError{
+						FromError: fmt.Errorf(
+							"failed to prepare a canonical reference without a tag for image '%s' and digest '%s': %w",
+							imageRef, digest.String(), err),
+					},
+				}
+			}
+			imageRef = canonicalRefWithoutTag.String()
+		}
 	}
 
 	metadata := engine.client.PullImage(engine.ctx, imageRef, container.RegistryAuthentication, engine.cfg.ImagePullTimeout)
@@ -1591,7 +1614,7 @@ func (engine *DockerTaskEngine) pullAndUpdateContainerReference(task *apitask.Ta
 	}
 	pullSucceeded := metadata.Error == nil
 
-	if pullSucceeded && imageRef != container.Image {
+	if pullSucceeded && imageRef != container.Image && !referenceutil.DigestExists(container.Image) {
 		// Resolved image manifest digest was used to pull the image.
 		// Tag the pulled image so that it can be found using the image reference in the task.
 		ctx, cancel := context.WithTimeout(engine.ctx, tagImageTimeout)

diff --git a/agent/utils/reference/reference.go b/agent/utils/reference/reference.go
@@ -93,7 +93,6 @@ func GetCanonicalRef(imageRef string, manifestDigest string) (reference.Canonica
 		return nil, fmt.Errorf("image reference '%s' is not a named reference, parsed as: %v",
 			imageRef, parsedImageRef)
 	}
-
 	canonicalRef, err := reference.WithDigest(namedImageRef, parsedDigest)
 	if err != nil {
 		return nil, fmt.Errorf(
@@ -103,3 +102,54 @@ func GetCanonicalRef(imageRef string, manifestDigest string) (reference.Canonica
 
 	return canonicalRef, nil
 }
+
+// check if image reference contains tag
+func TagExists(imageRef string) bool {
+
+	parsedImageRef, err := reference.Parse(imageRef)
+	if err != nil {
+		return false
+	}
+	_, ok := parsedImageRef.(reference.Tagged)
+	return ok
+}
+
+// check if image reference contains digest
+func DigestExists(imageRef string) bool {
+
+	parsedImageRef, err := reference.Parse(imageRef)
+	if err != nil {
+		return false
+	}
+	_, ok := parsedImageRef.(reference.Digested)
+	return ok
+}
+
+// Given an image reference and a manifest digest, returns a canonical reference without tag
+// for the image.
+// If the image reference has a digest then the canonical reference will still use the provided
+// manifest digest overwriting the existing digest in the image reference.
+func GetCanonicalRefWithoutTag(imageRef string, manifestDigest digest.Digest) (reference.Canonical, error) {
+
+	parsedImageRef, err := reference.Parse(imageRef)
+	if err != nil {
+		return nil, fmt.Errorf(
+			"failed to parse image reference '%s': %w", imageRef, err)
+	}
+	namedImageRef, ok := parsedImageRef.(reference.Named)
+	if !ok {
+		return nil, fmt.Errorf("image reference '%s' is not a named reference, parsed as: %v",
+			imageRef, parsedImageRef)
+	}
+
+	trimNamedImageRef := reference.TrimNamed(namedImageRef)
+
+	canonicalRef, err := reference.WithDigest(trimNamedImageRef, manifestDigest)
+	if err != nil {
+		return nil, fmt.Errorf(
+			"failed to produce a canonical reference using named reference '%v' and digest '%v': %w",
+			namedImageRef, manifestDigest, err)
+	}
+
+	return canonicalRef, nil
+}
diff --git a/agent/utils/reference/reference_test.go b/agent/utils/reference/reference_test.go
@@ -233,6 +233,12 @@ func TestGetCanonicalRef(t *testing.T) {
 			manifestDigest: "sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
 			expected:       "public.ecr.aws/library/alpine:latest@sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
 		},
+		{
+			name:           "has tag and digest ecr",
+			imageRef:       "public.ecr.aws/library/alpine:latest@sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+			manifestDigest: "sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+			expected:       "public.ecr.aws/library/alpine:latest@sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+		},
 	}
 	for _, tc := range tcs {
 		t.Run(tc.name, func(t *testing.T) {
@@ -246,3 +252,153 @@ func TestGetCanonicalRef(t *testing.T) {
 		})
 	}
 }
+
+func TestGetCanonicalRefWithoutTag(t *testing.T) {
+	tcs := []struct {
+		name           string
+		imageRef       string
+		manifestDigest digest.Digest
+		expected       string
+		expectedError  string
+	}{
+		{
+			name:           "invalid digest",
+			imageRef:       "alpine",
+			manifestDigest: "invalid digest",
+			expectedError:  "failed to produce a canonical reference using named reference 'alpine' and digest 'invalid digest': invalid digest format",
+		},
+		{
+			name:           "invalid image reference format",
+			imageRef:       "invalid reference",
+			manifestDigest: "sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+			expectedError:  "failed to parse image reference 'invalid reference': invalid reference format",
+		},
+		{
+			name:           "no tag",
+			imageRef:       "alpine",
+			manifestDigest: "sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+			expected:       "alpine@sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+		},
+		{
+			name:           "has tag",
+			imageRef:       "alpine:latest",
+			manifestDigest: "sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+			expected:       "alpine@sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+		},
+		{
+			name:           "image reference's digest is overwritten",
+			imageRef:       "alpine@sha256:c5b1261d6d3e43071626931fc004f70149baeba2c8ec672bd4f27761f8e1ad6b",
+			manifestDigest: "sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+			expected:       "alpine@sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+		},
+		{
+			name:           "no tag ecr",
+			imageRef:       "public.ecr.aws/library/alpine",
+			manifestDigest: "sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+			expected:       "public.ecr.aws/library/alpine@sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+		},
+		{
+			name:           "has tag ecr",
+			imageRef:       "public.ecr.aws/library/alpine:latest",
+			manifestDigest: "sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+			expected:       "public.ecr.aws/library/alpine@sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+		},
+		{
+			name:           "has tag and digest ecr",
+			imageRef:       "public.ecr.aws/library/alpine:latest@sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+			manifestDigest: "sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+			expected:       "public.ecr.aws/library/alpine@sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+		},
+	}
+	for _, tc := range tcs {
+		t.Run(tc.name, func(t *testing.T) {
+			canonicalRef, err := GetCanonicalRefWithoutTag(tc.imageRef, tc.manifestDigest)
+			if tc.expectedError == "" {
+				require.NoError(t, err)
+				assert.Equal(t, tc.expected, canonicalRef.String())
+			} else {
+				assert.EqualError(t, err, tc.expectedError)
+			}
+		})
+	}
+}
+
+func TestTagExists(t *testing.T) {
+	tcs := []struct {
+		name     string
+		imageRef string
+		expected bool
+	}{
+		{
+			name:     "invalid imageRef",
+			imageRef: "invalid imageRef",
+			expected: false,
+		},
+		{
+			name:     "no tag",
+			imageRef: "public.ecr.aws/library/alpine",
+			expected: false,
+		},
+		{
+			name:     "has tag",
+			imageRef: "public.ecr.aws/library/alpine:latest",
+			expected: true,
+		},
+		{
+			name:     "has tag and digest",
+			imageRef: "public.ecr.aws/library/alpine:latest@sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+			expected: true,
+		},
+		{
+			name:     "has digest no tag",
+			imageRef: "public.ecr.aws/library/alpine@sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+			expected: false,
+		},
+	}
+	for _, tc := range tcs {
+		t.Run(tc.name, func(t *testing.T) {
+			ok := TagExists(tc.imageRef)
+			assert.Equal(t, ok, tc.expected)
+		})
+	}
+}
+
+func TestDigestExists(t *testing.T) {
+	tcs := []struct {
+		name     string
+		imageRef string
+		expected bool
+	}{
+		{
+			name:     "invalid imageRef",
+			imageRef: "invalid imageRef",
+			expected: false,
+		},
+		{
+			name:     "no tag no digest",
+			imageRef: "public.ecr.aws/library/alpine",
+			expected: false,
+		},
+		{
+			name:     "has tag",
+			imageRef: "public.ecr.aws/library/alpine:latest",
+			expected: false,
+		},
+		{
+			name:     "has tag and digest",
+			imageRef: "public.ecr.aws/library/alpine:latest@sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+			expected: true,
+		},
+		{
+			name:     "has digest no tag",
+			imageRef: "public.ecr.aws/library/alpine@sha256:c3839dd800b9eb7603340509769c43e146a74c63dca3045a8e7dc8ee07e53966",
+			expected: true,
+		},
+	}
+	for _, tc := range tcs {
+		t.Run(tc.name, func(t *testing.T) {
+			ok := DigestExists(tc.imageRef)
+			assert.Equal(t, ok, tc.expected)
+		})
+	}
+}