test: get token audience from props

mongodb · May 25, 2023 · 5c6d31c · 5c6d31c
1 parent 65c10db
commit 5c6d31c
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 9 deletions.
diff --git a/src/cmap/auth/mongo_credentials.ts b/src/cmap/auth/mongo_credentials.ts
@@ -2,6 +2,7 @@
 import type { Document } from '../../bson';
 import {
   MongoAPIError,
+  MongoAzureError,
   MongoInvalidArgumentError,
   MongoMissingCredentialsError
 } from '../../error';
@@ -43,6 +44,10 @@ export const DEFAULT_ALLOWED_HOSTS = [
   '::1'
 ];
 
+/** Error for when the token audience is missing in the environment. */
+const TOKEN_AUDIENCE_MISSING_ERROR =
+  'TOKEN_AUDIENCE must be set in the auth mechanism properties when PROVIDER_NAME is azure.';
+
 /** @public */
 export interface AuthMechanismProperties extends Document {
   SERVICE_HOST?: string;
@@ -55,9 +60,11 @@ export interface AuthMechanismProperties extends Document {
   /** @experimental */
   REFRESH_TOKEN_CALLBACK?: OIDCRefreshFunction;
   /** @experimental */
-  PROVIDER_NAME?: 'aws';
+  PROVIDER_NAME?: 'aws' | 'azure';
   /** @experimental */
   ALLOWED_HOSTS?: string[];
+  /** @experimental */
+  TOKEN_AUDIENCE?: string;
 }
 
 /** @public */
@@ -177,6 +184,13 @@ export class MongoCredentials {
         );
       }
 
+      if (
+        this.mechanismProperties.PROVIDER_NAME === 'azure' &&
+        !this.mechanismProperties.TOKEN_AUDIENCE
+      ) {
+        throw new MongoAzureError(TOKEN_AUDIENCE_MISSING_ERROR);
+      }
+
       if (
         this.mechanismProperties.PROVIDER_NAME &&
         !ALLOWED_PROVIDER_NAMES.includes(this.mechanismProperties.PROVIDER_NAME)

diff --git a/src/cmap/auth/mongodb_oidc/azure_service_workflow.ts b/src/cmap/auth/mongodb_oidc/azure_service_workflow.ts
@@ -1,11 +1,9 @@
 import { MongoAzureError } from '../../../error';
 import { request } from '../../../utils';
+import type { MongoCredentials } from '../mongo_credentials';
 import { AzureTokenCache } from './azure_token_cache';
 import { ServiceWorkflow } from './service_workflow';
 
-/** Error for when the token audience is missing in the environment. */
-const TOKEN_AUDIENCE_MISSING_ERROR = 'TOKEN_AUDIENCE must be set in the environment.';
-
 /** Base URL for getting Azure tokens. */
 const AZURE_BASE_URL =
   'http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01';
@@ -20,6 +18,10 @@ const RESULT_PROPERTIES = ['access_token', 'expires_in'];
 const ENDPOINT_RESULT_ERROR =
   'Azure endpoint did not return a value with only access_token and expires_in properties';
 
+/** Error for when the token audience is missing in the environment. */
+const TOKEN_AUDIENCE_MISSING_ERROR =
+  'TOKEN_AUDIENCE must be set in the auth mechanism properties when PROVIDER_NAME is azure.';
+
 /**
  * The Azure access token format.
  * @internal
@@ -48,8 +50,8 @@ export class AzureServiceWorkflow extends ServiceWorkflow {
   /**
    * Get the token from the environment.
    */
-  async getToken(): Promise<string> {
-    const tokenAudience = process.env.TOKEN_AUDIENCE;
+  async getToken(credentials?: MongoCredentials): Promise<string> {
+    const tokenAudience = credentials?.mechanismProperties.TOKEN_AUDIENCE;
     if (!tokenAudience) {
       throw new MongoAzureError(TOKEN_AUDIENCE_MISSING_ERROR);
     }

diff --git a/src/cmap/auth/mongodb_oidc/service_workflow.ts b/src/cmap/auth/mongodb_oidc/service_workflow.ts
@@ -16,7 +16,7 @@ export abstract class ServiceWorkflow implements Workflow {
    * and then attempts to read the token from that path.
    */
   async execute(connection: Connection, credentials: MongoCredentials): Promise<Document> {
-    const token = await this.getToken();
+    const token = await this.getToken(credentials);
     const command = commandDocument(token);
     return connection.commandAsync(ns(credentials.source), command, undefined);
   }
@@ -25,7 +25,7 @@ export abstract class ServiceWorkflow implements Workflow {
    * Get the document to add for speculative authentication.
    */
   async speculativeAuth(credentials: MongoCredentials): Promise<Document> {
-    const token = await this.getToken();
+    const token = await this.getToken(credentials);
     const document = commandDocument(token);
     document.db = credentials.source;
     return { speculativeAuthenticate: document };
@@ -34,7 +34,7 @@ export abstract class ServiceWorkflow implements Workflow {
   /**
    * Get the token from the environment or endpoint.
    */
-  abstract getToken(): Promise<string>;
+  abstract getToken(credentials: MongoCredentials): Promise<string>;
 }
 
 /**