chore(util-endpoints): internal partitions api (#4455)

* chore(util-endpoints): internal partitions api

* chore(middleware-user-agent): add md/internal prefix when using custom partitions

* chore(util-endpoint): add dependency

* chore(util-endpoint): user input UA prefix

packages/middleware-retry/src/AdaptiveRetryStrategy.ts

@@ -11,7 +11,7 @@ export interface AdaptiveRetryStrategyOptions extends StandardRetryStrategyOptio
 }
 
 /**
- * @deprected use AdaptiveRetryStrategy from @aws-sdk/util-retry
+ * @deprecated use AdaptiveRetryStrategy from @aws-sdk/util-retry
  */
 export class AdaptiveRetryStrategy extends StandardRetryStrategy {
   private rateLimiter: RateLimiter;

packages/middleware-retry/src/StandardRetryStrategy.ts

@@ -29,7 +29,7 @@ export interface StandardRetryStrategyOptions {
 }
 
 /**
- * @deprected use StandardRetryStrategy from @aws-sdk/util-retry
+ * @deprecated use StandardRetryStrategy from @aws-sdk/util-retry
  */
 export class StandardRetryStrategy implements RetryStrategy {
   private retryDecider: RetryDecider;

packages/middleware-user-agent/package.json

@@ -22,6 +22,7 @@
   "dependencies": {
     "@aws-sdk/protocol-http": "*",
     "@aws-sdk/types": "*",
+    "@aws-sdk/util-endpoints": "*",
     "tslib": "^2.3.1"
   },
   "devDependencies": {

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

@@ -1,16 +1,22 @@
 import { HttpRequest } from "@aws-sdk/protocol-http";
 import { UserAgentPair } from "@aws-sdk/types";
+import { setPartitionInfo, useDefaultPartitionInfo } from "@aws-sdk/util-endpoints";
 
 import { USER_AGENT, X_AMZ_USER_AGENT } from "./constants";
 import { userAgentMiddleware } from "./user-agent-middleware";
 
 describe("userAgentMiddleware", () => {
   const mockNextHandler = jest.fn();
+  const mockInternalNextHandler = jest.fn();
 
   beforeEach(() => {
     jest.clearAllMocks();
   });
 
+  afterEach(() => {
+    useDefaultPartitionInfo();
+  });
+
   describe("should collect user agent pair from default, custom-supplied, and handler context", () => {
     [
       { runtime: "node", sdkUserAgentKey: USER_AGENT, userAgentKey: X_AMZ_USER_AGENT },
@@ -71,6 +77,21 @@ describe("userAgentMiddleware", () => {
               expect.stringContaining(expected)
             );
           });
+
+          it(`should include internal metadata, user agent ${ua} customization: ${expected}`, async () => {
+            const middleware = userAgentMiddleware({
+              defaultUserAgentProvider: async () => [ua],
+              runtime,
+            });
+
+            // internal variant
+            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("a-test-prefix " + expected)
+            );
+          });
         }
       })
     );

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

@@ -10,6 +10,7 @@ import {
   Pluggable,
   UserAgentPair,
 } from "@aws-sdk/types";
+import { getUserAgentPrefix } from "@aws-sdk/util-endpoints";
 
 import { UserAgentResolvedConfig } from "./configurations";
 import { SPACE, UA_ESCAPE_REGEX, USER_AGENT, X_AMZ_USER_AGENT } from "./constants";
@@ -39,9 +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 prefix = getUserAgentPrefix();
 
     // Set value to AWS-specific user agent header
-    const sdkUserAgentValue = [...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-")),

packages/util-endpoints/src/lib/aws/partition.spec.ts

@@ -1,4 +1,5 @@
-import { partitions } from "./partitions.json";
+import { getUserAgentPrefix, partition, setPartitionInfo, useDefaultPartitionInfo } from "./partition";
+import partitions from "./partitions.json";
 
 const MOCK_DEFAULT_PARTITION = {
   id: "aws",
@@ -33,57 +34,91 @@ const MOCK_PARTITION = {
 };
 
 describe("partition", () => {
-  describe("should reuturn data when default partition exists", () => {
-    jest.isolateModules(() => {
-      jest.mock("./partitions.json", () => ({
-        partitions: [MOCK_DEFAULT_PARTITION, MOCK_PARTITION],
-      }));
-      const { partition } = require("./partition");
-
-      describe("should return the data when region is found", () => {
-        it("returns region data if it exists", () => {
-          const regionWithRegionData = "mock-region-1";
-          expect(partition(regionWithRegionData)).toEqual({
-            ...MOCK_DEFAULT_PARTITION.outputs,
-            ...MOCK_DEFAULT_PARTITION.regions[regionWithRegionData],
-          });
-        });
+  afterEach(() => {
+    useDefaultPartitionInfo();
+  });
 
-        it("returns partition data if region data does not exist", () => {
-          const regionWithoutRegionData = "mock-region-2";
-          expect(partition(regionWithoutRegionData)).toEqual({
-            ...MOCK_DEFAULT_PARTITION.outputs,
-          });
-        });
+  describe("should return data when default partition exists", () => {
+    beforeEach(() => {
+      setPartitionInfo({
+        partitions: [MOCK_DEFAULT_PARTITION, MOCK_PARTITION],
       });
+    });
 
-      it("should return the partition data when region is matched with regionRegex", () => {
-        expect(partition(MOCK_DEFAULT_PARTITION.regionRegex)).toEqual({
+    describe("should return the data when region is found", () => {
+      it("returns region data if it exists", () => {
+        const regionWithRegionData = "mock-region-1";
+        expect(partition(regionWithRegionData)).toEqual({
           ...MOCK_DEFAULT_PARTITION.outputs,
-        });
-        expect(partition(MOCK_PARTITION.regionRegex)).toEqual({
-          ...MOCK_PARTITION.outputs,
+          ...MOCK_DEFAULT_PARTITION.regions[regionWithRegionData],
         });
       });
 
-      it("should return the default partition when the region is not found", () => {
-        expect(partition("non-existant-region")).toEqual({
+      it("returns partition data if region data does not exist", () => {
+        const regionWithoutRegionData = "mock-region-2";
+        expect(partition(regionWithoutRegionData)).toEqual({
           ...MOCK_DEFAULT_PARTITION.outputs,
         });
       });
     });
+
+    it("should return the partition data when region is matched with regionRegex", () => {
+      expect(partition(MOCK_DEFAULT_PARTITION.regionRegex)).toEqual({
+        ...MOCK_DEFAULT_PARTITION.outputs,
+      });
+      expect(partition(MOCK_PARTITION.regionRegex)).toEqual({
+        ...MOCK_PARTITION.outputs,
+      });
+    });
+
+    it("should return the default partition when the region is not found", () => {
+      expect(partition("non-existant-region")).toEqual({
+        ...MOCK_DEFAULT_PARTITION.outputs,
+      });
+    });
   });
 
   it("should throw an error when the default partition is not found, and region doesn't match in partition array or regex", () => {
-    jest.isolateModules(() => {
-      jest.mock("./partitions.json", () => ({
-        partitions: [MOCK_PARTITION],
-      }));
-      const { partition } = require("./partition");
-      expect(() => partition("non-existant-region")).toThrow(
-        "Provided region was not found in the partition array or regex," +
-          " and default partition with id 'aws' doesn't exist."
-      );
+    setPartitionInfo({
+      partitions: [MOCK_PARTITION],
     });
+    expect(() => partition("non-existant-region")).toThrow(
+      "Provided region was not found in the partition array or regex," +
+        " and default partition with id 'aws' doesn't exist."
+    );
+  });
+
+  it("should allow setting a custom partitions file", async () => {
+    const copy = JSON.parse(JSON.stringify(partitions));
+    setPartitionInfo(copy);
+    const testRegion = "us-test-135";
+    copy.partitions[0].regions[testRegion] = {
+      description: "not a real region",
+    };
+    const result = partition(testRegion);
+    expect(result).toEqual({
+      description: "not a real region",
+      dnsSuffix: "amazonaws.com",
+      dualStackDnsSuffix: "api.aws",
+      name: "aws",
+      supportsDualStack: true,
+      supportsFIPS: true,
+    });
+
+    useDefaultPartitionInfo();
+    // result is matched by regex, but customization is no longer present.
+    expect(partition(testRegion)).toEqual({
+      description: void 0,
+      dnsSuffix: "amazonaws.com",
+      dualStackDnsSuffix: "api.aws",
+      name: "aws",
+      supportsDualStack: true,
+      supportsFIPS: true,
+    });
+  });
+
+  it("should optionally set a user agent prefix", async () => {
+    setPartitionInfo(null as any, "a-string-prefix");
+    expect(getUserAgentPrefix()).toBe("a-string-prefix");
   });
 });

packages/util-endpoints/src/lib/aws/partition.ts

@@ -2,8 +2,37 @@ import { EndpointPartition } from "@aws-sdk/types";
 
 import partitionsInfo from "./partitions.json";
 
-const { partitions } = partitionsInfo;
-const DEFAULT_PARTITION = partitions.find((partition) => partition.id === "aws");
+export type PartitionsInfo = {
+  partitions: Array<{
+    id: string;
+    outputs: {
+      dnsSuffix: string;
+      dualStackDnsSuffix: string;
+      name: string;
+      supportsDualStack: boolean;
+      supportsFIPS: boolean;
+    };
+    regionRegex: string;
+    regions: Record<
+      string,
+      | {
+          description?: string;
+        }
+      | undefined
+    >;
+  }>;
+};
+
+/**
+ * The partitions.json data to be used in resolving endpoints.
+ * @internal
+ */
+let selectedPartitionsInfo: PartitionsInfo = partitionsInfo;
+
+/**
+ * @internal
+ */
+let selectedUserAgentPrefix = "";
 
 /**
  * Evaluates a single string argument value as a region, and matches the
@@ -12,6 +41,7 @@ const DEFAULT_PARTITION = partitions.find((partition) => partition.id === "aws")
  * that the region has been determined to be a part of.
  */
 export const partition = (value: string): EndpointPartition => {
+  const { partitions } = selectedPartitionsInfo;
   // Check for explicit region listed in the regions array.
   for (const partition of partitions) {
     const { regions, outputs } = partition;
@@ -35,6 +65,8 @@ export const partition = (value: string): EndpointPartition => {
     }
   }
 
+  const DEFAULT_PARTITION = partitions.find((partition) => partition.id === "aws");
+
   if (!DEFAULT_PARTITION) {
     throw new Error(
       "Provided region was not found in the partition array or regex," +
@@ -47,3 +79,25 @@ export const partition = (value: string): EndpointPartition => {
     ...DEFAULT_PARTITION.outputs,
   };
 };
+
+/**
+ * Set custom partitions.json data.
+ * @internal
+ */
+export const setPartitionInfo = (partitionsInfo: PartitionsInfo, userAgentPrefix = "") => {
+  selectedPartitionsInfo = partitionsInfo;
+  selectedUserAgentPrefix = userAgentPrefix;
+};
+
+/**
+ * Reset to the default partitions.json data.
+ * @internal
+ */
+export const useDefaultPartitionInfo = () => {
+  setPartitionInfo(partitionsInfo, "");
+};
+
+/**
+ * @internal
+ */
+export const getUserAgentPrefix = () => selectedUserAgentPrefix;