From fcb8c7e518906fcff21624bc0d02c1ebd09601e2 Mon Sep 17 00:00:00 2001
From: Pradithya Aria <pradithya.pura@go-jek.com>
Date: Mon, 17 Apr 2023 16:28:54 +0800
Subject: [PATCH] Implement ability to specify additional annotations when
 using Vault secret manager

Signed-off-by: Pradithya Aria <pradithya.pura@go-jek.com>
---
 pkg/webhook/config/config.go             |  5 ++--
 pkg/webhook/vault_secret_manager.go      |  5 ++--
 pkg/webhook/vault_secret_manager_test.go | 32 ++++++++++++++++++++++++
 3 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/pkg/webhook/config/config.go b/pkg/webhook/config/config.go
index c3cf9f2d8..ab6e8722b 100644
--- a/pkg/webhook/config/config.go
+++ b/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."`
+	ExtraAnnotations map[string]string `json:"extraAnnotations" pflag:"-,Additional annotation to be added to the pod. Useful to further customize Vault integration (https://developer.hashicorp.com/vault/docs/platform/k8s/injector/annotations)"`
 }
 
 func GetConfig() *Config {
diff --git a/pkg/webhook/vault_secret_manager.go b/pkg/webhook/vault_secret_manager.go
index 841e3f45c..2ff9ff8f6 100644
--- a/pkg/webhook/vault_secret_manager.go
+++ b/pkg/webhook/vault_secret_manager.go
@@ -74,8 +74,9 @@ 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(i.cfg.ExtraAnnotations, p.ObjectMeta.Annotations)
+		p.ObjectMeta.Annotations = utils.UnionMaps(commonVaultAnnotations, p.ObjectMeta.Annotations)
+		p.ObjectMeta.Annotations = utils.UnionMaps(secretVaultAnnotations, 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/pkg/webhook/vault_secret_manager_test.go b/pkg/webhook/vault_secret_manager_test.go
index 9af9b6dd4..2e8b80a01 100644
--- a/pkg/webhook/vault_secret_manager_test.go
+++ b/pkg/webhook/vault_secret_manager_test.go
@@ -82,6 +82,26 @@ func ExpectedKVv2(uuid string) *corev1.Pod {
 	return expected
 }
 
+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 NewInputPod() *corev1.Pod {
 	// Need to create a new Pod for every test since annotations are otherwise appended to original reference object
 	p := &corev1.Pod{
@@ -137,6 +157,18 @@ func TestVaultSecretManagerInjector_Inject(t *testing.T) {
 			want:    ExpectedKVv2,
 			wantErr: false,
 		},
+		{
+			name: "KVv3 Secret - with extra annotations",
+			args: args{
+				cfg: config.VaultSecretManagerConfig{Role: "flyte", KVVersion: config.KVVersion2, ExtraAnnotations: map[string]string{
+					"vault.hashicorp.com/auth-config-type": "gce",
+				}},
+				secret: inputSecret,
+				p:      NewInputPod(),
+			},
+			want:    ExpectedKVv3,
+			wantErr: false,
+		},
 		{
 			name: "Unsupported KV version",
 			args: args{