feat(middleware-user-agent): add injectUserAgent plugin

Author: AllanZhengYP

packages/middleware-user-agent/jest.config.js

@@ -0,0 +1,5 @@
+const base = require("../../jest.config.base.js");
+
+module.exports = {
+  ...base,
+};

packages/middleware-user-agent/package.json

@@ -22,6 +22,7 @@
     "tslib": "^1.8.0"
   },
   "devDependencies": {
+    "@aws-sdk/middleware-stack": "3.0.0",
     "@aws-sdk/types": "3.0.0",
     "@types/jest": "^26.0.4",
     "jest": "^26.1.0",

packages/middleware-user-agent/src/constants.ts

@@ -0,0 +1,7 @@
+export const USER_AGENT = "user-agent";
+
+export const X_AMZ_USER_AGENT = "x-amz-user-agent";
+
+export const SPACE = " ";
+
+export const UA_ESCAPE_REGEX = /[^\!\#\$\%\&\'\*\+\-\.\^\_\`\|\~\d\w]/g;

packages/middleware-user-agent/src/index.ts

@@ -1,2 +1,3 @@
 export * from "./configurations";
 export * from "./user-agent-middleware";
+export * from "./inject-ua-plugin";

packages/middleware-user-agent/src/inject-ua-plugin.spec.ts

@@ -0,0 +1,38 @@
+import { constructStack } from "@aws-sdk/middleware-stack";
+import { Handler } from "@aws-sdk/types";
+
+import { injectUserAgent } from "./inject-ua-plugin";
+
+describe("injectUserAgent", () => {
+  it("should inject user agent to handler execution context", async () => {
+    const stack = constructStack();
+    const addSpy = jest.spyOn(stack, "add");
+    stack.use(
+      injectUserAgent({
+        name: "foo-api",
+        version: "1.0.0",
+        userAgentIdentifier: "FooApi",
+      })
+    );
+    stack.add(
+      (next, context) => (args) => {
+        // @ts-ignore
+        args.input.context = context;
+        return next(args);
+      },
+      { step: "build" }
+    );
+    const mockCoreHandler: Handler<object, object> = jest.fn();
+    const handler = stack.resolve(mockCoreHandler, {});
+    await handler({ input: {} });
+    expect(mockCoreHandler).toHaveBeenCalled();
+    expect((mockCoreHandler as jest.Mock).mock.calls[0][0]?.input?.context).toMatchObject({
+      userAgent: [["foo-api", "1.0.0"]],
+    });
+    expect(addSpy.mock.calls[0][1]).toMatchObject({
+      name: "injectFooApiUAMiddleware",
+      step: "initialize",
+      tags: ["SET_USER_AGENT", "USER_AGENT"],
+    });
+  });
+});

packages/middleware-user-agent/src/inject-ua-plugin.ts

@@ -0,0 +1,38 @@
+import { MiddlewareStack, Pluggable } from "@aws-sdk/types";
+
+export interface InjectUserAgentOptions {
+  /**
+   * The aame of the user agent to be injected to the request, e.g. a dependency name.
+   */
+  name: string;
+  /**
+   * The value of the user agent to be injected to the request, e.g. a dependency version.
+   */
+  version?: string;
+  /**
+   * The identifier of the middleware that injecting this user agent pair to the handler execution context.
+   * The middleware name will be `inject${userAgentIdentifier}UAMiddleware`. You can use this name to remove
+   * this middleware.
+   */
+  userAgentIdentifier?: string;
+}
+
+/**
+ * A helper plugin that inject user agent pair.
+ */
+export const injectUserAgent = (options: InjectUserAgentOptions): Pluggable<any, any> => ({
+  applyToStack: (stack: MiddlewareStack<any, any>) => {
+    stack.add(
+      (next, context) => (args) => {
+        if (!context.userAgent) context.userAgent = [];
+        context.userAgent.push([options.name, options.version]);
+        return next(args);
+      },
+      {
+        step: "initialize",
+        name: options.userAgentIdentifier ? `inject${options.userAgentIdentifier}UAMiddleware` : undefined,
+        tags: ["SET_USER_AGENT", "USER_AGENT"],
+      }
+    );
+  },
+});

packages/middleware-user-agent/src/user-agent-middleware.spec.ts

@@ -0,0 +1,70 @@
+import { HttpRequest } from "@aws-sdk/protocol-http";
+import { UserAgentPair } from "@aws-sdk/types";
+
+import { USER_AGENT, X_AMZ_USER_AGENT } from "./constants";
+import { userAgentMiddleware } from "./user-agent-middleware";
+
+describe("userAgentMiddleware", () => {
+  const mockNextHandler = jest.fn();
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it("should collect user agent pair from default, custom-supplied, and handler context", async () => {
+    const middleware = userAgentMiddleware({
+      defaultUserAgent: async () => [
+        ["default_agent", "1.0.0"],
+        ["aws-sdk-js", "1.0.0"],
+      ],
+      customUserAgent: [["custom_ua/abc"]],
+      runtime: "node",
+    });
+    const handler = middleware(mockNextHandler, { userAgent: [["cfg/retry-mode", "standard"]] });
+    await handler({ input: {}, request: new HttpRequest({ headers: {} }) });
+    expect(mockNextHandler.mock.calls.length).toEqual(1);
+    const sdkUserAgent = mockNextHandler.mock.calls[0][0].request.headers[X_AMZ_USER_AGENT];
+    expect(sdkUserAgent).toEqual(expect.stringContaining("aws-sdk-js/1.0.0"));
+    expect(sdkUserAgent).toEqual(expect.stringContaining("default_agent/1.0.0"));
+    expect(sdkUserAgent).toEqual(expect.stringContaining("custom_ua/abc"));
+    expect(sdkUserAgent).toEqual(expect.stringContaining("cfg/retry-mode/standard"));
+    expect(mockNextHandler.mock.calls[0][0].request.headers[USER_AGENT]).toEqual(
+      expect.stringContaining("aws-sdk-js/1.0.0")
+    );
+  });
+
+  it(`should not set ${USER_AGENT} header in browser`, async () => {
+    const middleware = userAgentMiddleware({
+      defaultUserAgent: async () => [["aws-sdk-js", "1.0.0"]],
+      runtime: "browser",
+    });
+    const handler = middleware(mockNextHandler, {});
+    await handler({ input: {}, request: new HttpRequest({ headers: {} }) });
+    expect(mockNextHandler.mock.calls.length).toEqual(1);
+    expect(mockNextHandler.mock.calls[0][0].request.headers[USER_AGENT]).toBeUndefined();
+  });
+
+  describe("should sanitize the user agent string", () => {
+    const cases: { ua: UserAgentPair; expected: string }[] = [
+      { ua: ["/name", "1.0.0"], expected: "name/1.0.0" },
+      { ua: ["Name", "1.0.0"], expected: "Name/1.0.0" },
+      { ua: ["md/name", "1.0.0"], expected: "md/name/1.0.0" },
+      { ua: ["$prefix/&name", "1.0.0"], expected: "$prefix/&name/1.0.0" },
+      { ua: ["name(or not)", "1.0.0"], expected: "name_or_not_/1.0.0" },
+      { ua: ["name", "1.0.0(test_version)"], expected: "name/1.0.0_test_version" },
+    ];
+    for (const { ua, expected } of cases) {
+      it(`should sanitize user agent ${ua} to ${expected}`, async () => {
+        const middleware = userAgentMiddleware({
+          defaultUserAgent: async () => [ua],
+          runtime: "browser",
+        });
+        const handler = middleware(mockNextHandler, {});
+        await handler({ input: {}, request: new HttpRequest({ headers: {} }) });
+        expect(mockNextHandler.mock.calls[0][0].request.headers[X_AMZ_USER_AGENT]).toEqual(
+          expect.stringContaining(expected)
+        );
+      });
+    }
+  });
+});

packages/middleware-user-agent/src/user-agent-middleware.ts

@@ -11,12 +11,7 @@ import {
 } from "@aws-sdk/types";
 
 import { UserAgentResolvedConfig } from "./configurations";
-
-const USER_AGENT = "user-agent";
-const X_AMZ_USER_AGENT = "x-amz-user-agent";
-const SPACE = " ";
-
-const UA_ESCAPE_REGEX = /[^\!\#\$\%\&\'\*\+\-\.\^\_\`\|\~\d\w]/g;
+import { SPACE, UA_ESCAPE_REGEX, USER_AGENT, X_AMZ_USER_AGENT } from "./constants";
 
 /**
  * Build user agent header sections from:
@@ -59,13 +54,18 @@ export const userAgentMiddleware = (options: UserAgentResolvedConfig) => <Output
 
 /**
  * Escape the each pair according to https://tools.ietf.org/html/rfc5234 and join the pair with pattern `name/version`.
+ * User agent name may include prefix like `md/`, `api/`, `os/` etc., we should not escape the `/` after the prefix.
  * @private
  */
-export const escapeUserAgent = (pair: UserAgentPair): string =>
-  pair
-    .filter((item) => item !== undefined)
-    .map((item) => item.replace(UA_ESCAPE_REGEX, "_"))
+export const escapeUserAgent = ([name, version]: UserAgentPair): string => {
+  const prefixSeparatorIndex = name.indexOf("/");
+  const prefix = name.substring(0, prefixSeparatorIndex); // If no prefix, prefix is just ""
+  const uaName = name.substring(prefixSeparatorIndex + 1);
+  return [prefix, uaName, version]
+    .filter((item) => item && item.length > 0)
+    .map((item) => item?.replace(UA_ESCAPE_REGEX, "_"))
     .join("/");
+};
 
 export const getUserAgentMiddlewareOptions: BuildHandlerOptions = {
   name: "getUserAgentMiddleware",