wip

Author: nbbeeken

src/client-side-encryption/state_machine.ts

@@ -2,7 +2,6 @@ import { type MongoCryptContext, type MongoCryptKMSRequest } from 'mongodb-clien
 import * as net from 'net';
 import * as tls from 'tls';
 
-import { type AutoEncrypter } from '..';
 import {
   type BSONSerializeOptions,
   deserialize,
@@ -14,7 +13,7 @@ import { type ProxyOptions } from '../cmap/connection';
 import { CursorTimeoutContext } from '../cursor/abstract_cursor';
 import { getSocks, type SocksLib } from '../deps';
 import { MongoOperationTimeoutError } from '../error';
-import { type MongoClient, type MongoClientOptions } from '../mongo_client';
+import { type IO, type MongoClient, type MongoClientOptions } from '../mongo_client';
 import { type Abortable } from '../mongo_types';
 import { type CollectionInfo } from '../operations/list_collections';
 import { Timeout, type TimeoutContext, TimeoutError } from '../timeout';
@@ -25,7 +24,7 @@ import {
   MongoDBCollectionNamespace,
   promiseWithResolvers
 } from '../utils';
-import { autoSelectSocketOptions, type ClientEncryption, type DataKey } from './client_encryption';
+import { autoSelectSocketOptions, type DataKey } from './client_encryption';
 import { MongoCryptError } from './errors';
 import { type MongocryptdManager } from './mongocryptd_manager';
 import { type KMSProviders } from './providers';
@@ -186,10 +185,10 @@ export type StateMachineOptions = {
  */
 // TODO(DRIVERS-2671): clarify CSOT behavior for FLE APIs
 export class StateMachine {
-  private parent: AutoEncrypter | ClientEncryption;
+  private parent: { _client: { io: IO } };
 
   constructor(
-    parent: AutoEncrypter | ClientEncryption,
+    parent: { _client: { io: IO } },
     private options: StateMachineOptions,
     private bsonOptions = pluckBSONSerializeOptions(options)
   ) {

src/cmap/auth/mongodb_oidc/machine_workflow.ts

@@ -1,7 +1,7 @@
 import { setTimeout } from 'timers/promises';
 
 import { type Document } from '../../../bson';
-import { type MongoClient } from '../../../mongo_client';
+import { type IO } from '../../../mongo_client';
 import { ns } from '../../../utils';
 import type { Connection } from '../../connection';
 import type { MongoCredentials } from '../mongo_credentials';
@@ -24,7 +24,7 @@ export interface AccessToken {
 /** @internal */
 export type OIDCTokenFunction = (
   credentials: MongoCredentials,
-  client: MongoClient
+  client: { io: IO }
 ) => Promise<AccessToken>;
 
 /**
@@ -35,12 +35,12 @@ export abstract class MachineWorkflow implements Workflow {
   cache: TokenCache;
   callback: OIDCTokenFunction;
   lastExecutionTime: number;
-  client: MongoClient;
+  client: { io: IO };
 
   /**
    * Instantiate the machine workflow.
    */
-  constructor(client: MongoClient, cache: TokenCache) {
+  constructor(client: { io: IO }, cache: TokenCache) {
     this.client = client;
     this.cache = cache;
     this.callback = this.withLock(this.getToken.bind(this));
@@ -144,5 +144,5 @@ export abstract class MachineWorkflow implements Workflow {
   /**
    * Get the token from the environment or endpoint.
    */
-  abstract getToken(credentials: MongoCredentials, client: MongoClient): Promise<AccessToken>;
+  abstract getToken(credentials: MongoCredentials, client: { io: IO }): Promise<AccessToken>;
 }

src/cmap/connect.ts

@@ -16,6 +16,7 @@ import {
   MongoRuntimeError,
   needsRetryableWriteLabel
 } from '../error';
+import { type IO } from '../mongo_client';
 import { type Monitor, type RTTPinger } from '../sdam/monitor';
 import { HostAddress, ns, promiseWithResolvers } from '../utils';
 import { AuthContext } from './auth/auth_provider';
@@ -38,7 +39,7 @@ import {
 export type Stream = Socket | TLSSocket;
 
 export async function connect(
-  parent: Monitor | RTTPinger | ConnectionPool,
+  parent: { client: { io: IO } },
   options: ConnectionOptions
 ): Promise<Connection> {
   let connection: Connection | null = null;
@@ -54,7 +55,7 @@ export async function connect(
 }
 
 export function makeConnection(
-  parent: Monitor | RTTPinger | ConnectionPool,
+  parent: { client: { io: IO } },
   options: ConnectionOptions,
   socket: Stream
 ): Connection {

src/cmap/connection.ts

@@ -30,7 +30,7 @@ import {
   MongoServerError,
   MongoUnexpectedServerResponseError
 } from '../error';
-import type { MongoClient, ServerApi, SupportedNodeConnectionOptions } from '../mongo_client';
+import type { IO, ServerApi, SupportedNodeConnectionOptions } from '../mongo_client';
 import { type MongoClientAuthProviders } from '../mongo_client_auth_providers';
 import { MongoLoggableComponent, type MongoLogger, SeverityLevel } from '../mongo_logger';
 import { type Abortable, type CancellationToken, TypedEventEmitter } from '../mongo_types';
@@ -209,6 +209,7 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
   private clusterTime: Document | null = null;
   private error: Error | null = null;
   private dataEvents: AsyncGenerator<Buffer, void, void> | null = null;
+  private parent: { client: { io: IO } };
 
   private readonly socketTimeoutMS: number;
   private readonly monitorCommands: boolean;
@@ -230,17 +231,11 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
   /** @event */
   static readonly UNPINNED = UNPINNED;
 
-  private parent: Monitor | ConnectionPool | RTTPinger;
-
-  get client(): MongoClient {
+  get client(): { io: IO } {
     return this.parent.client;
   }
 
-  constructor(
-    parent: Monitor | ConnectionPool | RTTPinger,
-    stream: Stream,
-    options: ConnectionOptions
-  ) {
+  constructor(parent: { client: { io: IO } }, stream: Stream, options: ConnectionOptions) {
     super();
     this.on('error', noop);
 

src/cmap/handshake/client_metadata.ts

@@ -3,7 +3,7 @@ import * as process from 'process';
 
 import { BSON, type Document, Int32 } from '../../bson';
 import { MongoInvalidArgumentError } from '../../error';
-import type { MongoClient, MongoOptions } from '../../mongo_client';
+import type { IO, MongoOptions } from '../../mongo_client';
 
 // eslint-disable-next-line @typescript-eslint/no-require-imports
 const NODE_DRIVER_VERSION = require('../../../package.json').version;
@@ -157,7 +157,7 @@ export function makeClientMetadata(options: MakeClientMetadataOptions): ClientMe
 
 let dockerPromise: Promise<boolean>;
 /** @internal */
-async function getContainerMetadata(client: MongoClient) {
+async function getContainerMetadata(client: { io: IO }) {
   const containerMetadata: Record<string, any> = {};
   dockerPromise ??= client.io.fs.access('/.dockerenv');
   const isDocker = await dockerPromise;
@@ -176,7 +176,7 @@ async function getContainerMetadata(client: MongoClient) {
  * Re-add each metadata value.
  * Attempt to add new env container metadata, but keep old data if it does not fit.
  */
-export async function addContainerMetadata(client: MongoClient, originalMetadata: ClientMetadata) {
+export async function addContainerMetadata(client: { io: IO }, originalMetadata: ClientMetadata) {
   const containerMetadata = await getContainerMetadata(client);
   if (Object.keys(containerMetadata).length === 0) return originalMetadata;
 

test/mongodb.ts

@@ -123,6 +123,7 @@ export * from '../src/cmap/auth/mongodb_oidc/azure_machine_workflow';
 export * from '../src/cmap/auth/mongodb_oidc/callback_workflow';
 export * from '../src/cmap/auth/mongodb_oidc/gcp_machine_workflow';
 export * from '../src/cmap/auth/mongodb_oidc/machine_workflow';
+export * from '../src/cmap/auth/mongodb_oidc/token_cache';
 export * from '../src/cmap/auth/mongodb_oidc/token_machine_workflow';
 export * from '../src/cmap/auth/plain';
 export * from '../src/cmap/auth/providers';

test/unit/client-side-encryption/providers/credentialsProvider.test.ts

@@ -59,20 +59,26 @@ describe('#refreshKMSCredentials', function () {
     const secretKey = 'example';
     const sessionToken = 'example';
 
-    after(function () {
+    beforeEach(function () {
+      process.env.AWS_ACCESS_KEY_ID = accessKey;
+      process.env.AWS_SECRET_ACCESS_KEY = secretKey;
+      process.env.AWS_SESSION_TOKEN = sessionToken;
+    });
+
+    afterEach(function () {
       // After the entire suite runs, set the env back for the rest of the test run.
-      process.env.AWS_ACCESS_KEY_ID = originalAccessKeyId;
-      process.env.AWS_SECRET_ACCESS_KEY = originalSecretAccessKey;
-      process.env.AWS_SESSION_TOKEN = originalSessionToken;
+      if (typeof originalAccessKeyId === 'string') {
+        process.env.AWS_ACCESS_KEY_ID = originalAccessKeyId;
+      }
+      if (typeof originalSecretAccessKey === 'string') {
+        process.env.AWS_SECRET_ACCESS_KEY = originalSecretAccessKey;
+      }
+      if (typeof originalSessionToken === 'string') {
+        process.env.AWS_SESSION_TOKEN = originalSessionToken;
+      }
     });
 
     context('when the credential provider finds credentials', function () {
-      before(function () {
-        process.env.AWS_ACCESS_KEY_ID = accessKey;
-        process.env.AWS_SECRET_ACCESS_KEY = secretKey;
-        process.env.AWS_SESSION_TOKEN = sessionToken;
-      });
-
       context('when the credentials are empty', function () {
         const kmsProviders = { aws: {} };