Adding support for cloud-based auth

api/datadog/client.go

@@ -57,10 +57,29 @@ type Service struct {
 // SetAuthKeys sets the appropriate values in the headers parameter.
 func SetAuthKeys(ctx context.Context, headerParams *map[string]string, keys ...[2]string) {
 	if ctx != nil {
-		for _, key := range keys {
-			if auth, ok := ctx.Value(ContextAPIKeys).(map[string]APIKey); ok {
-				if apiKey, ok := auth[key[0]]; ok {
-					(*headerParams)[key[1]] = apiKey.Key
+		// If we have a cloud provider auth token and the route uses appKeyAuth, use that over the app key auth headers
+		hasCloudProviderAuth := false
+		if cloudProviderAuth, ok := ctx.Value(ContextCloudProvider).(*CloudProviderConfig); ok {
+			hasAppKeyAuth := false
+			for _, key := range keys {
+				if key[0] == "appKeyAuth" {
+					hasAppKeyAuth = true
+					break
+				}
+			}
+
+			if hasAppKeyAuth && cloudProviderAuth.DatadogToken != "" {
+				(*headerParams)["dd-auth-token"] = cloudProviderAuth.DatadogToken
+				hasCloudProviderAuth = true
+			}
+		}
+
+		if !hasCloudProviderAuth {
+			for _, key := range keys {
+				if auth, ok := ctx.Value(ContextAPIKeys).(map[string]APIKey); ok {
+					if apiKey, ok := auth[key[0]]; ok {
+						(*headerParams)[key[1]] = apiKey.Key
+					}
 				}
 			}
 		}
@@ -441,6 +460,20 @@ func (c *APIClient) Decode(v interface{}, b []byte, contentType string) (err err
 	return nil
 }
 
+func (c *APIClient) GetTokenUsingCloudProviderAuth(ctx context.Context) (*CloudProviderConfig, error) {
+	if cloudProvider, ok := ctx.Value(ContextCloudProvider).(*CloudProviderConfig); ok {
+		auth, err := cloudProvider.ProviderAuth.Authenticate()
+		if err != nil {
+			return nil, err
+		}
+		cloudProvider.OrgUUID = auth.OrgUUID
+		cloudProvider.CloudProviderProof = auth.CloudProviderProof
+		cloudProvider.DatadogToken = auth.DatadogToken
+		return cloudProvider, nil
+	}
+	return nil, fmt.Errorf("cloud provider auth not found in context")
+}
+
 // Add a file to the multipart request.
 func addFile(w *multipart.Writer, fieldName, path string) error {
 	file, err := os.Open(path)

api/datadog/configuration.go

@@ -40,6 +40,9 @@ var (
 	// ContextAPIKeys takes a string apikey as authentication for the request
 	ContextAPIKeys = contextKey("apiKeys")
 
+	// ContextCloudProvider takes a string cloud provider proof as authentication for the request
+	ContextCloudProvider = contextKey("cloudProvider")
+
 	// ContextHttpSignatureAuth takes HttpSignatureAuth as authentication for the request.
 	ContextHttpSignatureAuth = contextKey("httpsignature")
 
@@ -68,6 +71,24 @@ type APIKey struct {
 	Prefix string
 }
 
+type CloudProviderCredentials struct {
+	OrgUUID            string
+	DatadogToken       string
+	CloudProviderProof string
+}
+
+// CloudProviderConfig provides cloud provider based authentication to a request passed via context using ContextCloudProvider.
+type CloudProviderConfig struct {
+	CloudProviderCredentials
+	ProviderAuth CloudProviderAuth
+	Provider     string
+}
+
+// CloudProviderAuth is an interface for cloud provider authentication.
+type CloudProviderAuth interface {
+	Authenticate() (*CloudProviderCredentials, error)
+}
+
 // ServerVariable stores the information about a server variable.
 type ServerVariable struct {
 	Description  string

tests/api/client_test.go

@@ -0,0 +1,54 @@
+package api
+
+import (
+	"context"
+	"github.com/DataDog/datadog-api-client-go/v2/api/datadog"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+const FAKE_TOKEN = "fake-token"
+
+type FakeAWSAuth struct {
+	orgUUID string
+}
+
+func (f *FakeAWSAuth) Authenticate() (*datadog.CloudProviderCredentials, error) {
+	return &datadog.CloudProviderCredentials{
+		OrgUUID:            f.orgUUID,
+		DatadogToken:       FAKE_TOKEN,
+		CloudProviderProof: "proof",
+	}, nil
+}
+
+func TestAuthenticate(t *testing.T) {
+	awsAuth := FakeAWSAuth{
+		orgUUID: "1234",
+	}
+
+	testAuthCtx := context.WithValue(
+		context.Background(),
+		datadog.ContextCloudProvider,
+		&datadog.CloudProviderConfig{
+			ProviderAuth: &awsAuth,
+			Provider:     "aws",
+		},
+	)
+	config := datadog.NewConfiguration()
+	testAPIClient := datadog.NewAPIClient(config)
+
+	token, err := testAPIClient.GetTokenUsingCloudProviderAuth(testAuthCtx)
+	assert.Nil(t, err)
+	assert.Equal(t, token.DatadogToken, FAKE_TOKEN)
+
+	headers := map[string]string{}
+	datadog.SetAuthKeys(
+		testAuthCtx,
+		&headers,
+		[2]string{"apiKeyAuth", "DD-API-KEY"},
+		[2]string{"appKeyAuth", "DD-APPLICATION-KEY"},
+	)
+	assert.NotEmpty(t, headers)
+	assert.Equal(t, headers["dd-auth-token"], FAKE_TOKEN)
+
+}