Fix Error 18456 with Entra ID auth when re-authenticating with stale tokens (#18456)

extensions/mssql/src/azure/constants.ts

@@ -65,6 +65,12 @@ export const AADSTS50173 = "AADSTS50173";
  * We have the user sign in again when this error occurs.
  */
 export const AADSTS50020 = "AADSTS50020";
+/**
+ * The presented token is not valid. This often occurs when the token has been sent
+ * to a different tenant than it was intended for (multi-tenant scenario).
+ * Have the user sign in again with the correct tenant.
+ */
+export const AADSTS50078 = "AADSTS50078";
 /**
  * Error thrown from STS - indicates user account not found in MSAL cache.
  * We request user to sign in again.

extensions/mssql/src/azure/msal/msalAzureAuth.ts

@@ -131,7 +131,10 @@ export abstract class MsalAzureAuth {
         return account;
     }
 
-    protected abstract login(tenant: ITenant): Promise<{
+    protected abstract login(
+        tenant: ITenant,
+        scopes?: string[],
+    ): Promise<{
         response: AuthenticationResult | null;
         authComplete: IDeferred<void, Error>;
     }>;
@@ -227,7 +230,8 @@ export abstract class MsalAzureAuth {
             error instanceof InteractionRequiredAuthError ||
             error.errorMessage.includes(Constants.AADSTS70043) ||
             error.errorMessage.includes(Constants.AADSTS50020) ||
-            error.errorMessage.includes(Constants.AADSTS50173)
+            error.errorMessage.includes(Constants.AADSTS50173) ||
+            error.errorMessage.includes(Constants.AADSTS50078)
         );
     }
 
@@ -254,7 +258,9 @@ export abstract class MsalAzureAuth {
                     key: tokenResult.account!.homeAccountId,
                     token: tokenResult.accessToken,
                     tokenType: tokenResult.tokenType,
-                    expiresOn: tokenResult.account!.idTokenClaims!.exp,
+                    expiresOn: tokenResult.expiresOn
+                        ? Math.floor(tokenResult.expiresOn.getTime() / 1000)
+                        : undefined,
                 };
 
                 return await this.hydrateAccount(token, tokenClaims);
@@ -358,16 +364,32 @@ export abstract class MsalAzureAuth {
         settings: IAADResource,
         promptUser: boolean = true,
     ): Promise<AuthenticationResult | null> {
+        // Compute the correct scopes for the target resource so that
+        // re-auth acquires a token for the right audience.  Using the
+        // default this.scopes (ARM-only for the public cloud) when the
+        // caller needs a SQL token causes Error 18456 in multi-tenant
+        // scenarios.
+        const endpoint = settings.endpoint.endsWith("/")
+            ? settings.endpoint
+            : settings.endpoint + "/";
+        const resourceScopes =
+            settings.id === this.providerSettings.settings.windowsManagementResource.id
+                ? [`${endpoint}user_impersonation`]
+                : [`${endpoint}.default`];
+        // Always include OpenID Connect scopes so the login result
+        // contains an id_token that can be used for account hydration.
+        const loginScopes = [...resourceScopes, "openid", "email", "profile", "offline_access"];
+
         let shouldOpen: boolean;
         if (promptUser) {
             shouldOpen = await this.askUserForInteraction(tenant, settings);
             if (shouldOpen) {
-                const result = await this.login(tenant);
+                const result = await this.login(tenant, loginScopes);
                 result?.authComplete?.resolve();
                 return result?.response;
             }
         } else {
-            const result = await this.login(tenant);
+            const result = await this.login(tenant, loginScopes);
             result?.authComplete?.resolve();
             return result?.response;
         }

extensions/mssql/src/azure/msal/msalAzureCodeGrant.ts

@@ -59,7 +59,10 @@ export class MsalAzureCodeGrant extends MsalAzureAuth {
         };
     }
 
-    protected async login(tenant: ITenant): Promise<{
+    protected async login(
+        tenant: ITenant,
+        scopes?: string[],
+    ): Promise<{
         response: AuthenticationResult;
         authComplete: IDeferred<void, Error>;
     }> {
@@ -80,13 +83,15 @@ export class MsalAzureCodeGrant extends MsalAzureAuth {
         const state = `${serverPort},${this.pkceCodes.nonce}`;
         let authCodeRequest: AuthorizationCodeRequest;
 
+        const effectiveScopes = scopes ?? this.scopes;
+
         let authority = this.loginEndpointUrl + tenant.id;
         this.logger.info(`Authority URL set to: ${authority}`);
 
         try {
             let authUrlRequest: AuthorizationUrlRequest;
             authUrlRequest = {
-                scopes: this.scopes,
+                scopes: effectiveScopes,
                 redirectUri: `${this.redirectUri}:${serverPort}/redirect`,
                 codeChallenge: this.pkceCodes.codeChallenge,
                 codeChallengeMethod: this.pkceCodes.challengeMethod,
@@ -95,7 +100,7 @@ export class MsalAzureCodeGrant extends MsalAzureAuth {
                 state: state,
             };
             authCodeRequest = {
-                scopes: this.scopes,
+                scopes: effectiveScopes,
                 redirectUri: `${this.redirectUri}:${serverPort}/redirect`,
                 codeVerifier: this.pkceCodes.codeVerifier,
                 authority: authority,

extensions/mssql/src/azure/msal/msalAzureController.ts

@@ -207,7 +207,9 @@ export class MsalAzureController extends AzureController {
                     key: result.account.homeAccountId,
                     token: result.accessToken,
                     tokenType: result.tokenType,
-                    expiresOn: result.account.idTokenClaims.exp,
+                    expiresOn: result.expiresOn
+                        ? Math.floor(result.expiresOn.getTime() / 1000)
+                        : undefined,
                 };
                 return token;
             }
@@ -264,8 +266,8 @@ export class MsalAzureController extends AzureController {
 
             await accountStore.addAccount(newAccount!);
             return await this.getAccountSecurityToken(
-                account,
-                tenantId ?? account.properties.owningTenant.id,
+                newAccount!,
+                tenantId ?? newAccount!.properties.owningTenant.id,
                 settings,
             );
         } catch (ex) {
@@ -279,13 +281,16 @@ export class MsalAzureController extends AzureController {
                         account.properties.azureAuthType,
                         getCloudId(account.key.providerId),
                     );
-                    if (newAccount!.isStale === true) {
+                    if (!newAccount) {
+                        return undefined;
+                    }
+                    if (newAccount.isStale === true) {
                         return undefined;
                     }
                     await accountStore.addAccount(newAccount!);
                     return await this.getAccountSecurityToken(
-                        account,
-                        tenantId ?? account.properties.owningTenant.id,
+                        newAccount!,
+                        tenantId ?? newAccount!.properties.owningTenant.id,
                         settings,
                     );
                 } catch (ex) {

extensions/mssql/src/azure/msal/msalAzureDeviceCode.ts

@@ -30,7 +30,10 @@ export class MsalAzureDeviceCode extends MsalAzureAuth {
         );
     }
 
-    protected async login(tenant: ITenant): Promise<{
+    protected async login(
+        tenant: ITenant,
+        scopes?: string[],
+    ): Promise<{
         response: AuthenticationResult;
         authComplete: IDeferred<void, Error>;
     }> {
@@ -42,8 +45,10 @@ export class MsalAzureDeviceCode extends MsalAzureAuth {
         let authority = this.loginEndpointUrl + tenant.id;
         this.logger.info(`Authority URL set to: ${authority}`);
 
+        const effectiveScopes = scopes ?? this.scopes;
+
         const deviceCodeRequest: DeviceCodeRequest = {
-            scopes: this.scopes,
+            scopes: effectiveScopes,
             authority: authority,
             deviceCodeCallback: async (response) => {
                 await this.displayDeviceCodeScreen(

extensions/mssql/src/controllers/connectionManager.ts

@@ -2408,6 +2408,7 @@ function needsAccountRefresh(errorMessage: string, username: string): boolean {
         errorMessage.includes(AzureConstants.AADSTS70043) ||
         errorMessage.includes(AzureConstants.AADSTS50173) ||
         errorMessage.includes(AzureConstants.AADSTS50020) ||
+        errorMessage.includes(AzureConstants.AADSTS50078) ||
         errorMessage.includes(AzureConstants.mdsUserAccountNotReceived) ||
         errorMessage.includes(Utils.formatString(AzureConstants.mdsUserAccountNotFound, email))
     );