mongodb 5.2.0 → 5.4.0

package/src/cmap/auth/mongodb_oidc/callback_workflow.ts

@@ -1,16 +1,33 @@
 import { Binary, BSON, type Document } from 'bson';
 
-import { MongoInvalidArgumentError, MongoMissingCredentialsError } from '../../../error';
+import { MONGODB_ERROR_CODES, MongoError, MongoMissingCredentialsError } from '../../../error';
 import { ns } from '../../../utils';
 import type { Connection } from '../../connection';
 import type { MongoCredentials } from '../mongo_credentials';
-import type { OIDCMechanismServerStep1, OIDCRequestTokenResult } from '../mongodb_oidc';
+import type {
+  IdPServerInfo,
+  IdPServerResponse,
+  OIDCCallbackContext,
+  OIDCRefreshFunction,
+  OIDCRequestFunction,
+  Workflow
+} from '../mongodb_oidc';
 import { AuthMechanism } from '../providers';
+import { CallbackLockCache } from './callback_lock_cache';
 import { TokenEntryCache } from './token_entry_cache';
-import type { Workflow } from './workflow';
 
-/* 5 minutes in milliseconds */
-const TIMEOUT_MS = 300000;
+/** The current version of OIDC implementation. */
+const OIDC_VERSION = 0;
+
+/** 5 minutes in seconds */
+const TIMEOUT_S = 300;
+
+/** Properties allowed on results of callbacks. */
+const RESULT_PROPERTIES = ['accessToken', 'expiresInSeconds', 'refreshToken'];
+
+/** Error message when the callback result is invalid. */
+const CALLBACK_RESULT_ERROR =
+  'User provided OIDC callbacks must return a valid object with an accessToken.';
 
 /**
  * OIDC implementation of a callback based workflow.
@@ -18,206 +35,248 @@ const TIMEOUT_MS = 300000;
  */
 export class CallbackWorkflow implements Workflow {
   cache: TokenEntryCache;
+  callbackCache: CallbackLockCache;
 
   /**
    * Instantiate the workflow
    */
   constructor() {
     this.cache = new TokenEntryCache();
+    this.callbackCache = new CallbackLockCache();
   }
 
   /**
-   * Get the document to add for speculative authentication. Is empty when
-   * callbacks are in play.
+   * Get the document to add for speculative authentication. This also needs
+   * to add a db field from the credentials source.
    */
-  speculativeAuth(): Promise<Document> {
-    return Promise.resolve({});
+  async speculativeAuth(credentials: MongoCredentials): Promise<Document> {
+    const document = startCommandDocument(credentials);
+    document.db = credentials.source;
+    return { speculativeAuthenticate: document };
   }
 
   /**
-   * Execute the workflow.
-   *
-   * Steps:
-   * - If an entry is in the cache
-   *   - If it is not expired
-   *     - Skip step one and use the entry to execute step two.
-   *   - If it is expired
-   *     - If the refresh callback exists
-   *       - remove expired entry from cache
-   *       - call the refresh callback.
-   *       - put the new entry in the cache.
-   *       - execute step two.
-   *     - If the refresh callback does not exist.
-   *       - remove expired entry from cache
-   *       - call the request callback.
-   *       - put the new entry in the cache.
-   *       - execute step two.
-   * - If no entry is in the cache.
-   *   - execute step one.
-   *   - call the refresh callback.
-   *   - put the new entry in the cache.
-   *   - execute step two.
+   * Execute the OIDC callback workflow.
    */
   async execute(
     connection: Connection,
     credentials: MongoCredentials,
-    reauthenticate = false
+    reauthenticating: boolean,
+    response?: Document
   ): Promise<Document> {
-    const request = credentials.mechanismProperties.REQUEST_TOKEN_CALLBACK;
-    const refresh = credentials.mechanismProperties.REFRESH_TOKEN_CALLBACK;
-
-    const entry = this.cache.getEntry(
-      connection.address,
-      credentials.username,
-      request || null,
-      refresh || null
+    // Get the callbacks with locks from the callback lock cache.
+    const { requestCallback, refreshCallback, callbackHash } = this.callbackCache.getCallbacks(
+      connection,
+      credentials
     );
+    // Look for an existing entry in the cache.
+    const entry = this.cache.getEntry(connection.address, credentials.username, callbackHash);
+    let result;
     if (entry) {
-      // Check if the entry is not expired and if we are reauthenticating.
-      if (!reauthenticate && entry.isValid()) {
-        // Skip step one and execute the step two saslContinue.
-        try {
-          const result = await finishAuth(entry.tokenResult, undefined, connection, credentials);
-          return result;
-        } catch (error) {
-          // If authentication errors when using a cached token we remove it from
-          // the cache.
-          this.cache.deleteEntry(
-            connection.address,
-            credentials.username || '',
-            request || null,
-            refresh || null
-          );
-          throw error;
-        }
-      } else {
-        // Remove the expired entry from the cache.
-        this.cache.deleteEntry(
-          connection.address,
-          credentials.username || '',
-          request || null,
-          refresh || null
+      // Reauthentication cannot use a token from the cache since the server has
+      // stated it is invalid by the request for reauthentication.
+      if (entry.isValid() && !reauthenticating) {
+        // Presence of a valid cache entry means we can skip to the finishing step.
+        result = await this.finishAuthentication(
+          connection,
+          credentials,
+          entry.tokenResult,
+          response?.speculativeAuthenticate?.conversationId
         );
-        // Execute a refresh of the token and finish auth.
-        return this.refreshAndFinish(
+      } else {
+        // Presence of an expired cache entry means we must fetch a new one and
+        // then execute the final step.
+        const tokenResult = await this.fetchAccessToken(
           connection,
           credentials,
-          entry.serverResult,
-          entry.tokenResult
+          entry.serverInfo,
+          reauthenticating,
+          callbackHash,
+          requestCallback,
+          refreshCallback
         );
+        try {
+          result = await this.finishAuthentication(
+            connection,
+            credentials,
+            tokenResult,
+            reauthenticating ? undefined : response?.speculativeAuthenticate?.conversationId
+          );
+        } catch (error) {
+          // If we are reauthenticating and this errors with reauthentication
+          // required, we need to do the entire process over again and clear
+          // the cache entry.
+          if (
+            reauthenticating &&
+            error instanceof MongoError &&
+            error.code === MONGODB_ERROR_CODES.Reauthenticate
+          ) {
+            this.cache.deleteEntry(connection.address, credentials.username, callbackHash);
+            result = await this.execute(connection, credentials, reauthenticating);
+          } else {
+            throw error;
+          }
+        }
       }
     } else {
-      // No entry means to start with the step one saslStart.
-      const result = await connection.commandAsync(
-        ns(credentials.source),
-        startCommandDocument(credentials),
-        undefined
+      // No entry in the cache requires us to do all authentication steps
+      // from start to finish, including getting a fresh token for the cache.
+      const startDocument = await this.startAuthentication(
+        connection,
+        credentials,
+        reauthenticating,
+        response
+      );
+      const conversationId = startDocument.conversationId;
+      const serverResult = BSON.deserialize(startDocument.payload.buffer) as IdPServerInfo;
+      const tokenResult = await this.fetchAccessToken(
+        connection,
+        credentials,
+        serverResult,
+        reauthenticating,
+        callbackHash,
+        requestCallback,
+        refreshCallback
+      );
+      result = await this.finishAuthentication(
+        connection,
+        credentials,
+        tokenResult,
+        conversationId
       );
-      const stepOne = BSON.deserialize(result.payload.buffer) as OIDCMechanismServerStep1;
-      // Call the request callback and finish auth.
-      return this.requestAndFinish(connection, credentials, stepOne, result.conversationId);
     }
+    return result;
   }
 
   /**
-   * Execute the refresh callback if it exists, otherwise the request callback, then
-   * finish the authentication.
+   * Starts the callback authentication process. If there is a speculative
+   * authentication document from the initial handshake, then we will use that
+   * value to get the issuer, otherwise we will send the saslStart command.
    */
-  private async refreshAndFinish(
+  private async startAuthentication(
     connection: Connection,
     credentials: MongoCredentials,
-    stepOneResult: OIDCMechanismServerStep1,
-    tokenResult: OIDCRequestTokenResult,
-    conversationId?: number
+    reauthenticating: boolean,
+    response?: Document
   ): Promise<Document> {
-    const request = credentials.mechanismProperties.REQUEST_TOKEN_CALLBACK;
-    const refresh = credentials.mechanismProperties.REFRESH_TOKEN_CALLBACK;
-    // If a refresh callback exists, use it. Otherwise use the request callback.
-    if (refresh) {
-      const result: OIDCRequestTokenResult = await refresh(
-        credentials.username,
-        stepOneResult,
-        tokenResult,
-        TIMEOUT_MS
-      );
-      // Validate the result.
-      if (!result || !result.accessToken) {
-        throw new MongoMissingCredentialsError(
-          'REFRESH_TOKEN_CALLBACK must return a valid object with an accessToken'
-        );
-      }
-      // Cache a new entry and continue with the saslContinue.
-      this.cache.addEntry(
-        connection.address,
-        credentials.username || '',
-        request || null,
-        refresh,
-        result,
-        stepOneResult
-      );
-      return finishAuth(result, conversationId, connection, credentials);
+    let result;
+    if (!reauthenticating && response?.speculativeAuthenticate) {
+      result = response.speculativeAuthenticate;
     } else {
-      // Fallback to using the request callback.
-      return this.requestAndFinish(connection, credentials, stepOneResult, conversationId);
+      result = await connection.commandAsync(
+        ns(credentials.source),
+        startCommandDocument(credentials),
+        undefined
+      );
     }
+    return result;
   }
 
   /**
-   * Execute the request callback and finish authentication.
+   * Finishes the callback authentication process.
    */
-  private async requestAndFinish(
+  private async finishAuthentication(
     connection: Connection,
     credentials: MongoCredentials,
-    stepOneResult: OIDCMechanismServerStep1,
+    tokenResult: IdPServerResponse,
     conversationId?: number
   ): Promise<Document> {
-    // Call the request callback.
-    const request = credentials.mechanismProperties.REQUEST_TOKEN_CALLBACK;
-    const refresh = credentials.mechanismProperties.REFRESH_TOKEN_CALLBACK;
-    // Always clear expired entries from the cache on each finish as cleanup.
-    this.cache.deleteExpiredEntries();
-    if (!request) {
-      // Request callback must be present.
-      throw new MongoInvalidArgumentError(
-        'Auth mechanism property REQUEST_TOKEN_CALLBACK is required.'
-      );
+    const result = await connection.commandAsync(
+      ns(credentials.source),
+      finishCommandDocument(tokenResult.accessToken, conversationId),
+      undefined
+    );
+    return result;
+  }
+
+  /**
+   * Fetches an access token using either the request or refresh callbacks and
+   * puts it in the cache.
+   */
+  private async fetchAccessToken(
+    connection: Connection,
+    credentials: MongoCredentials,
+    serverInfo: IdPServerInfo,
+    reauthenticating: boolean,
+    callbackHash: string,
+    requestCallback: OIDCRequestFunction,
+    refreshCallback?: OIDCRefreshFunction
+  ): Promise<IdPServerResponse> {
+    // Get the token from the cache.
+    const entry = this.cache.getEntry(connection.address, credentials.username, callbackHash);
+    let result;
+    const context: OIDCCallbackContext = { timeoutSeconds: TIMEOUT_S, version: OIDC_VERSION };
+    // Check if there's a token in the cache.
+    if (entry) {
+      // If the cache entry is valid, return the token result.
+      if (entry.isValid() && !reauthenticating) {
+        return entry.tokenResult;
+      }
+      // If the cache entry is not valid, remove it from the cache and first attempt
+      // to use the refresh callback to get a new token. If no refresh callback
+      // exists, then fallback to the request callback.
+      if (refreshCallback) {
+        context.refreshToken = entry.tokenResult.refreshToken;
+        result = await refreshCallback(serverInfo, context);
+      } else {
+        result = await requestCallback(serverInfo, context);
+      }
+    } else {
+      // With no token in the cache we use the request callback.
+      result = await requestCallback(serverInfo, context);
     }
-    const tokenResult = await request(credentials.username, stepOneResult, TIMEOUT_MS);
-    // Validate the result.
-    if (!tokenResult || !tokenResult.accessToken) {
-      throw new MongoMissingCredentialsError(
-        'REQUEST_TOKEN_CALLBACK must return a valid object with an accessToken'
-      );
+    // Validate that the result returned by the callback is acceptable. If it is not
+    // we must clear the token result from the cache.
+    if (isCallbackResultInvalid(result)) {
+      this.cache.deleteEntry(connection.address, credentials.username, callbackHash);
+      throw new MongoMissingCredentialsError(CALLBACK_RESULT_ERROR);
     }
-    // Cache a new entry and continue with the saslContinue.
+    // Cleanup the cache.
+    this.cache.deleteExpiredEntries();
+    // Put the new entry into the cache.
     this.cache.addEntry(
       connection.address,
       credentials.username || '',
-      request,
-      refresh || null,
-      tokenResult,
-      stepOneResult
+      callbackHash,
+      result,
+      serverInfo
     );
-    return finishAuth(tokenResult, conversationId, connection, credentials);
+    return result;
   }
 }
 
 /**
- * Cache the result of the user supplied callback and execute the
- * step two saslContinue.
+ * Generate the finishing command document for authentication. Will be a
+ * saslStart or saslContinue depending on the presence of a conversation id.
  */
-async function finishAuth(
-  result: OIDCRequestTokenResult,
-  conversationId: number | undefined,
-  connection: Connection,
-  credentials: MongoCredentials
-): Promise<Document> {
-  // Execute the step two saslContinue.
-  return connection.commandAsync(
-    ns(credentials.source),
-    continueCommandDocument(result.accessToken, conversationId),
-    undefined
-  );
+function finishCommandDocument(token: string, conversationId?: number): Document {
+  if (conversationId != null && typeof conversationId === 'number') {
+    return {
+      saslContinue: 1,
+      conversationId: conversationId,
+      payload: new Binary(BSON.serialize({ jwt: token }))
+    };
+  }
+  // saslContinue requires a conversationId in the command to be valid so in this
+  // case the server allows "step two" to actually be a saslStart with the token
+  // as the jwt since the use of the cached value has no correlating conversating
+  // on the particular connection.
+  return {
+    saslStart: 1,
+    mechanism: AuthMechanism.MONGODB_OIDC,
+    payload: new Binary(BSON.serialize({ jwt: token }))
+  };
+}
+
+/**
+ * Determines if a result returned from a request or refresh callback
+ * function is invalid. This means the result is nullish, doesn't contain
+ * the accessToken required field, and does not contain extra fields.
+ */
+function isCallbackResultInvalid(tokenResult: unknown): boolean {
+  if (tokenResult == null || typeof tokenResult !== 'object') return true;
+  if (!('accessToken' in tokenResult)) return true;
+  return !Object.getOwnPropertyNames(tokenResult).every(prop => RESULT_PROPERTIES.includes(prop));
 }
 
 /**
@@ -235,25 +294,3 @@ function startCommandDocument(credentials: MongoCredentials): Document {
     payload: new Binary(BSON.serialize(payload))
   };
 }
-
-/**
- * Generate the saslContinue command document.
- */
-function continueCommandDocument(token: string, conversationId?: number): Document {
-  if (conversationId) {
-    return {
-      saslContinue: 1,
-      conversationId: conversationId,
-      payload: new Binary(BSON.serialize({ jwt: token }))
-    };
-  }
-  // saslContinue requires a conversationId in the command to be valid so in this
-  // case the server allows "step two" to actually be a saslStart with the token
-  // as the jwt since the use of the cached value has no correlating conversating
-  // on the particular connection.
-  return {
-    saslStart: 1,
-    mechanism: AuthMechanism.MONGODB_OIDC,
-    payload: new Binary(BSON.serialize({ jwt: token }))
-  };
-}

package/src/cmap/auth/mongodb_oidc/service_workflow.ts

@@ -3,8 +3,8 @@ import { BSON, type Document } from 'bson';
 import { ns } from '../../../utils';
 import type { Connection } from '../../connection';
 import type { MongoCredentials } from '../mongo_credentials';
+import type { Workflow } from '../mongodb_oidc';
 import { AuthMechanism } from '../providers';
-import type { Workflow } from './workflow';
 
 /**
  * Common behaviour for OIDC device workflows.
@@ -24,9 +24,11 @@ export abstract class ServiceWorkflow implements Workflow {
   /**
    * Get the document to add for speculative authentication.
    */
-  async speculativeAuth(): Promise<Document> {
+  async speculativeAuth(credentials: MongoCredentials): Promise<Document> {
     const token = await this.getToken();
-    return { speculativeAuthenticate: commandDocument(token) };
+    const document = commandDocument(token);
+    document.db = credentials.source;
+    return { speculativeAuthenticate: document };
   }
 
   /**

package/src/cmap/auth/mongodb_oidc/token_entry_cache.ts

@@ -1,41 +1,22 @@
-import type {
-  OIDCMechanismServerStep1,
-  OIDCRefreshFunction,
-  OIDCRequestFunction,
-  OIDCRequestTokenResult
-} from '../mongodb_oidc';
+import type { IdPServerInfo, IdPServerResponse } from '../mongodb_oidc';
+import { Cache } from './cache';
 
-/* 5 minutes in milliseonds */
+/* 5 minutes in milliseconds */
 const EXPIRATION_BUFFER_MS = 300000;
 /* Default expiration is now for when no expiration provided */
 const DEFAULT_EXPIRATION_SECS = 0;
-/* Counter for function "hashes".*/
-let FN_HASH_COUNTER = 0;
-/* No function present function */
-const NO_FUNCTION: OIDCRequestFunction = () => {
-  return Promise.resolve({ accessToken: 'test' });
-};
-/* The map of function hashes */
-const FN_HASHES = new WeakMap<OIDCRequestFunction | OIDCRefreshFunction, number>();
-/* Put the no function hash in the map. */
-FN_HASHES.set(NO_FUNCTION, FN_HASH_COUNTER);
-
 /** @internal */
 export class TokenEntry {
-  tokenResult: OIDCRequestTokenResult;
-  serverResult: OIDCMechanismServerStep1;
+  tokenResult: IdPServerResponse;
+  serverInfo: IdPServerInfo;
   expiration: number;
 
   /**
    * Instantiate the entry.
    */
-  constructor(
-    tokenResult: OIDCRequestTokenResult,
-    serverResult: OIDCMechanismServerStep1,
-    expiration: number
-  ) {
+  constructor(tokenResult: IdPServerResponse, serverInfo: IdPServerInfo, expiration: number) {
     this.tokenResult = tokenResult;
-    this.serverResult = serverResult;
+    this.serverInfo = serverInfo;
     this.expiration = expiration;
   }
 
@@ -52,62 +33,38 @@ export class TokenEntry {
  * Cache of OIDC token entries.
  * @internal
  */
-export class TokenEntryCache {
-  entries: Map<string, TokenEntry>;
-
-  constructor() {
-    this.entries = new Map();
-  }
-
+export class TokenEntryCache extends Cache<TokenEntry> {
   /**
    * Set an entry in the token cache.
    */
   addEntry(
     address: string,
     username: string,
-    requestFn: OIDCRequestFunction | null,
-    refreshFn: OIDCRefreshFunction | null,
-    tokenResult: OIDCRequestTokenResult,
-    serverResult: OIDCMechanismServerStep1
+    callbackHash: string,
+    tokenResult: IdPServerResponse,
+    serverInfo: IdPServerInfo
   ): TokenEntry {
     const entry = new TokenEntry(
       tokenResult,
-      serverResult,
+      serverInfo,
       expirationTime(tokenResult.expiresInSeconds)
     );
-    this.entries.set(cacheKey(address, username, requestFn, refreshFn), entry);
+    this.entries.set(this.cacheKey(address, username, callbackHash), entry);
     return entry;
   }
 
-  /**
-   * Clear the cache.
-   */
-  clear(): void {
-    this.entries.clear();
-  }
-
   /**
    * Delete an entry from the cache.
    */
-  deleteEntry(
-    address: string,
-    username: string,
-    requestFn: OIDCRequestFunction | null,
-    refreshFn: OIDCRefreshFunction | null
-  ): void {
-    this.entries.delete(cacheKey(address, username, requestFn, refreshFn));
+  deleteEntry(address: string, username: string, callbackHash: string): void {
+    this.entries.delete(this.cacheKey(address, username, callbackHash));
   }
 
   /**
    * Get an entry from the cache.
    */
-  getEntry(
-    address: string,
-    username: string,
-    requestFn: OIDCRequestFunction | null,
-    refreshFn: OIDCRefreshFunction | null
-  ): TokenEntry | undefined {
-    return this.entries.get(cacheKey(address, username, requestFn, refreshFn));
+  getEntry(address: string, username: string, callbackHash: string): TokenEntry | undefined {
+    return this.entries.get(this.cacheKey(address, username, callbackHash));
   }
 
   /**
@@ -128,39 +85,3 @@ export class TokenEntryCache {
 function expirationTime(expiresInSeconds?: number): number {
   return Date.now() + (expiresInSeconds ?? DEFAULT_EXPIRATION_SECS) * 1000;
 }
-
-/**
- * Create a cache key from the address and username.
- */
-function cacheKey(
-  address: string,
-  username: string,
-  requestFn: OIDCRequestFunction | null,
-  refreshFn: OIDCRefreshFunction | null
-): string {
-  return `${address}-${username}-${hashFunctions(requestFn, refreshFn)}`;
-}
-
-/**
- * Get the hash string for the request and refresh functions.
- */
-function hashFunctions(
-  requestFn: OIDCRequestFunction | null,
-  refreshFn: OIDCRefreshFunction | null
-): string {
-  let requestHash = FN_HASHES.get(requestFn || NO_FUNCTION);
-  let refreshHash = FN_HASHES.get(refreshFn || NO_FUNCTION);
-  if (!requestHash && requestFn) {
-    // Create a new one for the function and put it in the map.
-    FN_HASH_COUNTER++;
-    requestHash = FN_HASH_COUNTER;
-    FN_HASHES.set(requestFn, FN_HASH_COUNTER);
-  }
-  if (!refreshHash && refreshFn) {
-    // Create a new one for the function and put it in the map.
-    FN_HASH_COUNTER++;
-    refreshHash = FN_HASH_COUNTER;
-    FN_HASHES.set(refreshFn, FN_HASH_COUNTER);
-  }
-  return `${requestHash}-${refreshHash}`;
-}