Extract arn and account access key

blairhyy-amazon · blairhyy-amazon · commit 0c2c969ead42 · 2025-07-02T10:48:29.000-07:00
diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/src/aws-attribute-keys.ts b/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_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_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',
diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/src/patches/instrumentation-patch.ts b/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');
@@ -328,6 +392,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) {
+                  const credentials = await credsProvider();
+                  if (credentials.accessKeyId) {
+                    span.setAttribute(AWS_ATTRIBUTE_KEYS.AWS_AUTH_ACCESS_KEY, credentials.accessKeyId);
+                  }
+                }
+                if (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,
+          }
+        );
+
         command[V3_CLIENT_CONFIG_KEY] = this.config;
         return original.apply(this, [command, ...args]);
       };
diff --git a/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/instrumentation-patch.test.ts b/aws-distro-opentelemetry-node-autoinstrumentation/test/patches/instrumentation-patch.test.ts
@@ -58,6 +58,8 @@ const _BEDROCK_GUARDRAIL_ARN: string = 'arn:aws:bedrock:us-east-1:123456789012:g
 const _BEDROCK_KNOWLEDGEBASE_ID: string = 'KnowledgeBaseId';
 const _GEN_AI_SYSTEM: string = 'aws.bedrock';
 const _GEN_AI_REQUEST_MODEL: string = 'genAiReuqestModelId';
+const _STREAM_ARN: string = 'arn:aws:kinesis:us-west-2:123456789012:stream/testStream';
+const _TABLE_ARN: string = 'arn:aws:dynamodb:us-west-2:123456789012:table/testTable';
 
 const mockHeaders = {
   'x-test-header': 'test-value',
@@ -90,6 +92,8 @@ describe('InstrumentationPatchTest', () => {
     expect(services.get('Lambda').requestPreSpanHook).toBeTruthy();
     expect(services.get('SQS')._requestPreSpanHook).toBeFalsy();
     expect(services.get('SQS').requestPreSpanHook).toBeTruthy();
+    expect(services.get('Kinesis')._requestPreSpanHook).toBeFalsy();
+    expect(services.get('Kinesis').requestPreSpanHook).toBeTruthy();
     expect(services.has('Bedrock')).toBeFalsy();
     expect(services.has('BedrockAgent')).toBeFalsy();
     expect(services.get('BedrockAgentRuntime')).toBeFalsy();
@@ -119,6 +123,8 @@ describe('InstrumentationPatchTest', () => {
     expect(services.get('Lambda').requestPreSpanHook).toBeTruthy();
     expect(services.get('SQS')._requestPreSpanHook).toBeTruthy();
     expect(services.get('SQS').requestPreSpanHook).toBeTruthy();
+    expect(services.get('Kinesis')._requestPreSpanHook).toBeTruthy();
+    expect(services.get('Kinesis').requestPreSpanHook).toBeTruthy();
     expect(services.has('Bedrock')).toBeTruthy();
     expect(services.has('BedrockAgent')).toBeTruthy();
     expect(services.get('BedrockAgentRuntime')).toBeTruthy();
@@ -179,6 +185,14 @@ describe('InstrumentationPatchTest', () => {
     expect(() => doExtractBedrockAttributes(services, 'Bedrock')).toThrow();
   });
 
+  it('Kinesis without patching', () => {
+    const unpatchedAwsSdkInstrumentation: AwsInstrumentation = extractAwsSdkInstrumentation(UNPATCHED_INSTRUMENTATIONS);
+    const services: Map<string, any> = extractServicesFromAwsSdkInstrumentation(unpatchedAwsSdkInstrumentation);
+    expect(() => doExtractKinesisAttributes(services)).not.toThrow();
+    const kinesisAttributes: Attributes = doExtractKinesisAttributes(services);
+    expect(kinesisAttributes[AWS_ATTRIBUTE_KEYS.AWS_KINESIS_STREAM_ARN]).toBeUndefined();
+  });
+
   it('SNS with patching', () => {
     const patchedAwsSdkInstrumentation: AwsInstrumentation = extractAwsSdkInstrumentation(PATCHED_INSTRUMENTATIONS);
     const services: Map<string, any> = extractServicesFromAwsSdkInstrumentation(patchedAwsSdkInstrumentation);
@@ -233,6 +247,20 @@ describe('InstrumentationPatchTest', () => {
     expect(responseHookSecretsManagerAttributes[AWS_ATTRIBUTE_KEYS.AWS_SECRETSMANAGER_SECRET_ARN]).toBe(_SECRETS_ARN);
   });
 
+  it('Kinesis with patching', () => {
+    const patchedAwsSdkInstrumentation: AwsInstrumentation = extractAwsSdkInstrumentation(PATCHED_INSTRUMENTATIONS);
+    const services: Map<string, any> = extractServicesFromAwsSdkInstrumentation(patchedAwsSdkInstrumentation);
+    const requestKinesisAttributes: Attributes = doExtractKinesisAttributes(services);
+    expect(requestKinesisAttributes[AWS_ATTRIBUTE_KEYS.AWS_KINESIS_STREAM_ARN]).toEqual(_STREAM_ARN);
+  });
+
+  it('DynamoDB with patching', () => {
+    const patchedAwsSdkInstrumentation: AwsInstrumentation = extractAwsSdkInstrumentation(PATCHED_INSTRUMENTATIONS);
+    const services: Map<string, any> = extractServicesFromAwsSdkInstrumentation(patchedAwsSdkInstrumentation);
+    const responseDynamoDbAttributes: Attributes = doResponseHookDynamoDb(services);
+    expect(responseDynamoDbAttributes[AWS_ATTRIBUTE_KEYS.AWS_DYNAMODB_TABLE_ARN]).toEqual(_TABLE_ARN);
+  });
+
   it('Bedrock with patching', () => {
     const patchedAwsSdkInstrumentation: AwsInstrumentation = extractAwsSdkInstrumentation(PATCHED_INSTRUMENTATIONS);
     const services: Map<string, any> = extractServicesFromAwsSdkInstrumentation(patchedAwsSdkInstrumentation);
@@ -446,6 +474,18 @@ describe('InstrumentationPatchTest', () => {
     return doExtractAttributes(services, serviceName, params);
   }
 
+  function doExtractKinesisAttributes(services: Map<string, ServiceExtension>): Attributes {
+    const serviceName: string = 'Kinesis';
+    const params: NormalizedRequest = {
+      serviceName: serviceName,
+      commandName: 'mockCommandName',
+      commandInput: {
+        StreamARN: _STREAM_ARN,
+      },
+    };
+    return doExtractAttributes(services, serviceName, params);
+  }
+
   function doExtractAttributes(
     services: Map<string, ServiceExtension>,
     serviceName: string,
@@ -492,6 +532,23 @@ describe('InstrumentationPatchTest', () => {
     return doResponseHook(services, 'Lambda', results as NormalizedResponse);
   }
 
+  function doResponseHookDynamoDb(services: Map<string, ServiceExtension>): Attributes {
+    const results: Partial<NormalizedResponse> = {
+      data: {
+        Table: {
+          TableArn: _TABLE_ARN,
+        },
+      },
+      request: {
+        commandInput: {},
+        commandName: 'dummy_operation',
+        serviceName: 'DynamoDB',
+      },
+    };
+
+    return doResponseHook(services, 'DynamoDB', results as NormalizedResponse);
+  }
+
   function doResponseHookBedrock(
     services: Map<string, ServiceExtension>,
     serviceName: string,
@@ -563,9 +620,13 @@ describe('InstrumentationPatchTest', () => {
         mockedMiddlewareStack = {
           add: (arg1: any, arg2: any) => mockedMiddlewareStackInternal.push([arg1, arg2]),
         };
+        const mockConfig = {
+          credentials: () => Promise.resolve({ accessKeyId: 'test-access-key' }),
+          region: () => Promise.resolve('us-west-2'),
+        };
         const send = extractAwsSdkInstrumentation(PATCHED_INSTRUMENTATIONS)
           ['_getV3SmithyClientSendPatch']((...args: unknown[]) => Promise.resolve())
-          .bind({ middlewareStack: mockedMiddlewareStack });
+          .bind({ middlewareStack: mockedMiddlewareStack, config: mockConfig });
 
         middlewareArgsHeader = {
           request: {
@@ -617,6 +678,18 @@ describe('InstrumentationPatchTest', () => {
 
         expect(mockedMiddlewareStackInternal[0][1].name).toEqual('_adotInjectXrayContextMiddleware');
       });
+
+      it('Add cross account information span attributes from STS credentials', async () => {
+        const mockSpan = { setAttribute: sinon.stub() };
+        sinon.stub(trace, 'getSpan').returns(mockSpan as unknown as Span);
+        const credentialsMiddlewareArgs: any = {};
+        await mockedMiddlewareStackInternal[1][0]((arg: any) => Promise.resolve(arg), null)(credentialsMiddlewareArgs);
+        expect(mockedMiddlewareStackInternal[1][1].name).toEqual('_extractSignerCredentials');
+        expect(
+          mockSpan.setAttribute.calledWith(AWS_ATTRIBUTE_KEYS.AWS_AUTH_ACCESS_KEY, 'test-access-key')
+        ).toBeTruthy();
+        expect(mockSpan.setAttribute.calledWith(AWS_ATTRIBUTE_KEYS.AWS_AUTH_REGION, 'us-west-2')).toBeTruthy();
+      });
     });
 
     it('injects trace context header into request via propagator', async () => {
