From 364a0ceef29eb18c675301bc51128d6dcf5b4266 Mon Sep 17 00:00:00 2001
From: James Stoker <jstoker.github@jcs.me.uk>
Date: Fri, 4 Oct 2019 11:47:24 +0100
Subject: [PATCH] Add config parameter to Azure storage backend to allow
 specifying the ARM endpoint to support Azure Stack.

---
 physical/azure/azure.go                       | 28 ++++++++++++++---
 physical/azure/azure_test.go                  | 31 ++++++++++++-------
 .../docs/configuration/storage/azure.html.md  |  4 +++
 3 files changed, 46 insertions(+), 17 deletions(-)

diff --git a/physical/azure/azure.go b/physical/azure/azure.go
index 1f418891c9b8..4782e4425a8a 100644
--- a/physical/azure/azure.go
+++ b/physical/azure/azure.go
@@ -74,11 +74,29 @@ func NewAzureBackend(conf map[string]string, logger log.Logger) (physical.Backen
 			environmentName = "AzurePublicCloud"
 		}
 	}
-	environment, err := azure.EnvironmentFromName(environmentName)
-	if err != nil {
-		errorMsg := fmt.Sprintf("failed to look up Azure environment descriptor for name %q: {{err}}",
-			environmentName)
-		return nil, errwrap.Wrapf(errorMsg, err)
+	
+	environmentUrl := os.Getenv("AZURE_ARM_ENDPOINT")
+	if environmentUrl == "" {
+		environmentUrl = conf["arm_endpoint"]
+	}
+
+	var environment azure.Environment
+	var err error
+
+	if environmentUrl != "" {
+		environment, err = azure.EnvironmentFromURL(environmentUrl)
+		if err != nil {
+			errorMsg := fmt.Sprintf("failed to look up Azure environment descriptor for URL %q: {{err}}",
+			environmentUrl)
+			return nil, errwrap.Wrapf(errorMsg, err)
+		}
+	} else {
+		environment, err = azure.EnvironmentFromName(environmentName)
+		if err != nil {
+			errorMsg := fmt.Sprintf("failed to look up Azure environment descriptor for name %q: {{err}}",
+				environmentName)
+			return nil, errwrap.Wrapf(errorMsg, err)
+		}
 	}
 
 	client, err := storage.NewBasicClientOnSovereignCloud(accountName, accountKey, environment)
diff --git a/physical/azure/azure_test.go b/physical/azure/azure_test.go
index 8ec90b5d6bc3..6953091ab83f 100644
--- a/physical/azure/azure_test.go
+++ b/physical/azure/azure_test.go
@@ -16,9 +16,12 @@ import (
 	"github.com/hashicorp/vault/sdk/physical"
 )
 
-func environmentForCleanupClient(name string) (azure.Environment, error) {
+func environmentForCleanupClient(name string, armUrl string) (azure.Environment, error) {
+	if armUrl != "" {
+		return azure.EnvironmentFromURL(armUrl)
+	}
 	if name == "" {
-		return azure.EnvironmentFromName("AzurePublicCloud")
+		name = "AzurePublicCloud"
 	}
 	return azure.EnvironmentFromName(name)
 }
@@ -32,11 +35,12 @@ func TestAzureBackend(t *testing.T) {
 	accountName := os.Getenv("AZURE_ACCOUNT_NAME")
 	accountKey := os.Getenv("AZURE_ACCOUNT_KEY")
 	environmentName := os.Getenv("AZURE_ENVIRONMENT")
+	environmentUrl := os.Getenv("AZURE_ARM_ENDPOINT")
 
 	ts := time.Now().UnixNano()
 	name := fmt.Sprintf("vault-test-%d", ts)
 
-	cleanupEnvironment, err := environmentForCleanupClient(environmentName)
+	cleanupEnvironment, err := environmentForCleanupClient(environmentName, environmentUrl)
 	if err != nil {
 		t.Fatalf("err: %s", err)
 	}
@@ -46,10 +50,11 @@ func TestAzureBackend(t *testing.T) {
 	logger := logging.NewVaultLogger(log.Debug)
 
 	backend, err := NewAzureBackend(map[string]string{
-		"container":   name,
-		"accountName": accountName,
-		"accountKey":  accountKey,
-		"environment": environmentName,
+		"container":    name,
+		"accountName":  accountName,
+		"accountKey":   accountKey,
+		"environment":  environmentName,
+		"arm_endpoint": environmentUrl,
 	}, logger)
 
 	defer func() {
@@ -75,11 +80,12 @@ func TestAzureBackend_ListPaging(t *testing.T) {
 	accountName := os.Getenv("AZURE_ACCOUNT_NAME")
 	accountKey := os.Getenv("AZURE_ACCOUNT_KEY")
 	environmentName := os.Getenv("AZURE_ENVIRONMENT")
+	environmentUrl := os.Getenv("AZURE_ARM_ENDPOINT")
 
 	ts := time.Now().UnixNano()
 	name := fmt.Sprintf("vault-test-%d", ts)
 
-	cleanupEnvironment, err := environmentForCleanupClient(environmentName)
+	cleanupEnvironment, err := environmentForCleanupClient(environmentName, environmentUrl)
 	if err != nil {
 		t.Fatalf("err: %s", err)
 	}
@@ -89,10 +95,11 @@ func TestAzureBackend_ListPaging(t *testing.T) {
 	logger := logging.NewVaultLogger(log.Debug)
 
 	backend, err := NewAzureBackend(map[string]string{
-		"container":   name,
-		"accountName": accountName,
-		"accountKey":  accountKey,
-		"environment": environmentName,
+		"container":    name,
+		"accountName":  accountName,
+		"accountKey":   accountKey,
+		"environment":  environmentName,
+		"arm_endpoint": environmentUrl,
 	}, logger)
 
 	defer func() {
diff --git a/website/source/docs/configuration/storage/azure.html.md b/website/source/docs/configuration/storage/azure.html.md
index a3c4cf0432b8..150c7c6538c5 100644
--- a/website/source/docs/configuration/storage/azure.html.md
+++ b/website/source/docs/configuration/storage/azure.html.md
@@ -49,6 +49,10 @@ The current implementation is limited to a maximum of 4 megabytes per blob.
    environment the storage account belongs to by way of the case-insensitive
    name defined in the [Azure Go SDK][azure-environment].
 
+- `arm_endpoint` `(string: "")` - Specifies the cloud environment
+  the storage account belongs to by way of the Azure Resource Manager endpoint
+  URL.
+
 - `max_parallel` `(string: "128")` – Specifies The maximum number of concurrent
   requests to Azure.