Check for undefined XRay Trace Id in AwsSdk Instrumentation

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

@@ -38,6 +38,8 @@ import { AwsLambdaInstrumentation } from '@opentelemetry/instrumentation-aws-lam
 import type { Command as AwsV3Command } from '@aws-sdk/types';
 
 export const traceContextEnvironmentKey = '_X_AMZN_TRACE_ID';
+export const AWSXRAY_TRACE_ID_HEADER_CAPITALIZED = 'X-Amzn-Trace-Id';
+
 const awsPropagator = new AWSXRayPropagator();
 export const headerGetter: TextMapGetter<APIGatewayProxyEventHeaders> = {
   keys(carrier: any): string[] {
@@ -294,7 +296,6 @@ function patchAwsLambdaInstrumentation(instrumentation: Instrumentation): void {
 // Override the upstream private _getV3SmithyClientSendPatch method to add middleware to inject X-Ray Trace Context into HTTP Headers
 // 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 AWSXRAY_TRACE_ID_HEADER_CAPITALIZED = 'X-Amzn-Trace-Id';
 const V3_CLIENT_CONFIG_KEY = Symbol('opentelemetry.instrumentation.aws-sdk.client.config');
 type V3PluginCommand = AwsV3Command<any, any, any, any, any> & {
   [V3_CLIENT_CONFIG_KEY]?: any;
@@ -311,9 +312,12 @@ function patchAwsSdkInstrumentation(instrumentation: Instrumentation): void {
             // Need to set capitalized version of the trace id to ensure that the Recursion Detection Middleware
             // of aws-sdk-js-v3 will detect the propagated X-Ray Context
             // See: https://github.com/aws/aws-sdk-js-v3/blob/v3.768.0/packages/middleware-recursion-detection/src/index.ts#L13
-            middlewareArgs.request.headers[AWSXRAY_TRACE_ID_HEADER_CAPITALIZED] =
-              middlewareArgs.request.headers[AWSXRAY_TRACE_ID_HEADER];
-            delete middlewareArgs.request.headers[AWSXRAY_TRACE_ID_HEADER];
+            const xrayTraceId = middlewareArgs.request.headers[AWSXRAY_TRACE_ID_HEADER];
+
+            if (xrayTraceId) {
+              middlewareArgs.request.headers[AWSXRAY_TRACE_ID_HEADER_CAPITALIZED] = xrayTraceId;
+              delete middlewareArgs.request.headers[AWSXRAY_TRACE_ID_HEADER];
+            }
             const result = await next(middlewareArgs);
             return result;
           },

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

@@ -6,11 +6,13 @@ import {
   diag,
   Context as OtelContext,
   trace,
+  context,
   propagation,
   Span,
   Tracer,
   AttributeValue,
   TextMapSetter,
+  INVALID_SPAN_CONTEXT,
 } from '@opentelemetry/api';
 import { getNodeAutoInstrumentations } from '@opentelemetry/auto-instrumentations-node';
 import { Instrumentation } from '@opentelemetry/instrumentation';
@@ -19,7 +21,12 @@ import { AwsLambdaInstrumentation, AwsLambdaInstrumentationConfig } from '@opent
 import { expect } from 'expect';
 import { AWS_ATTRIBUTE_KEYS } from '../../src/aws-attribute-keys';
 import { RequestMetadata, ServiceExtension } from '../../src/third-party/otel/aws/services/ServiceExtension';
-import { applyInstrumentationPatches, customExtractor, headerGetter } from './../../src/patches/instrumentation-patch';
+import {
+  applyInstrumentationPatches,
+  AWSXRAY_TRACE_ID_HEADER_CAPITALIZED,
+  customExtractor,
+  headerGetter,
+} from './../../src/patches/instrumentation-patch';
 import * as sinon from 'sinon';
 import { AWSXRAY_TRACE_ID_HEADER, AWSXRayPropagator } from '@opentelemetry/propagator-aws-xray';
 import { Context } from 'aws-lambda';
@@ -535,34 +542,72 @@ describe('InstrumentationPatchTest', () => {
     let lambda: Lambda;
     const region = 'us-east-1';
 
-    it('overridden _getV3SmithyClientSendPatch updates MiddlewareStack', async () => {
-      const mockedMiddlewareStackInternal: any = [];
-      const mockedMiddlewareStack = {
-        add: (arg1: any, arg2: any) => mockedMiddlewareStackInternal.push([arg1, arg2]),
-      };
-      const send = extractAwsSdkInstrumentation(PATCHED_INSTRUMENTATIONS)
-        ['_getV3SmithyClientSendPatch']((...args: unknown[]) => Promise.resolve())
-        .bind({ middlewareStack: mockedMiddlewareStack });
-      sinon
-        .stub(AWSXRayPropagator.prototype, 'inject')
-        .callsFake((context: OtelContext, carrier: unknown, setter: TextMapSetter) => {
-          (carrier as any)['isCarrierModified'] = 'carrierIsModified';
-        });
+    describe('overridden _getV3SmithyClientSendPatch updates MiddlewareStack', async () => {
+      let mockedMiddlewareStackInternal: any;
+      let mockedMiddlewareStack;
+      let middlewareArgsHeader: any;
+      const testXrayTraceHeader = 'test-xray-trace-header';
+
+      beforeEach(async () => {
+        // Clear environment variables before each test
+        mockedMiddlewareStackInternal = [];
+        mockedMiddlewareStack = {
+          add: (arg1: any, arg2: any) => mockedMiddlewareStackInternal.push([arg1, arg2]),
+        };
+        const send = extractAwsSdkInstrumentation(PATCHED_INSTRUMENTATIONS)
+          ['_getV3SmithyClientSendPatch']((...args: unknown[]) => Promise.resolve())
+          .bind({ middlewareStack: mockedMiddlewareStack });
+
+        middlewareArgsHeader = {
+          request: {
+            headers: {},
+          },
+        };
+
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-ignore
+        await send({}, null);
+      });
 
-      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-      // @ts-ignore
-      await send({}, null);
+      afterEach(() => {
+        sinon.restore();
+      });
 
-      const middlewareArgs: any = {
-        request: {
-          headers: {},
-        },
-      };
-      await mockedMiddlewareStackInternal[0][0]((arg: any) => Promise.resolve(), null)(middlewareArgs);
+      it('Updates trace header casing when AWSXRayPropagator injects trace header successfully', async () => {
+        sinon
+          .stub(AWSXRayPropagator.prototype, 'inject')
+          .callsFake((context: OtelContext, carrier: unknown, setter: TextMapSetter) => {
+            (carrier as any)['isCarrierModified'] = 'carrierIsModified';
+            (carrier as any)[AWSXRAY_TRACE_ID_HEADER] = testXrayTraceHeader;
+          });
+        await mockedMiddlewareStackInternal[0][0]((arg: any) => Promise.resolve(), null)(middlewareArgsHeader);
+
+        expect(middlewareArgsHeader.request.headers['isCarrierModified']).toEqual('carrierIsModified');
+        expect(middlewareArgsHeader.request.headers).not.toHaveProperty(AWSXRAY_TRACE_ID_HEADER);
+        expect(middlewareArgsHeader.request.headers).toHaveProperty(AWSXRAY_TRACE_ID_HEADER_CAPITALIZED);
+        expect(middlewareArgsHeader.request.headers[AWSXRAY_TRACE_ID_HEADER_CAPITALIZED]).toEqual(testXrayTraceHeader);
+
+        expect(mockedMiddlewareStackInternal[0][1].name).toEqual('_adotInjectXrayContextMiddleware');
+      });
 
-      sinon.restore();
-      expect(middlewareArgs.request.headers['isCarrierModified']).toEqual('carrierIsModified');
-      expect(mockedMiddlewareStackInternal[0][1].name).toEqual('_adotInjectXrayContextMiddleware');
+      it('Does not set trace header when AWSXRayPropagator does not inject trace header', async () => {
+        const invalidContext: OtelContext = {
+          getValue: (key: symbol) => ({
+            spanContext: () => INVALID_SPAN_CONTEXT,
+          }),
+          setValue: (key: symbol, value: unknown) => invalidContext,
+          deleteValue: (key: symbol) => invalidContext,
+        };
+
+        sinon.stub(context, 'active').returns(invalidContext);
+
+        await mockedMiddlewareStackInternal[0][0]((arg: any) => Promise.resolve(), null)(middlewareArgsHeader);
+
+        expect(middlewareArgsHeader.request.headers).not.toHaveProperty(AWSXRAY_TRACE_ID_HEADER);
+        expect(middlewareArgsHeader.request.headers).not.toHaveProperty(AWSXRAY_TRACE_ID_HEADER_CAPITALIZED);
+
+        expect(mockedMiddlewareStackInternal[0][1].name).toEqual('_adotInjectXrayContextMiddleware');
+      });
     });
 
     it('injects trace context header into request via propagator', async () => {