From 7b0243db8e60bfb7f98dede0da77802149ab125c Mon Sep 17 00:00:00 2001
From: aria <pradithya.pura@go-jek.com>
Date: Wed, 19 Apr 2023 21:38:35 +0800
Subject: [PATCH] Implement ability to specify additional/override annotations
 when using Vault Secret Manager (#556)

* Implement ability to specify additional annotations when using Vault secret manager

Signed-off-by: Pradithya Aria <pradithya.pura@go-jek.com>

* Infer GOOS and GOARCH from environment (#552)

Signed-off-by: Jeev B <jeevb@users.noreply.github.com>
Signed-off-by: Pradithya Aria <pradithya.pura@go-jek.com>

* fix makefile to read variables from environment and overrides (#554)

Signed-off-by: Jeev B <jeevb@users.noreply.github.com>
Signed-off-by: Pradithya Aria <pradithya.pura@go-jek.com>

* Remove BarrierTick (#545)

* removed barrier logic

Signed-off-by: Daniel Rammer <daniel@union.ai>

* deprecated TransitionTypeBarrier

Signed-off-by: Daniel Rammer <daniel@union.ai>

* removed barrier tests

Signed-off-by: Daniel Rammer <daniel@union.ai>

* bumping flyteplugins

Signed-off-by: Daniel Rammer <daniel@union.ai>

---------

Signed-off-by: Daniel Rammer <daniel@union.ai>
Signed-off-by: Pradithya Aria <pradithya.pura@go-jek.com>

* Check for TerminateExecution error and eat Precondition status (#553)

* Check for TerminateExecution error and eat Precondition status

Signed-off-by: Haytham Abuelfutuh <haytham@afutuh.com>

* lint

Signed-off-by: Haytham Abuelfutuh <haytham@afutuh.com>

---------

Signed-off-by: Haytham Abuelfutuh <haytham@afutuh.com>
Signed-off-by: Pradithya Aria <pradithya.pura@go-jek.com>

* Rename to annotation

Signed-off-by: Pradithya Aria <pradithya.pura@go-jek.com>

* Inline merging annotations

Signed-off-by: Pradithya Aria <pradithya.pura@go-jek.com>

---------

Signed-off-by: Pradithya Aria <pradithya.pura@go-jek.com>
Signed-off-by: Jeev B <jeevb@users.noreply.github.com>
Signed-off-by: Daniel Rammer <daniel@union.ai>
Signed-off-by: Haytham Abuelfutuh <haytham@afutuh.com>
Co-authored-by: Jeev B <jeevb@users.noreply.github.com>
Co-authored-by: Dan Rammer <daniel@union.ai>
Co-authored-by: Haytham Abuelfutuh <haytham@afutuh.com>
---
 flytepropeller/pkg/webhook/config/config.go   |   5 +-
 .../pkg/webhook/vault_secret_manager.go       |   3 +-
 .../pkg/webhook/vault_secret_manager_test.go  | 104 +++++++++++++++++-
 3 files changed, 103 insertions(+), 9 deletions(-)

diff --git a/flytepropeller/pkg/webhook/config/config.go b/flytepropeller/pkg/webhook/config/config.go
index c3cf9f2d8..b4a68702b 100644
--- a/flytepropeller/pkg/webhook/config/config.go
+++ b/flytepropeller/pkg/webhook/config/config.go
@@ -113,8 +113,9 @@ type GCPSecretManagerConfig struct {
 }
 
 type VaultSecretManagerConfig struct {
-	Role      string    `json:"role" pflag:",Specifies the vault role to use"`
-	KVVersion KVVersion `json:"kvVersion" pflag:"-,The KV Engine Version. Defaults to 2. Use 1 for unversioned secrets. Refer to - https://www.vaultproject.io/docs/secrets/kv#kv-secrets-engine."`
+	Role        string            `json:"role" pflag:",Specifies the vault role to use"`
+	KVVersion   KVVersion         `json:"kvVersion" pflag:"-,The KV Engine Version. Defaults to 2. Use 1 for unversioned secrets. Refer to - https://www.vaultproject.io/docs/secrets/kv#kv-secrets-engine."`
+	Annotations map[string]string `json:"annotations" pflag:"-,Annotation to be added to user task pod. The annotation can also be used to override default annotations added by Flyte. Useful to customize Vault integration (https://developer.hashicorp.com/vault/docs/platform/k8s/injector/annotations)"`
 }
 
 func GetConfig() *Config {
diff --git a/flytepropeller/pkg/webhook/vault_secret_manager.go b/flytepropeller/pkg/webhook/vault_secret_manager.go
index 841e3f45c..6ac2281f8 100644
--- a/flytepropeller/pkg/webhook/vault_secret_manager.go
+++ b/flytepropeller/pkg/webhook/vault_secret_manager.go
@@ -74,8 +74,7 @@ func (i VaultSecretManagerInjector) Inject(ctx context.Context, secret *coreIdl.
 			return p, false, err
 		}
 
-		p.ObjectMeta.Annotations = utils.UnionMaps(p.ObjectMeta.Annotations, commonVaultAnnotations)
-		p.ObjectMeta.Annotations = utils.UnionMaps(p.ObjectMeta.Annotations, secretVaultAnnotations)
+		p.ObjectMeta.Annotations = utils.UnionMaps(secretVaultAnnotations, commonVaultAnnotations, i.cfg.Annotations, p.ObjectMeta.Annotations)
 
 	case coreIdl.Secret_ENV_VAR:
 		return p, false, fmt.Errorf("Env_Var is not a supported mount requirement for Vault Secret Manager")
diff --git a/flytepropeller/pkg/webhook/vault_secret_manager_test.go b/flytepropeller/pkg/webhook/vault_secret_manager_test.go
index 9af9b6dd4..fb4148aa9 100644
--- a/flytepropeller/pkg/webhook/vault_secret_manager_test.go
+++ b/flytepropeller/pkg/webhook/vault_secret_manager_test.go
@@ -82,11 +82,69 @@ func ExpectedKVv2(uuid string) *corev1.Pod {
 	return expected
 }
 
-func NewInputPod() *corev1.Pod {
+func ExpectedKVv3(uuid string) *corev1.Pod {
+	// Injects uuid into expected output for KV v2 secrets
+	expected := &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Annotations: map[string]string{
+				"vault.hashicorp.com/agent-inject":                                "true",
+				"vault.hashicorp.com/secret-volume-path":                          "/etc/flyte/secrets",
+				"vault.hashicorp.com/role":                                        "flyte",
+				"vault.hashicorp.com/agent-pre-populate-only":                     "true",
+				fmt.Sprintf("vault.hashicorp.com/agent-inject-secret-%s", uuid):   "foo",
+				fmt.Sprintf("vault.hashicorp.com/agent-inject-file-%s", uuid):     "foo/bar",
+				fmt.Sprintf("vault.hashicorp.com/agent-inject-template-%s", uuid): `{{- with secret "foo" -}}{{ .Data.data.bar }}{{- end -}}`,
+				"vault.hashicorp.com/auth-config-type":                            "gce",
+			},
+		},
+		Spec: PodSpec,
+	}
+	return expected
+}
+
+func ExpectedKVv4(uuid string) *corev1.Pod {
+	// Injects uuid into expected output for KV v2 secrets
+	expected := &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Annotations: map[string]string{
+				"vault.hashicorp.com/agent-inject":                                "true",
+				"vault.hashicorp.com/secret-volume-path":                          "/etc/flyte/secrets",
+				"vault.hashicorp.com/role":                                        "my-role",
+				"vault.hashicorp.com/agent-pre-populate-only":                     "true",
+				fmt.Sprintf("vault.hashicorp.com/agent-inject-secret-%s", uuid):   "foo",
+				fmt.Sprintf("vault.hashicorp.com/agent-inject-file-%s", uuid):     "foo/bar",
+				fmt.Sprintf("vault.hashicorp.com/agent-inject-template-%s", uuid): `{{- with secret "foo" -}}{{ .Data.data.bar }}{{- end -}}`,
+			},
+		},
+		Spec: PodSpec,
+	}
+	return expected
+}
+
+func ExpectedKVv5(uuid string) *corev1.Pod {
+	// Injects uuid into expected output for KV v2 secrets
+	expected := &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Annotations: map[string]string{
+				"vault.hashicorp.com/agent-inject":                                "true",
+				"vault.hashicorp.com/secret-volume-path":                          "/etc/flyte/secrets",
+				"vault.hashicorp.com/role":                                        "flyte",
+				"vault.hashicorp.com/agent-pre-populate-only":                     "false",
+				fmt.Sprintf("vault.hashicorp.com/agent-inject-secret-%s", uuid):   "foo",
+				fmt.Sprintf("vault.hashicorp.com/agent-inject-file-%s", uuid):     "foo/bar",
+				fmt.Sprintf("vault.hashicorp.com/agent-inject-template-%s", uuid): `{{- with secret "foo" -}}{{ .Data.data.bar }}{{- end -}}`,
+			},
+		},
+		Spec: PodSpec,
+	}
+	return expected
+}
+
+func NewInputPod(annotations map[string]string) *corev1.Pod {
 	// Need to create a new Pod for every test since annotations are otherwise appended to original reference object
 	p := &corev1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
-			Annotations: map[string]string{},
+			Annotations: annotations,
 		},
 		Spec: corev1.PodSpec{
 			Containers: []corev1.Container{
@@ -122,7 +180,7 @@ func TestVaultSecretManagerInjector_Inject(t *testing.T) {
 			args: args{
 				cfg:    config.VaultSecretManagerConfig{Role: "flyte", KVVersion: config.KVVersion1},
 				secret: inputSecret,
-				p:      NewInputPod(),
+				p:      NewInputPod(map[string]string{}),
 			},
 			want:    ExpectedKVv1,
 			wantErr: false,
@@ -132,17 +190,53 @@ func TestVaultSecretManagerInjector_Inject(t *testing.T) {
 			args: args{
 				cfg:    config.VaultSecretManagerConfig{Role: "flyte", KVVersion: config.KVVersion2},
 				secret: inputSecret,
-				p:      NewInputPod(),
+				p:      NewInputPod(map[string]string{}),
 			},
 			want:    ExpectedKVv2,
 			wantErr: false,
 		},
+		{
+			name: "KVv3 Secret - extra annotations",
+			args: args{
+				cfg: config.VaultSecretManagerConfig{Role: "flyte", KVVersion: config.KVVersion2, Annotations: map[string]string{
+					"vault.hashicorp.com/auth-config-type": "gce",
+				}},
+				secret: inputSecret,
+				p:      NewInputPod(map[string]string{}),
+			},
+			want:    ExpectedKVv3,
+			wantErr: false,
+		},
+		{
+			name: "KVv4 Secret - user override annotation",
+			args: args{
+				cfg:    config.VaultSecretManagerConfig{Role: "flyte", KVVersion: config.KVVersion2, Annotations: map[string]string{}},
+				secret: inputSecret,
+				p: NewInputPod(map[string]string{
+					"vault.hashicorp.com/role": "my-role",
+				}),
+			},
+			want:    ExpectedKVv4,
+			wantErr: false,
+		},
+		{
+			name: "KVv5 Secret - system override annotation",
+			args: args{
+				cfg: config.VaultSecretManagerConfig{Role: "flyte", KVVersion: config.KVVersion2, Annotations: map[string]string{
+					"vault.hashicorp.com/agent-pre-populate-only": "false", // override vault.hashicorp.com/agent-pre-populate-only
+				}},
+				secret: inputSecret,
+				p:      NewInputPod(map[string]string{}),
+			},
+			want:    ExpectedKVv5,
+			wantErr: false,
+		},
 		{
 			name: "Unsupported KV version",
 			args: args{
 				cfg:    config.VaultSecretManagerConfig{Role: "flyte", KVVersion: 3},
 				secret: inputSecret,
-				p:      NewInputPod(),
+				p:      NewInputPod(map[string]string{}),
 			},
 			want:    nil,
 			wantErr: true,