From c6397ee107889730761da0aa91e75ec9d09df94e Mon Sep 17 00:00:00 2001
From: Neal Beeken <neal.beeken@mongodb.com>
Date: Mon, 21 Aug 2023 14:27:09 -0400
Subject: [PATCH 1/3] fix(NODE-5477): set AWS region from environment variable
 for STSClient

---
 src/cmap/auth/mongodb_aws.ts              |  56 +++++++-
 src/deps.ts                               |   4 +
 test/integration/auth/mongodb_aws.test.ts | 152 +++++++++++++++++++++-
 3 files changed, 207 insertions(+), 5 deletions(-)

diff --git a/src/cmap/auth/mongodb_aws.ts b/src/cmap/auth/mongodb_aws.ts
index 0001c080475..cfaf8e6f9b3 100644
--- a/src/cmap/auth/mongodb_aws.ts
+++ b/src/cmap/auth/mongodb_aws.ts
@@ -1,4 +1,5 @@
 import * as crypto from 'crypto';
+import * as process from 'process';
 import { promisify } from 'util';
 
 import type { Binary, BSONSerializeOptions } from '../../bson';
@@ -15,6 +16,28 @@ import { type AuthContext, AuthProvider } from './auth_provider';
 import { MongoCredentials } from './mongo_credentials';
 import { AuthMechanism } from './providers';
 
+/**
+ * The following regions use the global AWS STS endpoint, sts.amazonaws.com, by default
+ * https://docs.aws.amazon.com/sdkref/latest/guide/feature-sts-regionalized-endpoints.html
+ */
+const LEGACY_REGIONS = new Set([
+  'ap-northeast-1',
+  'ap-south-1',
+  'ap-southeast-1',
+  'ap-southeast-2',
+  'aws-global',
+  'ca-central-1',
+  'eu-central-1',
+  'eu-north-1',
+  'eu-west-1',
+  'eu-west-2',
+  'eu-west-3',
+  'sa-east-1',
+  'us-east-1',
+  'us-east-2',
+  'us-west-1',
+  'us-west-2'
+]);
 const ASCII_N = 110;
 const AWS_RELATIVE_URI = 'http://169.254.170.2';
 const AWS_EC2_URI = 'http://169.254.169.254';
@@ -34,6 +57,7 @@ interface AWSSaslContinuePayload {
 }
 
 export class MongoDBAWS extends AuthProvider {
+  static credentialProvider: ReturnType<typeof getAwsCredentialProvider> | null = null;
   randomBytesAsync: (size: number) => Promise<Buffer>;
 
   constructor() {
@@ -174,11 +198,11 @@ async function makeTempCredentials(credentials: MongoCredentials): Promise<Mongo
     });
   }
 
-  const credentialProvider = getAwsCredentialProvider();
+  MongoDBAWS.credentialProvider ??= getAwsCredentialProvider();
 
   // Check if the AWS credential provider from the SDK is present. If not,
   // use the old method.
-  if ('kModuleError' in credentialProvider) {
+  if ('kModuleError' in MongoDBAWS.credentialProvider) {
     // If the environment variable AWS_CONTAINER_CREDENTIALS_RELATIVE_URI
     // is set then drivers MUST assume that it was set by an AWS ECS agent
     if (process.env.AWS_CONTAINER_CREDENTIALS_RELATIVE_URI) {
@@ -209,6 +233,32 @@ async function makeTempCredentials(credentials: MongoCredentials): Promise<Mongo
 
     return makeMongoCredentialsFromAWSTemp(creds);
   } else {
+    let { AWS_STS_REGIONAL_ENDPOINTS = '', AWS_REGION = '' } = process.env;
+    AWS_STS_REGIONAL_ENDPOINTS = AWS_STS_REGIONAL_ENDPOINTS.toLowerCase();
+    AWS_REGION = AWS_REGION.toLowerCase();
+
+    /** The option setting should work only for users who have explicit settings in their environment, the driver should not encode "defaults" */
+    const awsRegionSettingsExist =
+      AWS_REGION.length !== 0 && AWS_STS_REGIONAL_ENDPOINTS.length !== 0;
+
+    /**
+     * If AWS_STS_REGIONAL_ENDPOINTS is set to regional, users are opting into the new behavior of respecting the region settings
+     *
+     * If AWS_STS_REGIONAL_ENDPOINTS is set to legacy, then "old" regions need to keep using the global setting.
+     * Technically the SDK gets this wrong, it reaches out to 'sts.us-east-1.amazonaws.com' when it should be 'sts.amazonaws.com'.
+     * That is not our bug to fix here. We leave that up to the SDK.
+     */
+    const useRegionalSts =
+      AWS_STS_REGIONAL_ENDPOINTS === 'regional' ||
+      (AWS_STS_REGIONAL_ENDPOINTS === 'legacy' && !LEGACY_REGIONS.has(AWS_REGION));
+
+    const provider =
+      awsRegionSettingsExist && useRegionalSts
+        ? MongoDBAWS.credentialProvider.fromNodeProviderChain({
+            clientConfig: { region: AWS_REGION }
+          })
+        : MongoDBAWS.credentialProvider.fromNodeProviderChain();
+
     /*
      * Creates a credential provider that will attempt to find credentials from the
      * following sources (listed in order of precedence):
@@ -219,8 +269,6 @@ async function makeTempCredentials(credentials: MongoCredentials): Promise<Mongo
      * - Shared credentials and config ini files
      * - The EC2/ECS Instance Metadata Service
      */
-    const { fromNodeProviderChain } = credentialProvider;
-    const provider = fromNodeProviderChain();
     try {
       const creds = await provider();
       return makeMongoCredentialsFromAWSTemp({
diff --git a/src/deps.ts b/src/deps.ts
index 131146c0479..7a3c121f689 100644
--- a/src/deps.ts
+++ b/src/deps.ts
@@ -86,6 +86,10 @@ export interface AWSCredentials {
 }
 
 type CredentialProvider = {
+  fromNodeProviderChain(
+    this: void,
+    options: { clientConfig: { region: string } }
+  ): () => Promise<AWSCredentials>;
   fromNodeProviderChain(this: void): () => Promise<AWSCredentials>;
 };
 
diff --git a/test/integration/auth/mongodb_aws.test.ts b/test/integration/auth/mongodb_aws.test.ts
index 81eb0373c56..608971232ca 100644
--- a/test/integration/auth/mongodb_aws.test.ts
+++ b/test/integration/auth/mongodb_aws.test.ts
@@ -3,7 +3,7 @@ import * as http from 'http';
 import { performance } from 'perf_hooks';
 import * as sinon from 'sinon';
 
-import { MongoAWSError, type MongoClient, MongoServerError } from '../../mongodb';
+import { MongoAWSError, type MongoClient, MongoDBAWS, MongoServerError } from '../../mongodb';
 
 describe('MONGODB-AWS', function () {
   let client: MongoClient;
@@ -88,4 +88,154 @@ describe('MONGODB-AWS', function () {
       expect(timeTaken).to.be.below(12000);
     });
   });
+
+  describe('when using AssumeRoleWithWebIdentity', () => {
+    const tests = [
+      {
+        ctx: 'when no AWS region settings are set',
+        title: 'uses the default region',
+        env: {
+          AWS_STS_REGIONAL_ENDPOINTS: undefined,
+          AWS_REGION: undefined
+        },
+        outcome: null
+      },
+      {
+        ctx: 'when only AWS_STS_REGIONAL_ENDPOINTS is set',
+        title: 'uses the default region',
+        env: {
+          AWS_STS_REGIONAL_ENDPOINTS: 'regional',
+          AWS_REGION: undefined
+        },
+        outcome: null
+      },
+      {
+        ctx: 'when only AWS_REGION is set',
+        title: 'uses the default region',
+        env: {
+          AWS_STS_REGIONAL_ENDPOINTS: undefined,
+          AWS_REGION: 'us-west-2'
+        },
+        outcome: null
+      },
+
+      {
+        ctx: 'when AWS_STS_REGIONAL_ENDPOINTS is set to regional and region is legacy',
+        title: 'uses the region from the environment',
+        env: {
+          AWS_STS_REGIONAL_ENDPOINTS: 'regional',
+          AWS_REGION: 'us-west-2'
+        },
+        outcome: 'us-west-2'
+      },
+      {
+        ctx: 'when AWS_STS_REGIONAL_ENDPOINTS is set to regional and region is new',
+        title: 'uses the region from the environment',
+        env: {
+          AWS_STS_REGIONAL_ENDPOINTS: 'regional',
+          AWS_REGION: 'sa-east-1'
+        },
+        outcome: 'sa-east-1'
+      },
+
+      {
+        ctx: 'when AWS_STS_REGIONAL_ENDPOINTS is set to legacy and region is legacy',
+        title: 'uses the region from the environment',
+        env: {
+          AWS_STS_REGIONAL_ENDPOINTS: 'legacy',
+          AWS_REGION: 'us-west-2'
+        },
+        outcome: null
+      },
+      {
+        ctx: 'when AWS_STS_REGIONAL_ENDPOINTS is set to legacy and region is new',
+        title: 'uses the region from the environment',
+        env: {
+          AWS_STS_REGIONAL_ENDPOINTS: 'legacy',
+          AWS_REGION: 'sa-east-1'
+        },
+        outcome: null
+      }
+    ];
+
+    for (const test of tests) {
+      context(test.ctx, () => {
+        let credentialProvider;
+        let storedEnv;
+        let calledArguments;
+
+        const shouldSkip = () => {
+          const { AWS_WEB_IDENTITY_TOKEN_FILE = '' } = process.env;
+          credentialProvider = (() => {
+            try {
+              return require('@aws-sdk/credential-providers');
+            } catch {
+              return null;
+            }
+          })();
+          return AWS_WEB_IDENTITY_TOKEN_FILE.length === 0 || credentialProvider == null;
+        };
+
+        beforeEach(function () {
+          if (shouldSkip()) {
+            this.skipReason = 'only relevant to AssumeRoleWithWebIdentity with SDK installed';
+            return this.skip();
+          }
+
+          client = this.configuration.newClient(process.env.MONGODB_URI);
+
+          storedEnv = process.env;
+          if (test.env.AWS_STS_REGIONAL_ENDPOINTS === undefined) {
+            delete process.env.AWS_STS_REGIONAL_ENDPOINTS;
+          } else {
+            process.env.AWS_STS_REGIONAL_ENDPOINTS = test.env.AWS_STS_REGIONAL_ENDPOINTS;
+          }
+          if (test.env.AWS_REGION === undefined) {
+            delete process.env.AWS_REGION;
+          } else {
+            process.env.AWS_REGION = test.env.AWS_REGION;
+          }
+
+          calledArguments = [];
+          MongoDBAWS.credentialProvider = {
+            fromNodeProviderChain(...args) {
+              calledArguments = args;
+              return credentialProvider.fromNodeProviderChain(...args);
+            }
+          };
+        });
+
+        afterEach(() => {
+          if (shouldSkip()) {
+            return;
+          }
+          if (typeof storedEnv.AWS_STS_REGIONAL_ENDPOINTS === 'string') {
+            process.env.AWS_STS_REGIONAL_ENDPOINTS = storedEnv.AWS_STS_REGIONAL_ENDPOINTS;
+          }
+          if (typeof storedEnv.AWS_STS_REGIONAL_ENDPOINTS === 'string') {
+            process.env.AWS_REGION = storedEnv.AWS_REGION;
+          }
+          MongoDBAWS.credentialProvider = credentialProvider;
+          calledArguments = [];
+        });
+
+        it(test.title, async function () {
+          const result = await client
+            .db('aws')
+            .collection('aws_test')
+            .estimatedDocumentCount()
+            .catch(error => error);
+
+          expect(result).to.not.be.instanceOf(MongoServerError);
+          expect(result).to.be.a('number');
+
+          if (test.outcome != null) {
+            expect(calledArguments).to.deep.equal([{ clientConfig: { region: test.outcome } }]);
+          } else {
+            expect(calledArguments).to.deep.equal([]);
+          }
+        });
+      });
+    }
+  });
 });

From a89a5ff9f75d3ba88c1afa0ef3af7fee65d6bf8e Mon Sep 17 00:00:00 2001
From: Neal Beeken <neal.beeken@mongodb.com>
Date: Fri, 25 Aug 2023 17:34:33 -0400
Subject: [PATCH 2/3] test: cleanup

---
 test/integration/auth/mongodb_aws.test.ts | 22 +++++++++-------------
 1 file changed, 9 insertions(+), 13 deletions(-)

diff --git a/test/integration/auth/mongodb_aws.test.ts b/test/integration/auth/mongodb_aws.test.ts
index 608971232ca..b88e4445fb5 100644
--- a/test/integration/auth/mongodb_aws.test.ts
+++ b/test/integration/auth/mongodb_aws.test.ts
@@ -98,7 +98,7 @@ describe('MONGODB-AWS', function () {
           AWS_STS_REGIONAL_ENDPOINTS: undefined,
           AWS_REGION: undefined
         },
-        outcome: null
+        calledWith: []
       },
       {
         ctx: 'when only AWS_STS_REGIONAL_ENDPOINTS is set',
@@ -107,7 +107,7 @@ describe('MONGODB-AWS', function () {
           AWS_STS_REGIONAL_ENDPOINTS: 'regional',
           AWS_REGION: undefined
         },
-        outcome: null
+        calledWith: []
       },
       {
         ctx: 'when only AWS_REGION is set',
@@ -116,7 +116,7 @@ describe('MONGODB-AWS', function () {
           AWS_STS_REGIONAL_ENDPOINTS: undefined,
           AWS_REGION: 'us-west-2'
         },
-        outcome: null
+        calledWith: []
       },
 
       {
@@ -126,7 +126,7 @@ describe('MONGODB-AWS', function () {
           AWS_STS_REGIONAL_ENDPOINTS: 'regional',
           AWS_REGION: 'us-west-2'
         },
-        outcome: 'us-west-2'
+        calledWith: [{ clientConfig: { region: 'us-west-2' } }]
       },
       {
         ctx: 'when AWS_STS_REGIONAL_ENDPOINTS is set to regional and region is new',
@@ -135,7 +135,7 @@ describe('MONGODB-AWS', function () {
           AWS_STS_REGIONAL_ENDPOINTS: 'regional',
           AWS_REGION: 'sa-east-1'
         },
-        outcome: 'sa-east-1'
+        calledWith: [{ clientConfig: { region: 'sa-east-1' } }]
       },
 
       {
@@ -145,16 +145,16 @@ describe('MONGODB-AWS', function () {
           AWS_STS_REGIONAL_ENDPOINTS: 'legacy',
           AWS_REGION: 'us-west-2'
         },
-        outcome: null
+        calledWith: []
       },
       {
         ctx: 'when AWS_STS_REGIONAL_ENDPOINTS is set to legacy and region is new',
-        title: 'uses the region from the environment',
+        title: 'uses the default region',
         env: {
           AWS_STS_REGIONAL_ENDPOINTS: 'legacy',
           AWS_REGION: 'sa-east-1'
         },
-        outcome: null
+        calledWith: []
       }
     ];
 
@@ -229,11 +229,7 @@ describe('MONGODB-AWS', function () {
           expect(result).to.not.be.instanceOf(MongoServerError);
           expect(result).to.be.a('number');
 
-          if (test.outcome != null) {
-            expect(calledArguments).to.deep.equal([{ clientConfig: { region: test.outcome } }]);
-          } else {
-            expect(calledArguments).to.deep.equal([]);
-          }
+          expect(calledArguments).to.deep.equal(test.calledWith);
         });
       });
     }

From ded1a060994cfa4867262785ee0a2bba42e06e03 Mon Sep 17 00:00:00 2001
From: Neal Beeken <neal.beeken@mongodb.com>
Date: Fri, 25 Aug 2023 17:36:34 -0400
Subject: [PATCH 3/3] test: cleanup

---
 test/integration/auth/mongodb_aws.test.ts | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/test/integration/auth/mongodb_aws.test.ts b/test/integration/auth/mongodb_aws.test.ts
index b88e4445fb5..287c9b50c66 100644
--- a/test/integration/auth/mongodb_aws.test.ts
+++ b/test/integration/auth/mongodb_aws.test.ts
@@ -163,8 +163,9 @@ describe('MONGODB-AWS', function () {
         let credentialProvider;
         let storedEnv;
         let calledArguments;
+        let shouldSkip = false;
 
-        const shouldSkip = () => {
+        const envCheck = () => {
           const { AWS_WEB_IDENTITY_TOKEN_FILE = '' } = process.env;
           credentialProvider = (() => {
             try {
@@ -177,7 +178,8 @@ describe('MONGODB-AWS', function () {
         };
 
         beforeEach(function () {
-          if (shouldSkip()) {
+          shouldSkip = envCheck();
+          if (shouldSkip) {
             this.skipReason = 'only relevant to AssumeRoleWithWebIdentity with SDK installed';
             return this.skip();
           }
@@ -206,7 +208,7 @@ describe('MONGODB-AWS', function () {
         });
 
         afterEach(() => {
-          if (shouldSkip()) {
+          if (shouldSkip) {
             return;
           }
           if (typeof storedEnv.AWS_STS_REGIONAL_ENDPOINTS === 'string') {