feat(server-utils): Port SQS, SNS and Lambda aws-sdk extensions with trace propagation

packages/server-utils/src/integrations/tracing-channel/aws-sdk/aws-sdk.types.ts

@@ -0,0 +1,59 @@
+/*
+ * AWS SDK for JavaScript
+ * Copyright 2012-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * This product includes software developed at
+ * Amazon Web Services, Inc. (http://aws.amazon.com/).
+ */
+
+/*
+  These are slightly modified and simplified versions of the actual SQS/SNS types included
+  in the official distribution:
+  https://github.com/aws/aws-sdk-js/blob/master/clients/sqs.d.ts
+  These are brought here to avoid having users install the `aws-sdk` whenever they
+  require this instrumentation.
+*/
+
+// eslint-disable-next-line @typescript-eslint/no-empty-object-type
+interface Blob {}
+type Binary = Buffer | Uint8Array | Blob | string;
+
+// eslint-disable-next-line @typescript-eslint/no-namespace -- Prefer to contain the types copied over in one location
+export namespace SNS {
+  interface MessageAttributeValue {
+    DataType: string;
+    StringValue?: string;
+    BinaryValue?: Binary;
+  }
+
+  export type MessageAttributeMap = { [key: string]: MessageAttributeValue };
+}
+
+// eslint-disable-next-line @typescript-eslint/no-namespace -- Prefer to contain the types copied over in one location
+export namespace SQS {
+  type StringList = string[];
+  type BinaryList = Binary[];
+  interface MessageAttributeValue {
+    StringValue?: string;
+    BinaryValue?: Binary;
+    StringListValues?: StringList;
+    BinaryListValues?: BinaryList;
+    DataType: string;
+  }
+
+  export type MessageBodyAttributeMap = {
+    [key: string]: MessageAttributeValue;
+  };
+
+  type MessageSystemAttributeMap = { [key: string]: string };
+
+  export interface Message {
+    MessageId?: string;
+    ReceiptHandle?: string;
+    MD5OfBody?: string;
+    Body?: string;
+    Attributes?: MessageSystemAttributeMap;
+    MD5OfMessageAttributes?: string;
+    MessageAttributes?: MessageBodyAttributeMap;
+  }
+}

packages/server-utils/src/integrations/tracing-channel/aws-sdk/constants.ts

@@ -46,3 +46,17 @@ export const DB_SYSTEM_VALUE_DYNAMODB = 'dynamodb';
 export const ATTR_AWS_SECRETSMANAGER_SECRET_ARN = 'aws.secretsmanager.secret.arn';
 export const ATTR_AWS_STEP_FUNCTIONS_ACTIVITY_ARN = 'aws.step_functions.activity.arn';
 export const ATTR_AWS_STEP_FUNCTIONS_STATE_MACHINE_ARN = 'aws.step_functions.state_machine.arn';
+
+// SNS
+export const ATTR_AWS_SNS_TOPIC_ARN = 'aws.sns.topic.arn';
+
+// Lambda (faas)
+export const ATTR_FAAS_INVOKED_NAME = 'faas.invoked_name';
+export const ATTR_FAAS_INVOKED_PROVIDER = 'faas.invoked_provider';
+export const ATTR_FAAS_INVOKED_REGION = 'faas.invoked_region';
+export const ATTR_FAAS_EXECUTION = 'faas.execution';
+
+// Messaging (obsolete OTel conventions kept for parity with the OTel integration)
+export const ATTR_MESSAGING_DESTINATION = 'messaging.destination';
+export const ATTR_MESSAGING_DESTINATION_KIND = 'messaging.destination_kind';
+export const MESSAGING_DESTINATION_KIND_VALUE_TOPIC = 'topic';

packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/MessageAttributes.ts

@@ -0,0 +1,85 @@
+import type { Span } from '@sentry/core';
+import {
+  debug,
+  dynamicSamplingContextToSentryBaggageHeader,
+  getDynamicSamplingContextFromSpan,
+  spanToTraceHeader,
+} from '@sentry/core';
+import { DEBUG_BUILD } from '../../../../debug-build';
+import type { SNS, SQS } from '../aws-sdk.types';
+
+// https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/sqs-quotas.html
+export const MAX_MESSAGE_ATTRIBUTES = 10;
+
+// Sentry trace-propagation headers written into / read from AWS message attributes.
+const SENTRY_TRACE_HEADER = 'sentry-trace';
+const BAGGAGE_HEADER = 'baggage';
+const PROPAGATION_FIELDS = [SENTRY_TRACE_HEADER, BAGGAGE_HEADER];
+
+export interface AwsSdkContextObject {
+  [key: string]: {
+    StringValue?: string;
+    Value?: string;
+  };
+}
+
+/** Build the `sentry-trace`/`baggage` header pair carrying the span's trace context. */
+export function getPropagationHeaders(span: Span): Record<string, string> {
+  const headers: Record<string, string> = {
+    [SENTRY_TRACE_HEADER]: spanToTraceHeader(span),
+  };
+  const baggage = dynamicSamplingContextToSentryBaggageHeader(getDynamicSamplingContextFromSpan(span));
+  if (baggage) {
+    headers[BAGGAGE_HEADER] = baggage;
+  }
+  return headers;
+}
+
+/**
+ * Inject the span's trace-propagation headers into an SQS/SNS message-attribute map, so the consumer
+ * can continue the trace. Respects the SQS 10-attribute quota. Mirrors the OTel integration's
+ * `injectPropagationContext`, but writes Sentry's `sentry-trace`/`baggage` instead of W3C headers.
+ */
+export function injectPropagationContext(
+  attributesMap: SQS.MessageBodyAttributeMap | SNS.MessageAttributeMap | undefined,
+  span: Span,
+): SQS.MessageBodyAttributeMap | SNS.MessageAttributeMap {
+  const attributes = attributesMap ?? {};
+  const headers = getPropagationHeaders(span);
+  const headerKeys = Object.keys(headers);
+
+  if (Object.keys(attributes).length + headerKeys.length <= MAX_MESSAGE_ATTRIBUTES) {
+    for (const key of headerKeys) {
+      // Index-assigning into the SQS/SNS map union needs one concrete map type; the written value
+      // shape is valid for both.
+      (attributes as SQS.MessageBodyAttributeMap)[key] = { DataType: 'String', StringValue: headers[key] };
+    }
+  } else {
+    DEBUG_BUILD &&
+      debug.warn(
+        '[orchestrion:aws-sdk] cannot set trace propagation on SQS/SNS message due to maximum amount of MessageAttributes',
+      );
+  }
+  return attributes;
+}
+
+/** Read the propagation headers back off a received SQS message, if present. */
+export function extractPropagationHeaders(
+  message: SQS.Message,
+): { sentryTrace?: string; baggage?: string } | undefined {
+  const carrier = (message.MessageAttributes ?? {}) as AwsSdkContextObject;
+  const sentryTrace = carrier[SENTRY_TRACE_HEADER]?.StringValue ?? carrier[SENTRY_TRACE_HEADER]?.Value;
+  if (!sentryTrace) {
+    return undefined;
+  }
+  return {
+    sentryTrace,
+    baggage: carrier[BAGGAGE_HEADER]?.StringValue ?? carrier[BAGGAGE_HEADER]?.Value,
+  };
+}
+
+export function addPropagationFieldsToAttributeNames(messageAttributeNames: string[] = []): string[] {
+  return messageAttributeNames.length
+    ? Array.from(new Set([...messageAttributeNames, ...PROPAGATION_FIELDS]))
+    : [...PROPAGATION_FIELDS];
+}

packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/ServicesExtensions.ts

@@ -2,9 +2,12 @@ import type { Span } from '@sentry/core';
 import type { NormalizedRequest, NormalizedResponse, RequestMetadata } from '../types';
 import { DynamodbServiceExtension } from './dynamodb';
 import { KinesisServiceExtension } from './kinesis';
+import { LambdaServiceExtension } from './lambda';
 import { S3ServiceExtension } from './s3';
 import { SecretsManagerServiceExtension } from './secretsmanager';
 import type { ServiceExtension } from './ServiceExtension';
+import { SnsServiceExtension } from './sns';
+import { SqsServiceExtension } from './sqs';
 import { StepFunctionsServiceExtension } from './stepfunctions';
 
 export class ServicesExtensions implements ServiceExtension {
@@ -16,7 +19,10 @@ export class ServicesExtensions implements ServiceExtension {
     this._services = new Map();
     this._services.set('SecretsManager', new SecretsManagerServiceExtension());
     this._services.set('SFN', new StepFunctionsServiceExtension());
+    this._services.set('SQS', new SqsServiceExtension());
+    this._services.set('SNS', new SnsServiceExtension());
     this._services.set('DynamoDB', new DynamodbServiceExtension());
+    this._services.set('Lambda', new LambdaServiceExtension());
     this._services.set('S3', new S3ServiceExtension());
     this._services.set('Kinesis', new KinesisServiceExtension());
   }

packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/lambda.ts

@@ -0,0 +1,86 @@
+import type { Span } from '@sentry/core';
+import { debug, SPAN_KIND } from '@sentry/core';
+import { DEBUG_BUILD } from '../../../../debug-build';
+import {
+  ATTR_FAAS_EXECUTION,
+  ATTR_FAAS_INVOKED_NAME,
+  ATTR_FAAS_INVOKED_PROVIDER,
+  ATTR_FAAS_INVOKED_REGION,
+} from '../constants';
+import type { NormalizedRequest, NormalizedResponse } from '../types';
+import { getPropagationHeaders } from './MessageAttributes';
+import type { RequestMetadata, ServiceExtension } from './ServiceExtension';
+
+const INVOKE_COMMAND = 'Invoke';
+
+export class LambdaServiceExtension implements ServiceExtension {
+  public requestPreSpanHook(request: NormalizedRequest): RequestMetadata {
+    const functionName = request.commandInput?.FunctionName;
+
+    let spanAttributes: Record<string, unknown> = {};
+    let spanName: string | undefined;
+
+    switch (request.commandName) {
+      case INVOKE_COMMAND:
+        spanAttributes = {
+          [ATTR_FAAS_INVOKED_NAME]: functionName,
+          [ATTR_FAAS_INVOKED_PROVIDER]: 'aws',
+        };
+        if (request.region) {
+          spanAttributes[ATTR_FAAS_INVOKED_REGION] = request.region;
+        }
+        spanName = `${functionName} ${INVOKE_COMMAND}`;
+        break;
+    }
+    return {
+      spanAttributes,
+      spanKind: SPAN_KIND.CLIENT,
+      spanName,
+    };
+  }
+
+  public requestPostSpanHook(request: NormalizedRequest, span: Span): void {
+    if (request.commandName === INVOKE_COMMAND && request.commandInput) {
+      request.commandInput.ClientContext = injectLambdaPropagationContext(request.commandInput.ClientContext, span);
+    }
+  }
+
+  public responseHook(response: NormalizedResponse, span: Span): void {
+    if (response.request.commandName === INVOKE_COMMAND) {
+      span.setAttribute(ATTR_FAAS_EXECUTION, response.requestId);
+    }
+  }
+}
+
+function injectLambdaPropagationContext(clientContext: string | undefined, span: Span): string | undefined {
+  try {
+    const propagatedContext = getPropagationHeaders(span);
+
+    const parsedClientContext = clientContext ? JSON.parse(Buffer.from(clientContext, 'base64').toString('utf8')) : {};
+
+    const updatedClientContext = {
+      ...parsedClientContext,
+      custom: {
+        ...parsedClientContext.custom,
+        ...propagatedContext,
+      },
+    };
+
+    const encodedClientContext = Buffer.from(JSON.stringify(updatedClientContext)).toString('base64');
+
+    // The length of client context is capped at 3583 bytes of base64 encoded data
+    // (https://docs.aws.amazon.com/lambda/latest/dg/API_Invoke.html#API_Invoke_RequestSyntax)
+    if (encodedClientContext.length > 3583) {
+      DEBUG_BUILD &&
+        debug.warn(
+          '[orchestrion:aws-sdk] cannot set trace propagation on lambda invoke parameters due to ClientContext length limitations.',
+        );
+      return clientContext;
+    }
+
+    return encodedClientContext;
+  } catch (e) {
+    DEBUG_BUILD && debug.log('[orchestrion:aws-sdk] failed to set trace propagation on lambda ClientContext', e);
+    return clientContext;
+  }
+}

packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/sns.ts

@@ -0,0 +1,74 @@
+import type { Span, SpanKindValue } from '@sentry/core';
+import { SPAN_KIND } from '@sentry/core';
+import { MESSAGING_DESTINATION_NAME, MESSAGING_SYSTEM } from '@sentry/conventions/attributes';
+import {
+  ATTR_AWS_SNS_TOPIC_ARN,
+  ATTR_MESSAGING_DESTINATION,
+  ATTR_MESSAGING_DESTINATION_KIND,
+  MESSAGING_DESTINATION_KIND_VALUE_TOPIC,
+} from '../constants';
+import type { NormalizedRequest, NormalizedResponse } from '../types';
+import { injectPropagationContext } from './MessageAttributes';
+import type { RequestMetadata, ServiceExtension } from './ServiceExtension';
+
+export class SnsServiceExtension implements ServiceExtension {
+  public requestPreSpanHook(request: NormalizedRequest): RequestMetadata {
+    let spanKind: SpanKindValue = SPAN_KIND.CLIENT;
+    let spanName = `SNS ${request.commandName}`;
+    const spanAttributes: Record<string, unknown> = {
+      [MESSAGING_SYSTEM]: 'aws.sns',
+    };
+
+    if (request.commandName === 'Publish') {
+      spanKind = SPAN_KIND.PRODUCER;
+
+      spanAttributes[ATTR_MESSAGING_DESTINATION_KIND] = MESSAGING_DESTINATION_KIND_VALUE_TOPIC;
+      const { TopicArn, TargetArn, PhoneNumber } = request.commandInput;
+      const destinationName = extractDestinationName(TopicArn, TargetArn, PhoneNumber);
+      spanAttributes[ATTR_MESSAGING_DESTINATION] = destinationName;
+      spanAttributes[MESSAGING_DESTINATION_NAME] = TopicArn || TargetArn || PhoneNumber || 'unknown';
+
+      spanName = `${PhoneNumber ? 'phone_number' : destinationName} send`;
+    }
+
+    const topicArn = request.commandInput?.TopicArn;
+    if (topicArn) {
+      spanAttributes[ATTR_AWS_SNS_TOPIC_ARN] = topicArn;
+    }
+
+    return {
+      spanAttributes,
+      spanKind,
+      spanName,
+    };
+  }
+
+  public requestPostSpanHook(request: NormalizedRequest, span: Span): void {
+    if (request.commandName === 'Publish') {
+      const origMessageAttributes = request.commandInput.MessageAttributes ?? {};
+      request.commandInput.MessageAttributes = injectPropagationContext(origMessageAttributes, span);
+    }
+  }
+
+  public responseHook(response: NormalizedResponse, span: Span): void {
+    const topicArn = response.data?.TopicArn;
+    if (topicArn) {
+      span.setAttribute(ATTR_AWS_SNS_TOPIC_ARN, topicArn);
+    }
+  }
+}
+
+function extractDestinationName(topicArn: string, targetArn: string, phoneNumber: string): string {
+  if (topicArn || targetArn) {
+    const arn = topicArn ?? targetArn;
+    try {
+      return arn.substring(arn.lastIndexOf(':') + 1);
+    } catch {
+      return arn;
+    }
+  } else if (phoneNumber) {
+    return phoneNumber;
+  } else {
+    return 'unknown';
+  }
+}