Mount multiple file secrets using single init container (#338)

## Overview
Follow up to #293, in which each secret mounted as file uses a separate init container. This change utilizes a single init container.

## Test Plan
1. Deploy to dogfood.
2. Create multiple secrets using unionai's `-f` flag and file.
3. Run a workflow using the secrets created in previous step.
4. Inspect pod spec.
5. Verify that secret values were read correctly.

## Rollout Plan (if applicable)
Run `managed-cluster-<env>-sync-all`.

## Upstream Changes
Should this change be upstreamed to OSS (flyteorg/flyte)? If not, please uncheck this box, which is used for auditing. Note, it is the responsibility of each developer to actually upstream their changes. See [this guide](https://unionai.atlassian.net/wiki/spaces/ENG/pages/447610883/Flyte+-+Union+Cloud+Development+Runbook/#When-are-versions-updated%3F).
- [ ] To be upstreamed to OSS

## Issue
ref COR-811

## Checklist
* [x] Added tests
* [ ] Ran a deploy dry run and shared the terraform plan
* [ ] Added logging and metrics
* [ ] Updated [dashboards](https://unionai.grafana.net/dashboards) and [alerts](https://unionai.grafana.net/alerting/list)
* [ ] Updated documentation

.gitignore

@@ -1,4 +1,5 @@
 .idea
+.vscode/launch.json
 _build/
 _bin/
 build/

flytepropeller/go.mod

@@ -126,6 +126,7 @@ require (
 	github.com/prometheus/common v0.44.0 // indirect
 	github.com/prometheus/procfs v0.10.1 // indirect
 	github.com/ray-project/kuberay/ray-operator v1.1.0-rc.1 // indirect
+	github.com/samber/lo v1.39.0 // indirect
 	github.com/spf13/afero v1.8.2 // indirect
 	github.com/spf13/cast v1.4.1 // indirect
 	github.com/spf13/jwalterweatherman v1.1.0 // indirect

flytepropeller/go.sum

@@ -391,6 +391,8 @@ github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFR
 github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjRBZyWFQ=
 github.com/rogpeppe/go-internal v1.10.0/go.mod h1:UQnix2H7Ngw/k4C5ijL5+65zddjncjaFoBhdsK/akog=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
+github.com/samber/lo v1.39.0 h1:4gTz1wUhNYLhFSKl6O+8peW0v2F4BCY034GRpU9WnuA=
+github.com/samber/lo v1.39.0/go.mod h1:+m/ZKRl6ClXCE2Lgf3MsQlWfh4bn1bz6CXEOxnEXnEA=
 github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6MwdIuYE2rE=
 github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
 github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=

flytepropeller/pkg/secret/embedded_secret_manager.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/base64"
 	"fmt"
+	"slices"
 	"strings"
 	"unicode/utf8"
 
@@ -19,8 +20,18 @@ import (
 const (
 	UnionSecretEnvVarPrefix = "_UNION_"
 	// Static name of the volume used for mounting secrets with file mount requirement.
-	EmbeddedSecretsVolumeName = "embedded-secret-vol" // #nosec G101
-	EmbeddedSecretsMountPath  = "/etc/flyte/secrets"  // #nosec G101
+	EmbeddedSecretsFileMountVolumeName        = "embedded-secret-vol"  // #nosec G101
+	EmbeddedSecretsFileMountPath              = "/etc/flyte/secrets"   // #nosec G101
+	EmbeddedSecretsFileMountInitContainerName = "init-embedded-secret" // #nosec G101
+
+	// Name of the environment variable in the init container used for mounting secrets as files.
+	// This environment variable is used to pass secret names and values to the init container.
+	// The init container then reads its value and writes secrets to files.
+	// Format of this environment variable's value:
+	// 		secret_name1=base64_encoded_secret_value1
+	// 		secret_name2=base64_encoded_secret_value2
+	//		...
+	EmbeddedSecretsFileMountInitContainerEnvVariableName = "SECRETS" // #nosec G101
 
 	SecretFieldSeparator              = "__"
 	ValueFormatter                    = "%s"
@@ -197,9 +208,15 @@ func (i EmbeddedSecretManagerInjector) injectAsEnvVar(secretKey string, secretVa
 }
 
 func (i EmbeddedSecretManagerInjector) injectAsFile(secretKey string, secretValue []byte, pod *corev1.Pod) {
-	// A volume with a static name so that if we try to inject multiple secrets, we won't mount multiple volumes.
+	initContainer, exists := i.getOrAppendFileMountInitContainer(pod)
+	appendSecretToFileMountInitContainer(initContainer, secretKey, secretValue)
+
+	if exists {
+		return
+	}
+
 	volume := corev1.Volume{
-		Name: EmbeddedSecretsVolumeName,
+		Name: EmbeddedSecretsFileMountVolumeName,
 		VolumeSource: corev1.VolumeSource{
 			EmptyDir: &corev1.EmptyDirVolumeSource{
 				Medium: corev1.StorageMediumMemory,
@@ -208,46 +225,19 @@ func (i EmbeddedSecretManagerInjector) injectAsFile(secretKey string, secretValu
 	}
 	pod.Spec.Volumes = appendVolumeIfNotExists(pod.Spec.Volumes, volume)
 
-	secretFilePath := EmbeddedSecretsMountPath + "/" + secretKey
-	secretInitContainer := corev1.Container{
-		Name:  "init-embedded-secret-" + secretKey,
-		Image: i.cfg.FileMountInitContainer.Image,
-		Env: []corev1.EnvVar{
-			{
-				Name:  "SECRET_VALUE",
-				Value: base64.StdEncoding.EncodeToString(secretValue),
-			},
-		},
-		Command: []string{
-			"sh",
-			"-c",
-			fmt.Sprintf("printf \"%%s\" \"$SECRET_VALUE\" | base64 -d > \"%s\"", secretFilePath),
-		},
-		Resources: i.cfg.FileMountInitContainer.Resources,
-		VolumeMounts: []corev1.VolumeMount{
-			{
-				Name:      EmbeddedSecretsVolumeName,
-				ReadOnly:  false,
-				MountPath: EmbeddedSecretsMountPath,
-			},
-		},
-	}
-	pod.Spec.InitContainers = append(pod.Spec.InitContainers, secretInitContainer)
-
-	secretVolumeMount := corev1.VolumeMount{
-		Name:      EmbeddedSecretsVolumeName,
+	volumeMount := corev1.VolumeMount{
+		Name:      EmbeddedSecretsFileMountVolumeName,
 		ReadOnly:  true,
-		MountPath: EmbeddedSecretsMountPath,
+		MountPath: EmbeddedSecretsFileMountPath,
 	}
-	pod.Spec.InitContainers = AppendVolumeMounts(pod.Spec.InitContainers, secretVolumeMount)
-	pod.Spec.Containers = AppendVolumeMounts(pod.Spec.Containers, secretVolumeMount)
+	pod.Spec.InitContainers = AppendVolumeMounts(pod.Spec.InitContainers, volumeMount)
+	pod.Spec.Containers = AppendVolumeMounts(pod.Spec.Containers, volumeMount)
 
-	// Inject AWS secret-inject webhook annotations to mount the secret in a predictable location.
 	envVars := []corev1.EnvVar{
 		// Set environment variable to let the containers know where to find the mounted files.
 		{
 			Name:  SecretPathDefaultDirEnvVar,
-			Value: EmbeddedSecretsMountPath,
+			Value: EmbeddedSecretsFileMountPath,
 		},
 		// Sets an empty prefix to let the containers know the file names will match the secret keys as-is.
 		{
@@ -259,6 +249,71 @@ func (i EmbeddedSecretManagerInjector) injectAsFile(secretKey string, secretValu
 	pod.Spec.Containers = AppendEnvVars(pod.Spec.Containers, envVars...)
 }
 
+func (i EmbeddedSecretManagerInjector) getOrAppendFileMountInitContainer(pod *corev1.Pod) (*corev1.Container, bool /*exists*/) {
+	index := slices.IndexFunc(
+		pod.Spec.InitContainers,
+		func(c corev1.Container) bool { return c.Name == EmbeddedSecretsFileMountInitContainerName })
+	if index != -1 {
+		return &pod.Spec.InitContainers[index], true
+	}
+
+	pod.Spec.InitContainers = append(pod.Spec.InitContainers, corev1.Container{
+		Name:  EmbeddedSecretsFileMountInitContainerName,
+		Image: i.cfg.FileMountInitContainer.Image,
+		Command: []string{
+			"sh",
+			"-c",
+			// Script below expects an environment variable EmbeddedSecretsFileMountInitContainerEnvVariableName
+			// with contents in the following format:
+			// 		secret_name1=base64_encoded_secret_value1
+			// 		secret_name2=base64_encoded_secret_value2
+			//		...
+			//
+			// It base64-decodes each secret value and writes it to a separate
+			// file in the EmbeddedSecretsFileMountPath directory.
+			fmt.Sprintf(`
+				printf "%%s" "$%s" \
+				| awk '/^.+=/ {
+					i = index($0, "=");
+					name = substr($0, 0, i - 1);
+					value = substr($0, i + 1);
+					output_file = "%s/" name;
+					print value | "base64 -d > " output_file;
+				}'
+				`,
+				EmbeddedSecretsFileMountInitContainerEnvVariableName,
+				EmbeddedSecretsFileMountPath),
+		},
+		Resources: i.cfg.FileMountInitContainer.Resources,
+		VolumeMounts: []corev1.VolumeMount{
+			{
+				Name:      EmbeddedSecretsFileMountVolumeName,
+				ReadOnly:  false,
+				MountPath: EmbeddedSecretsFileMountPath,
+			},
+		},
+	})
+
+	return &pod.Spec.InitContainers[len(pod.Spec.InitContainers)-1], false
+}
+
+func appendSecretToFileMountInitContainer(initContainer *corev1.Container, secretKey string, secretValue []byte) {
+	var envVar *corev1.EnvVar
+	index := slices.IndexFunc(
+		initContainer.Env,
+		func(env corev1.EnvVar) bool { return env.Name == EmbeddedSecretsFileMountInitContainerEnvVariableName })
+	if index != -1 {
+		envVar = &initContainer.Env[index]
+	} else {
+		initContainer.Env = append(initContainer.Env, corev1.EnvVar{
+			Name: EmbeddedSecretsFileMountInitContainerEnvVariableName,
+		})
+		envVar = &initContainer.Env[len(initContainer.Env)-1]
+	}
+
+	envVar.Value += fmt.Sprintf("%s=%s\n", secretKey, base64.StdEncoding.EncodeToString(secretValue))
+}
+
 func NewEmbeddedSecretManagerInjector(cfg config.EmbeddedSecretManagerConfig, secretFetcher SecretFetcher) SecretsInjector {
 	return EmbeddedSecretManagerInjector{
 		cfg:           cfg,

flytepropeller/pkg/secret/embedded_secret_manager_test.go

@@ -7,6 +7,7 @@ import (
 
 	"cloud.google.com/go/secretmanager/apiv1/secretmanagerpb"
 	"github.com/go-test/deep"
+	"github.com/samber/lo"
 	"github.com/stretchr/testify/assert"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -178,3 +179,57 @@ func TestEmbeddedSecretManagerInjector_Inject(t *testing.T) {
 		})
 	}
 }
+
+func TestEmbeddedSecretManagerInjector_InjectAsFile(t *testing.T) {
+	ctx = context.Background()
+
+	secret := &core.Secret{
+		Key:              "secret1",
+		MountRequirement: core.Secret_FILE,
+	}
+
+	pod := &corev1.Pod{
+		Spec: corev1.PodSpec{},
+		ObjectMeta: metav1.ObjectMeta{
+			Labels: map[string]string{
+				"organization": "organization",
+				"project":      "project",
+				"domain":       "domain",
+			},
+		},
+	}
+
+	injector := NewEmbeddedSecretManagerInjector(
+		config.EmbeddedSecretManagerConfig{},
+		secretFetcherMock{
+			Secrets: map[string]SecretValue{
+				"u__org__organization__domain__domain__project__project__key__secret1": {
+					BinaryValue: []byte("banana"),
+				},
+			},
+		})
+
+	pod, injected, err := injector.Inject(ctx, secret, pod)
+	assert.NoError(t, err)
+	assert.True(t, injected)
+	assert.Len(t, pod.Spec.InitContainers, 1)
+
+	env, found := lo.Find(
+		pod.Spec.InitContainers[0].Env,
+		func(env corev1.EnvVar) bool { return env.Name == "SECRETS" })
+	assert.True(t, found)
+	assert.Equal(t, "secret1=YmFuYW5h\n", env.Value)
+}
+
+type secretFetcherMock struct {
+	Secrets map[string]SecretValue
+}
+
+func (f secretFetcherMock) GetSecretValue(ctx context.Context, secretID string) (*SecretValue, error) {
+	v, ok := f.Secrets[secretID]
+	if !ok {
+		return nil, fmt.Errorf("secret %q not found", secretID)
+	}
+
+	return &v, nil
+}