feat: Extract resource arn and remote resource access key for cross-account support (#192)

**Description of changes:**

1. Add auto-instrumentation support for the following AWS resources.
- Populate `aws.stream.arn` in Span by extracting StreamArn from the
request body.
- Populate `aws.table.arn` in Span by extracting TableArn from the
response body.

2. Add auto-instrumentation support for remote resource access key.
- Populate `aws.auth.account.access_key` and `aws.auth.region` in Span
from STS credentials in client config
 
3. Generate cross-account metrics attributes when remote resource
identifier is present
- If remote resource arn is available, extract account id and region
from arn.
- Otherwise, pass account access key id and region from span to metric
if available.

4. Add unit tests, contract tests. Done E2E tests with CW agent. 

By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.

Author: blairhyy-amazon

aws-distro-opentelemetry-node-autoinstrumentation/src/aws-attribute-keys.ts

@@ -5,12 +5,17 @@ import { SEMATTRS_AWS_DYNAMODB_TABLE_NAMES } from '@opentelemetry/semantic-conve
 
 // Utility class holding attribute keys with special meaning to AWS components
 export const AWS_ATTRIBUTE_KEYS = {
+  AWS_AUTH_ACCOUNT_ACCESS_KEY: 'aws.auth.account.access_key',
+  AWS_AUTH_REGION: 'aws.auth.region',
   AWS_SPAN_KIND: 'aws.span.kind',
   AWS_LOCAL_SERVICE: 'aws.local.service',
   AWS_LOCAL_OPERATION: 'aws.local.operation',
   AWS_REMOTE_SERVICE: 'aws.remote.service',
   AWS_REMOTE_ENVIRONMENT: 'aws.remote.environment',
   AWS_REMOTE_OPERATION: 'aws.remote.operation',
+  AWS_REMOTE_RESOURCE_ACCOUNT_ACCESS_KEY: 'aws.remote.resource.account.access_key',
+  AWS_REMOTE_RESOURCE_ACCOUNT_ID: 'aws.remote.resource.account.id',
+  AWS_REMOTE_RESOURCE_REGION: 'aws.remote.resource.region',
   AWS_REMOTE_RESOURCE_TYPE: 'aws.remote.resource.type',
   AWS_REMOTE_RESOURCE_IDENTIFIER: 'aws.remote.resource.identifier',
   AWS_SDK_DESCENDANT: 'aws.sdk.descendant',
@@ -31,7 +36,9 @@ export const AWS_ATTRIBUTE_KEYS = {
   AWS_S3_BUCKET: 'aws.s3.bucket',
   AWS_SQS_QUEUE_URL: 'aws.sqs.queue.url',
   AWS_SQS_QUEUE_NAME: 'aws.sqs.queue.name',
+  AWS_KINESIS_STREAM_ARN: 'aws.kinesis.stream.arn',
   AWS_KINESIS_STREAM_NAME: 'aws.kinesis.stream.name',
+  AWS_DYNAMODB_TABLE_ARN: 'aws.dynamodb.table.arn',
   AWS_DYNAMODB_TABLE_NAMES: SEMATTRS_AWS_DYNAMODB_TABLE_NAMES,
   AWS_BEDROCK_DATA_SOURCE_ID: 'aws.bedrock.data_source.id',
   AWS_BEDROCK_KNOWLEDGE_BASE_ID: 'aws.bedrock.knowledge_base.id',

aws-distro-opentelemetry-node-autoinstrumentation/src/aws-metric-attribute-generator.ts

@@ -32,6 +32,7 @@ import {
   MetricAttributeGenerator,
   SERVICE_METRIC,
 } from './metric-attribute-generator';
+import { RegionalResourceArnParser } from './regional-resource-arn-parser';
 import { SqsUrlParser } from './sqs-url-parser';
 import { LAMBDA_APPLICATION_SIGNALS_REMOTE_ENVIRONMENT } from './aws-opentelemetry-configurator';
 
@@ -112,8 +113,20 @@ export class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
     AwsMetricAttributeGenerator.setService(resource, span, attributes);
     AwsMetricAttributeGenerator.setEgressOperation(span, attributes);
     AwsMetricAttributeGenerator.setRemoteServiceAndOperation(span, attributes);
-    AwsMetricAttributeGenerator.setRemoteResourceTypeAndIdentifier(span, attributes);
+    const isRemoteResourceIdentifierPresent = AwsMetricAttributeGenerator.setRemoteResourceTypeAndIdentifier(
+      span,
+      attributes
+    );
     AwsMetricAttributeGenerator.setRemoteEnvironment(span, attributes);
+    if (isRemoteResourceIdentifierPresent) {
+      const isAccountIdAndRegionPresent = AwsMetricAttributeGenerator.setRemoteResourceAccountIdAndRegion(
+        span,
+        attributes
+      );
+      if (!isAccountIdAndRegionPresent) {
+        AwsMetricAttributeGenerator.setRemoteResourceAccessKeyAndRegion(span, attributes);
+      }
+    }
     AwsMetricAttributeGenerator.setSpanKindForDependency(span, attributes);
     AwsMetricAttributeGenerator.setRemoteDbUser(span, attributes);
 
@@ -369,7 +382,7 @@ export class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
    * href="https://docs.aws.amazon.com/cloudcontrolapi/latest/userguide/supported-resources.html">AWS
    * Cloud Control resource format</a>.
    */
-  private static setRemoteResourceTypeAndIdentifier(span: ReadableSpan, attributes: Attributes): void {
+  private static setRemoteResourceTypeAndIdentifier(span: ReadableSpan, attributes: Attributes): boolean {
     let remoteResourceType: AttributeValue | undefined;
     let remoteResourceIdentifier: AttributeValue | undefined;
     let cloudFormationIdentifier: AttributeValue | undefined;
@@ -383,11 +396,25 @@ export class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
       ) {
         remoteResourceType = NORMALIZED_DYNAMO_DB_SERVICE_NAME + '::Table';
         remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(awsTableNames[0]);
+      } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_DYNAMODB_TABLE_ARN)) {
+        remoteResourceType = NORMALIZED_DYNAMO_DB_SERVICE_NAME + '::Table';
+        remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
+          RegionalResourceArnParser.extractDynamoDbTableNameFromArn(
+            span.attributes[AWS_ATTRIBUTE_KEYS.AWS_DYNAMODB_TABLE_ARN]
+          )
+        );
       } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_KINESIS_STREAM_NAME)) {
         remoteResourceType = NORMALIZED_KINESIS_SERVICE_NAME + '::Stream';
         remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
           span.attributes[AWS_ATTRIBUTE_KEYS.AWS_KINESIS_STREAM_NAME]
         );
+      } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_KINESIS_STREAM_ARN)) {
+        remoteResourceType = NORMALIZED_KINESIS_SERVICE_NAME + '::Stream';
+        remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
+          RegionalResourceArnParser.extractKinesisStreamNameFromArn(
+            span.attributes[AWS_ATTRIBUTE_KEYS.AWS_KINESIS_STREAM_ARN]
+          )
+        );
       } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_S3_BUCKET)) {
         remoteResourceType = NORMALIZED_S3_SERVICE_NAME + '::Bucket';
         remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
@@ -398,31 +425,31 @@ export class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
 
         remoteResourceType = NORMALIZED_SNS_SERVICE_NAME + '::Topic';
         remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
-          this.extractResourceNameFromArn(snsArn)
+          RegionalResourceArnParser.extractResourceNameFromArn(snsArn)
         );
         cloudFormationIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(snsArn);
       } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_SECRETSMANAGER_SECRET_ARN)) {
         const secretsArn = span.attributes[AWS_ATTRIBUTE_KEYS.AWS_SECRETSMANAGER_SECRET_ARN];
 
         remoteResourceType = NORMALIZED_SECRETSMANAGER_SERVICE_NAME + '::Secret';
         remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
-          this.extractResourceNameFromArn(secretsArn)
+          RegionalResourceArnParser.extractResourceNameFromArn(secretsArn)
         );
         cloudFormationIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(secretsArn);
       } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_STEPFUNCTIONS_STATEMACHINE_ARN)) {
         const stateMachineArn = span.attributes[AWS_ATTRIBUTE_KEYS.AWS_STEPFUNCTIONS_STATEMACHINE_ARN];
 
         remoteResourceType = NORMALIZED_STEPFUNCTIONS_SERVICE_NAME + '::StateMachine';
         remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
-          this.extractResourceNameFromArn(stateMachineArn)
+          RegionalResourceArnParser.extractResourceNameFromArn(stateMachineArn)
         );
         cloudFormationIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(stateMachineArn);
       } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_STEPFUNCTIONS_ACTIVITY_ARN)) {
         const activityArn = span.attributes[AWS_ATTRIBUTE_KEYS.AWS_STEPFUNCTIONS_ACTIVITY_ARN];
 
         remoteResourceType = NORMALIZED_STEPFUNCTIONS_SERVICE_NAME + '::Activity';
         remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
-          this.extractResourceNameFromArn(activityArn)
+          RegionalResourceArnParser.extractResourceNameFromArn(activityArn)
         );
         cloudFormationIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(activityArn);
       } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_FUNCTION_NAME)) {
@@ -500,7 +527,10 @@ export class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
       attributes[AWS_ATTRIBUTE_KEYS.AWS_REMOTE_RESOURCE_TYPE] = remoteResourceType;
       attributes[AWS_ATTRIBUTE_KEYS.AWS_REMOTE_RESOURCE_IDENTIFIER] = remoteResourceIdentifier;
       attributes[AWS_ATTRIBUTE_KEYS.AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER] = cloudFormationIdentifier;
+      return true;
     }
+
+    return false;
   }
 
   /**
@@ -522,6 +552,56 @@ export class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
     }
   }
 
+  private static setRemoteResourceAccountIdAndRegion(span: ReadableSpan, attributes: Attributes): boolean {
+    const ARN_ATTRIBUTES: string[] = [
+      AWS_ATTRIBUTE_KEYS.AWS_DYNAMODB_TABLE_ARN,
+      AWS_ATTRIBUTE_KEYS.AWS_KINESIS_STREAM_ARN,
+      AWS_ATTRIBUTE_KEYS.AWS_SNS_TOPIC_ARN,
+      AWS_ATTRIBUTE_KEYS.AWS_SECRETSMANAGER_SECRET_ARN,
+      AWS_ATTRIBUTE_KEYS.AWS_STEPFUNCTIONS_STATEMACHINE_ARN,
+      AWS_ATTRIBUTE_KEYS.AWS_STEPFUNCTIONS_ACTIVITY_ARN,
+      AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_FUNCTION_ARN,
+      AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ARN,
+    ];
+    let remoteResourceAccountId: string | undefined = undefined;
+    let remoteResourceRegion: string | undefined = undefined;
+
+    if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_SQS_QUEUE_URL)) {
+      const sqsQueueUrl = AwsMetricAttributeGenerator.escapeDelimiters(
+        span.attributes[AWS_ATTRIBUTE_KEYS.AWS_SQS_QUEUE_URL]
+      );
+      remoteResourceAccountId = SqsUrlParser.getAccountId(sqsQueueUrl);
+      remoteResourceRegion = SqsUrlParser.getRegion(sqsQueueUrl);
+    } else {
+      for (const attributeKey of ARN_ATTRIBUTES) {
+        if (AwsSpanProcessingUtil.isKeyPresent(span, attributeKey)) {
+          const arn = span.attributes[attributeKey];
+          remoteResourceAccountId = RegionalResourceArnParser.getAccountId(arn);
+          remoteResourceRegion = RegionalResourceArnParser.getRegion(arn);
+          break;
+        }
+      }
+    }
+
+    if (remoteResourceAccountId !== undefined && remoteResourceRegion !== undefined) {
+      attributes[AWS_ATTRIBUTE_KEYS.AWS_REMOTE_RESOURCE_ACCOUNT_ID] = remoteResourceAccountId;
+      attributes[AWS_ATTRIBUTE_KEYS.AWS_REMOTE_RESOURCE_REGION] = remoteResourceRegion;
+      return true;
+    }
+
+    return false;
+  }
+
+  private static setRemoteResourceAccessKeyAndRegion(span: ReadableSpan, attributes: Attributes): void {
+    if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_AUTH_ACCOUNT_ACCESS_KEY)) {
+      attributes[AWS_ATTRIBUTE_KEYS.AWS_REMOTE_RESOURCE_ACCOUNT_ACCESS_KEY] =
+        span.attributes[AWS_ATTRIBUTE_KEYS.AWS_AUTH_ACCOUNT_ACCESS_KEY];
+    }
+    if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_AUTH_REGION)) {
+      attributes[AWS_ATTRIBUTE_KEYS.AWS_REMOTE_RESOURCE_REGION] = span.attributes[AWS_ATTRIBUTE_KEYS.AWS_AUTH_REGION];
+    }
+  }
+
   /**
    * RemoteResourceIdentifier is populated with rule <code>
    *     ^[{db.name}|]?{address}[|{port}]?
@@ -649,16 +729,6 @@ export class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
     return rpcService === 'Lambda' && span.attributes[SEMATTRS_RPC_METHOD] === LAMBDA_INVOKE_OPERATION;
   }
 
-  // Extracts the name of the resource from an arn
-  private static extractResourceNameFromArn(attribute: AttributeValue | undefined): string | undefined {
-    if (typeof attribute === 'string' && attribute.startsWith('arn:aws:')) {
-      const split = attribute.split(':');
-      return split[split.length - 1];
-    }
-
-    return undefined;
-  }
-
   /** Span kind is needed for differentiating metrics in the EMF exporter */
   private static setSpanKindForService(span: ReadableSpan, attributes: Attributes): void {
     let spanKind: string = SpanKind[span.kind];

aws-distro-opentelemetry-node-autoinstrumentation/src/patches/instrumentation-patch.ts

@@ -81,6 +81,8 @@ export function applyInstrumentationPatches(instrumentations: Instrumentation[])
         patchSqsServiceExtension(services.get('SQS'));
         patchSnsServiceExtension(services.get('SNS'));
         patchLambdaServiceExtension(services.get('Lambda'));
+        patchKinesisServiceExtension(services.get('Kinesis'));
+        patchDynamoDbServiceExtension(services.get('DynamoDB'));
       }
     } else if (instrumentation.instrumentationName === '@opentelemetry/instrumentation-aws-lambda') {
       diag.debug('Patching aws lambda instrumentation');
@@ -190,6 +192,69 @@ function patchSnsServiceExtension(snsServiceExtension: any): void {
   }
 }
 
+/*
+ * This patch extends the existing upstream extension for Kinesis. Extensions allow for custom logic for adding
+ * service-specific information to spans, such as attributes. Specifically, we are adding logic to add
+ * `aws.kinesis.stream.arn` attribute, to be used to generate RemoteTarget and achieve parity with the Java/Python instrumentation.
+ *
+ *
+ * @param kinesisServiceExtension Kinesis Service Extension obtained the service extension list from the AWS SDK OTel Instrumentation
+ */
+function patchKinesisServiceExtension(kinesisServiceExtension: any): void {
+  if (kinesisServiceExtension) {
+    const requestPreSpanHook = kinesisServiceExtension.requestPreSpanHook;
+    kinesisServiceExtension._requestPreSpanHook = requestPreSpanHook;
+
+    const patchedRequestPreSpanHook = (
+      request: NormalizedRequest,
+      _config: AwsSdkInstrumentationConfig
+    ): RequestMetadata => {
+      const requestMetadata: RequestMetadata = kinesisServiceExtension._requestPreSpanHook(request, _config);
+      if (requestMetadata.spanAttributes) {
+        const streamArn = request.commandInput?.StreamARN;
+        if (streamArn) {
+          requestMetadata.spanAttributes[AWS_ATTRIBUTE_KEYS.AWS_KINESIS_STREAM_ARN] = streamArn;
+        }
+      }
+      return requestMetadata;
+    };
+
+    kinesisServiceExtension.requestPreSpanHook = patchedRequestPreSpanHook;
+  }
+}
+
+/*
+ * This patch extends the existing upstream extension for DynamoDB. Extensions allow for custom logic for adding
+ * service-specific information to spans, such as attributes. Specifically, we are adding logic to add
+ * `aws.dynamodb.table.arn` attribute, to be used to generate RemoteTarget and achieve parity with the Java/Python instrumentation.
+ *
+ *
+ * @param dynamoDbServiceExtension DynamoDB Service Extension obtained the service extension list from the AWS SDK OTel Instrumentation
+ */
+function patchDynamoDbServiceExtension(dynamoDbServiceExtension: any): void {
+  if (dynamoDbServiceExtension) {
+    if (typeof dynamoDbServiceExtension.responseHook === 'function') {
+      const responseHook = dynamoDbServiceExtension.responseHook;
+
+      const patchedResponseHook = (
+        response: NormalizedResponse,
+        span: Span,
+        tracer: Tracer,
+        config: AwsSdkInstrumentationConfig
+      ): void => {
+        responseHook.call(dynamoDbServiceExtension, response, span, tracer, config);
+
+        const tableArn = response?.data?.Table?.TableArn;
+        if (tableArn) {
+          span.setAttribute(AWS_ATTRIBUTE_KEYS.AWS_DYNAMODB_TABLE_ARN, tableArn);
+        }
+      };
+
+      dynamoDbServiceExtension.responseHook = patchedResponseHook;
+    }
+  }
+}
+
 /*
  * This patch extends the existing upstream extension for Lambda. Extensions allow for custom logic for adding
  * service-specific information to spans, such as attributes. Specifically, we are adding logic to add
@@ -346,7 +411,7 @@ function patchAwsLambdaInstrumentation(instrumentation: Instrumentation): void {
   }
 }
 
-// Override the upstream private _getV3SmithyClientSendPatch method to add middleware to inject X-Ray Trace Context into HTTP Headers
+// Override the upstream private _getV3SmithyClientSendPatch method to add middlewares to inject X-Ray Trace Context into HTTP Headers and to extract account access key id and region for cross-account support
 // https://github.com/open-telemetry/opentelemetry-js-contrib/blob/instrumentation-aws-sdk-v0.48.0/plugins/node/opentelemetry-instrumentation-aws-sdk/src/aws-sdk.ts#L373-L384
 const V3_CLIENT_CONFIG_KEY = Symbol('opentelemetry.instrumentation.aws-sdk.client.config');
 type V3PluginCommand = AwsV3Command<any, any, any, any, any> & {
@@ -380,6 +445,40 @@ function patchAwsSdkInstrumentation(instrumentation: Instrumentation): void {
           }
         );
 
+        this.middlewareStack?.add(
+          (next: any, context: any) => async (middlewareArgs: any) => {
+            const activeContext = otelContext.active();
+            const span = trace.getSpan(activeContext);
+
+            if (span) {
+              try {
+                const credsProvider = this.config.credentials;
+                if (credsProvider instanceof Function) {
+                  const credentials = await credsProvider();
+                  if (credentials?.accessKeyId) {
+                    span.setAttribute(AWS_ATTRIBUTE_KEYS.AWS_AUTH_ACCOUNT_ACCESS_KEY, credentials.accessKeyId);
+                  }
+                }
+                if (this.config.region instanceof Function) {
+                  const region = await this.config.region();
+                  if (region) {
+                    span.setAttribute(AWS_ATTRIBUTE_KEYS.AWS_AUTH_REGION, region);
+                  }
+                }
+              } catch (err) {
+                diag.debug('Failed to get auth account access key and region:', err);
+              }
+            }
+
+            return await next(middlewareArgs);
+          },
+          {
+            step: 'build',
+            name: '_adotExtractSignerCredentials',
+            override: true,
+          }
+        );
+
         command[V3_CLIENT_CONFIG_KEY] = this.config;
         return original.apply(this, [command, ...args]);
       };

aws-distro-opentelemetry-node-autoinstrumentation/src/regional-resource-arn-parser.ts

@@ -0,0 +1,38 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { AttributeValue } from '@opentelemetry/api';
+import { isAccountId } from './utils';
+
+export class RegionalResourceArnParser {
+  /** Parses ARN with formats:
+   * arn:partition:service:region:account-id:resource-type/resource-id or
+   * arn:partition:service:region:account-id:resource-type:resource-id
+   */
+  private static parseArn(arn: AttributeValue | undefined): string[] | undefined {
+    if (typeof arn !== 'string') return undefined;
+    const parts = arn.split(':');
+    return parts.length >= 6 && parts[0] === 'arn' && isAccountId(parts[4]) ? parts : undefined;
+  }
+
+  public static getAccountId(arn: AttributeValue | undefined): string | undefined {
+    return this.parseArn(arn)?.[4];
+  }
+
+  public static getRegion(arn: AttributeValue | undefined): string | undefined {
+    return this.parseArn(arn)?.[3];
+  }
+
+  public static extractDynamoDbTableNameFromArn(arn: AttributeValue | undefined): string | undefined {
+    return this.extractResourceNameFromArn(arn)?.replace('table/', '');
+  }
+
+  public static extractKinesisStreamNameFromArn(arn: AttributeValue | undefined): string | undefined {
+    return this.extractResourceNameFromArn(arn)?.replace('stream/', '');
+  }
+
+  public static extractResourceNameFromArn(arn: AttributeValue | undefined): string | undefined {
+    const parts = this.parseArn(arn);
+    return parts?.[parts.length - 1];
+  }
+}