From d82a1f8ae5329767871a314738589ac4a63861ca Mon Sep 17 00:00:00 2001
From: Stephen Lewis <stephenrlewis@google.com>
Date: Fri, 21 Jan 2022 14:07:14 -0800
Subject: [PATCH] Allow extension of default user agent in GetConfig

---
 mmv1/third_party/validator/getconfig.go      | 18 +++++++++++++++++-
 mmv1/third_party/validator/getconfig_test.go | 18 +++++++++++++++++-
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/mmv1/third_party/validator/getconfig.go b/mmv1/third_party/validator/getconfig.go
index bd9fb3d56feb..a4a4755c8b78 100644
--- a/mmv1/third_party/validator/getconfig.go
+++ b/mmv1/third_party/validator/getconfig.go
@@ -2,13 +2,23 @@ package google
 
 import (
 	"context"
+	"fmt"
+	"os"
 
 	"github.com/pkg/errors"
+
+	"github.com/GoogleCloudPlatform/terraform-validator/version"
 )
 
+// Return the value of the private userAgent field
+func (c *Config) UserAgent() string {
+	return c.userAgent
+}
+
 func GetConfig(ctx context.Context, project string, offline bool) (*Config, error) {
 	cfg := &Config{
-		Project: project,
+		Project:   project,
+		userAgent: fmt.Sprintf("config-validator-tf/%s", version.BuildVersion()),
 	}
 
 	// Search for default credentials
@@ -26,6 +36,12 @@ func GetConfig(ctx context.Context, project string, offline bool) (*Config, erro
 		"GOOGLE_IMPERSONATE_SERVICE_ACCOUNT",
 	})
 
+	// opt in extension for adding to the User-Agent header
+	if ext := os.Getenv("GOOGLE_TERRAFORM_VALIDATOR_USERAGENT_EXTENSION"); ext != "" {
+		ua := cfg.userAgent
+		cfg.userAgent = fmt.Sprintf("%s %s", ua, ext)
+	}
+
 	if !offline {
 		ConfigureBasePaths(cfg)
 		if err := cfg.LoadAndValidate(ctx); err != nil {
diff --git a/mmv1/third_party/validator/getconfig_test.go b/mmv1/third_party/validator/getconfig_test.go
index 662f23176991..911b4abddb60 100644
--- a/mmv1/third_party/validator/getconfig_test.go
+++ b/mmv1/third_party/validator/getconfig_test.go
@@ -19,6 +19,9 @@ func getAccessToken(cfg *Config) string {
 func getImpersonateServiceAccount(cfg *Config) string {
 	return cfg.ImpersonateServiceAccount
 }
+func getUserAgent(cfg *Config) string {
+	return cfg.UserAgent()
+}
 
 func TestGetConfigExtractsEnvVars(t *testing.T) {
 	ctx := context.Background()
@@ -27,38 +30,51 @@ func TestGetConfigExtractsEnvVars(t *testing.T) {
 		name           string
 		envKey         string
 		envValue       string
+		expected       string
 		getConfigValue configAttrGetter
 	}{
 		{
 			name:           "GOOGLE_CREDENTIALS",
 			envKey:         "GOOGLE_CREDENTIALS",
 			envValue:       "whatever",
+			expected:       "whatever",
 			getConfigValue: getCredentials,
 		},
 		{
 			name:           "GOOGLE_CLOUD_KEYFILE_JSON",
 			envKey:         "GOOGLE_CLOUD_KEYFILE_JSON",
 			envValue:       "whatever",
+			expected:       "whatever",
 			getConfigValue: getCredentials,
 		},
 		{
 			name:           "GCLOUD_KEYFILE_JSON",
 			envKey:         "GCLOUD_KEYFILE_JSON",
 			envValue:       "whatever",
+			expected:       "whatever",
 			getConfigValue: getCredentials,
 		},
 		{
 			name:           "GOOGLE_OAUTH_ACCESS_TOKEN",
 			envKey:         "GOOGLE_OAUTH_ACCESS_TOKEN",
 			envValue:       "whatever",
+			expected:       "whatever",
 			getConfigValue: getAccessToken,
 		},
 		{
 			name:           "GOOGLE_IMPERSONATE_SERVICE_ACCOUNT",
 			envKey:         "GOOGLE_IMPERSONATE_SERVICE_ACCOUNT",
 			envValue:       "whatever",
+			expected:       "whatever",
 			getConfigValue: getImpersonateServiceAccount,
 		},
+		{
+			name:           "GOOGLE_TERRAFORM_VALIDATOR_USERAGENT_EXTENSION",
+			envKey:         "GOOGLE_TERRAFORM_VALIDATOR_USERAGENT_EXTENSION",
+			envValue:       "whatever",
+			expected:       "config-validator-tf/dev whatever",
+			getConfigValue: getUserAgent,
+		},
 	}
 
 	for _, c := range cases {
@@ -74,7 +90,7 @@ func TestGetConfigExtractsEnvVars(t *testing.T) {
 				t.Fatalf("error building converter: %s", err)
 			}
 
-			assert.EqualValues(t, c.getConfigValue(cfg), c.envValue)
+			assert.Equal(t, c.expected, c.getConfigValue(cfg))
 
 			if isSet {
 				err = os.Setenv(c.envKey, originalValue)