secrets/hashivault: Add support for VAULT_ADDR and VAULT_TOKEN as aliases for existing env variables

Author: stack72

secrets/hashivault/vault.go

@@ -19,8 +19,8 @@
 // URLs
 //
 // For secrets.OpenKeeper, hashivault registers for the scheme "hashivault".
-// The default URL opener will dial a Vault server using the environment
-// variables "VAULT_SERVER_URL" and "VAULT_SERVER_TOKEN".
+// The default URL opener will dial a Vault server using the environment variables
+// "VAULT_SERVER_URL" (or "VAULT_ADDR") and "VAULT_SERVER_TOKEN" (or "VAULT_TOKEN").
 // To customize the URL opener, or for more details on the URL format,
 // see URLOpener.
 // See https://gocloud.dev/concepts/urls/ for background information.
@@ -74,8 +74,41 @@ func init() {
 	secrets.DefaultURLMux().RegisterKeeper(Scheme, new(defaultDialer))
 }
 
+// getVaultURL ensures that we check both VAULT_SERVER_URL and VAULT_ADDR environment
+// variables for the API address for vault. VAULT_SERVER_URL takes precedence over VAULT_ADDR.
+func getVaultURL() (string, error) {
+	serverURL := os.Getenv("VAULT_SERVER_URL")
+	if serverURL != "" {
+		return serverURL, nil
+	}
+
+	vaultAddr := os.Getenv("VAULT_ADDR")
+	if vaultAddr != "" {
+		return vaultAddr, nil
+	}
+
+	return "", errors.New("neither VAULT_SERVER_URL nor VAULT_ADDR environment variables are set")
+}
+
+// getVaultToken ensures that we check both VAULT_SERVER_TOKEN and VAULT_TOKEN environment
+// variables for the API token for vault. VAULT_SERVER_TOKEN takes precedence over VAULT_TOKEN.
+// If neither environment variables are found, then we return an empty string as token is not required.
+func getVaultToken() string {
+	serverToken := os.Getenv("VAULT_SERVER_TOKEN")
+	if serverToken != "" {
+		return serverToken
+	}
+
+	vaultToken := os.Getenv("VAULT_TOKEN")
+	if vaultToken != "" {
+		return vaultToken
+	}
+
+	return ""
+}
+
 // defaultDialer dials a default Vault server based on the environment variables
-// VAULT_SERVER_URL and VAULT_SERVER_TOKEN.
+// VAULT_SERVER_URL / VAULT_ADDR and VAULT_SERVER_TOKEN / VAULT_TOKEN
 type defaultDialer struct {
 	init   sync.Once
 	opener *URLOpener
@@ -84,12 +117,12 @@ type defaultDialer struct {
 
 func (o *defaultDialer) OpenKeeperURL(ctx context.Context, u *url.URL) (*secrets.Keeper, error) {
 	o.init.Do(func() {
-		serverURL := os.Getenv("VAULT_SERVER_URL")
-		if serverURL == "" {
-			o.err = errors.New("VAULT_SERVER_URL environment variable is not set")
+		serverURL, err := getVaultURL()
+		if err != nil {
+			o.err = err
 			return
 		}
-		token := os.Getenv("VAULT_SERVER_TOKEN") // token is not required
+		token := getVaultToken()
 		cfg := Config{Token: token, APIConfig: api.Config{Address: serverURL}}
 		client, err := Dial(ctx, &cfg)
 		if err != nil {

secrets/hashivault/vault_test.go

@@ -71,11 +71,11 @@ func newHarness(ctx context.Context, t *testing.T) (drivertest.Harness, error) {
 	enableTransit.Do(func() {
 		s, err := c.Logical().Read("sys/mounts")
 		if err != nil {
-			t.Fatal(err, "; run secrets/vault/localvault.sh to start a dev vault container")
+			t.Fatal(err, "; run secrets/hashivault/localvault.sh to start a dev vault container")
 		}
 		if _, ok := s.Data["transit/"]; !ok {
 			if _, err := c.Logical().Write("sys/mounts/transit", map[string]interface{}{"type": "transit"}); err != nil {
-				t.Fatal(err, "; run secrets/vault/localvault.sh to start a dev vault container")
+				t.Fatal(err, "; run secrets/hashivault/localvault.sh to start a dev vault container")
 			}
 		}
 	})
@@ -144,6 +144,34 @@ func fakeConnectionStringInEnv() func() {
 	}
 }
 
+func alternativeConnectionStringEnvVars() func() {
+	oldURLVal := os.Getenv("VAULT_ADDR")
+	oldTokenVal := os.Getenv("VAULT_TOKEN")
+	os.Setenv("VAULT_ADDR", "http://myalternativevaultserver")
+	os.Setenv("VAULT_TOKEN", "faketoken2")
+	return func() {
+		os.Setenv("VAULT_ADDR", oldURLVal)
+		os.Setenv("VAULT_TOKEN", oldTokenVal)
+	}
+}
+
+func unsetConnectionStringEnvVars() func() {
+	oldURLVal := os.Getenv("VAULT_ADDR")
+	oldTokenVal := os.Getenv("VAULT_TOKEN")
+	oldServerURLVal := os.Getenv("VAULT_SERVER_URL")
+	oldServerTokenVal := os.Getenv("VAULT_SERVER_TOKEN")
+	os.Unsetenv("VAULT_ADDR")
+	os.Unsetenv("VAULT_TOKEN")
+	os.Unsetenv("VAULT_SERVER_URL")
+	os.Unsetenv("VAULT_SERVER_TOKEN")
+	return func() {
+		os.Setenv("VAULT_ADDR", oldURLVal)
+		os.Setenv("VAULT_SERVER_URL", oldServerURLVal)
+		os.Setenv("VAULT_TOKEN", oldTokenVal)
+		os.Setenv("VAULT_SERVER_TOKEN", oldServerTokenVal)
+	}
+}
+
 func TestOpenKeeper(t *testing.T) {
 	cleanup := fakeConnectionStringInEnv()
 	defer cleanup()
@@ -171,3 +199,55 @@ func TestOpenKeeper(t *testing.T) {
 		}
 	}
 }
+
+func TestGetVaultConnectionDetails(t *testing.T) {
+	t.Run("Test Current Env Vars", func(t *testing.T) {
+		cleanup := fakeConnectionStringInEnv()
+		defer cleanup()
+
+		serverUrl, err := getVaultURL()
+		if err != nil {
+			t.Errorf("got unexpected error: %v", err)
+		}
+		if serverUrl != "http://myvaultserver" {
+			t.Errorf("expected 'http://myvaultserver': got %q", serverUrl)
+		}
+
+		vaultToken := getVaultToken()
+		if vaultToken != "faketoken" {
+			t.Errorf("export 'faketoken': got %q", vaultToken)
+		}
+	})
+
+	t.Run("Test Alternative Env Vars", func(t *testing.T) {
+		cleanup := alternativeConnectionStringEnvVars()
+		defer cleanup()
+
+		serverUrl, err := getVaultURL()
+		if err != nil {
+			t.Errorf("got unexpected error: %v", err)
+		}
+		if serverUrl != "http://myalternativevaultserver" {
+			t.Errorf("export '': got %q", serverUrl)
+		}
+
+		vaultToken := getVaultToken()
+		if vaultToken != "faketoken2" {
+			t.Errorf("export 'faketoken2': got %q", vaultToken)
+		}
+	})
+	t.Run("Test Unset Env Vars Throws Error", func(t *testing.T) {
+		cleanup := unsetConnectionStringEnvVars()
+		defer cleanup()
+
+		serverUrl, err := getVaultURL()
+		if err == nil {
+			t.Errorf("expected error but got a url: %s", serverUrl)
+		}
+
+		vaultToken := getVaultToken()
+		if vaultToken != "" {
+			t.Errorf("export '': got %q", vaultToken)
+		}
+	})
+}