Fix regex and match group for env lookup akvs

cmd/azure-keyvault-env/environment.go

@@ -0,0 +1,67 @@
+package main
+
+import (
+	"fmt"
+	"regexp"
+	"strings"
+
+	"k8s.io/klog/v2"
+)
+
+const (
+	envLookupRegex = `^([a-zA-Z0-9]([a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?)@azurekeyvault(\?([a-zA-Z_][a-zA-Z0-9_\.]*)?)?$`
+)
+
+type EnvSecret struct {
+	akvsName string
+	query    string
+	index    int
+}
+
+func parseEnvVars(envVars []string) (map[string]EnvSecret, error) {
+	re := regexp.MustCompile(envLookupRegex)
+
+	result := make(map[string]EnvSecret)
+	for index, envVar := range envVars {
+		parts := strings.SplitN(envVar, "=", 2)
+		if len(parts) != 2 {
+			klog.ErrorS(fmt.Errorf("error splitting env var"), "env variable not properly formatted", "env", envVar)
+			continue
+		}
+
+		name := parts[0]
+		value := parts[1]
+		match := re.FindStringSubmatch(value)
+		if len(match) == 0 {
+			klog.V(4).InfoS("env variable not an azure key vault reference", "env", envVar)
+			continue
+		}
+
+		klog.V(4).InfoS("found env var to get azure key vault secret for", "env", name)
+
+		akvsName := match[1]
+		klog.V(4).InfoS("azure key vault secret name found", "akvsName", akvsName)
+
+		if akvsName == "" {
+			err := fmt.Errorf("error extracting secret name")
+			klog.ErrorS(err, "env variable not properly formatted", "env", name, "value", value)
+			return nil, err
+		}
+
+		var query string
+		if len(match) == 5 {
+			klog.V(4).InfoS("found query in env var", "env", name, "value", value, "query", query)
+			query = match[4]
+		} else {
+			query = ""
+		}
+
+		result[name] = EnvSecret{
+			akvsName: akvsName,
+			query:    query,
+			index:    index,
+		}
+	}
+
+	return result, nil
+}

cmd/azure-keyvault-env/environment_test.go

@@ -0,0 +1,70 @@
+package main
+
+import (
+	"testing"
+)
+
+func TestParseEnvVars(t *testing.T) {
+	envVars := []string{
+		"TEST_CERTIFICATE=certificate-inject@azurekeyvault?tls.crt",
+		"TEST_KEY=key-inject@azurekeyvault?tls.key",
+		"TEST_EXAMPLE=example-inject@azurekeyvault?example",
+		"TEST_EMPTY=empty-inject@azurekeyvault?",
+		"TEST_SECRET=secret-inject@azurekeyvault",
+		"TEST_SECRET_REGULAR=secret-no-inject",
+	}
+	expectedResult := map[string]EnvSecret{
+		"TEST_CERTIFICATE": {
+			akvsName: "certificate-inject",
+			query:    "tls.crt",
+			index:    0,
+		},
+		"TEST_KEY": {
+			akvsName: "key-inject",
+			query:    "tls.key",
+			index:    1,
+		},
+		"TEST_EXAMPLE": {
+			akvsName: "example-inject",
+			query:    "example",
+			index:    2,
+		},
+		"TEST_EMPTY": {
+			akvsName: "empty-inject",
+			query:    "",
+			index:    3,
+		},
+		"TEST_SECRET": {
+			akvsName: "secret-inject",
+			query:    "",
+			index:    4,
+		},
+	}
+
+	result, err := parseEnvVars(envVars)
+	if err != nil {
+		t.Errorf("Expected no error, but got %s", err)
+	}
+
+	if len(result) != len(expectedResult) {
+		t.Errorf("Expected length of result to be %d, but got %d", len(expectedResult), len(result))
+	}
+
+	for key, value := range result {
+		expectedValue, exists := expectedResult[key]
+		if !exists {
+			t.Errorf("Expected result to contain key %s", key)
+			continue
+		}
+
+		if value.akvsName != expectedValue.akvsName {
+			t.Errorf("Expected akvsName for key %s to be %s, but got %s", key, expectedValue.akvsName, value.akvsName)
+		}
+		if value.query != expectedValue.query {
+			t.Errorf("Expected query for key %s to be %s, but got %s", key, expectedValue.query, value.query)
+		}
+		if value.index != expectedValue.index {
+			t.Errorf("Expected index for key %s to be %d, but got %d", key, expectedValue.index, value.index)
+		}
+	}
+}

cmd/azure-keyvault-env/main.go

@@ -23,7 +23,6 @@ import (
 	"fmt"
 	"os"
 	"os/exec"
-	"regexp"
 	"strings"
 	"syscall"
 	"time"
@@ -42,10 +41,6 @@ import (
 	"k8s.io/klog/v2"
 )
 
-const (
-	envLookupRegex = `^([a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*)@azurekeyvault([\?]?[a-z-\d\.]*)$`
-)
-
 type injectorConfig struct {
 	namespace                    string
 	podName                      string
@@ -286,73 +281,51 @@ func main() {
 
 	environ := os.Environ()
 
-	re := regexp.MustCompile(envLookupRegex)
-
-	for i, env := range environ {
-		split := strings.SplitN(env, "=", 2)
-		name := split[0]
-		value := split[1]
-
-		regexMatches := re.FindAllStringSubmatch(value, -1)
-
-		// e.g. my-akv-secret-name@azurekeyvault?some-sub-key
-		for _, match := range regexMatches {
-			// e.g. my-akv-secret-name?some-sub-key
-			klog.V(4).InfoS("found env var to get azure key vault secret for", "env", name)
-
-			akvsName := strings.Join(match[1:], "")
+	variables, err := parseEnvVars(environ)
+	if err != nil {
+		klog.ErrorS(err, "error extracting environment variables to inject")
+		os.Exit(1)
+	}
 
-			if akvsName == "" {
-				klog.ErrorS(fmt.Errorf("error extracting secret name"), "env variable not properly formatted", "env", name, "value", value)
-				os.Exit(1)
-			}
+	for name, values := range variables {
+		akvsName := values.akvsName
+		secretQuery := values.query
+		index := values.index
 
-			var secretQuery string
-			if query := strings.Split(akvsName, "?"); len(query) > 1 {
-				if len(query) > 2 {
-					klog.ErrorS(fmt.Errorf("error extracting secret query"), "multiple query elements defined with '?' - only one supported", "secret", akvsName)
-					os.Exit(1)
-				}
-				akvsName = query[0]
-				secretQuery = query[1]
-				klog.V(4).InfoS("found query in env var", "env", name, "value", value, "query", secretQuery)
-			}
+		klog.V(4).InfoS("getting azurekeyvaultsecret", "azurekeyvaultsecret", klog.KRef(config.namespace, akvsName))
+		akvs, err := azureKeyVaultSecretClient.AzureKeyVaultV2beta1().AzureKeyVaultSecrets(config.namespace).Get(context.TODO(), akvsName, v1.GetOptions{})
+		if err != nil {
+			klog.ErrorS(err, "failed to get azurekeyvaultsecret", "azurekeyvaultsecret", klog.KRef(config.namespace, akvsName))
+			klog.InfoS("will retry getting azurekeyvaultsecret", "azurekeyvaultsecret", klog.KRef(config.namespace, akvsName), "retryTimes", config.retryTimes, "delay", config.waitTimeBetweenRetries)
 
-			klog.V(4).InfoS("getting azurekeyvaultsecret", "azurekeyvaultsecret", klog.KRef(config.namespace, akvsName))
-			akvs, err := azureKeyVaultSecretClient.AzureKeyVaultV2beta1().AzureKeyVaultSecrets(config.namespace).Get(context.TODO(), akvsName, v1.GetOptions{})
-			if err != nil {
-				klog.ErrorS(err, "failed to get azurekeyvaultsecret", "azurekeyvaultsecret", klog.KRef(config.namespace, akvsName))
-				klog.InfoS("will retry getting azurekeyvaultsecret", "azurekeyvaultsecret", klog.KRef(config.namespace, akvsName), "retryTimes", config.retryTimes, "delay", config.waitTimeBetweenRetries)
-
-				err = retry(config.retryTimes, time.Second*time.Duration(config.waitTimeBetweenRetries), func() error {
-					akvs, err = azureKeyVaultSecretClient.AzureKeyVaultV2beta1().AzureKeyVaultSecrets(config.namespace).Get(context.TODO(), akvsName, v1.GetOptions{})
-					if err != nil {
-						klog.V(4).ErrorS(err, "error getting azurekeyvaultsecret", "azurekeyvaultsecret", klog.KRef(config.namespace, akvsName))
-						return err
-					}
-					klog.InfoS("succeeded getting azurekeyvaultsecret", "azurekeyvaultsecret", klog.KObj(akvs))
-					return nil
-				})
+			err = retry(config.retryTimes, time.Second*time.Duration(config.waitTimeBetweenRetries), func() error {
+				akvs, err = azureKeyVaultSecretClient.AzureKeyVaultV2beta1().AzureKeyVaultSecrets(config.namespace).Get(context.TODO(), akvsName, v1.GetOptions{})
 				if err != nil {
-					klog.ErrorS(err, "error getting azurekeyvaultsecret", "azurekeyvaultsecret", klog.KRef(config.namespace, akvsName))
-					os.Exit(1)
+					klog.V(4).ErrorS(err, "error getting azurekeyvaultsecret", "azurekeyvaultsecret", klog.KRef(config.namespace, akvsName))
+					return err
 				}
-			}
-
-			klog.V(4).InfoS("getting secret value for from azure key vault, to inject into env var", "azurekeyvaultsecret", klog.KObj(akvs), "env", name)
-			secret, err := getSecretFromKeyVault(akvs, secretQuery, vaultService)
+				klog.InfoS("succeeded getting azurekeyvaultsecret", "azurekeyvaultsecret", klog.KObj(akvs))
+				return nil
+			})
 			if err != nil {
-				klog.ErrorS(err, "failed to read secret from azure key vault", "azurekeyvaultsecret", klog.KObj(akvs))
+				klog.ErrorS(err, "error getting azurekeyvaultsecret", "azurekeyvaultsecret", klog.KRef(config.namespace, akvsName))
 				os.Exit(1)
 			}
+		}
 
-			if secret == "" {
-				klog.ErrorS(fmt.Errorf("secret value empty"), "secret not found in azure key vault", "azurekeyvaultsecret", klog.KObj(akvs))
-				os.Exit(1)
-			} else {
-				klog.InfoS("secret injected into env var", "azurekeyvaultsecret", klog.KObj(akvs), "env", name)
-				environ[i] = fmt.Sprintf("%s=%s", name, secret)
-			}
+		klog.V(4).InfoS("getting secret value for from azure key vault, to inject into env var", "azurekeyvaultsecret", klog.KObj(akvs), "env", name)
+		secret, err := getSecretFromKeyVault(akvs, secretQuery, vaultService)
+		if err != nil {
+			klog.ErrorS(err, "failed to read secret from azure key vault", "azurekeyvaultsecret", klog.KObj(akvs))
+			os.Exit(1)
+		}
+
+		if secret == "" {
+			klog.ErrorS(fmt.Errorf("secret value empty"), "secret not found in azure key vault", "azurekeyvaultsecret", klog.KObj(akvs))
+			os.Exit(1)
+		} else {
+			klog.InfoS("secret injected into env var", "azurekeyvaultsecret", klog.KObj(akvs), "env", name)
+			environ[index] = fmt.Sprintf("%s=%s", name, secret)
 		}
 	}