fix: migrate machine workflows to callbacks

Author: durran

src/cmap/auth/mongodb_oidc.ts

@@ -4,11 +4,12 @@ import type { HandshakeDocument } from '../connect';
 import type { Connection } from '../connection';
 import { type AuthContext, AuthProvider } from './auth_provider';
 import type { MongoCredentials } from './mongo_credentials';
-import { AzureMachineWorkflow } from './mongodb_oidc/azure_machine_workflow';
-import { GCPMachineWorkflow } from './mongodb_oidc/gcp_machine_workflow';
-import { K8SMachineWorkflow } from './mongodb_oidc/k8s_machine_workflow';
+import { AutomatedCallbackWorkflow } from './mongodb_oidc/automated_callback_workflow';
+import { callback as azureCallback } from './mongodb_oidc/azure_machine_workflow';
+import { callback as gcpCallback } from './mongodb_oidc/gcp_machine_workflow';
+import { callback as k8sCallback } from './mongodb_oidc/k8s_machine_workflow';
 import { TokenCache } from './mongodb_oidc/token_cache';
-import { TokenMachineWorkflow } from './mongodb_oidc/token_machine_workflow';
+import { callback as testCallback } from './mongodb_oidc/token_machine_workflow';
 
 /** Error when credentials are missing. */
 const MISSING_CREDENTIALS_ERROR = 'AuthContext must provide credentials.';
@@ -78,6 +79,8 @@ export interface OIDCCallbackParams {
   idpInfo?: IdPInfo;
   /** The refresh token, if applicable, to be used by the callback to request a new token from the issuer. */
   refreshToken?: string;
+  /** The token audience for GCP and Azure. */
+  tokenAudience?: string;
 }
 
 /**
@@ -118,10 +121,10 @@ export interface Workflow {
 
 /** @internal */
 export const OIDC_WORKFLOWS: Map<EnvironmentName, () => Workflow> = new Map();
-OIDC_WORKFLOWS.set('test', () => new TokenMachineWorkflow(new TokenCache()));
-OIDC_WORKFLOWS.set('azure', () => new AzureMachineWorkflow(new TokenCache()));
-OIDC_WORKFLOWS.set('gcp', () => new GCPMachineWorkflow(new TokenCache()));
-OIDC_WORKFLOWS.set('k8s', () => new K8SMachineWorkflow(new TokenCache()));
+OIDC_WORKFLOWS.set('test', () => new AutomatedCallbackWorkflow(new TokenCache(), testCallback));
+OIDC_WORKFLOWS.set('azure', () => new AutomatedCallbackWorkflow(new TokenCache(), azureCallback));
+OIDC_WORKFLOWS.set('gcp', () => new AutomatedCallbackWorkflow(new TokenCache(), gcpCallback));
+OIDC_WORKFLOWS.set('k8s', () => new AutomatedCallbackWorkflow(new TokenCache(), k8sCallback));
 
 /**
  * OIDC auth provider.

src/cmap/auth/mongodb_oidc/automated_callback_workflow.ts

@@ -66,6 +66,9 @@ export class AutomatedCallbackWorkflow extends CallbackWorkflow {
     if (credentials.username) {
       params.username = credentials.username;
     }
+    if (credentials.mechanismProperties.TOKEN_RESOURCE) {
+      params.tokenAudience = credentials.mechanismProperties.TOKEN_RESOURCE;
+    }
     const timeout = Timeout.expires(AUTOMATED_TIMEOUT_MS);
     try {
       return await Promise.race([this.executeAndValidateCallback(params), timeout]);

src/cmap/auth/mongodb_oidc/azure_machine_workflow.ts

@@ -1,9 +1,7 @@
 import { addAzureParams, AZURE_BASE_URL } from '../../../client-side-encryption/providers/azure';
 import { MongoAzureError } from '../../../error';
 import { get } from '../../../utils';
-import type { MongoCredentials } from '../mongo_credentials';
-import { type AccessToken, MachineWorkflow } from './machine_workflow';
-import { type TokenCache } from './token_cache';
+import type { OIDCCallbackParams, OIDCResponse } from '../mongodb_oidc';
 
 /** Azure request headers. */
 const AZURE_HEADERS = Object.freeze({ Metadata: 'true', Accept: 'application/json' });
@@ -17,39 +15,27 @@ const TOKEN_RESOURCE_MISSING_ERROR =
   'TOKEN_RESOURCE must be set in the auth mechanism properties when ENVIRONMENT is azure.';
 
 /**
- * Device workflow implementation for Azure.
- *
- * @internal
+ * The callback function to be used in the automated callback workflow.
+ * @param params - The OIDC callback parameters.
+ * @returns The OIDC response.
  */
-export class AzureMachineWorkflow extends MachineWorkflow {
-  /**
-   * Instantiate the machine workflow.
-   */
-  constructor(cache: TokenCache) {
-    super(cache);
+export async function callback(params: OIDCCallbackParams): Promise<OIDCResponse> {
+  const tokenAudience = params.tokenAudience;
+  const username = params.username;
+  if (!tokenAudience) {
+    throw new MongoAzureError(TOKEN_RESOURCE_MISSING_ERROR);
   }
-
-  /**
-   * Get the token from the environment.
-   */
-  async getToken(credentials?: MongoCredentials): Promise<AccessToken> {
-    const tokenAudience = credentials?.mechanismProperties.TOKEN_RESOURCE;
-    const username = credentials?.username;
-    if (!tokenAudience) {
-      throw new MongoAzureError(TOKEN_RESOURCE_MISSING_ERROR);
-    }
-    const response = await getAzureTokenData(tokenAudience, username);
-    if (!isEndpointResultValid(response)) {
-      throw new MongoAzureError(ENDPOINT_RESULT_ERROR);
-    }
-    return response;
+  const response = await getAzureTokenData(tokenAudience, username);
+  if (!isEndpointResultValid(response)) {
+    throw new MongoAzureError(ENDPOINT_RESULT_ERROR);
   }
+  return response;
 }
 
 /**
  * Hit the Azure endpoint to get the token data.
  */
-async function getAzureTokenData(tokenAudience: string, username?: string): Promise<AccessToken> {
+async function getAzureTokenData(tokenAudience: string, username?: string): Promise<OIDCResponse> {
   const url = new URL(AZURE_BASE_URL);
   addAzureParams(url, tokenAudience, username);
   const response = await get(url, {
@@ -62,8 +48,8 @@ async function getAzureTokenData(tokenAudience: string, username?: string): Prom
   }
   const result = JSON.parse(response.body);
   return {
-    access_token: result.access_token,
-    expires_in: Number(result.expires_in)
+    accessToken: result.access_token,
+    expiresInSeconds: Number(result.expires_in)
   };
 }
 

src/cmap/auth/mongodb_oidc/gcp_machine_workflow.ts

@@ -1,8 +1,6 @@
 import { MongoGCPError } from '../../../error';
 import { get } from '../../../utils';
-import { type MongoCredentials } from '../mongo_credentials';
-import { type AccessToken, MachineWorkflow } from './machine_workflow';
-import { type TokenCache } from './token_cache';
+import type { OIDCCallbackParams, OIDCResponse } from '../mongodb_oidc';
 
 /** GCP base URL. */
 const GCP_BASE_URL =
@@ -15,30 +13,23 @@ const GCP_HEADERS = Object.freeze({ 'Metadata-Flavor': 'Google' });
 const TOKEN_RESOURCE_MISSING_ERROR =
   'TOKEN_RESOURCE must be set in the auth mechanism properties when ENVIRONMENT is gcp.';
 
-export class GCPMachineWorkflow extends MachineWorkflow {
-  /**
-   * Instantiate the machine workflow.
-   */
-  constructor(cache: TokenCache) {
-    super(cache);
-  }
-
-  /**
-   * Get the token from the environment.
-   */
-  async getToken(credentials?: MongoCredentials): Promise<AccessToken> {
-    const tokenAudience = credentials?.mechanismProperties.TOKEN_RESOURCE;
-    if (!tokenAudience) {
-      throw new MongoGCPError(TOKEN_RESOURCE_MISSING_ERROR);
-    }
-    return await getGcpTokenData(tokenAudience);
+/**
+ * The callback function to be used in the automated callback workflow.
+ * @param params - The OIDC callback parameters.
+ * @returns The OIDC response.
+ */
+export async function callback(params: OIDCCallbackParams): Promise<OIDCResponse> {
+  const tokenAudience = params.tokenAudience;
+  if (!tokenAudience) {
+    throw new MongoGCPError(TOKEN_RESOURCE_MISSING_ERROR);
   }
+  return await getGcpTokenData(tokenAudience);
 }
 
 /**
  * Hit the GCP endpoint to get the token data.
  */
-async function getGcpTokenData(tokenAudience: string): Promise<AccessToken> {
+async function getGcpTokenData(tokenAudience: string): Promise<OIDCResponse> {
   const url = new URL(GCP_BASE_URL);
   url.searchParams.append('audience', tokenAudience);
   const response = await get(url, {
@@ -49,5 +40,5 @@ async function getGcpTokenData(tokenAudience: string): Promise<AccessToken> {
       `Status code ${response.status} returned from the GCP endpoint. Response body: ${response.body}`
     );
   }
-  return { access_token: response.body };
+  return { accessToken: response.body };
 }

src/cmap/auth/mongodb_oidc/k8s_machine_workflow.ts

@@ -1,7 +1,6 @@
 import { readFile } from 'fs/promises';
 
-import { type AccessToken, MachineWorkflow } from './machine_workflow';
-import { type TokenCache } from './token_cache';
+import type { OIDCResponse } from '../mongodb_oidc';
 
 /** The fallback file name */
 const FALLBACK_FILENAME = '/var/run/secrets/kubernetes.io/serviceaccount/token';
@@ -12,27 +11,20 @@ const AZURE_FILENAME = 'AZURE_FEDERATED_TOKEN_FILE';
 /** The AWS environment variable for the file name. */
 const AWS_FILENAME = 'AWS_WEB_IDENTITY_TOKEN_FILE';
 
-export class K8SMachineWorkflow extends MachineWorkflow {
-  /**
-   * Instantiate the machine workflow.
-   */
-  constructor(cache: TokenCache) {
-    super(cache);
-  }
-
-  /**
-   * Get the token from the environment.
-   */
-  async getToken(): Promise<AccessToken> {
-    let filename: string;
-    if (process.env[AZURE_FILENAME]) {
-      filename = process.env[AZURE_FILENAME];
-    } else if (process.env[AWS_FILENAME]) {
-      filename = process.env[AWS_FILENAME];
-    } else {
-      filename = FALLBACK_FILENAME;
-    }
-    const token = await readFile(filename, 'utf8');
-    return { access_token: token };
+/**
+ * The callback function to be used in the automated callback workflow.
+ * @param params - The OIDC callback parameters.
+ * @returns The OIDC response.
+ */
+export async function callback(): Promise<OIDCResponse> {
+  let filename: string;
+  if (process.env[AZURE_FILENAME]) {
+    filename = process.env[AZURE_FILENAME];
+  } else if (process.env[AWS_FILENAME]) {
+    filename = process.env[AWS_FILENAME];
+  } else {
+    filename = FALLBACK_FILENAME;
   }
+  const token = await readFile(filename, 'utf8');
+  return { accessToken: token };
 }