Add manual tenant sign-in fallback for Conditional Access tenants

src/commands/accounts/signInToTenant.ts

@@ -0,0 +1,82 @@
+/*---------------------------------------------------------------------------------------------
+*  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);
+
+    return unauthenticatedTenants
+        .sort((a, b) => (a.displayName ?? '').localeCompare(b.displayName ?? ''))
+        .map(tenant => ({
+            label: tenant.displayName ?? tenant.defaultDomain ?? tenant.tenantId,
+            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.txt 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);
+    }
+}

test/signInToTenant.test.ts

@@ -0,0 +1,183 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import assert from 'assert';
+import { AzureAccount, AzureSubscriptionProvider, AzureTenant, NotSignedInError, RefreshSuggestedEvent, TenantIdAndAccount, getConfiguredAzureEnv } from '@microsoft/vscode-azext-azureauth';
+import { IActionContext, IAzureQuickPickItem } from '@microsoft/vscode-azext-utils';
+import type * as vscode from 'vscode';
+import { signInToTenant } from '../src/commands/accounts/signInToTenant';
+import { getManualSignInTenant, resetManualSignInTenantForTests, setManualSignInTenant } from '../src/utils/manualSignInTenant';
+
+suite('signInToTenant', () => {
+    setup(async () => {
+        await resetManualSignInTenantForTests();
+    });
+
+    teardown(async () => {
+        await resetManualSignInTenantForTests();
+    });
+
+    test('signs in to selected discovered tenant', async () => {
+        const account = createAccount();
+        const discoveredTenant: AzureTenant = {
+            account,
+            tenantId: 'tenantId',
+            displayName: 'Tenant',
+        };
+        let signedInTenant: Partial<TenantIdAndAccount> | undefined;
+        let inputBoxShown = false;
+
+        const context = createTestActionContext({
+            showQuickPick: async picks => picks[0],
+            showInputBox: async () => {
+                inputBoxShown = true;
+                return undefined;
+            },
+        });
+        const provider = createTestSubscriptionProvider({
+            getAccounts: async () => [account],
+            getUnauthenticatedTenantsForAccount: async () => [discoveredTenant],
+            signIn: async tenant => {
+                signedInTenant = tenant;
+                return true;
+            },
+        });
+
+        await signInToTenant(context, provider);
+
+        assert.strictEqual(signedInTenant, discoveredTenant);
+        assert.strictEqual(inputBoxShown, false);
+    });
+
+    test('prompts for tenant when no account exists for tenant discovery', async () => {
+        let signedInTenantId: string | undefined;
+        let signedInTenantAccount: AzureAccount | undefined;
+        let quickPickShown = false;
+
+        const context = createTestActionContext({
+            showQuickPick: async picks => {
+                quickPickShown = true;
+                return picks[0];
+            },
+            showInputBox: async () => 'contoso.onmicrosoft.com',
+        });
+        const provider = createTestSubscriptionProvider({
+            getAccounts: async () => {
+                throw new NotSignedInError();
+            },
+            signIn: async tenant => {
+                signedInTenantId = tenant?.tenantId;
+                signedInTenantAccount = tenant?.account;
+                return true;
+            },
+        });
+
+        await signInToTenant(context, provider);
+
+        assert.strictEqual(signedInTenantId, 'contoso.onmicrosoft.com');
+        assert.strictEqual(signedInTenantAccount, undefined);
+        assert.strictEqual(quickPickShown, false);
+        assert.strictEqual(getManualSignInTenant(), 'contoso.onmicrosoft.com');
+    });
+
+    test('uses tenant domain as quick pick label when display name is missing', async () => {
+        const account = createAccount();
+        const discoveredTenant: AzureTenant = {
+            account,
+            tenantId: 'tenantId',
+            defaultDomain: 'contoso.onmicrosoft.com',
+        };
+        let shownLabel: string | undefined;
+
+        const context = createTestActionContext({
+            showQuickPick: async picks => {
+                shownLabel = picks[0].label;
+                return picks[0];
+            },
+            showInputBox: async () => undefined,
+        });
+        const provider = createTestSubscriptionProvider({
+            getAccounts: async () => [account],
+            getUnauthenticatedTenantsForAccount: async () => [discoveredTenant],
+            signIn: async () => true,
+        });
+
+        await signInToTenant(context, provider);
+
+        assert.strictEqual(shownLabel, 'contoso.onmicrosoft.com');
+    });
+
+    test('uses tenant id as quick pick label when display name and domain are missing', async () => {
+        const account = createAccount();
+        const discoveredTenant: AzureTenant = {
+            account,
+            tenantId: 'tenantId',
+        };
+        let shownLabel: string | undefined;
+
+        const context = createTestActionContext({
+            showQuickPick: async picks => {
+                shownLabel = picks[0].label;
+                return picks[0];
+            },
+            showInputBox: async () => undefined,
+        });
+        const provider = createTestSubscriptionProvider({
+            getAccounts: async () => [account],
+            getUnauthenticatedTenantsForAccount: async () => [discoveredTenant],
+            signIn: async () => true,
+        });
+
+        await signInToTenant(context, provider);
+
+        assert.strictEqual(shownLabel, 'tenantId');
+    });
+
+    test('resetManualSignInTenantForTests clears persisted tenant', async () => {
+        await setManualSignInTenant('contoso.onmicrosoft.com');
+
+        await resetManualSignInTenantForTests();
+
+        assert.strictEqual(getManualSignInTenant(), undefined);
+    });
+});
+
+function createAccount(): AzureAccount {
+    return {
+        id: 'accountId',
+        label: 'Account',
+        environment: getConfiguredAzureEnv(),
+    };
+}
+
+function createTestActionContext(options: {
+    showQuickPick: (picks: IAzureQuickPickItem<TenantIdAndAccount>[]) => Promise<IAzureQuickPickItem<TenantIdAndAccount>>;
+    showInputBox: () => Promise<string | undefined>;
+}): IActionContext {
+    return {
+        ui: {
+            showQuickPick: options.showQuickPick,
+            showInputBox: options.showInputBox,
+        },
+    } as unknown as IActionContext;
+}
+
+function createTestSubscriptionProvider(options: {
+    getAccounts?: () => Promise<AzureAccount[]>;
+    getUnauthenticatedTenantsForAccount?: (account: AzureAccount) => Promise<AzureTenant[]>;
+    signIn: (tenant?: Partial<TenantIdAndAccount>) => Promise<boolean>;
+}): AzureSubscriptionProvider {
+    const onRefreshSuggested: vscode.Event<RefreshSuggestedEvent> = (_listener, _thisArgs, _disposables) => ({ dispose: () => { /* no-op */ } });
+
+    return {
+        onRefreshSuggested,
+        getAccounts: options.getAccounts ?? (async () => []),
+        getUnauthenticatedTenantsForAccount: options.getUnauthenticatedTenantsForAccount ?? (async () => []),
+        signIn: options.signIn,
+        getAvailableSubscriptions: async () => [],
+        getTenantsForAccount: async () => [],
+        getSubscriptionsForTenant: async () => [],
+    };
+}