feat(NODE-6632): destroy async resources when client closes

Author: nbbeeken

.eslintrc.json

@@ -119,6 +119,14 @@
       {
         "selector": "BinaryExpression[operator=/[=!]==?/] Literal[value='undefined']",
         "message": "Do not strictly check typeof undefined (NOTE: currently this rule only detects the usage of 'undefined' string literal so this could be a misfire)"
+      },
+      {
+        "selector": "CallExpression[callee.name='setTimeout']",
+        "message": "setTimeout must be abortable"
+      },
+      {
+        "selector": "CallExpression[callee.name='clearTimeout']",
+        "message": "clearTimeout must remove abort listener"
       }
     ],
     "@typescript-eslint/no-unused-vars": "error",

src/change_stream.ts

@@ -680,7 +680,8 @@ export class ChangeStream<
     if (this.options.timeoutMS != null) {
       this.timeoutContext = new CSOTTimeoutContext({
         timeoutMS: this.options.timeoutMS,
-        serverSelectionTimeoutMS
+        serverSelectionTimeoutMS,
+        closeSignal: this.cursor.client.closeSignal
       });
     }
   }

src/client-side-encryption/auto_encrypter.ts

@@ -393,13 +393,17 @@ export class AutoEncrypter {
     context.ns = ns;
     context.document = cmd;
 
-    const stateMachine = new StateMachine({
-      promoteValues: false,
-      promoteLongs: false,
-      proxyOptions: this._proxyOptions,
-      tlsOptions: this._tlsOptions,
-      socketOptions: autoSelectSocketOptions(this._client.s.options)
-    });
+    const stateMachine = new StateMachine(
+      {
+        promoteValues: false,
+        promoteLongs: false,
+        proxyOptions: this._proxyOptions,
+        tlsOptions: this._tlsOptions,
+        socketOptions: autoSelectSocketOptions(this._client.s.options)
+      },
+      undefined,
+      this._client.closeSignal
+    );
 
     return deserialize(await stateMachine.execute(this, context, options), {
       promoteValues: false,
@@ -420,12 +424,16 @@ export class AutoEncrypter {
 
     context.id = this._contextCounter++;
 
-    const stateMachine = new StateMachine({
-      ...options,
-      proxyOptions: this._proxyOptions,
-      tlsOptions: this._tlsOptions,
-      socketOptions: autoSelectSocketOptions(this._client.s.options)
-    });
+    const stateMachine = new StateMachine(
+      {
+        ...options,
+        proxyOptions: this._proxyOptions,
+        tlsOptions: this._tlsOptions,
+        socketOptions: autoSelectSocketOptions(this._client.s.options)
+      },
+      undefined,
+      this._client.closeSignal
+    );
 
     return await stateMachine.execute(this, context, options);
   }
@@ -438,7 +446,7 @@ export class AutoEncrypter {
    * the original ones.
    */
   async askForKMSCredentials(): Promise<KMSProviders> {
-    return await refreshKMSCredentials(this._kmsProviders);
+    return await refreshKMSCredentials(this._kmsProviders, this._client.closeSignal);
   }
 
   /**

src/client-side-encryption/client_encryption.ts

@@ -214,11 +214,15 @@ export class ClientEncryption {
       keyMaterial
     });
 
-    const stateMachine = new StateMachine({
-      proxyOptions: this._proxyOptions,
-      tlsOptions: this._tlsOptions,
-      socketOptions: autoSelectSocketOptions(this._client.s.options)
-    });
+    const stateMachine = new StateMachine(
+      {
+        proxyOptions: this._proxyOptions,
+        tlsOptions: this._tlsOptions,
+        socketOptions: autoSelectSocketOptions(this._client.s.options)
+      },
+      undefined,
+      this._client.closeSignal
+    );
 
     const timeoutContext =
       options?.timeoutContext ??
@@ -283,11 +287,15 @@ export class ClientEncryption {
     }
     const filterBson = serialize(filter);
     const context = this._mongoCrypt.makeRewrapManyDataKeyContext(filterBson, keyEncryptionKeyBson);
-    const stateMachine = new StateMachine({
-      proxyOptions: this._proxyOptions,
-      tlsOptions: this._tlsOptions,
-      socketOptions: autoSelectSocketOptions(this._client.s.options)
-    });
+    const stateMachine = new StateMachine(
+      {
+        proxyOptions: this._proxyOptions,
+        tlsOptions: this._tlsOptions,
+        socketOptions: autoSelectSocketOptions(this._client.s.options)
+      },
+      undefined,
+      this._client.closeSignal
+    );
 
     const timeoutContext = TimeoutContext.create(
       resolveTimeoutOptions(this._client, { timeoutMS: this._timeoutMS })
@@ -687,11 +695,15 @@ export class ClientEncryption {
     const valueBuffer = serialize({ v: value });
     const context = this._mongoCrypt.makeExplicitDecryptionContext(valueBuffer);
 
-    const stateMachine = new StateMachine({
-      proxyOptions: this._proxyOptions,
-      tlsOptions: this._tlsOptions,
-      socketOptions: autoSelectSocketOptions(this._client.s.options)
-    });
+    const stateMachine = new StateMachine(
+      {
+        proxyOptions: this._proxyOptions,
+        tlsOptions: this._tlsOptions,
+        socketOptions: autoSelectSocketOptions(this._client.s.options)
+      },
+      undefined,
+      this._client.closeSignal
+    );
 
     const timeoutContext =
       this._timeoutMS != null
@@ -712,7 +724,7 @@ export class ClientEncryption {
    * the original ones.
    */
   async askForKMSCredentials(): Promise<KMSProviders> {
-    return await refreshKMSCredentials(this._kmsProviders);
+    return await refreshKMSCredentials(this._kmsProviders, this._client.closeSignal);
   }
 
   static get libmongocryptVersion() {
@@ -771,11 +783,15 @@ export class ClientEncryption {
     }
 
     const valueBuffer = serialize({ v: value });
-    const stateMachine = new StateMachine({
-      proxyOptions: this._proxyOptions,
-      tlsOptions: this._tlsOptions,
-      socketOptions: autoSelectSocketOptions(this._client.s.options)
-    });
+    const stateMachine = new StateMachine(
+      {
+        proxyOptions: this._proxyOptions,
+        tlsOptions: this._tlsOptions,
+        socketOptions: autoSelectSocketOptions(this._client.s.options)
+      },
+      undefined,
+      this._client.closeSignal
+    );
     const context = this._mongoCrypt.makeExplicitEncryptionContext(valueBuffer, contextOptions);
 
     const timeoutContext =

src/client-side-encryption/providers/azure.ts

@@ -30,9 +30,9 @@ interface AzureTokenCacheEntry extends AccessToken {
 export class AzureCredentialCache {
   cachedToken: AzureTokenCacheEntry | null = null;
 
-  async getToken(): Promise<AccessToken> {
+  async getToken(closeSignal: AbortSignal): Promise<AccessToken> {
     if (this.cachedToken == null || this.needsRefresh(this.cachedToken)) {
-      this.cachedToken = await this._getToken();
+      this.cachedToken = await this._getToken(closeSignal);
     }
 
     return { accessToken: this.cachedToken.accessToken };
@@ -53,8 +53,8 @@ export class AzureCredentialCache {
   /**
    * exposed for testing
    */
-  _getToken(): Promise<AzureTokenCacheEntry> {
-    return fetchAzureKMSToken();
+  _getToken(closeSignal: AbortSignal): Promise<AzureTokenCacheEntry> {
+    return fetchAzureKMSToken(undefined, closeSignal);
   }
 }
 
@@ -156,11 +156,12 @@ export function prepareRequest(options: AzureKMSRequestOptions): {
  * [prose test 18](https://github.com/mongodb/specifications/tree/master/source/client-side-encryption/tests#azure-imds-credentials)
  */
 export async function fetchAzureKMSToken(
-  options: AzureKMSRequestOptions = {}
+  options: AzureKMSRequestOptions = {},
+  closeSignal: AbortSignal
 ): Promise<AzureTokenCacheEntry> {
   const { headers, url } = prepareRequest(options);
   try {
-    const response = await get(url, { headers });
+    const response = await get(url, { headers }, closeSignal);
     return await parseResponse(response);
   } catch (error) {
     if (error instanceof MongoNetworkTimeoutError) {
@@ -175,7 +176,10 @@ export async function fetchAzureKMSToken(
  *
  * @throws Will reject with a `MongoCryptError` if the http request fails or the http response is malformed.
  */
-export async function loadAzureCredentials(kmsProviders: KMSProviders): Promise<KMSProviders> {
-  const azure = await tokenCache.getToken();
+export async function loadAzureCredentials(
+  kmsProviders: KMSProviders,
+  closeSignal: AbortSignal
+): Promise<KMSProviders> {
+  const azure = await tokenCache.getToken(closeSignal);
   return { ...kmsProviders, azure };
 }

src/client-side-encryption/providers/index.ts

@@ -176,7 +176,10 @@ export function isEmptyCredentials(
  *
  * @internal
  */
-export async function refreshKMSCredentials(kmsProviders: KMSProviders): Promise<KMSProviders> {
+export async function refreshKMSCredentials(
+  kmsProviders: KMSProviders,
+  closeSignal: AbortSignal
+): Promise<KMSProviders> {
   let finalKMSProviders = kmsProviders;
 
   if (isEmptyCredentials('aws', kmsProviders)) {
@@ -188,7 +191,7 @@ export async function refreshKMSCredentials(kmsProviders: KMSProviders): Promise
   }
 
   if (isEmptyCredentials('azure', kmsProviders)) {
-    finalKMSProviders = await loadAzureCredentials(finalKMSProviders);
+    finalKMSProviders = await loadAzureCredentials(finalKMSProviders, closeSignal);
   }
   return finalKMSProviders;
 }

src/client-side-encryption/state_machine.ts

@@ -19,6 +19,7 @@ import { type Abortable } from '../mongo_types';
 import { Timeout, type TimeoutContext, TimeoutError } from '../timeout';
 import {
   addAbortListener,
+  addAbortSignalToStream,
   BufferPool,
   kDispose,
   MongoDBCollectionNamespace,
@@ -185,10 +186,15 @@ export type StateMachineOptions = {
  */
 // TODO(DRIVERS-2671): clarify CSOT behavior for FLE APIs
 export class StateMachine {
+  closeSignal: AbortSignal;
+
   constructor(
     private options: StateMachineOptions,
-    private bsonOptions = pluckBSONSerializeOptions(options)
-  ) {}
+    private bsonOptions = pluckBSONSerializeOptions(options),
+    closeSignal: AbortSignal
+  ) {
+    this.closeSignal = closeSignal;
+  }
 
   /**
    * Executes the state machine according to the specification
@@ -339,6 +345,9 @@ export class StateMachine {
     const buffer = new BufferPool();
 
     const netSocket: net.Socket = new net.Socket();
+
+    addAbortSignalToStream(this.closeSignal, netSocket);
+
     let socket: tls.TLSSocket;
 
     function destroySockets() {
@@ -451,7 +460,7 @@ export class StateMachine {
       await (options?.timeoutContext?.csotEnabled()
         ? Promise.all([
             willResolveKmsRequest,
-            Timeout.expires(options.timeoutContext?.remainingTimeMS)
+            Timeout.expires(options.timeoutContext.remainingTimeMS, this.closeSignal)
           ])
         : willResolveKmsRequest);
     } catch (error) {

src/cmap/auth/auth_provider.ts

@@ -47,7 +47,8 @@ export abstract class AuthProvider {
    */
   async prepare(
     handshakeDoc: HandshakeDocument,
-    _authContext: AuthContext
+    _authContext: AuthContext,
+    _closeSignal: AbortSignal
   ): Promise<HandshakeDocument> {
     return handshakeDoc;
   }
@@ -57,19 +58,19 @@ export abstract class AuthProvider {
    *
    * @param context - A shared context for authentication flow
    */
-  abstract auth(context: AuthContext): Promise<void>;
+  abstract auth(context: AuthContext, closeSignal: AbortSignal): Promise<void>;
 
   /**
    * Reauthenticate.
    * @param context - The shared auth context.
    */
-  async reauth(context: AuthContext): Promise<void> {
+  async reauth(context: AuthContext, closeSignal: AbortSignal): Promise<void> {
     if (context.reauthenticating) {
       throw new MongoRuntimeError('Reauthentication already in progress.');
     }
     try {
       context.reauthenticating = true;
-      await this.auth(context);
+      await this.auth(context, closeSignal);
     } finally {
       context.reauthenticating = false;
     }

src/cmap/auth/mongodb_oidc.ts

@@ -106,12 +106,20 @@ export interface Workflow {
   /**
    * Each workflow should specify the correct custom behaviour for reauthentication.
    */
-  reauthenticate(connection: Connection, credentials: MongoCredentials): Promise<void>;
+  reauthenticate(
+    connection: Connection,
+    credentials: MongoCredentials,
+    closeSignal: AbortSignal
+  ): Promise<void>;
 
   /**
    * Get the document to add for speculative authentication.
    */
-  speculativeAuth(connection: Connection, credentials: MongoCredentials): Promise<Document>;
+  speculativeAuth(
+    connection: Connection,
+    credentials: MongoCredentials,
+    closeSignal: AbortSignal
+  ): Promise<Document>;
 }
 
 /** @internal */
@@ -141,14 +149,14 @@ export class MongoDBOIDC extends AuthProvider {
   /**
    * Authenticate using OIDC
    */
-  override async auth(authContext: AuthContext): Promise<void> {
+  override async auth(authContext: AuthContext, closeSignal: AbortSignal): Promise<void> {
     const { connection, reauthenticating, response } = authContext;
     if (response?.speculativeAuthenticate?.done && !reauthenticating) {
       return;
     }
     const credentials = getCredentials(authContext);
     if (reauthenticating) {
-      await this.workflow.reauthenticate(connection, credentials);
+      await this.workflow.reauthenticate(connection, credentials, closeSignal);
     } else {
       await this.workflow.execute(connection, credentials, response);
     }
@@ -159,11 +167,12 @@ export class MongoDBOIDC extends AuthProvider {
    */
   override async prepare(
     handshakeDoc: HandshakeDocument,
-    authContext: AuthContext
+    authContext: AuthContext,
+    closeSignal: AbortSignal
   ): Promise<HandshakeDocument> {
     const { connection } = authContext;
     const credentials = getCredentials(authContext);
-    const result = await this.workflow.speculativeAuth(connection, credentials);
+    const result = await this.workflow.speculativeAuth(connection, credentials, closeSignal);
     return { ...handshakeDoc, ...result };
   }
 }

src/cmap/auth/mongodb_oidc/automated_callback_workflow.ts

@@ -19,8 +19,8 @@ export class AutomatedCallbackWorkflow extends CallbackWorkflow {
   /**
    * Instantiate the human callback workflow.
    */
-  constructor(cache: TokenCache, callback: OIDCCallbackFunction) {
-    super(cache, callback);
+  constructor(cache: TokenCache, callback: OIDCCallbackFunction, closeSignal: AbortSignal) {
+    super(cache, callback, closeSignal);
   }
 
   /**
@@ -66,7 +66,7 @@ export class AutomatedCallbackWorkflow extends CallbackWorkflow {
     if (credentials.username) {
       params.username = credentials.username;
     }
-    const timeout = Timeout.expires(AUTOMATED_TIMEOUT_MS);
+    const timeout = Timeout.expires(AUTOMATED_TIMEOUT_MS, this.closeSignal);
     try {
       return await Promise.race([this.executeAndValidateCallback(params), timeout]);
     } catch (error) {