Skip to content

Commit

Permalink
refactor(NODE-5464): refactor reauth signature
Browse files Browse the repository at this point in the history
  • Loading branch information
durran committed Nov 26, 2023
1 parent f782829 commit ad4ada2
Show file tree
Hide file tree
Showing 6 changed files with 101 additions and 98 deletions.
12 changes: 10 additions & 2 deletions src/cmap/auth/mongodb_oidc.ts
Original file line number Diff line number Diff line change
Expand Up @@ -62,10 +62,14 @@ export interface Workflow {
execute(
connection: Connection,
credentials: MongoCredentials,
reauthenticating: boolean,
response?: Document
): Promise<Document>;

/**
* Each workflow should specify the correct custom behaviour for reauthentication.
*/
reauthenticate(connection: Connection, credentials: MongoCredentials): Promise<Document>;

/**
* Get the document to add for speculative authentication.
*/
Expand Down Expand Up @@ -97,7 +101,11 @@ export class MongoDBOIDC extends AuthProvider {
const { connection, reauthenticating, response } = authContext;
const credentials = getCredentials(authContext);
const workflow = getWorkflow(credentials);
await workflow.execute(connection, credentials, reauthenticating, response);
if (reauthenticating) {
await workflow.reauthenticate(connection, credentials);
} else {
await workflow.execute(connection, credentials, response);
}
}

/**
Expand Down
72 changes: 20 additions & 52 deletions src/cmap/auth/mongodb_oidc/callback_workflow.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { Binary, BSON, type Document } from 'bson';
import { BSON, type Document } from 'bson';

import { MongoMissingCredentialsError } from '../../../error';
import { ns } from '../../../utils';
Expand All @@ -11,7 +11,7 @@ import type {
OIDCRequestFunction,
Workflow
} from '../mongodb_oidc';
import { AuthMechanism } from '../providers';
import { finishCommandDocument, startCommandDocument } from './command_builders';

/** The current version of OIDC implementation. */
const OIDC_VERSION = 0;
Expand Down Expand Up @@ -43,27 +43,35 @@ export class CallbackWorkflow implements Workflow {
return { speculativeAuthenticate: document };
}

/**
* Reauthenticate the callback workflow.
* For reauthentication:
* - Check if the connection's accessToken is not equal to the token manager's.
* - If they are different, use the token from the manager and set it on the connection and finish auth.
* - On success return, on error continue.
* - start auth to update the IDP information
* - If the idp info has changed, clear access token and refresh token.
* - If the idp info has not changed, attempt to use the refresh token.
* - if there's still a refresh token at this point, attempt to finish auth with that.
* - Attempt the full auth run, on error, raise to user.
*/
async reauthenticate(connection: Connection, credentials: MongoCredentials): Promise<Document> {
return this.execute(connection, credentials);
}

/**
* Execute the OIDC callback workflow.
*/
async execute(
connection: Connection,
credentials: MongoCredentials,
reauthenticating: boolean,
response?: Document
): Promise<Document> {
const requestCallback = credentials.mechanismProperties.REQUEST_TOKEN_CALLBACK;
if (!requestCallback) {
throw new MongoMissingCredentialsError(NO_REQUEST_CALLBACK);
}
// 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 startDocument = await this.startAuthentication(connection, credentials, response);
const conversationId = startDocument.conversationId;
const serverResult = BSON.deserialize(startDocument.payload.buffer) as IdPServerInfo;
const tokenResult = await this.fetchAccessToken(
Expand All @@ -89,11 +97,10 @@ export class CallbackWorkflow implements Workflow {
private async startAuthentication(
connection: Connection,
credentials: MongoCredentials,
reauthenticating: boolean,
response?: Document
): Promise<Document> {
let result;
if (!reauthenticating && response?.speculativeAuthenticate) {
if (response?.speculativeAuthenticate) {
result = response.speculativeAuthenticate;
} else {
result = await connection.commandAsync(
Expand Down Expand Up @@ -144,29 +151,6 @@ export class CallbackWorkflow implements Workflow {
}
}

/**
* Generate the finishing command document for authentication. Will be a
* saslStart or saslContinue depending on the presence of a conversation id.
*/
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
Expand All @@ -177,19 +161,3 @@ function isCallbackResultInvalid(tokenResult: unknown): boolean {
if (!('accessToken' in tokenResult)) return true;
return !Object.getOwnPropertyNames(tokenResult).every(prop => RESULT_PROPERTIES.includes(prop));
}

/**
* Generate the saslStart command document.
*/
function startCommandDocument(credentials: MongoCredentials): Document {
const payload: Document = {};
if (credentials.username) {
payload.n = credentials.username;
}
return {
saslStart: 1,
autoAuthorize: 1,
mechanism: AuthMechanism.MONGODB_OIDC,
payload: new Binary(BSON.serialize(payload))
};
}
43 changes: 43 additions & 0 deletions src/cmap/auth/mongodb_oidc/command_builders.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
import { Binary, BSON, type Document } from 'bson';

import { type MongoCredentials } from '../mongo_credentials';
import { AuthMechanism } from '../providers';

/**
* Generate the finishing command document for authentication. Will be a
* saslStart or saslContinue depending on the presence of a conversation id.
*/
export 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 }))
};
}

/**
* Generate the saslStart command document.
*/
export function startCommandDocument(credentials: MongoCredentials): Document {
const payload: Document = {};
if (credentials.username) {
payload.n = credentials.username;
}
return {
saslStart: 1,
autoAuthorize: 1,
mechanism: AuthMechanism.MONGODB_OIDC,
payload: new Binary(BSON.serialize(payload))
};
}
32 changes: 14 additions & 18 deletions src/cmap/auth/mongodb_oidc/machine_workflow.ts
Original file line number Diff line number Diff line change
@@ -1,32 +1,39 @@
import { BSON, type Document } from 'bson';
import { 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 { finishCommandDocument } from './command_builders';

/**
* Common behaviour for OIDC device workflows.
* Common behaviour for OIDC machine workflows.
* @internal
*/
export abstract class MachineWorkflow implements Workflow {
/**
* Execute the workflow. Looks for AWS_WEB_IDENTITY_TOKEN_FILE in the environment
* and then attempts to read the token from that path.
* Execute the workflow. Gets the token from the subclass implementation.
*/
async execute(connection: Connection, credentials: MongoCredentials): Promise<Document> {
const token = await this.getToken(credentials);
const command = commandDocument(token);
const command = finishCommandDocument(token);
return connection.commandAsync(ns(credentials.source), command, undefined);
}

/**
* Reauthenticate on a machine workflow just grabs the token again since the server
* has said the current access token is invalid or expired.
*/
async reauthenticate(connection: Connection, credentials: MongoCredentials): Promise<Document> {
return this.execute(connection, credentials);
}

/**
* Get the document to add for speculative authentication.
*/
async speculativeAuth(credentials: MongoCredentials): Promise<Document> {
const token = await this.getToken(credentials);
const document = commandDocument(token);
const document = finishCommandDocument(token);
document.db = credentials.source;
return { speculativeAuthenticate: document };
}
Expand All @@ -36,14 +43,3 @@ export abstract class MachineWorkflow implements Workflow {
*/
abstract getToken(credentials: MongoCredentials): Promise<string>;
}

/**
* Create the saslStart command document.
*/
export function commandDocument(token: string): Document {
return {
saslStart: 1,
mechanism: AuthMechanism.MONGODB_OIDC,
payload: BSON.serialize({ jwt: token })
};
}
26 changes: 0 additions & 26 deletions src/cmap/auth/mongodb_oidc/token_manager.ts

This file was deleted.

14 changes: 14 additions & 0 deletions src/cmap/connection.ts
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,8 @@ const kHello = Symbol('hello');
const kAutoEncrypter = Symbol('autoEncrypter');
/** @internal */
const kDelayedTimeoutId = Symbol('delayedTimeoutId');
/** @internal */
const kAccessToken = Symbol('accessToken');

const INVALID_QUEUE_SIZE = 'Connection internal queue contains more than 1 operation description';

Expand Down Expand Up @@ -138,6 +140,7 @@ export interface ConnectionOptions
socketTimeoutMS?: number;
cancellationToken?: CancellationToken;
metadata: ClientMetadata;
accessToken?: string;
}

/** @internal */
Expand Down Expand Up @@ -195,6 +198,8 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
[kHello]: Document | null;
/** @internal */
[kClusterTime]: Document | null;
/** @internal */
[kAccessToken]?: string;

/** @event */
static readonly COMMAND_STARTED = COMMAND_STARTED;
Expand Down Expand Up @@ -237,6 +242,7 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
this[kDescription] = new StreamDescription(this.address, options);
this[kGeneration] = options.generation;
this[kLastUseTime] = now();
this[kAccessToken] = options.accessToken;

// setup parser stream and message handling
this[kQueue] = new Map();
Expand Down Expand Up @@ -278,6 +284,14 @@ export class Connection extends TypedEventEmitter<ConnectionEvents> {
this[kHello] = response;
}

get accessToken(): string | undefined {
return this[kAccessToken];
}

set accessToken(value: string | undefined) {
this[kAccessToken] = value;
}

// Set the whether the message stream is for a monitoring connection.
set isMonitoringConnection(value: boolean) {
this[kMessageStream].isMonitoringConnection = value;
Expand Down

0 comments on commit ad4ada2

Please sign in to comment.