fix: using also user-agent to identify sqs calls (#180)

src/instrumentations/https/http.ts

@@ -9,12 +9,7 @@ import { Span } from '@opentelemetry/api';
 import { InstrumentationIfc } from '../hooksIfc';
 import { logger } from '../../logging';
 import { getAwsServiceData } from '../../spans/awsSpan';
-import {
-  isAwsService,
-  runOneTimeWrapper,
-  safeExecute,
-  getSpanAttributeMaxLength,
-} from '../../utils';
+import { runOneTimeWrapper, safeExecute, getSpanAttributeMaxLength } from '../../utils';
 import { contentType, scrubHttpPayload } from '../../tools/payloads';
 
 // eslint-disable-next-line @typescript-eslint/no-empty-function
@@ -148,10 +143,15 @@ export class Http {
         ScrubContext.HTTP_RESPONSE_BODY
       );
       span.setAttribute('http.response.body', scrubbed);
+
       try {
-        if (isAwsService(requestRawData.request.host, requestRawData.response)) {
-          span.setAttributes(getAwsServiceData(requestRawData.request, requestRawData.response));
-          span.setAttribute('aws.region', span.attributes?.['http.host'].split('.')[1]);
+        const serviceAttributes = getAwsServiceData(
+          requestRawData.request,
+          requestRawData.response,
+          span
+        );
+        if (serviceAttributes) {
+          span.setAttributes(serviceAttributes);
         }
       } catch (e) {
         logger.debug('Failed to parse aws service data', e);

src/spans/awsSpan.test.ts

@@ -0,0 +1,74 @@
+import {
+  AwsOtherService,
+  AwsParsedService,
+  getAwsServiceData,
+  getAwsServiceFromHost,
+} from './awsSpan';
+import { BasicTracerProvider } from '@opentelemetry/sdk-trace-base';
+
+describe('awsSpan', () => {
+  describe('getAwsServiceFromHost', () => {
+    test('with an ApiGateway', () => {
+      expect(
+        getAwsServiceFromHost(
+          new URL('https://my_happy_api.execute-api.eu-central-1.amazonaws.com/production/')
+            .hostname
+        )
+      ).toBe(AwsOtherService.ApiGateway);
+    });
+
+    test('with an SQS queue URL', () => {
+      // Example from https://docs.aws.amazon.com/AWSSimpleQueueService/latest/APIReference/API_GetQueueUrl.html
+      expect(
+        getAwsServiceFromHost(
+          new URL('https://sqs.us-east-1.amazonaws.com/177715257436/MyQueue').hostname
+        )
+      ).toBe(AwsParsedService.SQS);
+    });
+  });
+
+  describe('getAwsServiceData', () => {
+    describe('SQS', () => {
+      test('SQS queue', () => {
+        const requestData = {
+          body: '',
+        };
+        const responseData = {
+          body: '<?xml version="1.0"?><SendMessageResponse xmlns="http://queue.amazonaws.com/doc/2012-11-05/"><SendMessageResult><MessageId>85dc3997-b060-47bc-9d89-c754d7260dbd</MessageId><MD5OfMessageBody>c5cb6abef11b88049177473a73ed662f</MD5OfMessageBody></SendMessageResult><ResponseMetadata><RequestId>b6b5a045-23c6-5e3a-a54f-f7dd99f7b379</RequestId></ResponseMetadata></SendMessageResponse>',
+        };
+        const provider = new BasicTracerProvider();
+        const root = provider.getTracer('default').startSpan('root');
+        root.setAttribute('http.user_agent', 'aws-sqsd/3.0.4');
+        root.end();
+        // @ts-ignore
+        const awsServiceData = getAwsServiceData(requestData, responseData, root);
+
+        expect(awsServiceData).toMatchObject({
+          messageId: '85dc3997-b060-47bc-9d89-c754d7260dbd',
+        });
+        expect(awsServiceData).not.toHaveProperty('aws.region');
+      });
+
+      test('Elastic Beanstalk SQS Daemon', () => {
+        const requestData = {
+          host: 'localhost',
+          body: '',
+        };
+        const responseData = {
+          body: '<?xml version="1.0"?><SendMessageResponse xmlns="http://queue.amazonaws.com/doc/2012-11-05/"><SendMessageResult><MessageId>85dc3997-b060-47bc-9d89-c754d7260dbd</MessageId><MD5OfMessageBody>c5cb6abef11b88049177473a73ed662f</MD5OfMessageBody></SendMessageResult><ResponseMetadata><RequestId>b6b5a045-23c6-5e3a-a54f-f7dd99f7b379</RequestId></ResponseMetadata></SendMessageResponse>',
+        };
+        const provider = new BasicTracerProvider();
+        const root = provider.getTracer('default').startSpan('root');
+        root.setAttribute('http.user_agent', 'aws-sqsd/3.0.4');
+        root.end();
+        // @ts-ignore
+        const awsServiceData = getAwsServiceData(requestData, responseData, root);
+
+        expect(awsServiceData).toMatchObject({
+          messageId: '85dc3997-b060-47bc-9d89-c754d7260dbd',
+        });
+        expect(awsServiceData).not.toHaveProperty('aws.region');
+      });
+    });
+  });
+});

src/spans/awsSpan.ts

@@ -1,3 +1,6 @@
+import { Span } from '@opentelemetry/api';
+import { SemanticAttributes } from '@opentelemetry/semantic-conventions';
+
 import {
   apigwParser,
   awsParser,
@@ -9,19 +12,42 @@ import {
   sqsParser,
 } from '../parsers/aws';
 
-export const EXTERNAL_SERVICE = 'external';
+export enum AwsOtherService {
+  ApiGateway,
+  ExternalService,
+  ElasticBeanstalkSqsDaemon,
+}
 
-export const AWS_PARSED_SERVICES = ['dynamodb', 'sns', 'lambda', 'sqs', 'kinesis', 'events'];
+export enum AwsParsedService {
+  DynamoDB = 'dynamodb',
+  EventBridge = 'events',
+  Kinesis = 'kinesis',
+  Lambda = 'lambda',
+  SNS = 'sns',
+  SQS = 'sqs',
+}
 
-export const getAwsServiceFromHost = (host = '') => {
-  const service = host.split('.')[0];
-  if (AWS_PARSED_SERVICES.includes(service)) {
-    return service;
+const AMAZON_REQUESTID_HEADER_NAME = 'x-amzn-requestid';
+
+export const getAwsServiceFromHost = (host = ''): AwsParsedService | AwsOtherService => {
+  if (host?.includes('.execute-api.')) {
+    // E.g. `my_happy_api.execute-api.eu-central-1.amazonaws.com`
+    return AwsOtherService.ApiGateway;
   }
 
-  if (host.includes('execute-api')) return 'apigw';
+  /*
+   * The AWS service name for the API is usually the first segment of the host, e.g.,
+   * `sqs.us-east-1.amazonaws.com`.
+   */
+  const service = host.split('.')[0];
+
+  for (const awsParsedService in AwsParsedService) {
+    if (AwsParsedService[awsParsedService].includes(service)) {
+      return AwsParsedService[awsParsedService];
+    }
+  }
 
-  return EXTERNAL_SERVICE;
+  return AwsOtherService.ExternalService;
 };
 
 export type AwsServiceData = {
@@ -34,24 +60,66 @@ export type AwsServiceData = {
 };
 
 // eslint-disable-next-line @typescript-eslint/no-unused-vars
-export const getAwsServiceData = (requestData, responseData): AwsServiceData => {
+export const getAwsServiceData = (
+  requestData,
+  responseData,
+  span: Span & { attributes: Record<string, string> }
+): AwsServiceData => {
+  let awsService: AwsParsedService | AwsOtherService;
+
   const { host } = requestData;
-  const awsService = getAwsServiceFromHost(host);
+  if (host?.includes('amazonaws.com')) {
+    awsService = getAwsServiceFromHost(host);
+  } else if (span.attributes[SemanticAttributes.HTTP_USER_AGENT]?.startsWith('aws-sqsd')) {
+    /*
+     * Workaround for Elastic Beanstalk, where a local proxy called "AWS SQS Daemon"
+     * is used to fetch SQS messages, causing the hostname to be `localhost`.
+     *
+     * See https://docs.aws.amazon.com/elasticbeanstalk/latest/dg/using-features-managing-env-tiers.html#worker-daemon
+     */
+    awsService = AwsOtherService.ElasticBeanstalkSqsDaemon;
+  } else if (
+    responseData?.headers[AMAZON_REQUESTID_HEADER_NAME] ||
+    responseData?.headers[AMAZON_REQUESTID_HEADER_NAME]
+  ) {
+    awsService = AwsOtherService.ExternalService;
+  } else {
+    // not an aws service
+    return {};
+  }
 
+  const serviceData = getServiceAttributes(awsService, requestData, responseData);
+
+  // If the service is one in the AwsParsedService enum, we also need to extract the region
+  if (Object.values(AwsParsedService).includes(awsService as AwsParsedService)) {
+    serviceData['aws.region'] = host.split('.')[1];
+  }
+
+  return serviceData;
+};
+
+const getServiceAttributes = (
+  awsService: AwsParsedService | AwsOtherService,
+  requestData,
+  responseData
+) => {
   switch (awsService) {
-    case 'dynamodb':
+    case AwsOtherService.ApiGateway:
+      return apigwParser(requestData, responseData);
+    case AwsParsedService.DynamoDB:
       return dynamodbParser(requestData);
-    case 'sns':
+    case AwsParsedService.SNS:
       return snsParser(requestData, responseData);
-    case 'lambda':
+    case AwsParsedService.Lambda:
       return lambdaParser(requestData, responseData);
-    case 'sqs':
+    case AwsOtherService.ElasticBeanstalkSqsDaemon:
+      // Same as AwsParsedService.SQS
       return sqsParser(requestData, responseData);
-    case 'kinesis':
+    case AwsParsedService.SQS:
+      return sqsParser(requestData, responseData);
+    case AwsParsedService.Kinesis:
       return kinesisParser(requestData, responseData);
-    case 'apigw':
-      return apigwParser(requestData, responseData);
-    case 'events':
+    case AwsParsedService.EventBridge:
       return eventBridgeParser(requestData, responseData);
     default:
       return awsParser(requestData, responseData);

src/utils.ts

@@ -4,6 +4,7 @@ import * as https from 'https';
 
 import { logger } from './logging';
 import { sortify } from './tools/jsonSortify';
+import { Span } from '@opentelemetry/api';
 
 export const DEFAULT_ATTRIBUTE_VALUE_LENGTH_LIMIT = 2048;
 export const DEFAULT_CONNECTION_TIMEOUT = 5000;
@@ -136,17 +137,6 @@ export const safeGet = (obj, arr, dflt = null) => {
   return current || dflt;
 };
 
-export const isAwsService = (host, responseData = undefined): boolean => {
-  if (host && host.includes('amazonaws.com')) {
-    return true;
-  }
-  return !!(
-    responseData &&
-    responseData.headers &&
-    (responseData.headers['x-amzn-requestid'] || responseData.headers['x-amz-request-id'])
-  );
-};
-
 export const parseQueryParams = (queryParams) => {
   return safeExecute(() => {
     if (typeof queryParams !== 'string') return {};