From 6980f3e9908fca66ae2fe032609a6342fbe909d2 Mon Sep 17 00:00:00 2001
From: Bradley Jones <bradley.jones@anchore.com>
Date: Sat, 20 Jul 2024 00:51:07 +0100
Subject: [PATCH] feat: if no tag can be found fallback to setting UNKNOWN

Signed-off-by: Bradley Jones <bradley.jones@anchore.com>
---
 pkg/inventory/ecs.go      | 23 ++++++++++-----
 pkg/inventory/ecs_test.go | 62 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 8 deletions(-)

diff --git a/pkg/inventory/ecs.go b/pkg/inventory/ecs.go
index a1e5bbb..c884192 100644
--- a/pkg/inventory/ecs.go
+++ b/pkg/inventory/ecs.go
@@ -89,16 +89,10 @@ func fetchContainersFromTasks(client ecsiface.ECSAPI, cluster string, tasks []*s
 				logger.Log.Warnf("No image digest found for container: %s", *container.ContainerArn)
 				logger.Log.Warn("Ensure all ECS container hosts are running at least ECS Agent 1.70.0, which fixed a bug where image digests were not returned in the DescribeTasks API response.")
 			}
-			// Fix container image tag if it contains an @ symbol
-			if strings.Contains(*container.Image, "@") {
-				if tag, ok := containerTagMap[digest]; ok {
-					// replace the image tag with the correct one
-					container.Image = &tag
-				}
-			}
+			containerImage := getContainerImageTag(containerTagMap, container)
 			containers = append(containers, reporter.Container{
 				ARN:         *container.ContainerArn,
-				ImageTag:    *container.Image,
+				ImageTag:    containerImage,
 				ImageDigest: digest,
 				TaskARN:     *task.TaskArn,
 			})
@@ -108,6 +102,19 @@ func fetchContainersFromTasks(client ecsiface.ECSAPI, cluster string, tasks []*s
 	return containers, nil
 }
 
+func getContainerImageTag(containerTagMap map[string]string, container *ecs.Container) string {
+	// Fix container image tag if it contains an @ symbol
+	if strings.Contains(*container.Image, "@") {
+		// replace the image tag with the correct one
+		if tag, ok := containerTagMap[*container.ImageDigest]; ok {
+			return tag
+		}
+		logger.Log.Warnf("No image tag found for container setting to UNKNOWN: %s", *container.Image)
+		return strings.Split(*container.Image, "@")[0] + ":UNKNOWN"
+	}
+	return *container.Image
+}
+
 // Build a map of container image digests to image tags
 func buildContainerTagMap(tasks []*ecs.Task) map[string]string {
 	containerMap := make(map[string]string)
diff --git a/pkg/inventory/ecs_test.go b/pkg/inventory/ecs_test.go
index 5401800..01249d6 100644
--- a/pkg/inventory/ecs_test.go
+++ b/pkg/inventory/ecs_test.go
@@ -3,6 +3,7 @@ package inventory
 import (
 	"testing"
 
+	"github.com/aws/aws-sdk-go/service/ecs"
 	"github.com/aws/aws-sdk-go/service/ecs/ecsiface"
 	"github.com/stretchr/testify/assert"
 
@@ -496,3 +497,64 @@ func Test_constructServiceARN(t *testing.T) {
 		})
 	}
 }
+
+func Test_getContainerImageTag(t *testing.T) {
+	type args struct {
+		containerTagMap map[string]string
+		container       *ecs.Container
+	}
+	tests := []struct {
+		name string
+		args args
+		want string
+	}{
+		{
+			name: "return container image tag when it does not contain @ symbol",
+			args: args{
+				containerTagMap: map[string]string{
+					"sha256:1234567890123456789012345678901234567890123456789012345678901111": "image-1:latest",
+					"sha256:1234567890123456789012345678901234567890123456789012345678902222": "image-2:latest",
+				},
+				container: &ecs.Container{
+					Image:       GetPointerToValue("image-1:latest"),
+					ImageDigest: GetPointerToValue("sha256:1234567890123456789012345678901234567890123456789012345678901111"),
+				},
+			},
+			want: "image-1:latest",
+		},
+		{
+			name: "return container image tag from map when it does contain @ symbol",
+			args: args{
+				containerTagMap: map[string]string{
+					"sha256:1234567890123456789012345678901234567890123456789012345678901111": "image-1:latest",
+					"sha256:1234567890123456789012345678901234567890123456789012345678902222": "image-2:latest",
+				},
+				container: &ecs.Container{
+					Image:       GetPointerToValue("image-1@sha256:1234567890123456789012345678901234567890123456789012345678901111"),
+					ImageDigest: GetPointerToValue("sha256:1234567890123456789012345678901234567890123456789012345678901111"),
+				},
+			},
+			want: "image-1:latest",
+		},
+		{
+			name: "return UNKNOWN as the tag when image tag is not found in the map",
+			args: args{
+				containerTagMap: map[string]string{
+					"sha256:1234567890123456789012345678901234567890123456789012345678901111": "image-1:latest",
+					"sha256:1234567890123456789012345678901234567890123456789012345678902222": "image-2:latest",
+				},
+				container: &ecs.Container{
+					Image:       GetPointerToValue("image-1@sha256:0000"),
+					ImageDigest: GetPointerToValue("sha256:11"),
+				},
+			},
+			want: "image-1:UNKNOWN",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := getContainerImageTag(tt.args.containerTagMap, tt.args.container)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}