Extract cross account information from STS credentials and arn

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: { [key: string]: string } = {
+  AWS_AUTH_ACCESS_KEY: 'aws.auth.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_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: { [key: string]: string } = {
   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';
 
 // Does not exist in @opentelemetry/semantic-conventions
@@ -69,6 +70,16 @@ const DB_CONNECTION_RESOURCE_TYPE: string = 'DB::Connection';
 // - `defaultServiceName()` returns `unknown_service:${process.argv0}`
 const OTEL_UNKNOWN_SERVICE: string = defaultServiceName();
 
+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,
+];
 /**
  * AwsMetricAttributeGenerator generates very specific metric attributes based on low-cardinality
  * span and resource attributes. If such attributes are not present, we fallback to default values.
@@ -108,7 +119,19 @@ 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
+    );
+    if (isRemoteResourceIdentifierPresent) {
+      const isAccountIdAndRegionPresent = AwsMetricAttributeGenerator.setRemoteResourceAccountIdAndRegion(
+        span,
+        attributes
+      );
+      if (!isAccountIdAndRegionPresent) {
+        AwsMetricAttributeGenerator.setRemoteResourceAccessKeyAndRegion(span, attributes);
+      }
+    }
     AwsMetricAttributeGenerator.setSpanKindForDependency(span, attributes);
     AwsMetricAttributeGenerator.setRemoteDbUser(span, attributes);
 
@@ -353,7 +376,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;
@@ -367,11 +390,27 @@ 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(
+          this.extractResourceNameFromArn(span.attributes[AWS_ATTRIBUTE_KEYS.AWS_DYNAMODB_TABLE_ARN])?.replace(
+            'table/',
+            ''
+          )
+        );
       } 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(
+          this.extractResourceNameFromArn(span.attributes[AWS_ATTRIBUTE_KEYS.AWS_KINESIS_STREAM_ARN])?.replace(
+            'stream/',
+            ''
+          )
+        );
       } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_S3_BUCKET)) {
         remoteResourceType = NORMALIZED_S3_SERVICE_NAME + '::Bucket';
         remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
@@ -502,6 +541,48 @@ export class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
 
         attributes[AWS_ATTRIBUTE_KEYS.AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER] = cloudFormationIdentifier;
       }
+      return true;
+    }
+    return false;
+  }
+
+  private static setRemoteResourceAccountIdAndRegion(span: ReadableSpan, attributes: Attributes): boolean {
+    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_ACCESS_KEY)) {
+      attributes[AWS_ATTRIBUTE_KEYS.AWS_REMOTE_RESOURCE_ACCESS_KEY] =
+        span.attributes[AWS_ATTRIBUTE_KEYS.AWS_AUTH_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];
     }
   }
 

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

@@ -80,6 +80,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');
@@ -189,6 +191,68 @@ 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 originalResponseHook = dynamoDbServiceExtension.responseHook;
+      dynamoDbServiceExtension.responseHook = (
+        response: NormalizedResponse,
+        span: Span,
+        tracer: Tracer,
+        config: AwsSdkInstrumentationConfig
+      ): void => {
+        originalResponseHook.call(dynamoDbServiceExtension, response, span, tracer, config);
+
+        if (response.data && response.data.Table) {
+          const tableArn = response.data.Table.TableArn;
+          if (tableArn) {
+            span.setAttribute(AWS_ATTRIBUTE_KEYS.AWS_DYNAMODB_TABLE_ARN, tableArn);
+          }
+        }
+      };
+    }
+  }
+}
+
 /*
  * 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
@@ -293,7 +357,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 awsXrayPropagator = new AWSXRayPropagator();
 const V3_CLIENT_CONFIG_KEY = Symbol('opentelemetry.instrumentation.aws-sdk.client.config');
@@ -306,6 +370,38 @@ function patchAwsSdkInstrumentation(instrumentation: Instrumentation): void {
       original: (...args: unknown[]) => Promise<any>
     ) {
       return function send(this: any, command: V3PluginCommand, ...args: unknown[]): Promise<any> {
+        this.middlewareStack?.add(
+          (next: any, context: any) => async (middlewareArgs: any) => {
+            const activeContext = otelContext.active();
+            const span = trace.getSpan(activeContext);
+
+            try {
+              const credsProvider = this.config.credentials;
+              if (credsProvider && span) {
+                const credentials = await credsProvider();
+                if (credentials.accessKeyId) {
+                  span.setAttribute(AWS_ATTRIBUTE_KEYS.AWS_AUTH_ACCESS_KEY, credentials.accessKeyId);
+                }
+              }
+              if (span && this.config.region) {
+                const region = await this.config.region();
+                if (region) {
+                  span.setAttribute(AWS_ATTRIBUTE_KEYS.AWS_AUTH_REGION, region);
+                }
+              }
+            } catch (err) {
+              diag.debug('Fail to get auth account access key and region.');
+            }
+
+            return await next(middlewareArgs);
+          },
+          {
+            step: 'build',
+            name: '_extractSignerCredentials',
+            override: true,
+          }
+        );
+
         this.middlewareStack?.add(
           (next: any, context: any) => async (middlewareArgs: any) => {
             awsXrayPropagator.inject(otelContext.active(), middlewareArgs.request.headers, defaultTextMapSetter);

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

@@ -0,0 +1,52 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { AttributeValue } from '@opentelemetry/api';
+
+export class RegionalResourceArnParser {
+  public static getAccountId(arn: AttributeValue | undefined): string | undefined {
+    if (this.isArn(arn)) {
+      return (arn! as string).split(':')[4];
+    }
+    return undefined;
+  }
+
+  public static getRegion(arn: AttributeValue | undefined): string | undefined {
+    if (this.isArn(arn)) {
+      return (arn! as string).split(':')[3];
+    }
+    return undefined;
+  }
+
+  public static isArn(arn: AttributeValue | undefined): boolean {
+    // Check if arn follows the format:
+    // arn:partition:service:region:account-id:resource-type/resource-id or
+    // arn:partition:service:region:account-id:resource-type:resource-id
+    if (!arn || typeof arn !== 'string') {
+      return false;
+    }
+
+    if (!arn.startsWith('arn')) {
+      return false;
+    }
+
+    const arnParts = arn.split(':');
+    return arnParts.length >= 6 && this.isAccountId(arnParts[4]);
+  }
+
+  private static isAccountId(input: string): boolean {
+    if (input == null || input.length !== 12) {
+      return false;
+    }
+
+    if (!this._checkDigits(input)) {
+      return false;
+    }
+
+    return true;
+  }
+
+  private static _checkDigits(str: string): boolean {
+    return /^\d+$/.test(str);
+  }
+}

aws-distro-opentelemetry-node-autoinstrumentation/src/sqs-url-parser.ts

@@ -30,6 +30,57 @@ export class SqsUrlParser {
     return undefined;
   }
 
+  /**
+   * Extracts the account ID from an SQS URL.
+   */
+  public static getAccountId(url: AttributeValue | undefined): string | undefined {
+    if (typeof url !== 'string') {
+      return undefined;
+    }
+
+    url = url.replace(HTTP_SCHEMA, '').replace(HTTPS_SCHEMA, '');
+    if (this.isValidSqsUrl(url)) {
+      const splitUrl: string[] = url.split('/');
+      return splitUrl[1];
+    }
+
+    return undefined;
+  }
+
+  /**
+   * Extracts the region from an SQS URL.
+   */
+  public static getRegion(url: AttributeValue | undefined): string | undefined {
+    if (typeof url !== 'string') {
+      return undefined;
+    }
+
+    url = url.replace(HTTP_SCHEMA, '').replace(HTTPS_SCHEMA, '');
+    if (this.isValidSqsUrl(url)) {
+      const splitUrl: string[] = url.split('/');
+      const domain: string = splitUrl[0];
+      const domainParts: string[] = domain.split('.');
+      if (domainParts.length === 4) {
+        return domainParts[1];
+      }
+    }
+
+    return undefined;
+  }
+
+  /**
+   * Checks if the URL is a valid SQS URL.
+   */
+  private static isValidSqsUrl(url: string): boolean {
+    const splitUrl: string[] = url.split('/');
+    return (
+      splitUrl.length === 3 &&
+      splitUrl[0].toLowerCase().startsWith('sqs') &&
+      this.isAccountId(splitUrl[1]) &&
+      this.isValidQueueName(splitUrl[2])
+    );
+  }
+
   private static isAccountId(input: string): boolean {
     if (input == null || input.length !== 12) {
       return false;