add custom user agent middleware

Author: Alexander Melnyk

packages/commons/package.json

@@ -14,6 +14,7 @@
     "test:unit": "jest --group=unit --detectOpenHandles --coverage --verbose",
     "test:e2e": "echo 'Not Applicable'",
     "watch": "jest --watch",
+    "generateVersionFile": "echo \"// this file is auto generated, do not modify\nmodule.exports = { version: '$(jq -r '.version' package.json)' };\" > src/version.ts",
     "build": "tsc",
     "lint": "eslint --ext .ts,.js --no-error-on-unmatched-pattern .",
     "lint-fix": "eslint --fix --ext .ts,.js --no-error-on-unmatched-pattern .",

packages/commons/src/index.ts

@@ -5,3 +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 './middleware';

packages/commons/src/middleware/index.ts

@@ -1,2 +1,3 @@
 export * from './cleanupMiddlewares';
 export * from './constants';
+export * from './userAgentMiddleware';

packages/commons/src/middleware/userAgentMiddleware.ts

@@ -0,0 +1,39 @@
+// eslint-disable-next-line @typescript-eslint/ban-ts-comment
+// @ts-nocheck
+/**
+ * @internal
+ */
+import { version as PT_VERSION } from '../version';
+import { userAgentMiddleware } from '@aws-sdk/middleware-user-agent';
+
+const EXEC_ENV = process.env.AWS_EXECUTION_ENV || 'NA';
+const middlewareOptions = {
+  step: 'finalizeRequest',
+  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 (next, _context) => async (args) => {
+    const powertoolsUserAgent = `PT/${feature}/${PT_VERSION} PTEnv/${EXEC_ENV}`;
+    args.request.headers['user-agent'] = args.request.headers['user-agent']
+      ? `${args.request.headers['user-agent']} ${powertoolsUserAgent}`
+      : `${powertoolsUserAgent}`;
+
+    return await next(args);
+  };
+};
+
+const addUserAgentMiddleware = (client, feature): void => {
+  client.middlewareStack.add(
+    customUserAgentMiddleware(feature),
+    middlewareOptions
+  );
+};
+
+export { addUserAgentMiddleware, userAgentMiddleware };

packages/commons/src/version.ts

@@ -0,0 +1,2 @@
+// this file is auto generated, do not modify
+module.exports = { version: '1.10.0' };

packages/commons/tests/unit/userAgentMiddleware.test.ts

@@ -0,0 +1,49 @@
+import { addUserAgentMiddleware } from '../../src/middleware/userAgentMiddleware';
+import { InvokeCommand, LambdaClient } from '@aws-sdk/client-lambda';
+
+describe('Function: addUserAgentMiddleware', () => {
+  it('adds powertools user agent to request header at the end', async () => {
+    const lambdaClient = new LambdaClient({
+      logger: console,
+      region: 'us-east-1',
+      endpoint: 'http://localhost:9001',
+      maxAttempts: 1, // disable retry to have the correct number of assertions
+    });
+    addUserAgentMiddleware(lambdaClient, 'my-feature');
+
+    expect(lambdaClient.middlewareStack.identify()).toContain(
+      'addPowertoolsToUserAgent: POWERTOOLS,USER_AGENT'
+    );
+
+    lambdaClient.middlewareStack.add(
+      (next) => (args) => {
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-ignore
+        const userAgent = args?.request?.headers['user-agent'];
+        expect(userAgent).toContain('PT/my-feature/1.10.0 PTEnv/NA');
+        // make sure it's at the end of the user agent
+        expect(
+          // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+          // @ts-ignore
+          userAgent
+            ?.split(' ')
+            .slice(userAgent?.split(' ').length - 2)
+            .join(' ')
+        ).toEqual('PT/my-feature/1.10.0 PTEnv/NA');
+
+        return next(args);
+      },
+      {
+        step: 'finalizeRequest',
+      }
+    );
+    try {
+      await lambdaClient.send(
+        new InvokeCommand({
+          FunctionName: 'test',
+          Payload: new TextEncoder().encode(JSON.stringify('foo')),
+        })
+      );
+    } catch (e) {}
+  });
+});