Fixed image parsing logic (#470)

* fixed image parsing logic

* code review changes

docker/parse.go

@@ -5,8 +5,8 @@ import (
 	"io"
 	"os"
 	"sort"
-	"strings"
 
+	"github.com/docker/distribution/reference"
 	"github.com/moby/buildkit/frontend/dockerfile/command"
 	"github.com/moby/buildkit/frontend/dockerfile/parser"
 )
@@ -102,10 +102,26 @@ func GetImageTagFromParsedFile(cmds []Command) (baseImage, tag string) {
 	for _, cmd := range cmds {
 		if cmd.Cmd == command.From {
 			from := cmd.Value[0]
-			splittedFrom := strings.Split(from, ":")
-			baseImage, tag := splittedFrom[0], splittedFrom[1]
+			baseImage, tag := parseImageName(from)
 			return baseImage, tag
 		}
 	}
 	return "", ""
 }
+
+func parseImageName(imageName string) (baseImage, tag string) {
+	ref, err := reference.Parse(imageName)
+	if err != nil {
+		return baseImage, tag
+	}
+	parsedName, ok := ref.(reference.Named)
+	if ok {
+		baseImage = parsedName.Name()
+	}
+	tag = "latest"
+	parsedTag, ok := ref.(reference.Tagged)
+	if ok {
+		tag = parsedTag.Tag()
+	}
+	return baseImage, tag
+}

docker/parse_test.go

@@ -72,3 +72,23 @@ func TestGetImageTagFromParsedFile(t *testing.T) {
 	assert.Equal(t, "quay.io/astronomer/ap-airflow", image)
 	assert.Equal(t, "2.0.0-buster-onbuild", tag)
 }
+
+func TestParseImageName(t *testing.T) {
+	tests := []struct {
+		imageName         string
+		expectedBaseImage string
+		expectedTag       string
+	}{
+		{imageName: "localhost:5000/airflow:latest", expectedBaseImage: "localhost:5000/airflow", expectedTag: "latest"},
+		{imageName: "localhost/airflow:1.2.0", expectedBaseImage: "localhost/airflow", expectedTag: "1.2.0"},
+		{imageName: "quay.io:5000/airflow", expectedBaseImage: "quay.io:5000/airflow", expectedTag: "latest"},
+		{imageName: "airflow:latest", expectedBaseImage: "airflow", expectedTag: "latest"},
+		{imageName: "airflow", expectedBaseImage: "airflow", expectedTag: "latest"},
+	}
+
+	for _, tt := range tests {
+		baseImage, tag := parseImageName(tt.imageName)
+		assert.Equal(t, tt.expectedBaseImage, baseImage)
+		assert.Equal(t, tt.expectedTag, tag)
+	}
+}