feat(nrawssdk): Convert AccessKeyId to AccountID

Author: cade-conklin

v3/integrations/nrawssdk-v2/example/main.go

@@ -39,11 +39,17 @@ func main() {
 	txn := app.StartTransaction("My sample transaction")
 
 	ctx := context.Background()
+
 	awsConfig, err := config.LoadDefaultConfig(ctx, func(awsConfig *config.LoadOptions) error {
 		// Instrument all new AWS clients with New Relic
-		nrawssdk.AppendMiddlewares(&awsConfig.APIOptions, nil)
+
 		return nil
 	})
+	creds, err := awsConfig.Credentials.Retrieve(ctx)
+	if err != nil {
+		log.Println("Warning couldn't get flags")
+	}
+	nrawssdk.AppendMiddlewares(&awsConfig.APIOptions, nil, creds)
 	if err != nil {
 		log.Fatal(err)
 	}

v3/integrations/nrawssdk-v2/nrawssdk.go

@@ -28,6 +28,7 @@ package nrawssdk
 
 import (
 	"context"
+	"fmt"
 	"net/url"
 	"strconv"
 	"strings"
@@ -39,12 +40,14 @@ import (
 	"github.com/aws/smithy-go/middleware"
 	smithymiddle "github.com/aws/smithy-go/middleware"
 	smithyhttp "github.com/aws/smithy-go/transport/http"
+	"github.com/newrelic/go-agent/v3/internal/awssupport"
 	"github.com/newrelic/go-agent/v3/newrelic"
 	"github.com/newrelic/go-agent/v3/newrelic/integrationsupport"
 )
 
 type nrMiddleware struct {
-	txn *newrelic.Transaction
+	txn   *newrelic.Transaction
+	creds aws.Credentials
 }
 
 type contextKey string
@@ -69,13 +72,19 @@ func (m nrMiddleware) deserializeMiddleware(stack *smithymiddle.Stack) error {
 		}
 
 		smithyRequest := in.Request.(*smithyhttp.Request)
-
 		// The actual http.Request is inside the smithyhttp.Request
 		httpRequest := smithyRequest.Request
 		serviceName := awsmiddle.GetServiceID(ctx)
 		operation := awsmiddle.GetOperationName(ctx)
 		region := awsmiddle.GetRegion(ctx)
 
+		creds := awsmiddle.GetSigningCredentials(ctx)
+		accountID, err := awssupport.AWSAccountIdFromAWSAccessKey(creds)
+		if err != nil {
+			accountID = ""
+			fmt.Println(err.Error())
+		}
+
 		var segment endable
 
 		if serviceName == "dynamodb" || serviceName == "DynamoDB" {
@@ -129,6 +138,7 @@ func (m nrMiddleware) deserializeMiddleware(stack *smithymiddle.Stack) error {
 				integrationsupport.AddAgentSpanAttribute(txn, newrelic.AttributeAWSElastSearchDomainEndpoint, httpRequest.URL.String()) // this way I don't have to pull it out of context
 			}
 			// Set additional span attributes
+			integrationsupport.AddAgentSpanAttribute(txn, newrelic.AttributeCloudAccountID, accountID) // setting account ID here, why do we only do this if it is an SQS service?
 			integrationsupport.AddAgentSpanAttribute(txn,
 				newrelic.AttributeResponseCode, strconv.Itoa(response.StatusCode))
 			integrationsupport.AddAgentSpanAttribute(txn,
@@ -150,8 +160,8 @@ func (m nrMiddleware) serializeMiddleware(stack *middleware.Stack) error {
 	return stack.Initialize.Add(middleware.InitializeMiddlewareFunc("NRSerializeMiddleware", func(
 		ctx context.Context, in middleware.InitializeInput, next middleware.InitializeHandler) (
 		out middleware.InitializeOutput, metadata middleware.Metadata, err error) {
-
 		serviceName := awsmiddle.GetServiceID(ctx)
+		ctx = awsmiddle.SetSigningCredentials(ctx, m.creds)
 		switch serviceName {
 		case "dynamodb", "DynamoDB":
 			ctx = context.WithValue(ctx, dynamodbInputKey, dynamoDBInputFromMiddlewareInput(in))
@@ -219,8 +229,8 @@ func (m nrMiddleware) serializeMiddleware(stack *middleware.Stack) error {
 //	if err != nil {
 //		log.Fatal(err)
 //	}
-func AppendMiddlewares(apiOptions *[]func(*smithymiddle.Stack) error, txn *newrelic.Transaction) {
-	m := nrMiddleware{txn: txn}
+func AppendMiddlewares(apiOptions *[]func(*smithymiddle.Stack) error, txn *newrelic.Transaction, creds aws.Credentials) {
+	m := nrMiddleware{txn: txn, creds: creds}
 	*apiOptions = append(*apiOptions, m.deserializeMiddleware)
 	*apiOptions = append(*apiOptions, m.serializeMiddleware)
 }

v3/integrations/nrawssdk-v2/nrawssdk_test.go

@@ -64,7 +64,7 @@ var fakeCreds = func() interface{} {
 
 func newConfig(ctx context.Context, txn *newrelic.Transaction) aws.Config {
 	cfg, _ := config.LoadDefaultConfig(ctx, func(o *config.LoadOptions) error {
-		AppendMiddlewares(&o.APIOptions, txn)
+		AppendMiddlewares(&o.APIOptions, txn, aws.Credentials{})
 		return nil
 	})
 	cfg.Credentials = fakeCreds.(aws.CredentialsProvider)

v3/internal/awssupport/awssupport.go

@@ -8,10 +8,14 @@ package awssupport
 
 import (
 	"context"
-	"github.com/newrelic/go-agent/v3/newrelic/integrationsupport"
+	"encoding/base32"
+	"fmt"
 	"net/http"
 	"reflect"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/newrelic/go-agent/v3/newrelic/integrationsupport"
+
 	newrelic "github.com/newrelic/go-agent/v3/newrelic"
 )
 
@@ -114,3 +118,31 @@ func EndSegment(ctx context.Context, resp *http.Response) {
 		segment.End()
 	}
 }
+
+func AWSAccountIdFromAWSAccessKey(creds aws.Credentials) (string, error) {
+	if creds.AccountID != "" {
+		return creds.AccountID, nil
+	}
+	if creds.AccessKeyID == "" {
+		return "", fmt.Errorf("no access key id found")
+	}
+	if len(creds.AccessKeyID) < 16 {
+		return "", fmt.Errorf("improper access key id format")
+	}
+	trimmedAccessKey := creds.AccessKeyID[4:]
+	decoded, err := base32.StdEncoding.DecodeString(trimmedAccessKey)
+	if err != nil {
+		return "", fmt.Errorf("error decoding access keys")
+	}
+	var bigEndian uint64
+	for i := 0; i < 6; i++ {
+		bigEndian = bigEndian << 8      // shift 8 bits left.  Most significant byte read in first (decoded[i])
+		bigEndian |= uint64(decoded[i]) // apply OR for current byte
+	}
+
+	mask := uint64(0x7fffffffff80)
+
+	num := (bigEndian & mask) >> 7 // apply mask and get rid of last 7 bytes from mask
+
+	return fmt.Sprintf("%d", num), nil
+}

v3/internal/awssupport/awssupport_test.go

@@ -10,6 +10,8 @@ import (
 	"net/http"
 	"strings"
 	"testing"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
 )
 
 func TestGetTableName(t *testing.T) {
@@ -84,3 +86,133 @@ func TestGetRequestID(t *testing.T) {
 		}
 	}
 }
+
+func TestAWSAccountIdFromAWSAccessKey(t *testing.T) {
+	tests := []struct {
+		name string // description of this test case
+		// Named input parameters for target function.
+		creds      aws.Credentials
+		want       string
+		wantErr    bool
+		wantErrStr string // error message returned
+	}{
+		{
+			name: "first test",
+			creds: aws.Credentials{
+				AccountID:   "",
+				AccessKeyID: "AKIASAWSR23456AWS357",
+			},
+			want:    "138954266361",
+			wantErr: false,
+		},
+		{
+			name: "AccountID already exists and access key exists. Should return AccountID immediately",
+			creds: aws.Credentials{
+				AccountID:   "123451234512",
+				AccessKeyID: "ASKDHA123457AKJFHAKS",
+			},
+			want:    "123451234512",
+			wantErr: false,
+		},
+		{
+			name: "AccountID already exists and access key exists with too short of length. Should return AccountID immediately",
+			creds: aws.Credentials{
+				AccountID:   "123451234512",
+				AccessKeyID: "a",
+			},
+			want:    "123451234512",
+			wantErr: false,
+		},
+		{
+			name: "AccountID already exists and access key exists with improper format. Should return AccountID immediately",
+			creds: aws.Credentials{
+				AccountID:   "123451234512",
+				AccessKeyID: "a a a.                      ",
+			},
+			want:    "123451234512",
+			wantErr: false,
+		},
+		{
+			name: "AccountID already exists and access key does not exist. Should return AccountID immediately",
+			creds: aws.Credentials{
+				AccountID: "123451234512",
+			},
+			want:    "123451234512",
+			wantErr: false,
+		},
+		{
+			name:       "AccountID does not exist and access key does not exist.  Should return an error",
+			creds:      aws.Credentials{},
+			want:       "",
+			wantErr:    true,
+			wantErrStr: "no access key id found",
+		},
+		{
+			name: "AccountID does not exist and access key is in an improper format. Should return an error",
+			creds: aws.Credentials{
+				AccessKeyID: "123asdfas",
+			},
+			want:       "",
+			wantErr:    true,
+			wantErrStr: "improper access key id format",
+		},
+		{
+			name: "AccountID does not exist and access key is in an improper format with only one character. Should return an error",
+			creds: aws.Credentials{
+				AccessKeyID: "a",
+			},
+			want:       "",
+			wantErr:    true,
+			wantErrStr: "improper access key id format",
+		},
+		{
+			name: "AccountID does not exist and access key is in an improper format for decoding.",
+			creds: aws.Credentials{
+				AccessKeyID: "a a a.                      ",
+			},
+			want:       "",
+			wantErr:    true,
+			wantErrStr: "error decoding access keys",
+		},
+		{
+			name: "AccountID does not exist and access key contains non base32 characters",
+			creds: aws.Credentials{
+				AccessKeyID: "AKIA1234567899876541",
+			},
+			want:       "",
+			wantErr:    true,
+			wantErrStr: "error decoding access keys",
+		},
+		{
+			name: "AccountID does not exist and access key contains non base32 characters and is too short in length",
+			creds: aws.Credentials{
+				AccessKeyID: "AKIA1818",
+			},
+			want:       "",
+			wantErr:    true,
+			wantErrStr: "improper access key id format",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, gotErr := AWSAccountIdFromAWSAccessKey(tt.creds)
+			if gotErr != nil {
+				if !tt.wantErr {
+					t.Errorf("AWSAccountIdFromAWSAccessKey() failed: %v", gotErr)
+				} else {
+					if tt.wantErrStr != gotErr.Error() {
+						t.Errorf("AWSAccountIdFromAWSAccessKey() error = %v, want %v", gotErr.Error(), tt.wantErrStr)
+					}
+				}
+				return
+			}
+			if tt.wantErr {
+				t.Fatal("AWSAccountIdFromAWSAccessKey() succeeded unexpectedly")
+			}
+			// TODO: update the condition below to compare got with tt.want.
+			if tt.want != got {
+				t.Errorf("AWSAccountIdFromAWSAccessKey() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}