feat: if no tag can be found fallback to setting UNKNOWN

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)

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)
+		})
+	}
+}