From fa260c47f0f8fb85e7cf558477a5a4139458c85c Mon Sep 17 00:00:00 2001
From: shangruwu <shangruwu@tencent.com>
Date: Tue, 14 Jul 2020 14:58:08 +0800
Subject: [PATCH] fix(installer): use standard function to parse docker image
 name

---
 pkg/util/docker/docker.go      |  17 +++--
 pkg/util/docker/docker_test.go | 119 +++++++++++++++++++++++++++++++++
 2 files changed, 131 insertions(+), 5 deletions(-)
 create mode 100644 pkg/util/docker/docker_test.go

diff --git a/pkg/util/docker/docker.go b/pkg/util/docker/docker.go
index 843bd1d8d..30e5e8297 100644
--- a/pkg/util/docker/docker.go
+++ b/pkg/util/docker/docker.go
@@ -27,6 +27,7 @@ import (
 
 	"tkestack.io/tke/pkg/spec"
 
+	"github.com/docker/distribution/reference"
 	pkgerrors "github.com/pkg/errors"
 )
 
@@ -154,14 +155,20 @@ func (d *Docker) GetNameArchTag(image string) (name string, arch string, tag str
 // SplitImageNameAndTag returns the name & tag of the given image.
 // If the tag is <none>, return err.
 func (d *Docker) SplitImageNameAndTag(image string) (name string, tag string, err error) {
-	nameAndTag := strings.Split(image, ":")
-	if len(nameAndTag) != 2 {
+	imageRef, err := reference.Parse(image)
+	if err != nil {
 		return "", "", fmt.Errorf("fail to get name and tag for image: %v", image)
 	}
-	if nameAndTag[1] == "<none>" {
-		return "", "", fmt.Errorf("image %s is invalid", image)
+
+	switch v := imageRef.(type) {
+	case reference.NamedTagged:
+		if v.Tag() == "<none>" {
+			return "", "", fmt.Errorf("image %s is invalid", image)
+		}
+		return v.Name(), v.Tag(), nil
+	default:
+		return "", "", fmt.Errorf("image: %v does not have a name and tag", image)
 	}
-	return nameAndTag[0], nameAndTag[1], nil
 }
 
 // SplitNameAndArch returns the real name & arch of the given name.
diff --git a/pkg/util/docker/docker_test.go b/pkg/util/docker/docker_test.go
new file mode 100644
index 000000000..d44275706
--- /dev/null
+++ b/pkg/util/docker/docker_test.go
@@ -0,0 +1,119 @@
+package docker
+
+import (
+	"fmt"
+	"strconv"
+	"testing"
+)
+
+func TestDockerGetNameArchTag(t *testing.T) {
+	imageTestCases := []struct {
+		// image is the image name component of testcase
+		image string
+		// name is the string representation for the image name (without arch)
+		name string
+		// arch is the arch info for the image name
+		arch string
+		// tag is the tag for the image name
+		tag string
+		// err is the error expected from Parse, or nil
+		err error
+	}{
+		{
+			image: "localhost:5000/library/test-amd64:v1.2.3",
+			name:  "localhost:5000/library/test",
+			arch:  "amd64",
+			tag:   "v1.2.3",
+			err:   nil,
+		},
+		{
+			image: "127.0.0.1/tke/test-aaa:1.2.3",
+			name:  "127.0.0.1/tke/test-aaa",
+			arch:  "",
+			tag:   "1.2.3",
+			err:   nil,
+		},
+		{
+			image: "tke/test-arm64:1-2-3",
+			name:  "tke/test",
+			arch:  "arm64",
+			tag:   "1-2-3",
+			err:   nil,
+		},
+		{
+			image: "test-aaa:1.2.3-amd64",
+			name:  "test-aaa",
+			arch:  "",
+			tag:   "1.2.3-amd64",
+			err:   nil,
+		},
+		{
+			image: "test:5000/repo@sha256:ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff",
+			name:  "",
+			arch:  "",
+			tag:   "",
+			err:   fmt.Errorf("image: %v does not have a name and tag", "test:5000/repo@sha256:ffff"),
+		},
+		{
+			image: "tketest:<none>",
+			name:  "",
+			arch:  "",
+			tag:   "",
+			err:   fmt.Errorf("image %s is invalid", "tketest:<none>"),
+		},
+		{
+			image: "tketest",
+			name:  "",
+			arch:  "",
+			tag:   "",
+			err:   fmt.Errorf("image: %v does not have a name and tag", "tketest"),
+		},
+		{
+			image: "tketest/test/test",
+			name:  "",
+			arch:  "",
+			tag:   "",
+			err:   fmt.Errorf("image: %v does not have a name and tag", "tketest"),
+		},
+		{
+			image: "tketest/test1/test2/test3:v1.23",
+			name:  "tketest/test1/test2/test3",
+			arch:  "",
+			tag:   "v1.23",
+			err:   nil,
+		},
+	}
+
+	docker := New()
+	for _, testcase := range imageTestCases {
+		failf := func(format string, v ...interface{}) {
+			t.Logf(strconv.Quote(testcase.image)+": "+format, v...)
+			t.Fail()
+		}
+
+		name, arch, tag, err := docker.GetNameArchTag(testcase.image)
+		if testcase.err != nil {
+			if err == nil {
+				failf("missing expected error: %v", testcase.err)
+			}
+		}
+
+		if testcase.err == nil {
+			if err != nil {
+				failf("unexpected error: %v", err)
+			}
+		}
+
+		if name != testcase.name {
+			failf("mismatched name: got %q, expected %q", name, testcase.name)
+		}
+
+		if arch != testcase.arch {
+			failf("mismatched arch: got %q, expected %q", arch, testcase.arch)
+		}
+
+		if tag != testcase.tag {
+			failf("mismatched tag: got %q, expected %q", tag, testcase.tag)
+		}
+	}
+}