From 6284f99708f2ad92b2e41a120c174575a4fe1402 Mon Sep 17 00:00:00 2001
From: Paul Tyng <paul@paultyng.net>
Date: Thu, 15 Mar 2018 09:25:57 -0400
Subject: [PATCH] Allow callers to append to user agent

---
 httpclient/useragent.go      | 16 +++++++++++-
 httpclient/useragent_test.go | 47 ++++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 1 deletion(-)
 create mode 100644 httpclient/useragent_test.go

diff --git a/httpclient/useragent.go b/httpclient/useragent.go
index d8db087cf..5e2801768 100644
--- a/httpclient/useragent.go
+++ b/httpclient/useragent.go
@@ -2,15 +2,29 @@ package httpclient
 
 import (
 	"fmt"
+	"log"
 	"net/http"
+	"os"
+	"strings"
 
 	"github.com/hashicorp/terraform/version"
 )
 
 const userAgentFormat = "Terraform/%s"
+const uaEnvVar = "TF_APPEND_USER_AGENT"
 
 func UserAgentString() string {
-	return fmt.Sprintf(userAgentFormat, version.Version)
+	ua := fmt.Sprintf(userAgentFormat, version.Version)
+
+	if add := os.Getenv(uaEnvVar); add != "" {
+		add = strings.TrimSpace(add)
+		if len(add) > 0 {
+			ua += " " + add
+			log.Printf("[DEBUG] Using modified User-Agent: %s", ua)
+		}
+	}
+
+	return ua
 }
 
 type userAgentRoundTripper struct {
diff --git a/httpclient/useragent_test.go b/httpclient/useragent_test.go
new file mode 100644
index 000000000..9c577da72
--- /dev/null
+++ b/httpclient/useragent_test.go
@@ -0,0 +1,47 @@
+package httpclient
+
+import (
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/hashicorp/terraform/version"
+)
+
+func TestUserAgentString_env(t *testing.T) {
+	expectedBase := fmt.Sprintf(userAgentFormat, version.Version)
+	if oldenv, isSet := os.LookupEnv(uaEnvVar); isSet {
+		defer os.Setenv(uaEnvVar, oldenv)
+	} else {
+		defer os.Unsetenv(uaEnvVar)
+	}
+
+	for i, c := range []struct {
+		expected   string
+		additional string
+	}{
+		{expectedBase, ""},
+		{expectedBase, " "},
+		{expectedBase, " \n"},
+
+		{fmt.Sprintf("%s test/1", expectedBase), "test/1"},
+		{fmt.Sprintf("%s test/2", expectedBase), "test/2 "},
+		{fmt.Sprintf("%s test/3", expectedBase), " test/3 "},
+		{fmt.Sprintf("%s test/4", expectedBase), "test/4 \n"},
+	} {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			if c.additional == "" {
+				os.Unsetenv(uaEnvVar)
+			} else {
+				os.Setenv(uaEnvVar, c.additional)
+			}
+
+			actual := UserAgentString()
+
+			if c.expected != actual {
+				t.Fatalf("Expected User-Agent '%s' does not match '%s'", c.expected, actual)
+			}
+		})
+	}
+
+}