Check for undefined XRay Trace Id in AwsSdk Instrumentation

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

@@ -311,9 +311,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

@@ -11,6 +11,8 @@ import {
   Tracer,
   AttributeValue,
   TextMapSetter,
+  defaultTextMapSetter,
+  ROOT_CONTEXT,
 } from '@opentelemetry/api';
 import { getNodeAutoInstrumentations } from '@opentelemetry/auto-instrumentations-node';
 import { Instrumentation } from '@opentelemetry/instrumentation';
@@ -535,34 +537,71 @@ 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 send;
+      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]),
+        };
+        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);
+      });
+
+      it('propagator injects with valid context', 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);
+
+        sinon.restore();
+
+        expect(middlewareArgsHeader.request.headers['isCarrierModified']).toEqual('carrierIsModified');
+        expect(middlewareArgsHeader.request.headers).not.toHaveProperty(AWSXRAY_TRACE_ID_HEADER);
+        expect(middlewareArgsHeader.request.headers).toHaveProperty('X-Amzn-Trace-Id');
+        expect(middlewareArgsHeader.request.headers['X-Amzn-Trace-Id']).toEqual(testXrayTraceHeader);
+
+        expect(mockedMiddlewareStackInternal[0][1].name).toEqual('_adotInjectXrayContextMiddleware');
+      });
+
+      it('propagator does not inject with invalid context', async () => {
+        const invalidContext = trace.setSpanContext(ROOT_CONTEXT, {
+          traceId: 'invalid-trace-id',
+          spanId: 'invalid-span',
+          traceFlags: 0,
+          isRemote: false,
         });
 
-      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-      // @ts-ignore
-      await send({}, null);
+        const propagator = new AWSXRayPropagator();
 
-      const middlewareArgs: any = {
-        request: {
-          headers: {},
-        },
-      };
-      await mockedMiddlewareStackInternal[0][0]((arg: any) => Promise.resolve(), null)(middlewareArgs);
+        propagator.inject(invalidContext, middlewareArgsHeader.request.headers, defaultTextMapSetter);
+        sinon.restore();
 
-      sinon.restore();
-      expect(middlewareArgs.request.headers['isCarrierModified']).toEqual('carrierIsModified');
-      expect(mockedMiddlewareStackInternal[0][1].name).toEqual('_adotInjectXrayContextMiddleware');
+        expect(middlewareArgsHeader.request.headers).not.toHaveProperty(AWSXRAY_TRACE_ID_HEADER);
+        expect(middlewareArgsHeader.request.headers).not.toHaveProperty('X-Amzn-Trace-Id');
+
+        expect(mockedMiddlewareStackInternal[0][1].name).toEqual('_adotInjectXrayContextMiddleware');
+      });
     });
 
     it('injects trace context header into request via propagator', async () => {