Add manual tenant sign-in fallback for Conditional Access tenants

src/commands/accounts/signInToTenant.ts

@@ -0,0 +1,83 @@
+/*---------------------------------------------------------------------------------------------
+*  Copyright (c) Microsoft Corporation. All rights reserved.
+*  Licensed under the MIT License. See License.md in the project root for license information.
+*--------------------------------------------------------------------------------------------*/
+
+import { AzureAccount, AzureSubscriptionProvider, AzureTenant, isNotSignedInError, SignInOptions, TenantIdAndAccount } from '@microsoft/vscode-azext-azureauth';
+import { IActionContext, IAzureQuickPickItem } from '@microsoft/vscode-azext-utils';
+import { setManualSignInTenant } from '../../utils/manualSignInTenant';
+import { localize } from '../../utils/localize';
+
+type ManualTenantSignIn = Partial<TenantIdAndAccount> & Pick<TenantIdAndAccount, 'tenantId'>;
+type TenantSignIn = TenantIdAndAccount | ManualTenantSignIn;
+type ManualTenantSubscriptionProvider = Omit<AzureSubscriptionProvider, 'signIn'> & {
+    signIn(tenant?: Partial<TenantIdAndAccount>, options?: SignInOptions): Promise<boolean>;
+};
+
+export async function signInToTenant(context: IActionContext, subscriptionProvider: ManualTenantSubscriptionProvider, account?: AzureAccount): Promise<void> {
+    const tenantPick = await pickTenant(context, subscriptionProvider, account);
+    if (tenantPick) {
+        const signedIn = await subscriptionProvider.signIn(tenantPick.tenant);
+        if (signedIn && tenantPick.isManual) {
+            await setManualSignInTenant(tenantPick.tenant.tenantId);
+        }
+    }
+}
+
+interface TenantPick {
+    tenant: TenantSignIn;
+    isManual: boolean;
+}
+
+async function pickTenant(context: IActionContext, subscriptionProvider: ManualTenantSubscriptionProvider, account?: AzureAccount): Promise<TenantPick | undefined> {
+    try {
+        const picks = await getPicks(subscriptionProvider, account);
+        if (picks.length) {
+            const pick = await context.ui.showQuickPick(picks, {
+                placeHolder: localize('selectTenantPlaceholder', 'Select a Tenant (Directory) to Sign In To'),
+                matchOnDescription: true,
+            });
+            return { tenant: pick.data, isManual: false };
+        }
+    } catch (err) {
+        if (!isNotSignedInError(err)) {
+            throw err;
+        }
+    }
+
+    const tenant = await promptForTenantId(context);
+    return tenant ? { tenant, isManual: true } : undefined;
+}
+
+async function getPicks(subscriptionProvider: AzureSubscriptionProvider, account?: AzureAccount): Promise<IAzureQuickPickItem<TenantIdAndAccount>[]> {
+    const unauthenticatedTenants: AzureTenant[] = [];
+    const accounts = account ? [account] : await subscriptionProvider.getAccounts();
+    for (const account of accounts) {
+        unauthenticatedTenants.push(...await subscriptionProvider.getUnauthenticatedTenantsForAccount(account));
+    }
+
+    const duplicateTenants = new Set(unauthenticatedTenants
+        .filter((tenant, index, self) => index !== self.findIndex(t => t.tenantId === tenant.tenantId))
+        .map(tenant => tenant.tenantId));
+    const isDuplicate = (tenantId: string) => duplicateTenants.has(tenantId);
+    const tenantLabel = (tenant: AzureTenant) => tenant.displayName ?? tenant.defaultDomain ?? tenant.tenantId;
+
+    return unauthenticatedTenants
+        .sort((a, b) => tenantLabel(a).localeCompare(tenantLabel(b)))
+        .map(tenant => ({
+            label: tenantLabel(tenant),
+            description: `${tenant.tenantId}${isDuplicate(tenant.tenantId) ? ` (${tenant.account.label})` : ''}`,
+            detail: tenant.defaultDomain,
+            data: tenant,
+        }));
+}
+
+async function promptForTenantId(context: IActionContext): Promise<ManualTenantSignIn | undefined> {
+    const tenantId = await context.ui.showInputBox({
+        placeHolder: localize('enterTenantIdPlaceholder', 'Tenant ID or domain'),
+        prompt: localize('enterTenantIdPrompt', 'Enter the tenant ID or domain to sign in to.'),
+        validateInput: value => value?.trim() ? undefined : localize('enterTenantIdValidation', 'Enter a tenant ID or domain.'),
+    });
+
+    return tenantId ? { tenantId: tenantId.trim() } : undefined;
+}

src/commands/registerCommands.ts

@@ -3,7 +3,6 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { signInToTenant } from '@microsoft/vscode-azext-azureauth';
 import { AzExtTreeItem, IActionContext, isAzExtTreeItem, nonNullValue, openUrl, registerCommand, registerCommandWithTreeNodeUnwrapping, registerErrorHandler, registerReportIssueCommand } from '@microsoft/vscode-azext-utils';
 import { commands } from 'vscode';
 import { askAgentAboutActivityLog } from '../chat/askAgentAboutActivityLog/askAgentAboutActivityLog';
@@ -19,6 +18,7 @@ import { GroupingItem } from '../tree/azure/grouping/GroupingItem';
 import { TenantTreeItem } from '../tree/tenants/TenantTreeItem';
 import { logIn } from './accounts/logIn';
 import { SelectSubscriptionOptions, selectSubscriptions } from './accounts/selectSubscriptions';
+import { signInToTenant } from './accounts/signInToTenant';
 import { clearActivities } from './activities/clearActivities';
 import { createResource } from './createResource';
 import { createResourceGroup } from './createResourceGroup';
@@ -102,8 +102,8 @@ export function registerCommands(): void {
     registerCommand('azureResourceGroups.logIn', (context: IActionContext) => logIn(context, { clearSessionPreference: true }));
     registerCommand('azureTenantsView.addAccount', (context: IActionContext) => logIn(context, { clearSessionPreference: true }));
     registerCommand('azureResourceGroups.selectSubscriptions', (context: IActionContext, options: SelectSubscriptionOptions) => selectSubscriptions(context, options));
-    registerCommand('azureResourceGroups.signInToTenant', async () => {
-        await signInToTenant(await ext.subscriptionProviderFactory());
+    registerCommand('azureResourceGroups.signInToTenant', async (context: IActionContext) => {
+        await signInToTenant(context, await ext.subscriptionProviderFactory());
         ext.setClearCacheOnNextLoad('tenant');
         ext.actions.refreshTenantTree();
         ext.setClearCacheOnNextLoad('azure');

src/services/VSCodeAzureSubscriptionProvider.ts

@@ -3,10 +3,11 @@
 *  Licensed under the MIT License. See License.md in the project root for license information.
 *--------------------------------------------------------------------------------------------*/
 
-import { AzureAccount, AzureTenant, GetTenantsForAccountOptions, VSCodeAzureSubscriptionProvider } from '@microsoft/vscode-azext-azureauth';
-import { getSelectedTenantAndSubscriptionIds } from '../commands/accounts/selectSubscriptions';
-import { ext } from '../extensionVariables';
-import { isTenantFilteredOut } from '../utils/tenantSelection';
+import { AzureAccount, AzureTenant, GetTenantsForAccountOptions, TenantIdAndAccount, VSCodeAzureSubscriptionProvider } from '@microsoft/vscode-azext-azureauth';
+import { getSelectedTenantAndSubscriptionIds } from '../commands/accounts/selectSubscriptions';
+import { ext } from '../extensionVariables';
+import { getTenantIdForAuthentication } from '../utils/manualSignInTenant';
+import { isTenantFilteredOut } from '../utils/tenantSelection';
 
 /**
  * Extends {@link VSCodeAzureSubscriptionProvider} to additionally filter tenants based on the
@@ -29,9 +30,16 @@ class ResourceGroupsSubscriptionProvider extends VSCodeAzureSubscriptionProvider
             return filtered;
         }
 
-        return tenants;
-    }
-}
+        return tenants;
+    }
+
+    protected override getSubscriptionContext(tenant: Partial<TenantIdAndAccount>) {
+        return super.getSubscriptionContext({
+            ...tenant,
+            tenantId: getTenantIdForAuthentication(tenant.tenantId),
+        });
+    }
+}
 
 let vscodeAzureSubscriptionProvider: VSCodeAzureSubscriptionProvider | undefined;
 

src/utils/manualSignInTenant.ts

@@ -0,0 +1,46 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.md in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { ext } from '../extensionVariables';
+
+const ManualSignInTenantKey = 'manualSignInTenant';
+let testingManualSignInTenant: string | undefined;
+
+export function normalizeTenant(tenant: string | undefined): string | undefined {
+    const trimmedTenant = tenant?.trim();
+    return trimmedTenant ? trimmedTenant : undefined;
+}
+
+export function getManualSignInTenant(): string | undefined {
+    if (testingManualSignInTenant !== undefined) {
+        return testingManualSignInTenant;
+    }
+
+    if (!ext.context) {
+        return undefined;
+    }
+
+    return normalizeTenant(ext.context.globalState.get<string>(ManualSignInTenantKey));
+}
+
+export async function setManualSignInTenant(tenant: string): Promise<void> {
+    if (!ext.context) {
+        testingManualSignInTenant = normalizeTenant(tenant);
+        return;
+    }
+
+    await ext.context.globalState.update(ManualSignInTenantKey, normalizeTenant(tenant));
+}
+
+export function getTenantIdForAuthentication(tenantId: string | undefined): string | undefined {
+    return normalizeTenant(tenantId) ?? getManualSignInTenant();
+}
+
+export async function resetManualSignInTenantForTests(): Promise<void> {
+    testingManualSignInTenant = undefined;
+    if (ext.context) {
+        await ext.context.globalState.update(ManualSignInTenantKey, undefined);
+    }
+}