[Fix] Update AWS SDK Instrumentation to inject XRay trace context into HTTP Headers

aws-distro-opentelemetry-node-autoinstrumentation/.eslintignore

@@ -1,4 +1,5 @@
 build
 node_modules
 .eslintrc.js
-version.ts
+version.ts
+src/third-party

aws-distro-opentelemetry-node-autoinstrumentation/package.json

@@ -31,9 +31,18 @@
     "prepublishOnly": "npm run compile",
     "tdd": "yarn test -- --watch-extensions ts --watch",
     "test": "nyc ts-mocha --timeout 10000 -p tsconfig.json --require '@opentelemetry/contrib-test-utils' 'test/**/*.ts'",
-    "test:coverage": "nyc --all --check-coverage --functions 95 --lines 95 ts-mocha --timeout 10000 -p tsconfig.json --require '@opentelemetry/contrib-test-utils' 'test/**/*.ts'",
+    "test:coverage": "nyc --check-coverage --functions 95 --lines 95 ts-mocha --timeout 10000 -p tsconfig.json --require '@opentelemetry/contrib-test-utils' 'test/**/*.ts'",
     "watch": "tsc -w"
   },
+  "nyc": {
+    "all": true,
+    "include": [
+      "src/**/*.ts"
+    ],
+    "exclude": [
+      "src/third-party/**/*"
+    ]
+  },
   "bugs": {
     "url": "https://github.com/aws-observability/aws-otel-js-instrumentation/issues"
   },
@@ -65,6 +74,7 @@
     "@aws-sdk/client-bedrock-runtime": "3.632.0",
     "@aws-sdk/client-kinesis": "3.632.0",
     "@aws-sdk/client-s3": "3.632.0",
+    "@smithy/protocol-http": "4.1.8",
     "@opentelemetry/contrib-test-utils": "0.41.0",
     "@types/mocha": "7.0.2",
     "@types/node": "18.6.5",

aws-distro-opentelemetry-node-autoinstrumentation/src/patches/extended-instrumentations/aws-sdk-instrumentation-extended.ts

@@ -0,0 +1,42 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { AwsInstrumentation } from '@opentelemetry/instrumentation-aws-sdk';
+import { InstrumentationModuleDefinition, InstrumentationNodeModuleDefinition } from '@opentelemetry/instrumentation';
+import { context, defaultTextMapSetter } from '@opentelemetry/api';
+import { propwrap } from './../../third-party/otel/aws/propwrap';
+import { AWSXRayPropagator } from '@opentelemetry/propagator-aws-xray';
+
+const awsXrayPropagator = new AWSXRayPropagator();
+
+// This class extends the upstream AwsInstrumentation to add an additional
+// module instrumentation to patch `HttpRequest` of the `@smithy/protocol-http` module.
+// This additional module instrumentation will replace HttpRequest with an extended version
+// that injects the `X-Amzn-Trace-Id` HTTP header in the constructor so that aws-sdk-js-v3
+// client calls can propagate the X-Ray trace context
+export class AwsSdkInstrumentationExtended extends AwsInstrumentation {
+  protected override init(): InstrumentationModuleDefinition[] {
+    const instrumentationModuleDefinitions = super.init();
+
+    const v3SmithyProtocolHttp = new InstrumentationNodeModuleDefinition(
+      '@smithy/protocol-http',
+      ['>=2.0.0'],
+      (moduleExports: any) => {
+        const newExports = propwrap(moduleExports, 'HttpRequest', (origHttpRequest: any) => {
+          class ExtendedHttpRequest extends origHttpRequest {
+            constructor(...args: any[]) {
+              super(...args);
+              awsXrayPropagator.inject(context.active(), this.headers, defaultTextMapSetter);
+            }
+          }
+
+          return ExtendedHttpRequest;
+        });
+
+        return newExports;
+      }
+    );
+
+    return [...instrumentationModuleDefinitions, v3SmithyProtocolHttp];
+  }
+}

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

@@ -25,7 +25,9 @@ import {
 } from './aws/services/bedrock';
 import { KinesisServiceExtension } from './aws/services/kinesis';
 import { S3ServiceExtension } from './aws/services/s3';
-import { AwsLambdaInstrumentationPatch } from './aws/services/aws-lambda';
+import { AwsLambdaInstrumentationPatch } from './extended-instrumentations/aws-lambda';
+import { InstrumentationConfigMap } from '@opentelemetry/auto-instrumentations-node';
+import { AwsSdkInstrumentationExtended } from './extended-instrumentations/aws-sdk-instrumentation-extended';
 
 export const traceContextEnvironmentKey = '_X_AMZN_TRACE_ID';
 const awsPropagator = new AWSXRayPropagator();
@@ -38,7 +40,10 @@ export const headerGetter: TextMapGetter<APIGatewayProxyEventHeaders> = {
   },
 };
 
-export function applyInstrumentationPatches(instrumentations: Instrumentation[]): void {
+export function applyInstrumentationPatches(
+  instrumentations: Instrumentation[],
+  instrumentationConfigs?: InstrumentationConfigMap
+): void {
   /*
   Apply patches to upstream instrumentation libraries.
 
@@ -50,10 +55,16 @@ export function applyInstrumentationPatches(instrumentations: Instrumentation[])
   */
   instrumentations.forEach((instrumentation, index) => {
     if (instrumentation.instrumentationName === '@opentelemetry/instrumentation-aws-sdk') {
+      diag.debug('Overriding aws sdk instrumentation');
+      instrumentations[index] = new AwsSdkInstrumentationExtended(
+        instrumentationConfigs ? instrumentationConfigs['@opentelemetry/instrumentation-aws-sdk'] : undefined
+      );
+
       // Access private property servicesExtensions of AwsInstrumentation
       // eslint-disable-next-line @typescript-eslint/ban-ts-comment
       // @ts-ignore
-      const services: Map<string, ServiceExtension> | undefined = (instrumentation as any).servicesExtensions?.services;
+      const services: Map<string, ServiceExtension> | undefined = (instrumentations[index] as any).servicesExtensions
+        ?.services;
       if (services) {
         services.set('S3', new S3ServiceExtension());
         services.set('Kinesis', new KinesisServiceExtension());

aws-distro-opentelemetry-node-autoinstrumentation/src/register.ts

@@ -62,7 +62,7 @@ const instrumentationConfigs: InstrumentationConfigMap = {
 const instrumentations: Instrumentation[] = getNodeAutoInstrumentations(instrumentationConfigs);
 
 // Apply instrumentation patches
-applyInstrumentationPatches(instrumentations);
+applyInstrumentationPatches(instrumentations, instrumentationConfigs);
 
 const configurator: AwsOpentelemetryConfigurator = new AwsOpentelemetryConfigurator(instrumentations, useXraySampler);
 const configuration: Partial<opentelemetry.NodeSDKConfiguration> = configurator.configure();

aws-distro-opentelemetry-node-autoinstrumentation/src/third-party/otel/aws/propwrap.ts

@@ -0,0 +1,155 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*
+ * This block is derived from esbuild's bundling support.
+ * https://github.com/evanw/esbuild/blob/v0.14.42/internal/runtime/runtime.go#L22
+ *
+ * License:
+ * MIT License
+ *
+ * Copyright (c) 2020 Evan Wallace
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+const __defProp = Object.defineProperty;
+const __getOwnPropDesc = Object.getOwnPropertyDescriptor;
+const __hasOwnProp = Object.prototype.hasOwnProperty;
+const __getOwnPropNames = Object.getOwnPropertyNames;
+const __copyProps = (
+  to: any,
+  from: any,
+  except: string,
+  desc?: PropertyDescriptor | undefined
+) => {
+  if ((from && typeof from === 'object') || typeof from === 'function') {
+    for (const key of __getOwnPropNames(from)) {
+      if (!__hasOwnProp.call(to, key) && key !== except) {
+        __defProp(to, key, {
+          get: () => from[key] as any,
+          enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable,
+        });
+      }
+    }
+  }
+  return to;
+};
+
+/**
+ * Return a new object that is a copy of `obj`, with its `subpath` property
+ * replaced with the return value of `wrapper(original)`.
+ *
+ * This is similar to shimmer (i.e. `InstrumentationBase.prototype._wrap`).
+ * However, it uses a different technique to support wrapping properties that
+ * are only available via a getter (i.e. their property descriptor is `.writable
+ * === false`).
+ *
+ * For example:
+ *    var os = propwrap(require('os'), 'platform', (orig) => {
+ *      return function wrappedPlatform () {
+ *        return orig().toUpperCase()
+ *      }
+ *    })
+ *    console.log(os.platform()) // => DARWIN
+ *
+ * The subpath can indicate a nested property. Each property in that subpath,
+ * except the last, must identify an *Object*.
+ *
+ * Limitations:
+ * - This doesn't handle possible Symbol properties on the copied object(s).
+ * - This cannot wrap a property of a function, because we cannot create a
+ *   copy of the function.
+ *
+ * @param {object} obj
+ * @param {string} subpath - The property subpath on `obj` to wrap. This may
+ *    point to a nested property by using a '.' to separate levels. For example:
+ *        var fs = wrap(fs, 'promises.sync', (orig) => { ... })
+ * @param {Function} wrapper - A function of the form `function (orig)`, where
+ *    `orig` is the original property value. This must synchronously return the
+ *    new property value.
+ * @returns {object} A new object with the wrapped property.
+ * @throws {TypeError} if the subpath points to a non-existent property, or if
+ *    any but the last subpath part points to a non-Object.
+ */
+export const propwrap = (obj: any, subpath: string, wrapper: Function): any => {
+  const parts = subpath.split('.');
+  const namespaces = [obj];
+  let namespace = obj;
+  let key;
+  let val;
+
+  // 1. Traverse the subpath parts to sanity check and get references to the
+  //    Objects that we will be copying.
+  for (let i = 0; i < parts.length; i++) {
+    key = parts[i];
+    val = namespace[key];
+    if (!val) {
+      throw new TypeError(
+        `cannot wrap "${subpath}": "<obj>.${parts
+          .slice(0, i)
+          .join('.')}" is ${typeof val}`
+      );
+    } else if (i < parts.length - 1) {
+      if (typeof val !== 'object') {
+        throw new TypeError(
+          `cannot wrap "${subpath}": "<obj>.${parts
+            .slice(0, i)
+            .join('.')}" is not an Object`
+        );
+      }
+      namespace = val;
+      namespaces.push(namespace);
+    }
+  }
+
+  // 2. Now work backwards, wrapping each namespace with a new object that has a
+  //    copy of all the properties, except the one that we've wrapped.
+  for (let i = parts.length - 1; i >= 0; i--) {
+    key = parts[i];
+    namespace = namespaces[i];
+    if (i === parts.length - 1) {
+      const orig = namespace[key];
+      val = wrapper(orig);
+    } else {
+      val = namespaces[i + 1];
+    }
+    const desc = __getOwnPropDesc(namespace, key);
+    const wrappedNamespace = __defProp({}, key, {
+      value: val,
+      enumerable: !desc || desc.enumerable,
+    });
+    __copyProps(wrappedNamespace, namespace, key);
+    namespaces[i] = wrappedNamespace;
+  }
+
+  return namespaces[0];
+};

aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/aws-lambda.test.ts → aws-distro-opentelemetry-node-autoinstrumentation/test/patches/extended-instrumentations/aws-lambda.test.ts

@@ -6,7 +6,7 @@ import * as path from 'path';
 import * as fs from 'fs';
 import { diag } from '@opentelemetry/api';
 import { InstrumentationNodeModuleDefinition } from '@opentelemetry/instrumentation';
-import { AwsLambdaInstrumentationPatch } from '../../../../src/patches/aws/services/aws-lambda';
+import { AwsLambdaInstrumentationPatch } from '../../../src/patches/extended-instrumentations/aws-lambda';
 
 describe('AwsLambdaInstrumentationPatch', () => {
   let instrumentation: AwsLambdaInstrumentationPatch;

aws-distro-opentelemetry-node-autoinstrumentation/test/patches/extended-instrumentations/aws-sdk-instrumentation-extended.test.ts

@@ -0,0 +1,35 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+import * as sinon from 'sinon';
+import { AwsSdkInstrumentationExtended } from '../../../src/patches/extended-instrumentations/aws-sdk-instrumentation-extended';
+import { AwsInstrumentation } from '@opentelemetry/instrumentation-aws-sdk';
+import expect from 'expect';
+import { AWSXRayPropagator } from '@opentelemetry/propagator-aws-xray';
+import { Context, TextMapSetter } from '@opentelemetry/api';
+import { HttpRequest } from '@smithy/protocol-http';
+
+describe('AwsSdkInstrumentationExtended', () => {
+  let instrumentation: AwsSdkInstrumentationExtended;
+
+  beforeEach(() => {
+    instrumentation = new AwsSdkInstrumentationExtended({});
+  });
+
+  afterEach(() => {
+    sinon.restore();
+  });
+
+  it('overridden init patches smithy HttpRequest', () => {
+    sinon.stub(AwsInstrumentation.prototype as any, 'init').returns([]);
+    sinon
+      .stub(AWSXRayPropagator.prototype, 'inject')
+      .callsFake((context: Context, carrier: unknown, setter: TextMapSetter) => {
+        (carrier as any)['isCarrierModified'] = 'carrierIsModified';
+      });
+    const result = (instrumentation as any).init();
+    expect(result.length).toEqual(1);
+
+    const patchedHttpRequestObject = new HttpRequest({});
+    expect(patchedHttpRequestObject.headers['isCarrierModified']).toEqual('carrierIsModified');
+  });
+});