From 9be9e42166b9b5c0adcf6b3d8f055359e19a33a8 Mon Sep 17 00:00:00 2001
From: Chris Galvan <chgalvan@amazon.com>
Date: Wed, 17 Jul 2024 12:00:48 -0500
Subject: [PATCH] Fixed to account for GovCloud and some minor feedback fixes

---
 src/cognito/index.test.ts | 30 ++++++++++++++++++++++++++++++
 src/cognito/index.ts      |  4 ++--
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/src/cognito/index.test.ts b/src/cognito/index.test.ts
index aab2c6a..a37f332 100644
--- a/src/cognito/index.test.ts
+++ b/src/cognito/index.test.ts
@@ -13,6 +13,7 @@ describe("AuthHelper for Cognito", () => {
   const region = "us-west-2";
   const cognitoIdentityPoolId = `${region}:TEST-IDENTITY-POOL-ID`;
   const url = "https://maps.geo.us-west-2.amazonaws.com/";
+  const govCloudUrl = "https://maps.geo-fips.us-gov-west-1.amazonaws.com/";
   const nonAWSUrl = "https://example.com/";
   const nonLocationAWSUrl = "https://my.cool.service.us-west-2.amazonaws.com/";
   const mockedCredentials = {
@@ -163,6 +164,35 @@ describe("AuthHelper for Cognito", () => {
     expect(credential).toContain(mockedCredentials.accessKeyId);
   });
 
+  it("getMapAuthenticationOptions should contain transformRequest function to sign the AWS GovCloud Urls using our custom signer", async () => {
+    const authHelper = await withIdentityPoolId(cognitoIdentityPoolId);
+    const transformRequest = authHelper.getMapAuthenticationOptions().transformRequest;
+    const originalUrl = new URL(govCloudUrl);
+    const signedUrl = new URL(transformRequest(govCloudUrl).url);
+
+    // Host and pathname should still be the same
+    expect(signedUrl.hostname).toStrictEqual(originalUrl.hostname);
+    expect(signedUrl.pathname).toStrictEqual(originalUrl.pathname);
+
+    const searchParams = signedUrl.searchParams;
+    expect(searchParams.size).toStrictEqual(6);
+
+    // Verify these search params exist on the signed url
+    // We don't need to test the actual values since they are non-deterministic or constants
+    const expectedSearchParams = ["X-Amz-Algorithm", "X-Amz-Date", "X-Amz-SignedHeaders", "X-Amz-Signature"];
+    expectedSearchParams.forEach((value) => {
+      expect(searchParams.has(value)).toStrictEqual(true);
+    });
+
+    // We can expect the session token to match exactly as passed in
+    const securityToken = searchParams.get("X-Amz-Security-Token");
+    expect(securityToken).toStrictEqual(mockedCredentials.sessionToken);
+
+    // The credential starts with our access key, the rest is generated
+    const credential = searchParams.get("X-Amz-Credential");
+    expect(credential).toContain(mockedCredentials.accessKeyId);
+  });
+
   it("getMapAuthenticationOptions transformRequest function should pass-through non AWS Urls unchanged", async () => {
     const authHelper = await withIdentityPoolId(cognitoIdentityPoolId);
     const transformRequest = authHelper.getMapAuthenticationOptions().transformRequest;
diff --git a/src/cognito/index.ts b/src/cognito/index.ts
index 1e007be..e2da835 100644
--- a/src/cognito/index.ts
+++ b/src/cognito/index.ts
@@ -42,8 +42,8 @@ export async function withIdentityPoolId(
   return {
     getMapAuthenticationOptions: () => ({
       transformRequest: (url: string) => {
-        // Only sign aws location service URLs
-        if (url.match("(http|https)://(.*).geo.(.*).amazonaws.com")) {
+        // Only sign Amazon Location Service URLs
+        if (url.match("https://maps.(geo|geo-fips).(.*).amazonaws.com")) {
           return {
             url: Signer.signUrl(url, region, {
               access_key: credentials.accessKeyId,