feat: oidc updates

Author: durran

src/cmap/auth/mongodb_oidc.ts

@@ -1,16 +1,23 @@
 import { MongoInvalidArgumentError, MongoMissingCredentialsError } from '../../error';
+import { hostMatchesWildcards } from '../../utils';
 import type { HandshakeDocument } from '../connect';
 import { type AuthContext, AuthProvider } from './auth_provider';
-import type { MongoCredentials } from './mongo_credentials';
+import { DEFAULT_ALLOWED_HOSTS, MongoCredentials } from './mongo_credentials';
 import { AwsServiceWorkflow } from './mongodb_oidc/aws_service_workflow';
 import { CallbackWorkflow } from './mongodb_oidc/callback_workflow';
 import type { Workflow } from './mongodb_oidc/workflow';
 
+/**
+ * @internal
+ * The current version of OIDC implementation.
+ */
+export const OIDC_VERSION = 0;
+
 /**
  * @public
  * @experimental
  */
-export interface OIDCMechanismServerStep1 {
+export interface IdPServerInfo {
   issuer: string;
   clientId: string;
   requestScopes?: string[];
@@ -20,7 +27,7 @@ export interface OIDCMechanismServerStep1 {
  * @public
  * @experimental
  */
-export interface OIDCRequestTokenResult {
+export interface IdPServerResponse {
   accessToken: string;
   expiresInSeconds?: number;
   refreshToken?: string;
@@ -30,30 +37,30 @@ export interface OIDCRequestTokenResult {
  * @public
  * @experimental
  */
-export interface OIDCClientInfo {
-  principalName: string;
+export interface OIDCCallbackContext {
+  refreshToken?: string;
   timeoutSeconds?: number;
   timeoutContext?: AbortSignal;
+  version: number;
 }
 
 /**
  * @public
  * @experimental
  */
 export type OIDCRequestFunction = (
-  clientInfo: OIDCClientInfo,
-  serverInfo: OIDCMechanismServerStep1
-) => Promise<OIDCRequestTokenResult>;
+  info: IdPServerInfo,
+  context: OIDCCallbackContext
+) => Promise<IdPServerResponse>;
 
 /**
  * @public
  * @experimental
  */
 export type OIDCRefreshFunction = (
-  clientInfo: OIDCClientInfo,
-  serverInfo: OIDCMechanismServerStep1,
-  tokenResult: OIDCRequestTokenResult
-) => Promise<OIDCRequestTokenResult>;
+  info: IdPServerInfo,
+  context: OIDCCallbackContext
+) => Promise<IdPServerResponse>;
 
 type ProviderName = 'aws' | 'callback';
 
@@ -92,6 +99,11 @@ export class MongoDBOIDC extends AuthProvider {
     authContext: AuthContext
   ): Promise<HandshakeDocument> {
     const credentials = getCredentials(authContext);
+    const { connection } = authContext;
+    const allowedHosts = credentials.mechanismProperties.ALLOWED_HOSTS || DEFAULT_ALLOWED_HOSTS;
+    if (!hostMatchesWildcards(connection.address, allowedHosts)) {
+      throw new MongoInvalidArgumentError('Host does not match provided ALLOWED_HOSTS values');
+    }
     const workflow = getWorkflow(credentials);
     const result = await workflow.speculativeAuth(credentials);
     return { ...handshakeDoc, ...result };

src/cmap/auth/mongodb_oidc/callback_workflow.ts

@@ -1,14 +1,21 @@
 import { Binary, BSON, type Document } from 'bson';
 
-import { MongoInvalidArgumentError, MongoMissingCredentialsError } from '../../../error';
+import {
+  MONGODB_ERROR_CODES,
+  MongoError,
+  MongoInvalidArgumentError,
+  MongoMissingCredentialsError
+} from '../../../error';
 import { ns } from '../../../utils';
 import type { Connection } from '../../connection';
 import type { MongoCredentials } from '../mongo_credentials';
-import type {
-  OIDCMechanismServerStep1,
+import {
+  IdPServerInfo,
+  IdPServerResponse,
+  OIDC_VERSION,
+  OIDCCallbackContext,
   OIDCRefreshFunction,
-  OIDCRequestFunction,
-  OIDCRequestTokenResult
+  OIDCRequestFunction
 } from '../mongodb_oidc';
 import { AuthMechanism } from '../providers';
 import { TokenEntryCache } from './token_entry_cache';
@@ -71,8 +78,8 @@ export class CallbackWorkflow implements Workflow {
     let result;
     // Reauthentication must go through all the steps again regards of a cache entry
     // being present.
-    if (entry && !reauthenticating) {
-      if (entry.isValid()) {
+    if (entry) {
+      if (entry.isValid() && !reauthenticating) {
         // Presence of a valid cache entry means we can skip to the finishing step.
         result = await this.finishAuthentication(
           connection,
@@ -91,12 +98,33 @@ export class CallbackWorkflow implements Workflow {
           requestCallback,
           refreshCallback
         );
-        result = await this.finishAuthentication(
-          connection,
-          credentials,
-          tokenResult,
-          response?.speculativeAuthenticate?.conversationId
-        );
+        try {
+          result = await this.finishAuthentication(
+            connection,
+            credentials,
+            tokenResult,
+            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 || '',
+              requestCallback,
+              refreshCallback || null
+            );
+            result = await this.execute(connection, credentials, reauthenticating);
+          } else {
+            throw error;
+          }
+        }
       }
     } else {
       // No entry in the cache requires us to do all authentication steps
@@ -108,9 +136,7 @@ export class CallbackWorkflow implements Workflow {
         response
       );
       const conversationId = startDocument.conversationId;
-      const serverResult = BSON.deserialize(
-        startDocument.payload.buffer
-      ) as OIDCMechanismServerStep1;
+      const serverResult = BSON.deserialize(startDocument.payload.buffer) as IdPServerInfo;
       const tokenResult = await this.fetchAccessToken(
         connection,
         credentials,
@@ -159,7 +185,7 @@ export class CallbackWorkflow implements Workflow {
   private async finishAuthentication(
     connection: Connection,
     credentials: MongoCredentials,
-    tokenResult: OIDCRequestTokenResult,
+    tokenResult: IdPServerResponse,
     conversationId?: number
   ): Promise<Document> {
     const result = await connection.commandAsync(
@@ -177,11 +203,11 @@ export class CallbackWorkflow implements Workflow {
   private async fetchAccessToken(
     connection: Connection,
     credentials: MongoCredentials,
-    startResult: OIDCMechanismServerStep1,
+    serverInfo: IdPServerInfo,
     reauthenticating: boolean,
     requestCallback: OIDCRequestFunction,
     refreshCallback?: OIDCRefreshFunction
-  ): Promise<OIDCRequestTokenResult> {
+  ): Promise<IdPServerResponse> {
     // Get the token from the cache.
     const entry = this.cache.getEntry(
       connection.address,
@@ -190,7 +216,7 @@ export class CallbackWorkflow implements Workflow {
       refreshCallback || null
     );
     let result;
-    const clientInfo = { principalName: credentials.username, timeoutSeconds: TIMEOUT_S };
+    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.
@@ -201,13 +227,14 @@ export class CallbackWorkflow implements Workflow {
       // to use the refresh callback to get a new token. If no refresh callback
       // exists, then fallback to the request callback.
       if (refreshCallback) {
-        result = await refreshCallback(clientInfo, startResult, entry.tokenResult);
+        context.refreshToken = entry.tokenResult.refreshToken;
+        result = await refreshCallback(serverInfo, context);
       } else {
-        result = await requestCallback(clientInfo, startResult);
+        result = await requestCallback(serverInfo, context);
       }
     } else {
       // With no token in the cache we use the request callback.
-      result = await requestCallback(clientInfo, startResult);
+      result = await requestCallback(serverInfo, context);
     }
     // Validate that the result returned by the callback is acceptable.
     if (isCallbackResultInvalid(result)) {
@@ -224,7 +251,7 @@ export class CallbackWorkflow implements Workflow {
       requestCallback,
       refreshCallback || null,
       result,
-      startResult
+      serverInfo
     );
     return result;
   }

src/cmap/auth/mongodb_oidc/token_entry_cache.ts

@@ -1,8 +1,8 @@
 import type {
-  OIDCMechanismServerStep1,
+  PrincipalStepRequest,
   OIDCRefreshFunction,
   OIDCRequestFunction,
-  OIDCRequestTokenResult
+  IdPServerResponse
 } from '../mongodb_oidc';
 
 /* 5 minutes in milliseonds */
@@ -22,16 +22,16 @@ FN_HASHES.set(NO_FUNCTION, FN_HASH_COUNTER);
 
 /** @internal */
 export class TokenEntry {
-  tokenResult: OIDCRequestTokenResult;
-  serverResult: OIDCMechanismServerStep1;
+  tokenResult: IdPServerResponse;
+  serverResult: PrincipalStepRequest;
   expiration: number;
 
   /**
    * Instantiate the entry.
    */
   constructor(
-    tokenResult: OIDCRequestTokenResult,
-    serverResult: OIDCMechanismServerStep1,
+    tokenResult: IdPServerResponse,
+    serverResult: PrincipalStepRequest,
     expiration: number
   ) {
     this.tokenResult = tokenResult;
@@ -67,8 +67,8 @@ export class TokenEntryCache {
     username: string,
     requestFn: OIDCRequestFunction | null,
     refreshFn: OIDCRefreshFunction | null,
-    tokenResult: OIDCRequestTokenResult,
-    serverResult: OIDCMechanismServerStep1
+    tokenResult: IdPServerResponse,
+    serverResult: PrincipalStepRequest
   ): TokenEntry {
     const entry = new TokenEntry(
       tokenResult,

src/index.ts

@@ -204,11 +204,11 @@ export type {
   MongoCredentialsOptions
 } from './cmap/auth/mongo_credentials';
 export type {
-  OIDCClientInfo,
-  OIDCMechanismServerStep1,
+  OIDCCallbackContext as OIDCClientInfo,
+  PrincipalStepRequest as OIDCMechanismServerStep1,
   OIDCRefreshFunction,
   OIDCRequestFunction,
-  OIDCRequestTokenResult
+  IdPServerResponse as OIDCRequestTokenResult
 } from './cmap/auth/mongodb_oidc';
 export type {
   BinMsg,

src/utils.ts

@@ -58,6 +58,23 @@ export const ByteUtils = {
   }
 };
 
+/**
+ * Determines if a connection's address matches a user provided list
+ * of domain wildcards.
+ */
+export function hostMatchesWildcards(address: string, wildcards: string[]): boolean {
+  const host = HostAddress.fromString(address).host;
+  for (const wildcard of wildcards) {
+    if (
+      host === wildcard ||
+      (wildcard.startsWith('*.') && host?.endsWith(wildcard.substring(2, wildcard.length)))
+    ) {
+      return true;
+    }
+  }
+  return false;
+}
+
 /**
  * Throws if collectionName is not a valid mongodb collection namespace.
  * @internal