From 27dbebb7be20a6a8de13dd0d6d792b06b42cb926 Mon Sep 17 00:00:00 2001
From: Su Shi <1684739+metacpp@users.noreply.github.com>
Date: Mon, 15 Apr 2019 18:48:07 -0700
Subject: [PATCH] Enable correlation request ID header in ARM API.

---
 azurerm/config.go                     | 23 +--------------
 azurerm/config_test.go                | 16 -----------
 azurerm/helpers/azure/request.go      | 41 +++++++++++++++++++++++++++
 azurerm/helpers/azure/request_test.go | 31 ++++++++++++++++++++
 4 files changed, 73 insertions(+), 38 deletions(-)
 delete mode 100644 azurerm/config_test.go
 create mode 100644 azurerm/helpers/azure/request.go
 create mode 100644 azurerm/helpers/azure/request_test.go

diff --git a/azurerm/config.go b/azurerm/config.go
index 2ad8c521b2e3..5b952ce52e2a 100644
--- a/azurerm/config.go
+++ b/azurerm/config.go
@@ -74,7 +74,6 @@ import (
 	"github.com/Azure/go-autorest/autorest/adal"
 	az "github.com/Azure/go-autorest/autorest/azure"
 	"github.com/hashicorp/go-azure-helpers/authentication"
-	"github.com/hashicorp/go-uuid"
 	"github.com/hashicorp/terraform/httpclient"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/helpers/azure"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/utils"
@@ -369,30 +368,10 @@ type ArmClient struct {
 	policySetDefinitionsClient policy.SetDefinitionsClient
 }
 
-var (
-	msClientRequestIDOnce sync.Once
-	msClientRequestID     string
-)
-
-// clientRequestID generates a UUID to pass through `x-ms-client-request-id` header.
-func clientRequestID() string {
-	msClientRequestIDOnce.Do(func() {
-		var err error
-		msClientRequestID, err = uuid.GenerateUUID()
-
-		if err != nil {
-			log.Printf("[WARN] Fail to generate uuid for msClientRequestID: %+v", err)
-		}
-	})
-
-	log.Printf("[DEBUG] AzureRM Client Request Id: %s", msClientRequestID)
-	return msClientRequestID
-}
-
 func (c *ArmClient) configureClient(client *autorest.Client, auth autorest.Authorizer) {
 	setUserAgent(client, c.partnerId)
 	client.Authorizer = auth
-	//client.RequestInspector = azure.WithClientID(clientRequestID())
+	client.RequestInspector = azure.WithCorrelationRequestID(azure.CorrelationRequestID())
 	client.Sender = azure.BuildSender()
 	client.SkipResourceProviderRegistration = c.skipProviderRegistration
 	client.PollingDuration = 60 * time.Minute
diff --git a/azurerm/config_test.go b/azurerm/config_test.go
deleted file mode 100644
index 83d7d86dd7a0..000000000000
--- a/azurerm/config_test.go
+++ /dev/null
@@ -1,16 +0,0 @@
-package azurerm
-
-import "testing"
-
-func TestClientRequestID(t *testing.T) {
-	first := clientRequestID()
-
-	if first == "" {
-		t.Fatal("no client request ID generated")
-	}
-
-	second := clientRequestID()
-	if first != second {
-		t.Fatal("subsequent request ID not the same as the first")
-	}
-}
diff --git a/azurerm/helpers/azure/request.go b/azurerm/helpers/azure/request.go
new file mode 100644
index 000000000000..878047dc2eb3
--- /dev/null
+++ b/azurerm/helpers/azure/request.go
@@ -0,0 +1,41 @@
+package azure
+
+import (
+	"log"
+	"sync"
+
+	"github.com/Azure/go-autorest/autorest"
+	"github.com/hashicorp/go-uuid"
+)
+
+const (
+	// HeaderCorrelationRequestID is the Azure extension header to set a user-specified correlation request ID.
+	HeaderCorrelationRequestID = "x-ms-correlation-request-id"
+)
+
+var (
+	msCorrelationRequestIDOnce sync.Once
+	msCorrelationRequestID     string
+)
+
+// WithCorrelationRequestID returns a PrepareDecorator that adds an HTTP extension header of
+// `x-ms-correlation-request-id` whose value is passed, undecorated UUID (e.g.,
+// `7F5A6223-F475-4A9C-B9D5-12575AA6B11B`).
+func WithCorrelationRequestID(uuid string) autorest.PrepareDecorator {
+	return autorest.WithHeader(HeaderCorrelationRequestID, uuid)
+}
+
+// CorrelationRequestID generates an UUID to pass through `x-ms-correlation-request-id` header.
+func CorrelationRequestID() string {
+	msCorrelationRequestIDOnce.Do(func() {
+		var err error
+		msCorrelationRequestID, err = uuid.GenerateUUID()
+
+		if err != nil {
+			log.Printf("[WARN] Fail to generate uuid for msCorrelationRequestID: %+v", err)
+		}
+	})
+
+	log.Printf("[DEBUG] AzureRM Correlation Request Id: %s", msCorrelationRequestID)
+	return msCorrelationRequestID
+}
diff --git a/azurerm/helpers/azure/request_test.go b/azurerm/helpers/azure/request_test.go
new file mode 100644
index 000000000000..410bf5167eeb
--- /dev/null
+++ b/azurerm/helpers/azure/request_test.go
@@ -0,0 +1,31 @@
+package azure
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/Azure/go-autorest/autorest"
+)
+
+func TestCorrelationRequestID(t *testing.T) {
+	first := CorrelationRequestID()
+
+	if first == "" {
+		t.Fatal("no correlation request ID generated")
+	}
+
+	second := CorrelationRequestID()
+	if first != second {
+		t.Fatal("subsequent correlation request ID not the same as the first")
+	}
+}
+
+func TestWithCorrelationRequestID(t *testing.T) {
+	uuid := CorrelationRequestID()
+	req, _ := autorest.Prepare(&http.Request{}, WithCorrelationRequestID(uuid))
+
+	if req.Header.Get(HeaderCorrelationRequestID) != uuid {
+		t.Fatalf("azure: withCorrelationRequestID failed to set %s -- expected %s, received %s",
+			HeaderCorrelationRequestID, uuid, req.Header.Get(HeaderCorrelationRequestID))
+	}
+}