chore(util-endpoint): user input UA prefix

aws · Mar 14, 2023 · 28ab855 · 28ab855
1 parent 94a3480
commit 28ab855
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 11 deletions.
diff --git a/packages/middleware-user-agent/src/user-agent-middleware.spec.ts b/packages/middleware-user-agent/src/user-agent-middleware.spec.ts
@@ -78,18 +78,18 @@ describe("userAgentMiddleware", () => {
             );
           });
 
-          it(`should include internal metadata, user agent ${ua} to md/internal ${expected}`, async () => {
+          it(`should include internal metadata, user agent ${ua} customization: ${expected}`, async () => {
             const middleware = userAgentMiddleware({
               defaultUserAgentProvider: async () => [ua],
               runtime,
             });
 
             // internal variant
-            setPartitionInfo({} as any);
+            setPartitionInfo({} as any, "a-test-prefix");
             const handler = middleware(mockInternalNextHandler, {});
             await handler({ input: {}, request: new HttpRequest({ headers: {} }) });
             expect(mockInternalNextHandler.mock.calls[0][0].request.headers[sdkUserAgentKey]).toEqual(
-              expect.stringContaining("md/internal " + expected)
+              expect.stringContaining("a-test-prefix " + expected)
             );
           });
         }

diff --git a/packages/middleware-user-agent/src/user-agent-middleware.ts b/packages/middleware-user-agent/src/user-agent-middleware.ts
@@ -10,7 +10,7 @@ import {
   Pluggable,
   UserAgentPair,
 } from "@aws-sdk/types";
-import { isUsingDefaultPartitionInfo } from "@aws-sdk/util-endpoints";
+import { getUserAgentPrefix } from "@aws-sdk/util-endpoints";
 
 import { UserAgentResolvedConfig } from "./configurations";
 import { SPACE, UA_ESCAPE_REGEX, USER_AGENT, X_AMZ_USER_AGENT } from "./constants";
@@ -40,10 +40,13 @@ export const userAgentMiddleware =
     const userAgent = context?.userAgent?.map(escapeUserAgent) || [];
     const defaultUserAgent = (await options.defaultUserAgentProvider()).map(escapeUserAgent);
     const customUserAgent = options?.customUserAgent?.map(escapeUserAgent) || [];
-    const internalUserMetadata = !isUsingDefaultPartitionInfo() ? "md/internal" : "";
+    const prefix = getUserAgentPrefix();
 
     // Set value to AWS-specific user agent header
-    const sdkUserAgentValue = [internalUserMetadata, ...defaultUserAgent, ...userAgent, ...customUserAgent].join(SPACE);
+    const sdkUserAgentValue = (prefix ? [prefix] : [])
+      .concat([...defaultUserAgent, ...userAgent, ...customUserAgent])
+      .join(SPACE);
+
     // Get value to be sent with non-AWS-specific user agent header.
     const normalUAValue = [
       ...defaultUserAgent.filter((section) => section.startsWith("aws-sdk-")),

diff --git a/packages/util-endpoints/src/lib/aws/partition.spec.ts b/packages/util-endpoints/src/lib/aws/partition.spec.ts
@@ -1,4 +1,4 @@
-import { partition, setPartitionInfo, useDefaultPartitionInfo } from "./partition";
+import { getUserAgentPrefix, partition, setPartitionInfo, useDefaultPartitionInfo } from "./partition";
 import partitions from "./partitions.json";
 
 const MOCK_DEFAULT_PARTITION = {
@@ -116,4 +116,9 @@ describe("partition", () => {
       supportsFIPS: true,
     });
   });
+
+  it("should optionally set a user agent prefix", async () => {
+    setPartitionInfo(null as any, "a-string-prefix");
+    expect(getUserAgentPrefix()).toBe("a-string-prefix");
+  });
 });
diff --git a/packages/util-endpoints/src/lib/aws/partition.ts b/packages/util-endpoints/src/lib/aws/partition.ts
@@ -29,6 +29,11 @@ export type PartitionsInfo = {
  */
 let selectedPartitionsInfo: PartitionsInfo = partitionsInfo;
 
+/**
+ * @internal
+ */
+let selectedUserAgentPrefix = "";
+
 /**
  * Evaluates a single string argument value as a region, and matches the
  * string value to an AWS partition.
@@ -79,20 +84,20 @@ export const partition = (value: string): EndpointPartition => {
  * Set custom partitions.json data.
  * @internal
  */
-export const setPartitionInfo = (partitionsInfo: PartitionsInfo) => {
+export const setPartitionInfo = (partitionsInfo: PartitionsInfo, userAgentPrefix = "") => {
   selectedPartitionsInfo = partitionsInfo;
+  selectedUserAgentPrefix = userAgentPrefix;
 };
 
 /**
  * Reset to the default partitions.json data.
  * @internal
  */
 export const useDefaultPartitionInfo = () => {
-  setPartitionInfo(partitionsInfo);
+  setPartitionInfo(partitionsInfo, "");
 };
 
 /**
  * @internal
- * @returns whether the default partitions data are selected.
  */
-export const isUsingDefaultPartitionInfo = () => selectedPartitionsInfo === partitionsInfo;
+export const getUserAgentPrefix = () => selectedUserAgentPrefix;