From bc1e7aa603b96bdfd3b5fad762ccdee6b6f2f1ac Mon Sep 17 00:00:00 2001
From: Aaron Lehmann <aaron.lehmann@docker.com>
Date: Tue, 8 Nov 2016 14:47:33 -0800
Subject: [PATCH] digest: Preserve tag and digest in With* functions

When WithDigest is called on a reference that has a tag, it should
preserve the tag.

When WithTag is called on a reference that has digest, it should
preserve the digest.

Signed-off-by: Aaron Lehmann <aaron.lehmann@docker.com>
---
 reference/reference.go      | 14 ++++++++++++++
 reference/reference_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/reference/reference.go b/reference/reference.go
index dc3af1670..8f30f6b01 100644
--- a/reference/reference.go
+++ b/reference/reference.go
@@ -218,6 +218,13 @@ func WithTag(name Named, tag string) (NamedTagged, error) {
 	if !anchoredTagRegexp.MatchString(tag) {
 		return nil, ErrTagInvalidFormat
 	}
+	if canonical, ok := name.(Canonical); ok {
+		return reference{
+			name:   name.Name(),
+			tag:    tag,
+			digest: canonical.Digest(),
+		}, nil
+	}
 	return taggedReference{
 		name: name.Name(),
 		tag:  tag,
@@ -230,6 +237,13 @@ func WithDigest(name Named, digest digest.Digest) (Canonical, error) {
 	if !anchoredDigestRegexp.MatchString(digest.String()) {
 		return nil, ErrDigestInvalidFormat
 	}
+	if tagged, ok := name.(Tagged); ok {
+		return reference{
+			name:   name.Name(),
+			tag:    tagged.Tag(),
+			digest: digest,
+		}, nil
+	}
 	return canonicalReference{
 		name:   name.Name(),
 		digest: digest,
diff --git a/reference/reference_test.go b/reference/reference_test.go
index f60cf093e..be27a3628 100644
--- a/reference/reference_test.go
+++ b/reference/reference_test.go
@@ -467,6 +467,7 @@ func TestSerialization(t *testing.T) {
 func TestWithTag(t *testing.T) {
 	testcases := []struct {
 		name     string
+		digest   digest.Digest
 		tag      string
 		combined string
 	}{
@@ -490,6 +491,12 @@ func TestWithTag(t *testing.T) {
 			tag:      "TAG5",
 			combined: "test.com:8000/foo:TAG5",
 		},
+		{
+			name:     "test.com:8000/foo",
+			digest:   "sha256:1234567890098765432112345667890098765",
+			tag:      "TAG5",
+			combined: "test.com:8000/foo:TAG5@sha256:1234567890098765432112345667890098765",
+		},
 	}
 	for _, testcase := range testcases {
 		failf := func(format string, v ...interface{}) {
@@ -501,6 +508,14 @@ func TestWithTag(t *testing.T) {
 		if err != nil {
 			failf("error parsing name: %s", err)
 		}
+		if testcase.digest != "" {
+			canonical, err := WithDigest(named, testcase.digest)
+			if err != nil {
+				failf("error adding digest")
+			}
+			named = canonical
+		}
+
 		tagged, err := WithTag(named, testcase.tag)
 		if err != nil {
 			failf("WithTag failed: %s", err)
@@ -515,6 +530,7 @@ func TestWithDigest(t *testing.T) {
 	testcases := []struct {
 		name     string
 		digest   digest.Digest
+		tag      string
 		combined string
 	}{
 		{
@@ -532,6 +548,12 @@ func TestWithDigest(t *testing.T) {
 			digest:   "sha256:1234567890098765432112345667890098765",
 			combined: "test.com:8000/foo@sha256:1234567890098765432112345667890098765",
 		},
+		{
+			name:     "test.com:8000/foo",
+			digest:   "sha256:1234567890098765432112345667890098765",
+			tag:      "latest",
+			combined: "test.com:8000/foo:latest@sha256:1234567890098765432112345667890098765",
+		},
 	}
 	for _, testcase := range testcases {
 		failf := func(format string, v ...interface{}) {
@@ -543,6 +565,13 @@ func TestWithDigest(t *testing.T) {
 		if err != nil {
 			failf("error parsing name: %s", err)
 		}
+		if testcase.tag != "" {
+			tagged, err := WithTag(named, testcase.tag)
+			if err != nil {
+				failf("error adding tag")
+			}
+			named = tagged
+		}
 		digested, err := WithDigest(named, testcase.digest)
 		if err != nil {
 			failf("WithDigest failed: %s", err)