refactor(commons): aws sdk util + identity check (#1708)

Author: dreamorosi

package-lock.json

@@ -44,13 +44,5 @@
     "powertools",
     "serverless",
     "nodejs"
-  ],
-  "devDependencies": {
-    "@aws-sdk/client-appconfigdata": "^3.413.0",
-    "@aws-sdk/client-dynamodb": "^3.413.0",
-    "@aws-sdk/client-lambda": "^3.413.0",
-    "@aws-sdk/client-secrets-manager": "^3.413.0",
-    "@aws-sdk/client-ssm": "^3.413.0",
-    "@aws-sdk/util-utf8-node": "^3.259.0"
-  }
+  ]
 }

packages/commons/package.json

@@ -0,0 +1,2 @@
+export { addUserAgentMiddleware } from './userAgentMiddleware';
+export { isSdkClient } from './utils';

packages/commons/src/awsSdk/index.ts

@@ -0,0 +1,57 @@
+import { PT_VERSION } from '../version';
+import { isSdkClient } from './utils';
+import type { MiddlewareArgsLike } from '../types/awsSdk';
+
+/**
+ * @internal
+ */
+const EXEC_ENV = process.env.AWS_EXECUTION_ENV || 'NA';
+const middlewareOptions = {
+  relation: 'after',
+  toMiddleware: 'getUserAgentMiddleware',
+  name: 'addPowertoolsToUserAgent',
+  tags: ['POWERTOOLS', 'USER_AGENT'],
+};
+
+/**
+ * @internal
+ * returns a middleware function for the MiddlewareStack, that can be used for the SDK clients
+ * @param feature
+ */
+const customUserAgentMiddleware = (feature: string) => {
+  return <T extends MiddlewareArgsLike>(next: (arg0: T) => Promise<T>) =>
+    async (args: T) => {
+      const powertoolsUserAgent = `PT/${feature}/${PT_VERSION} PTEnv/${EXEC_ENV}`;
+      args.request.headers[
+        'user-agent'
+      ] = `${args.request.headers['user-agent']} ${powertoolsUserAgent}`;
+
+      return await next(args);
+    };
+};
+
+const addUserAgentMiddleware = (client: unknown, feature: string): void => {
+  try {
+    if (isSdkClient(client)) {
+      if (
+        client.middlewareStack
+          .identify()
+          .includes('addPowertoolsToUserAgent: POWERTOOLS,USER_AGENT')
+      ) {
+        return;
+      }
+      client.middlewareStack.addRelativeTo(
+        customUserAgentMiddleware(feature),
+        middlewareOptions
+      );
+    } else {
+      throw new Error(
+        `The client provided does not match the expected interface`
+      );
+    }
+  } catch (e) {
+    console.warn('Failed to add user agent middleware', e);
+  }
+};
+
+export { customUserAgentMiddleware, addUserAgentMiddleware };

packages/commons/src/awsSdk/userAgentMiddleware.ts

@@ -0,0 +1,25 @@
+import { SdkClient } from '../types/awsSdk';
+
+/**
+ * @internal
+ * Type guard to check if the client provided is a valid AWS SDK v3 client
+ */
+const isSdkClient = (client: unknown): client is SdkClient =>
+  typeof client === 'object' &&
+  client !== null &&
+  'send' in client &&
+  typeof client.send === 'function' &&
+  'config' in client &&
+  client.config !== undefined &&
+  typeof client.config === 'object' &&
+  client.config !== null &&
+  'middlewareStack' in client &&
+  client.middlewareStack !== undefined &&
+  typeof client.middlewareStack === 'object' &&
+  client.middlewareStack !== null &&
+  'identify' in client.middlewareStack &&
+  typeof client.middlewareStack.identify === 'function' &&
+  'addRelativeTo' in client.middlewareStack &&
+  typeof client.middlewareStack.addRelativeTo === 'function';
+
+export { isSdkClient };

packages/commons/src/awsSdk/utils.ts

@@ -5,4 +5,4 @@ export * as ContextExamples from './samples/resources/contexts';
 export * as Events from './samples/resources/events';
 export * from './types/middy';
 export * from './types/utils';
-export * from './userAgentMiddleware';
+export * from './awsSdk';

packages/commons/src/index.ts

@@ -0,0 +1,22 @@
+/**
+ * @internal
+ * Minimal interface for an AWS SDK v3 client
+ */
+interface SdkClient {
+  send: (args: unknown) => Promise<unknown>;
+  config: {
+    serviceId: string;
+  };
+  middlewareStack: {
+    identify: () => string[];
+    addRelativeTo: (middleware: unknown, options: unknown) => void;
+  };
+}
+
+/**
+ * @internal
+ * Minimal type for the arguments passed to a middleware function
+ */
+type MiddlewareArgsLike = { request: { headers: { [key: string]: string } } };
+
+export { SdkClient, MiddlewareArgsLike };

packages/commons/src/types/awsSdk.ts

@@ -0,0 +1,148 @@
+import { addUserAgentMiddleware, isSdkClient } from '../../src/awsSdk';
+import { PT_VERSION as version } from '../../src/version';
+import { customUserAgentMiddleware } from '../../src/awsSdk/userAgentMiddleware';
+
+describe('Helpers: awsSdk', () => {
+  describe('Function: userAgentMiddleware', () => {
+    beforeAll(() => {
+      jest.spyOn(console, 'warn').mockImplementation(() => ({}));
+    });
+
+    it('handles gracefully failures in adding a middleware and only log a warning', () => {
+      // Prepare
+      const client = {
+        middlewareStack: {
+          addRelativeTo: () => {
+            throw new Error('test');
+          },
+        },
+      };
+      const warningSpy = jest
+        .spyOn(console, 'warn')
+        .mockImplementation(() => ({}));
+
+      // Act & Assess
+      expect(() => addUserAgentMiddleware(client, 'my-feature')).not.toThrow();
+      expect(warningSpy).toHaveBeenCalledTimes(1);
+    });
+
+    it('should return and do nothing if the client already has a Powertools UA middleware', async () => {
+      // Prepare
+      const client = {
+        middlewareStack: {
+          identify: () => 'addPowertoolsToUserAgent: POWERTOOLS,USER_AGENT',
+          addRelativeTo: jest.fn(),
+        },
+        send: jest.fn(),
+        config: {
+          defaultSigningName: 'bar',
+        },
+      };
+      const feature = 'my-feature';
+
+      // Act
+      addUserAgentMiddleware(client, feature);
+
+      // Assess
+      expect(client.middlewareStack.addRelativeTo).toHaveBeenCalledTimes(0);
+    });
+
+    it('should call the function to add the middleware with the correct arguments', async () => {
+      // Prepare
+      const client = {
+        middlewareStack: {
+          identify: () => '',
+          addRelativeTo: jest.fn(),
+        },
+        send: jest.fn(),
+        config: {
+          defaultSigningName: 'bar',
+        },
+      };
+      const feature = 'my-feature';
+
+      // Act
+      addUserAgentMiddleware(client, feature);
+
+      // Assess
+      expect(client.middlewareStack.addRelativeTo).toHaveBeenNthCalledWith(
+        1,
+        expect.any(Function),
+        {
+          relation: 'after',
+          toMiddleware: 'getUserAgentMiddleware',
+          name: 'addPowertoolsToUserAgent',
+          tags: ['POWERTOOLS', 'USER_AGENT'],
+        }
+      );
+    });
+  });
+
+  describe('Function: customUserAgentMiddleware', () => {
+    it('returns a middleware function', () => {
+      // Prepare
+      const feature = 'my-feature';
+
+      // Act
+      const middleware = customUserAgentMiddleware(feature);
+
+      // Assess
+      expect(middleware).toBeInstanceOf(Function);
+    });
+
+    it('adds the Powertools UA to the request headers', async () => {
+      // Prepare
+      const feature = 'my-feature';
+      const middleware = customUserAgentMiddleware(feature);
+      const next = jest.fn();
+      const args = {
+        request: {
+          headers: {
+            'user-agent': 'foo',
+          },
+        },
+      };
+
+      // Act
+      await middleware(next)(args);
+
+      // Assess
+      expect(args.request.headers['user-agent']).toEqual(
+        `foo PT/my-feature/${version} PTEnv/NA`
+      );
+    });
+  });
+
+  describe('Function: isSdkClient', () => {
+    it('returns true if the client is a valid AWS SDK v3 client', () => {
+      // Prepare
+      const client = {
+        send: jest.fn(),
+        config: {
+          defaultSigningName: 'bar',
+        },
+        middlewareStack: {
+          identify: () => '',
+          addRelativeTo: jest.fn(),
+        },
+      };
+
+      // Act
+      const result = isSdkClient(client);
+
+      // Assess
+      expect(result).toEqual(true);
+    });
+
+    it('returns false if the client is not a valid AWS SDK v3 client', () => {
+      // Prepare
+      const client = {};
+
+      // Act
+      const result = isSdkClient(client);
+
+      // Assess
+      expect(result).toEqual(false);
+    });
+  });
+});