auth: Add BearerChallengePolicy for MFA step-up challenges during subscription listing (#2231)

* auth: Add BearerChallengePolicy to handle MFA step-up challenges

The auth package creates SubscriptionClient instances to list tenants and
subscriptions, but these clients had no challenge policy on their HTTP
pipeline. When users hit an MFA step-up challenge (401 with WWW-Authenticate
header) during tenant/subscription listing, the request simply fails.

Changes:
- Create BearerChallengePolicy in auth/src/utils/BearerChallengePolicy.ts
- Wire into SubscriptionClient pipeline in AzureSubscriptionProviderBase
- Export BearerChallengePolicy and getDefaultScopeFromEndpoint
- Add @azure/core-rest-pipeline as devDependency
- Bump version to 6.0.0-alpha.6

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address PR feedback: dep classification and retry tracking

- Move @azure/core-rest-pipeline back to regular dependencies since
  BearerChallengePolicy is exported and its .d.ts references pipeline types
- Replace HTTP header-based retry marker (x-azext-challenge-retry) with a
  WeakSet to track retried requests out-of-band, avoiding sending the
  marker over the wire

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: alexweininger

auth/package-lock.json

@@ -1,7 +1,7 @@
 {
     "name": "@microsoft/vscode-azext-azureauth",
     "author": "Microsoft Corporation",
-    "version": "6.0.0-alpha.5",
+    "version": "6.0.0-alpha.6",
     "description": "Azure authentication helpers for Visual Studio Code",
     "tags": [
         "azure",
@@ -55,6 +55,7 @@
     },
     "dependencies": {
         "@azure/arm-resources-subscriptions": "^2.1.0",
+        "@azure/core-rest-pipeline": "^1.22.2",
         "@azure/identity": "^4.13.0",
         "@azure/ms-rest-azure-env": "^2.0.0"
     },

auth/package.json

@@ -12,6 +12,7 @@ export * from './contracts/AzureTenant';
 // The `AzureDevOpsSubscriptionProvider` is intentionally not exported, it must be imported from `'@microsoft/vscode-azext-azureauth/azdo'`
 export * from './providers/AzureSubscriptionProviderBase';
 export * from './providers/VSCodeAzureSubscriptionProvider';
+export * from './utils/BearerChallengePolicy';
 export * from './utils/configuredAzureEnv';
 export * from './utils/dedupeSubscriptions';
 export * from './utils/getMetricsForTelemetry';

auth/src/index.ts

@@ -5,6 +5,7 @@
 
 import type { SubscriptionClient } from '@azure/arm-resources-subscriptions'; // Keep this as `import type` to avoid actually loading the package before necessary
 import type { GetTokenOptions, TokenCredential } from '@azure/core-auth'; // Keep this as `import type` to avoid actually loading the package (at all, this one is dev-only)
+import { BearerChallengePolicy } from '../utils/BearerChallengePolicy';
 import { inspect } from 'util';
 import * as vscode from 'vscode';
 import type { AzureAccount } from '../contracts/AzureAccount';
@@ -364,8 +365,23 @@ export abstract class AzureSubscriptionProviderBase implements AzureSubscription
 
         armSubs ??= await import('@azure/arm-resources-subscriptions');
 
+        const endpoint = getConfiguredAzureEnv().resourceManagerEndpointUrl;
+        const client = new armSubs.SubscriptionClient(credential, { endpoint });
+
+        client.pipeline.addPolicy(
+            new BearerChallengePolicy(
+                async (challenge) => {
+                    this.silenceRefreshEvents();
+                    const session = await getSessionFromVSCode(challenge, tenant.tenantId, { createIfNone: true, account: tenant.account });
+                    return session?.accessToken;
+                },
+                endpoint,
+            ),
+            { phase: 'Sign', afterPolicies: ['bearerTokenAuthenticationPolicy'] },
+        );
+
         return {
-            client: new armSubs.SubscriptionClient(credential, { endpoint: getConfiguredAzureEnv().resourceManagerEndpointUrl }),
+            client,
             credential: credential,
             authentication: {
                 getSession: async () => {

auth/src/providers/AzureSubscriptionProviderBase.ts

@@ -0,0 +1,44 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import type { PipelinePolicy, PipelineRequest, PipelineResponse, SendRequest } from '@azure/core-rest-pipeline';
+import type * as vscode from 'vscode';
+
+export function getDefaultScopeFromEndpoint(endpoint?: string): string {
+    let base = endpoint ?? 'https://management.azure.com/';
+    base = base.replace(/\/+$/, '');
+    return `${base}/.default`;
+}
+
+const challengeRetriedRequests = new WeakSet<PipelineRequest>();
+
+export class BearerChallengePolicy implements PipelinePolicy {
+    public readonly name = 'BearerChallengePolicy';
+
+    public constructor(
+        private readonly getTokenForChallenge: (request: vscode.AuthenticationWwwAuthenticateRequest) => Promise<string | undefined>,
+        private readonly endpoint?: string,
+    ) { }
+
+    public async sendRequest(request: PipelineRequest, next: SendRequest): Promise<PipelineResponse> {
+        const initial = await next(request);
+
+        if (initial.status === 401 && !challengeRetriedRequests.has(request)) {
+            const header = initial.headers.get('WWW-Authenticate') || initial.headers.get('www-authenticate');
+            if (header) {
+                const scopes = [getDefaultScopeFromEndpoint(this.endpoint)];
+                challengeRetriedRequests.add(request);
+
+                const token = await this.getTokenForChallenge({ wwwAuthenticate: header, fallbackScopes: scopes });
+                if (token) {
+                    request.headers.set('Authorization', `Bearer ${token}`);
+                    return await next(request);
+                }
+            }
+        }
+
+        return initial;
+    }
+}