From d7e0cfdef4141b8e2f4e2acdddd6a32ff582bfa7 Mon Sep 17 00:00:00 2001
From: Nick Wilburn <senior.crepe@gmail.com>
Date: Thu, 18 May 2023 16:08:29 -0500
Subject: [PATCH] feat: add scheduling logic to zarf injector (#1731)

## Description

Checks to ensure that a node does not have a `NoSchedule` taint before
choosing that node/image for the Zarf injector pod

## Related Issue

Fixes #1730
Fixes #905

## Type of change

- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Other (security config, docs update, etc)

## Checklist before merging

- [ ] Test, docs, adr added or updated as needed
- [ ] [Contributor Guide
Steps](https://github.com/defenseunicorns/zarf/blob/main/CONTRIBUTING.md#developer-workflow)
followed


Gave this a WAG. It seems to work locally, but am unsure if this is the
direction you guys want to go. Let me know!
---
 src/internal/cluster/injector.go |  7 +++++--
 src/pkg/k8s/images.go            | 26 +++++++++++++++++++-------
 src/pkg/k8s/nodes.go             |  6 ++++++
 3 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/src/internal/cluster/injector.go b/src/internal/cluster/injector.go
index 867c2469b3..c887cd0120 100644
--- a/src/internal/cluster/injector.go
+++ b/src/internal/cluster/injector.go
@@ -42,10 +42,13 @@ func (c *Cluster) StartInjectionMadness(tempPath types.TempPaths, injectorSeedTa
 	var seedImages []transform.Image
 
 	// Get all the images from the cluster
-	spinner.Updatef("Getting the list of existing cluster images")
-	if images, err = c.Kube.GetAllImages(); err != nil {
+	timeout := 5 * time.Minute
+	spinner.Updatef("Getting the list of existing cluster images (%s timeout)", timeout.String())
+	if images, err = c.Kube.GetAllImages(timeout); err != nil {
 		spinner.Fatalf(err, "Unable to generate a list of candidate images to perform the registry injection")
 	}
+	message.Debugf("Found %d images in the cluster", len(images))
+	message.Debugf("Images: %#v", images)
 
 	spinner.Updatef("Creating the injector configmap")
 	if err = c.createInjectorConfigmap(tempPath); err != nil {
diff --git a/src/pkg/k8s/images.go b/src/pkg/k8s/images.go
index 10bd84695c..753ac36ef6 100644
--- a/src/pkg/k8s/images.go
+++ b/src/pkg/k8s/images.go
@@ -19,8 +19,8 @@ type ImageMap map[string]bool
 type ImageNodeMap map[string][]string
 
 // GetAllImages returns a list of images and their nodes found in pods in the cluster.
-func (k *K8s) GetAllImages() (ImageNodeMap, error) {
-	timeout := time.After(5 * time.Minute)
+func (k *K8s) GetAllImages(timeoutDuration time.Duration) (ImageNodeMap, error) {
+	timeout := time.After(timeoutDuration)
 
 	for {
 		// Delay check 2 seconds.
@@ -44,7 +44,8 @@ func (k *K8s) GetAllImages() (ImageNodeMap, error) {
 	}
 }
 
-// GetImagesWithNodes returns all images and their nodes in a given namespace.
+// GetImagesWithNodes checks for images on schedulable nodes and returns
+// a map of these images and their nodes in a given namespace.
 func (k *K8s) GetImagesWithNodes(namespace string) (ImageNodeMap, error) {
 	result := make(ImageNodeMap)
 
@@ -53,16 +54,27 @@ func (k *K8s) GetImagesWithNodes(namespace string) (ImageNodeMap, error) {
 		return nil, fmt.Errorf("unable to get the list of pods in the cluster")
 	}
 
+	findImages:
 	for _, pod := range pods.Items {
-		node := pod.Spec.NodeName
+		nodeName := pod.Spec.NodeName
+		nodeDetails, err := k.GetNode(nodeName)
+		if err != nil {
+			return nil, fmt.Errorf("unable to get the node %s", pod.Spec.NodeName)
+		}
+
+		for _, taint := range nodeDetails.Spec.Taints {
+			if (taint.Effect == corev1.TaintEffectNoSchedule || taint.Effect == corev1.TaintEffectNoExecute) {
+				continue findImages
+			}
+		}
 		for _, container := range pod.Spec.InitContainers {
-			result[container.Image] = append(result[container.Image], node)
+			result[container.Image] = append(result[container.Image], nodeName)
 		}
 		for _, container := range pod.Spec.Containers {
-			result[container.Image] = append(result[container.Image], node)
+			result[container.Image] = append(result[container.Image], nodeName)
 		}
 		for _, container := range pod.Spec.EphemeralContainers {
-			result[container.Image] = append(result[container.Image], node)
+			result[container.Image] = append(result[container.Image], nodeName)
 		}
 	}
 
diff --git a/src/pkg/k8s/nodes.go b/src/pkg/k8s/nodes.go
index 1031c4ffe1..7df70c443e 100644
--- a/src/pkg/k8s/nodes.go
+++ b/src/pkg/k8s/nodes.go
@@ -16,3 +16,9 @@ func (k *K8s) GetNodes() (*corev1.NodeList, error) {
 	metaOptions := metav1.ListOptions{}
 	return k.Clientset.CoreV1().Nodes().List(context.TODO(), metaOptions)
 }
+
+// GetNode returns a node from the k8s cluster.
+func (k *K8s) GetNode(nodeName string) (*corev1.Node, error) {
+	return k.Clientset.CoreV1().Nodes().Get(context.TODO(), nodeName, metav1.GetOptions{})
+}
+