Add label that enables Azure workload ID (#265)

castai · Apr 30, 2024 · a5b43b7 · a5b43b7
1 parent a1d16ba
commit a5b43b7
Show file tree

Hide file tree

Showing 4 changed files with 227 additions and 185 deletions.
diff --git a/cmd/controller/main.go b/cmd/controller/main.go
@@ -39,6 +39,8 @@ var (
 
 	pyroscopeAddr = pflag.String("pyroscope-addr", "", "Enable pyroscope tracing")
 
+	cloudProvider = pflag.String("cloud-provider", "", "Cloud provider in which the cluster is running")
+
 	castaiSecretRefName      = pflag.String("castai-secret-ref-name", "castai-kvisor", "CASTAI k8s secret name")
 	castaiConfigSyncDuration = pflag.Duration("castai-config-sync-duration", 1*time.Minute, "CASTAI remote config sync duration")
 	castaiServerInsecure     = pflag.Bool("castai-server-insecure", false, "Use insecure connection to castai grpc server. Used for e2e.")
@@ -69,7 +71,8 @@ var (
 	kubeBenchScanInterval       = pflag.Duration("kube-bench-scan-interval", 5*time.Minute, "Kube bench scan interval")
 	kubeBenchForceScan          = pflag.Bool("kube-bench-force", false, "Kube Bench force scan")
 	kubeBenchJobImagePullPolicy = pflag.String("kube-bench-job-pull-policy", "IfNotPresent", "Kube bench job image pull policy")
-	kubeBenchCloudProvider      = pflag.String("kube-bench-cloud-provider", "", "Kube bench cloud provider")
+	// deprecated: use cloudProvider
+	kubeBenchCloudProvider = pflag.String("kube-bench-cloud-provider", "", "Kube bench cloud provider. Deprecated: use `cloud-provider` instead.")
 
 	kubeLinterEnabled      = pflag.Bool("kube-linter-enabled", false, "Kube linter enabled")
 	kubeLinterScanInterval = pflag.Duration("kube-linter-scan-interval", 60*time.Second, "Kube linter scan interval")
@@ -137,6 +140,14 @@ func main() {
 		os.Exit(1)
 	}
 
+	var cloudProviderVal string
+	if *cloudProvider != "" {
+		cloudProviderVal = *cloudProvider
+	} else {
+		slog.Warn(`--kube-bench-cloud-provider is deprecated, please use --cloud-provider instead.`)
+		cloudProviderVal = *kubeBenchCloudProvider
+	}
+
 	podNs := os.Getenv("POD_NAMESPACE")
 	appInstance := app.New(app.Config{
 		LogLevel:              *logLevel,
@@ -174,6 +185,7 @@ func main() {
 			CastaiClusterID:           castaiClusterID,
 			CastaiGrpcInsecure:        *castaiServerInsecure,
 			ImageScanBlobsCacheURL:    *imageScanBlobsCacheURL,
+			CloudProvider:             cloudProviderVal,
 		},
 		Linter: kubelinter.Config{
 			Enabled:      *kubeLinterEnabled,
@@ -185,7 +197,7 @@ func main() {
 			Force:              *kubeBenchForceScan,
 			ScanInterval:       *kubeBenchScanInterval,
 			JobImagePullPolicy: *kubeBenchJobImagePullPolicy,
-			CloudProvider:      *kubeBenchCloudProvider,
+			CloudProvider:      cloudProviderVal,
 			JobNamespace:       podNs,
 		},
 		Delta: delta.Config{

diff --git a/cmd/controller/state/imagescan/controller.go b/cmd/controller/state/imagescan/controller.go
@@ -46,6 +46,7 @@ type Config struct {
 	CastaiClusterID           string
 	CastaiGrpcInsecure        bool
 	ImageScanBlobsCacheURL    string
+	CloudProvider             string
 }
 
 type ImageScanImage struct {

diff --git a/cmd/controller/state/imagescan/scanner.go b/cmd/controller/state/imagescan/scanner.go
@@ -402,6 +402,12 @@ func scanJobSpec(
 	vol volumesAndMounts,
 	tolerations []corev1.Toleration,
 ) *batchv1.Job {
+
+	podLabels := map[string]string{}
+	if cfg.CloudProvider == "aks" {
+		podLabels["azure.workload.identity/use"] = "true"
+	}
+
 	job := &batchv1.Job{
 		TypeMeta: metav1.TypeMeta{
 			Kind:       "Job",
@@ -423,6 +429,7 @@ func scanJobSpec(
 			Template: corev1.PodTemplateSpec{
 				ObjectMeta: metav1.ObjectMeta{
 					Annotations: annotations,
+					Labels:      podLabels,
 				},
 				Spec: corev1.PodSpec{
 					RestartPolicy:      corev1.RestartPolicyNever,