Repository: newrelic/newrelic-client-go
Pull request #854 (2022-03-01 00:02:16)

feat(auth): Add X-Account-ID header if value exists in request context
Closes https://github.com/newrelic/newrelic-client-go/issues/697

This PR adds an X-Account-ID header on REST requests so the same API key can be used across all account it has access to (not just it's primary account id).

To add the header, users should leverage the request's `context` and add an `X-Account-ID` key in the request's context.  This is done with a `contextkeys` package.

This work is meant to address an issue creating channels for child accounts with New Relic's [terraform-provider-newrelic](https://github.com/newrelic/terraform-provider-newrelic/issues/1645).


CLAassistant: [![CLA assistant check](https://cla-assistant.io/pull/badge/not_signed)](https://cla-assistant.io/newrelic/newrelic-client-go?pullRequest=854) <br/>Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our [Contributor License Agreement](https://cla-assistant.io/newrelic/newrelic-client-go?pullRequest=854) before we can accept your contribution.<br/><sub>You have signed the CLA already but the status is still pending? Let us [recheck](https://cla-assistant.io/check/newrelic/newrelic-client-go?pullRequest=854) it.</sub>

Patch:
diff --git a/internal/http/auth.go b/internal/http/auth.go
index 20363c7..7c54b3a 100644
--- a/internal/http/auth.go
+++ b/internal/http/auth.go
@@ -2,6 +2,7 @@ package http
 
 import (
 	"github.com/newrelic/newrelic-client-go/pkg/config"
+	"github.com/newrelic/newrelic-client-go/pkg/contextkeys"
 )
 
 // RequestAuthorizer is an interface that allows customizatino of how a request is authorized.
@@ -27,6 +28,10 @@ func (a *PersonalAPIKeyCapableV2Authorizer) AuthorizeRequest(r *Request, c *conf
 	} else {
 		r.SetHeader("X-Api-Key", c.AdminAPIKey)
 	}
+
+	if accountID, ok := contextkeys.GetAccountID(r.request.Context()); ok {
+		r.SetHeader("X-Account-ID", accountID)
+	}
 }
 
 // ClassicV2Authorizer authorizes V2 endpoints that cannot use a personal API key.
diff --git a/pkg/contextkeys/contextkeys.go b/pkg/contextkeys/contextkeys.go
new file mode 100644
index 0000000..fbe4606
--- /dev/null
+++ b/pkg/contextkeys/contextkeys.go
@@ -0,0 +1,43 @@
+package contextkeys
+
+import "context"
+
+// contextKeys gets and sets context values from a context.
+type contextKeys struct {
+	accountID struct{}
+}
+
+var (
+	keys = contextKeys{
+		accountID: struct{}{},
+	}
+)
+
+// SetAccountID inserts the account ID value into context.
+func SetAccountID(ctx context.Context, value string) context.Context {
+	return keys.SetAccountID(ctx, value)
+}
+
+// GetAccountID returns the account ID from the context.
+func GetAccountID(ctx context.Context) (string, bool) {
+	return keys.GetAccountID(ctx)
+}
+
+// SetAccountID inserts the account ID value into context.
+func (c contextKeys) SetAccountID(ctx context.Context, value string) context.Context {
+	if nil == ctx {
+		ctx = context.Background()
+	}
+
+	return context.WithValue(ctx, c.accountID, value)
+}
+
+// GetAccountID returns the account ID from the context.
+func (c contextKeys) GetAccountID(ctx context.Context) (string, bool) {
+	if nil != ctx {
+		accountID, ok := ctx.Value(keys.accountID).(string)
+		return accountID, ok
+	}
+
+	return "", false
+}
diff --git a/pkg/entities/types.go b/pkg/entities/types.go
index 50bd228..1808b51 100644
--- a/pkg/entities/types.go
+++ b/pkg/entities/types.go
@@ -40,7 +40,7 @@ var AccountStatusTypes = struct {
 	UPGRADED:               "UPGRADED",
 }
 
-// AgentApplicationSettingsBrowserLoader - Determines which Browser Loader will be configured. Some allowed return values are specified for backwards-compatability and do not represent currently allowed values for new applications.
+// AgentApplicationSettingsBrowserLoader - Determines which Browser Loader will be configured. Some allowed return values are specified for backwards-compatibility and do not represent currently allowed values for new applications.
 // See [documentation](https://docs.newrelic.com/docs/browser/browser-monitoring/installation/install-browser-monitoring-agent/#agent-types) for further information.
 type AgentApplicationSettingsBrowserLoader string
 
@@ -913,7 +913,7 @@ var DashboardEntityAlertStatusTypes = struct {
 	YELLOW: "YELLOW",
 }
 
-// DashboardEntityPermissions - Permisions that represent visibility & editability
+// DashboardEntityPermissions - Permissions that represent visibility & editability
 type DashboardEntityPermissions string
 
 var DashboardEntityPermissionsTypes = struct {
@@ -1062,7 +1062,7 @@ var DashboardMetricLineChartWidgetVisualizationTypeTypes = struct {
 	METRIC_LINE_CHART: "METRIC_LINE_CHART",
 }
 
-// DashboardPermissions - Permisions that represent visibility & editability
+// DashboardPermissions - Permissions that represent visibility & editability
 type DashboardPermissions string
 
 var DashboardPermissionsTypes = struct {
@@ -3071,7 +3071,7 @@ type AgentTracesErrorTrace struct {
 	Path string `json:"path"`
 	// Error stack trace.
 	StackTrace []AgentTracesStackTraceFrame `json:"stackTrace,omitempty"`
-	// When the error occured.
+	// When the error occurred.
 	StartTime *nrtime.EpochMilliseconds `json:"startTime"`
 	// URI.
 	Uri string `json:"uri,omitempty"`
@@ -3147,7 +3147,7 @@ type AgentTracesSqlTrace struct {
 	Sql string `json:"sql,omitempty"`
 	// An agent generated `sql_id`.
 	SqlId string `json:"sqlId"`
-	// When the SQL query occured.
+	// When the SQL query occurred.
 	StartTime *nrtime.EpochMilliseconds `json:"startTime"`
 	// Call time, as added across all `call_count` traces.
 	TotalCallTime Milliseconds `json:"totalCallTime,omitempty"`
@@ -4706,7 +4706,7 @@ type ApmApplicationDeployment struct {
 	Permalink string `json:"permalink,omitempty"`
 	// The revision of the app that was deployed
 	Revision string `json:"revision,omitempty"`
-	// The moment the deployment occured
+	// The moment the deployment occurred
 	Timestamp *nrtime.EpochMilliseconds `json:"timestamp,omitempty"`
 	// The user who triggered the deployment
 	User string `json:"user,omitempty"`
diff --git a/pkg/installevents/types.go b/pkg/installevents/types.go
index 73a8088..be607cb 100644
--- a/pkg/installevents/types.go
+++ b/pkg/installevents/types.go
@@ -214,7 +214,7 @@ type InstallationRecipeStatus struct {
 	KernelVersion string `json:"kernelVersion"`
 	// The path to the log file on the customer's host.
 	LogFilePath string `json:"logFilePath"`
-	// Additional key:value data related to an error or related to the environment where the installation ocurred.
+	// Additional key:value data related to an error or related to the environment where the installation occurred.
 	Metadata map[string]interface{} `json:"metadata,omitempty"`
 	// The unique name for a given recipe.
 	Name string `json:"name"`


Test patch:
diff --git a/internal/http/client_test.go b/internal/http/client_test.go
index c0f8454..d55f7af 100644
--- a/internal/http/client_test.go
+++ b/internal/http/client_test.go
@@ -4,12 +4,15 @@
 package http
 
 import (
+	"context"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
 	"testing"
 	"time"
 
+	"github.com/newrelic/newrelic-client-go/pkg/contextkeys"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -311,6 +314,29 @@ func TestCustomRequestHeaders(t *testing.T) {
 	assert.Nil(t, err)
 }
 
+func TestAccountIDHeaderWithPersonalAPIKeyCapableV2Authorizer(t *testing.T) {
+	// Given mock server
+	t.Parallel()
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Then X-Account-ID should be set in the header
+		assert.Equal(t, "custom-account-id", r.Header.Get("X-Account-ID"))
+	}))
+	tc := mock.NewTestConfig(t, ts)
+
+	// Given a client with PersonalAPIKeyCapableV2Authorizer Auth Strategy
+	c := NewClient(tc)
+	c.SetAuthStrategy(&PersonalAPIKeyCapableV2Authorizer{})
+
+	// When a request is made with context
+	req, err := c.NewRequest("GET", c.config.Region().RestURL("path"), nil, nil, nil)
+	ctx := contextkeys.SetAccountID(context.Background(), "custom-account-id")
+	req.WithContext(ctx)
+
+	// Then there are no errors with the request
+	_, err = c.Do(req)
+	assert.Nil(t, err)
+}
+
 func TestErrNotFound(t *testing.T) {
 	t.Parallel()
 	c := NewTestAPIClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
diff --git a/pkg/contextkeys/contextkeys_test.go b/pkg/contextkeys/contextkeys_test.go
new file mode 100644
index 0000000..5e48a11
--- /dev/null
+++ b/pkg/contextkeys/contextkeys_test.go
@@ -0,0 +1,79 @@
+//go:build unit
+// +build unit
+
+package contextkeys
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetAccountID_ReturnsAccountIDIfFound(t *testing.T) {
+	testCtx := context.WithValue(context.Background(), keys.accountID, "some-account-id")
+	accountID, ok := GetAccountID(testCtx)
+
+	assert.Equal(t, true, ok)
+	assert.Equal(t, "some-account-id", accountID)
+}
+
+func TestGetAccountID_DefaultReturnValueWhenNoAccountIDFound(t *testing.T) {
+	defaultValue, ok := GetAccountID(context.Background())
+
+	assert.Equal(t, false, ok)
+	assert.Equal(t, "", defaultValue)
+}
+
+func TestSetAccountID_SetsTheAccountID(t *testing.T) {
+	ctx := SetAccountID(context.Background(), "some-account-id")
+	result, ok := GetAccountID(ctx)
+
+	assert.Equal(t, true, ok)
+	assert.Equal(t, "some-account-id", result)
+}
+
+func TestSetAccountID_CreatesContextIfNil(t *testing.T) {
+	ctx := SetAccountID(nil, "some-account-id")
+	assert.NotNil(t, ctx)
+
+	result, ok := GetAccountID(ctx)
+	assert.Equal(t, true, ok)
+	assert.Equal(t, "some-account-id", result)
+}
+
+func TestContextKeys_SetAccountID_SetsValues(t *testing.T) {
+	x := contextKeys{}
+	ctx := x.SetAccountID(context.Background(), "some-account-id")
+	result, ok := x.GetAccountID(ctx)
+
+	assert.Equal(t, true, ok)
+	assert.Equal(t, "some-account-id", result)
+}
+
+func TestContextKeys_SetAccountID_CreatesContextIfNil(t *testing.T) {
+	x := contextKeys{}
+	ctx := x.SetAccountID(nil, "some-account-id")
+	assert.NotNil(t, ctx)
+
+	result, ok := x.GetAccountID(ctx)
+	assert.Equal(t, true, ok)
+	assert.Equal(t, "some-account-id", result)
+}
+
+func TestContextKeys_GetAccountID_ReturnsAccountIDIfFound(t *testing.T) {
+	x := contextKeys{}
+	testCtx := context.WithValue(context.Background(), keys.accountID, "some-account-id")
+	accountID, ok := x.GetAccountID(testCtx)
+
+	assert.Equal(t, true, ok)
+	assert.Equal(t, "some-account-id", accountID)
+}
+
+func TestContextKeys_GetAccountID_DefaultReturnValueWhenNoAccountIDFound(t *testing.T) {
+	x := contextKeys{}
+	defaultValue, ok := x.GetAccountID(context.Background())
+
+	assert.Equal(t, false, ok)
+	assert.Equal(t, "", defaultValue)
+}