From 38b876f265923e1cc1739f1f48542c3f75e65118 Mon Sep 17 00:00:00 2001
From: Masashi Honma <masashi.honma@gmail.com>
Date: Tue, 18 Jun 2024 15:57:10 +0900
Subject: [PATCH] Add support for injecting tolerations to sonobuoy pod

Resolves #1973.

We can inject some tolerations to sonobuoy aggregator pod by adding trailing
description into sonobuoy config json.

{
  "AggregatorTolerations": [
    {
      "effect": "NoSchedule",
      "key": "key1",
      "operator": "Equal",
      "value": "value1"
    },
    {
      "effect": "NoSchedule",
      "key": "key2",
      "operator": "Equal",
      "value": "value2"
    }
  ]
}

Signed-off-by: Masashi Honma <masashi.honma@gmail.com>
---
 pkg/client/gen.go    | 32 ++++++++++++++++++++++++++++++++
 pkg/config/config.go | 21 +++++++++++----------
 2 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/pkg/client/gen.go b/pkg/client/gen.go
index 2367080ff..43c8c02f6 100644
--- a/pkg/client/gen.go
+++ b/pkg/client/gen.go
@@ -410,6 +410,38 @@ func generateAggregatorAndService(w io.Writer, cfg *GenConfig) error {
 	if len(cfg.Config.CustomAnnotations) > 0 {
 		p.ObjectMeta.Annotations = cfg.Config.CustomAnnotations
 	}
+	if len(cfg.Config.AggregatorTolerations) > 0 {
+		for _, t := range cfg.Config.AggregatorTolerations {
+			var toleration corev1.Toleration
+			if val, exists := t["key"]; exists {
+				toleration.Key = val
+			}
+			if val, exists := t["value"]; exists {
+				toleration.Value = val
+			}
+			if val, exists := t["effect"]; exists {
+				if val == "NoSchedule" {
+					toleration.Effect = corev1.TaintEffectNoSchedule
+				} else if val == "NoExecute" {
+					toleration.Effect = corev1.TaintEffectNoExecute
+				} else if val == "PreferNoSchedule" {
+					toleration.Effect = corev1.TaintEffectPreferNoSchedule
+				} else {
+					return errors.New("Invalid effect: " + val)
+				}
+			}
+			if val, exists := t["operator"]; exists {
+				if val == "Equal" {
+					toleration.Operator = corev1.TolerationOpEqual
+				} else if val == "Exists" {
+					toleration.Operator = corev1.TolerationOpExists
+				} else {
+					return errors.New("Invalid operator: " + val)
+				}
+			}
+			p.Spec.Tolerations = append(p.Spec.Tolerations, toleration)
+		}
+	}
 
 	switch cfg.Config.SecurityContextMode {
 	case "none":
diff --git a/pkg/config/config.go b/pkg/config/config.go
index 9861dd852..87df54539 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -147,16 +147,17 @@ type Config struct {
 	///////////////////////////////////////////////
 	// Sonobuoy configuration
 	///////////////////////////////////////////////
-	WorkerImage              string            `json:"WorkerImage" mapstructure:"WorkerImage"`
-	ImagePullPolicy          string            `json:"ImagePullPolicy" mapstructure:"ImagePullPolicy"`
-	ForceImagePullPolicy     bool              `json:"ForceImagePullPolicy,omitempty" mapstructure:"ForceImagePullPolicy"`
-	ImagePullSecrets         string            `json:"ImagePullSecrets" mapstructure:"ImagePullSecrets"`
-	CustomAnnotations        map[string]string `json:"CustomAnnotations,omitempty" mapstructure:"CustomAnnotations"`
-	AggregatorPermissions    string            `json:"AggregatorPermissions" mapstructure:"AggregatorPermissions"`
-	ServiceAccountName       string            `json:"ServiceAccountName" mapstructure:"ServiceAccountName"`
-	ExistingServiceAccount   bool              `json:"ExistingServiceAccount,omitempty" mapstructure:"ExistingServiceAccount,omitempty"`
-	E2EDockerConfigFile      string            `json:"E2EDockerConfigFile,omitempty" mapstructure:"E2EDockerConfigFile,omitempty"`
-	NamespacePSAEnforceLevel string            `json:"NamespacePSAEnforceLevel,omitempty" mapstructure:"NamespacePSAEnforceLevel,omitempty"`
+	WorkerImage              string              `json:"WorkerImage" mapstructure:"WorkerImage"`
+	ImagePullPolicy          string              `json:"ImagePullPolicy" mapstructure:"ImagePullPolicy"`
+	ForceImagePullPolicy     bool                `json:"ForceImagePullPolicy,omitempty" mapstructure:"ForceImagePullPolicy"`
+	ImagePullSecrets         string              `json:"ImagePullSecrets" mapstructure:"ImagePullSecrets"`
+	CustomAnnotations        map[string]string   `json:"CustomAnnotations,omitempty" mapstructure:"CustomAnnotations"`
+	AggregatorPermissions    string              `json:"AggregatorPermissions" mapstructure:"AggregatorPermissions"`
+	AggregatorTolerations    []map[string]string `json:"AggregatorTolerations,omitempty" mapstructure:"AggregatorTolerations"`
+	ServiceAccountName       string              `json:"ServiceAccountName" mapstructure:"ServiceAccountName"`
+	ExistingServiceAccount   bool                `json:"ExistingServiceAccount,omitempty" mapstructure:"ExistingServiceAccount,omitempty"`
+	E2EDockerConfigFile      string              `json:"E2EDockerConfigFile,omitempty" mapstructure:"E2EDockerConfigFile,omitempty"`
+	NamespacePSAEnforceLevel string              `json:"NamespacePSAEnforceLevel,omitempty" mapstructure:"NamespacePSAEnforceLevel,omitempty"`
 
 	// ProgressUpdatesPort is the port on which the Sonobuoy worker will listen for status updates from its plugin.
 	ProgressUpdatesPort string `json:"ProgressUpdatesPort,omitempty" mapstructure:"ProgressUpdatesPort"`