From 829d7bed8cd0d53d966367c8507837980530cf01 Mon Sep 17 00:00:00 2001
From: Daria Pardue <daria.pardue@mongodb.com>
Date: Mon, 2 May 2022 11:21:17 -0400
Subject: [PATCH] fix(NODE-4208): add aws http request timeout handler (#3225)

---
 .evergreen/config.in.yml                  |  1 +
 .evergreen/config.yml                     |  1 +
 src/cmap/auth/mongodb_aws.ts              |  5 +++
 src/error.ts                              | 17 +++++++++
 src/index.ts                              |  1 +
 test/integration/auth/mongodb_aws.test.js | 46 +++++++++++++++++++++++
 6 files changed, 71 insertions(+)

diff --git a/.evergreen/config.in.yml b/.evergreen/config.in.yml
index 699ab73ded..de48ac71b2 100644
--- a/.evergreen/config.in.yml
+++ b/.evergreen/config.in.yml
@@ -599,6 +599,7 @@ functions:
         working_dir: "src"
         script: |
           ${PREPARE_SHELL}
+          export IS_EC2=true
           ${PROJECT_DIRECTORY}/.evergreen/run-mongodb-aws-test.sh
 
   "run aws auth test with aws credentials as environment variables":
diff --git a/.evergreen/config.yml b/.evergreen/config.yml
index 1f35e694e8..1b2caba4d6 100644
--- a/.evergreen/config.yml
+++ b/.evergreen/config.yml
@@ -565,6 +565,7 @@ functions:
         working_dir: src
         script: |
           ${PREPARE_SHELL}
+          export IS_EC2=true
           ${PROJECT_DIRECTORY}/.evergreen/run-mongodb-aws-test.sh
   run aws auth test with aws credentials as environment variables:
     - command: shell.exec
diff --git a/src/cmap/auth/mongodb_aws.ts b/src/cmap/auth/mongodb_aws.ts
index ad2014c4fc..3136dc8794 100644
--- a/src/cmap/auth/mongodb_aws.ts
+++ b/src/cmap/auth/mongodb_aws.ts
@@ -6,6 +6,7 @@ import type { Binary, BSONSerializeOptions } from '../../bson';
 import * as BSON from '../../bson';
 import { aws4 } from '../../deps';
 import {
+  MongoAWSError,
   MongoCompatibilityError,
   MongoMissingCredentialsError,
   MongoRuntimeError
@@ -283,6 +284,10 @@ function request(uri: string, _options: RequestOptions | undefined, callback: Ca
     });
   });
 
+  req.on('timeout', () => {
+    req.destroy(new MongoAWSError(`AWS request to ${uri} timed out after ${options.timeout} ms`));
+  });
+
   req.on('error', err => callback(err));
   req.end();
 }
diff --git a/src/error.ts b/src/error.ts
index 7b1554c628..7a788b923b 100644
--- a/src/error.ts
+++ b/src/error.ts
@@ -347,6 +347,23 @@ export class MongoKerberosError extends MongoRuntimeError {
   }
 }
 
+/**
+ * A error generated when the user attempts to authenticate
+ * via AWS, but fails
+ *
+ * @public
+ * @category Error
+ */
+export class MongoAWSError extends MongoRuntimeError {
+  constructor(message: string) {
+    super(message);
+  }
+
+  override get name(): string {
+    return 'MongoAWSError';
+  }
+}
+
 /**
  * An error generated when a ChangeStream operation fails to execute.
  *
diff --git a/src/index.ts b/src/index.ts
index 9775f079b2..965f584521 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -38,6 +38,7 @@ export const ObjectID = ObjectId;
 export { AnyBulkWriteOperation, BulkWriteOptions, MongoBulkWriteError } from './bulk/common';
 export {
   MongoAPIError,
+  MongoAWSError,
   MongoBatchReExecutionError,
   MongoChangeStreamError,
   MongoCompatibilityError,
diff --git a/test/integration/auth/mongodb_aws.test.js b/test/integration/auth/mongodb_aws.test.js
index 6d2cc42fb8..688ad36dac 100644
--- a/test/integration/auth/mongodb_aws.test.js
+++ b/test/integration/auth/mongodb_aws.test.js
@@ -1,6 +1,10 @@
 'use strict';
 const { expect } = require('chai');
 const { removeAuthFromConnectionString } = require('../../tools/utils');
+const sinon = require('sinon');
+const http = require('http');
+const { performance } = require('perf_hooks');
+const { MongoAWSError } = require('../../../src');
 
 describe('MONGODB-AWS', function () {
   beforeEach(function () {
@@ -52,4 +56,46 @@ describe('MONGODB-AWS', function () {
       .to.have.nested.property('options.credentials.mechanismProperties.AWS_SESSION_TOKEN')
       .that.equals('');
   });
+
+  describe('EC2 with missing credentials', () => {
+    let client;
+
+    beforeEach(function () {
+      if (!process.env.IS_EC2) {
+        this.currentTest.skipReason = 'requires an AWS EC2 environment';
+        this.skip();
+      }
+      sinon.stub(http, 'request').callsFake(function () {
+        arguments[0].hostname = 'www.example.com';
+        arguments[0].port = 81;
+        return http.request.wrappedMethod.apply(this, arguments);
+      });
+    });
+
+    afterEach(async () => {
+      sinon.restore();
+      if (client) {
+        await client.close();
+      }
+    });
+
+    it('should respect the default timeout of 10000ms', async function () {
+      const config = this.configuration;
+      client = config.newClient(process.env.MONGODB_URI, { authMechanism: 'MONGODB-AWS' }); // use the URI built by the test environment
+      const startTime = performance.now();
+
+      let caughtError = null;
+      await client.connect().catch(err => {
+        caughtError = err;
+      });
+
+      const endTime = performance.now();
+      const timeTaken = endTime - startTime;
+      expect(caughtError).to.be.instanceOf(MongoAWSError);
+      expect(caughtError)
+        .property('message')
+        .match(/timed out after/);
+      expect(timeTaken).to.be.within(10000, 12000);
+    });
+  });
 });