fix SetContainerInfo for initcontainers

Signed-off-by: Matthias Bertschy <[email protected]>
kubescape · Jan 23, 2025 · ddd019f · ddd019f
1 parent 8f6f8b7
commit ddd019f
Show file tree

Hide file tree

Showing 6 changed files with 57 additions and 51 deletions.
diff --git a/pkg/applicationprofilemanager/v1/applicationprofile_manager_test.go b/pkg/applicationprofilemanager/v1/applicationprofile_manager_test.go
@@ -70,12 +70,14 @@ func ensureInstanceID(container *containercollection.Container, watchedContainer
 	if err != nil {
 		return fmt.Errorf("failed to generate instanceID: %w", err)
 	}
-	watchedContainer.InstanceID = instanceIDs[0]
 	for i := range instanceIDs {
 		if instanceIDs[i].GetContainerName() == container.K8s.ContainerName {
 			watchedContainer.InstanceID = instanceIDs[i]
 		}
 	}
+	if watchedContainer.InstanceID == nil {
+		return fmt.Errorf("failed to find instance id for container %s", container.K8s.ContainerName)
+	}
 	return nil
 }
 

diff --git a/pkg/containerwatcher/v1/container_watcher_private.go b/pkg/containerwatcher/v1/container_watcher_private.go
@@ -149,13 +149,14 @@ func (ch *IGContainerWatcher) getSharedWatchedContainerData(container *container
 	if err != nil {
 		return nil, fmt.Errorf("failed to generate instance id: %w", err)
 	}
-	watchedContainer.InstanceID = instanceIDs[0]
 	for i := range instanceIDs {
 		if instanceIDs[i].GetContainerName() == container.K8s.ContainerName {
 			watchedContainer.InstanceID = instanceIDs[i]
 		}
 	}
-
+	if watchedContainer.InstanceID == nil {
+		return nil, fmt.Errorf("failed to find instance id for container %s", container.K8s.ContainerName)
+	}
 	return &watchedContainer, nil
 }
 

diff --git a/pkg/k8sclient/k8sclient_mock.go b/pkg/k8sclient/k8sclient_mock.go
@@ -41,6 +41,10 @@ func (k *K8sClientMock) GetWorkload(namespace, _, name string) (k8sinterface.IWo
 		},
 		"status": map[string]interface{}{
 			"containerStatuses": []interface{}{
+				map[string]interface{}{
+					"name":    "log",
+					"imageID": storage.FluentBitImageID,
+				},
 				map[string]interface{}{
 					"name":    "cont",
 					"imageID": storage.NginxImageID,

diff --git a/pkg/networkmanager/v2/network_manager_test.go b/pkg/networkmanager/v2/network_manager_test.go
@@ -72,12 +72,14 @@ func ensureInstanceID(container *containercollection.Container, watchedContainer
 	if err != nil {
 		return fmt.Errorf("failed to generate instanceID: %w", err)
 	}
-	watchedContainer.InstanceID = instanceIDs[0]
 	for i := range instanceIDs {
 		if instanceIDs[i].GetContainerName() == container.K8s.ContainerName {
 			watchedContainer.InstanceID = instanceIDs[i]
 		}
 	}
+	if watchedContainer.InstanceID == nil {
+		return fmt.Errorf("failed to find instance id for container %s", container.K8s.ContainerName)
+	}
 	return nil
 }
 

diff --git a/pkg/storage/storage_mock.go b/pkg/storage/storage_mock.go
@@ -6,8 +6,9 @@ import (
 )
 
 const (
-	NginxKey     = "nginx-c9b3ae"
-	NginxImageID = "nginx@sha256:6a59f1cbb8d28ac484176d52c473494859a512ddba3ea62a547258cf16c9b3ae"
+	NginxKey         = "nginx-c9b3ae"
+	NginxImageID     = "nginx@sha256:6a59f1cbb8d28ac484176d52c473494859a512ddba3ea62a547258cf16c9b3ae"
+	FluentBitImageID = "fluentbit@sha256:236f7d961b0ba8b91796955f155819d64801e0d00fa666147502ab9b5b80f623"
 )
 
 type StorageHttpClientMock struct {

diff --git a/pkg/utils/utils.go b/pkg/utils/utils.go
@@ -9,6 +9,7 @@ import (
 	"fmt"
 	"hash"
 	"io"
+	"iter"
 	"math/rand"
 	"os"
 	"path/filepath"
@@ -239,9 +240,6 @@ func (watchedContainer *WatchedContainerData) StatusUpdated() bool {
 }
 
 func (watchedContainer *WatchedContainerData) SetContainerInfo(wl workloadinterface.IWorkload, containerName string) error {
-	if watchedContainer.ContainerInfos == nil {
-		watchedContainer.ContainerInfos = make(map[ContainerType][]ContainerInfo)
-	}
 	podSpec, err := wl.GetPodSpec()
 	if err != nil {
 		return fmt.Errorf("failed to get pod spec: %w", err)
@@ -254,67 +252,65 @@ func (watchedContainer *WatchedContainerData) SetContainerInfo(wl workloadinterf
 	if podSpec.SecurityContext != nil && podSpec.SecurityContext.SeccompProfile != nil {
 		watchedContainer.SeccompProfilePath = podSpec.SecurityContext.SeccompProfile.LocalhostProfile
 	}
-	// TODO rewrite with podutil.VisitContainers()
-	checkContainers := func(containers []v1.Container, ephemeralContainers []v1.EphemeralContainer, containerType ContainerType) error {
-		var imageID string
-		for _, c := range podStatus.ContainerStatuses {
-			if c.Name == containerName {
-				imageID = c.ImageID
-			}
-		}
-		if imageID == "" {
-			return fmt.Errorf("failed to get imageID for container %s", containerName)
-		}
+	// fill container infos
+	if watchedContainer.ContainerInfos == nil {
+		watchedContainer.ContainerInfos = make(map[ContainerType][]ContainerInfo)
+	}
+	checkContainers := func(containers iter.Seq2[int, v1.Container], containerStatuses []v1.ContainerStatus, containerType ContainerType) error {
 		var containersInfo []ContainerInfo
-		if containerType == EphemeralContainer {
-			for i, c := range ephemeralContainers {
-				containersInfo = append(containersInfo, ContainerInfo{
-					Name:     c.Name,
-					ImageTag: c.Image,
-					ImageID:  imageID,
-				})
-				if c.Name == containerName {
-					watchedContainer.ContainerIndex = i
-					watchedContainer.ContainerType = containerType
-					if c.SecurityContext != nil && c.SecurityContext.SeccompProfile != nil {
-						watchedContainer.SeccompProfilePath = c.SecurityContext.SeccompProfile.LocalhostProfile
-					}
-				}
-			}
-		} else {
-			for i, c := range containers {
-				containersInfo = append(containersInfo, ContainerInfo{
-					Name:     c.Name,
-					ImageTag: c.Image,
-					ImageID:  imageID,
-				})
-				if c.Name == containerName {
-					watchedContainer.ContainerIndex = i
-					watchedContainer.ContainerType = containerType
-					if c.SecurityContext != nil && c.SecurityContext.SeccompProfile != nil {
-						watchedContainer.SeccompProfilePath = c.SecurityContext.SeccompProfile.LocalhostProfile
-					}
+		for i, c := range containers {
+			containersInfo = append(containersInfo, ContainerInfo{
+				Name:     c.Name,
+				ImageTag: c.Image,
+				ImageID:  containerStatuses[i].ImageID,
+			})
+			if c.Name == containerName {
+				watchedContainer.ContainerIndex = i
+				watchedContainer.ContainerType = containerType
+				if c.SecurityContext != nil && c.SecurityContext.SeccompProfile != nil {
+					watchedContainer.SeccompProfilePath = c.SecurityContext.SeccompProfile.LocalhostProfile
 				}
 			}
 		}
 		watchedContainer.ContainerInfos[containerType] = containersInfo
 		return nil
 	}
 	// containers
-	if err := checkContainers(podSpec.Containers, nil, Container); err != nil {
+	if err := checkContainers(containersIterator(podSpec.Containers), podStatus.ContainerStatuses, Container); err != nil {
 		return err
 	}
 	// initContainers
-	if err := checkContainers(podSpec.InitContainers, nil, InitContainer); err != nil {
+	if err := checkContainers(containersIterator(podSpec.InitContainers), podStatus.InitContainerStatuses, InitContainer); err != nil {
 		return err
 	}
 	// ephemeralContainers
-	if err := checkContainers(nil, podSpec.EphemeralContainers, EphemeralContainer); err != nil {
+	if err := checkContainers(ephemeralContainersIterator(podSpec.EphemeralContainers), podStatus.EphemeralContainerStatuses, EphemeralContainer); err != nil {
 		return err
 	}
+	logger.L().Info("Matthias - SetContainerInfo", helpers.Interface("ContainerInfos", watchedContainer.ContainerInfos))
 	return nil
 }
 
+func containersIterator(c []v1.Container) iter.Seq2[int, v1.Container] {
+	return func(yield func(int, v1.Container) bool) {
+		for i := 0; i < len(c); i++ {
+			if !yield(i, c[i]) {
+				return
+			}
+		}
+	}
+}
+
+func ephemeralContainersIterator(c []v1.EphemeralContainer) iter.Seq2[int, v1.Container] {
+	return func(yield func(int, v1.Container) bool) {
+		for i := 0; i < len(c); i++ {
+			if !yield(i, v1.Container(c[i].EphemeralContainerCommon)) {
+				return
+			}
+		}
+	}
+}
+
 type PatchOperation struct {
 	Op    string      `json:"op"`
 	Path  string      `json:"path"`