feat: add device flow support for authorizing a new org

Codeneos · Dec 3, 2023 · 57326c1 · 57326c1
1 parent 882c55a
commit 57326c1
Show file tree

Hide file tree

Showing 2 changed files with 228 additions and 23 deletions.
diff --git a/packages/util/src/sfdx.ts b/packages/util/src/sfdx.ts
@@ -41,6 +41,48 @@ try {
     // Ignore errors while updating SFDX logger
 }
 
+export interface SalesforceOAuth2LoginOptions {
+    /**
+     * Login URL to use for the login request, defaults to {@link defaultLoginUrl} when not specified.
+     */
+    instanceUrl?: string;
+    /**
+     * When specified the alias is used to store the org details in the SFDX configuration store. Otherwise the new login is only
+     * stored under the username in the SFDX configuration store.
+     */
+    alias?: string;
+    /**
+     * OAuth2 client ID to use for the login request, defaults to the SFDX default client ID when not specified.
+     */
+    clientId?: string;
+}
+
+export interface DeviceLoginFlow {
+    /**
+     * User code to display to the user that must be entered on the device login page ({@link verificationUrl}) to approve the login request.
+     * The device login page {@link verificationUrl} can be opened automatically by setting the `openVerificationUrl` option to `true` when calling {@link awaitDeviceApproval}.
+     * @readonly
+     */
+    readonly userCode: string;
+    /**
+     * URL to open in the browser to approve the device login request.
+     * @readonly
+     */
+    readonly verificationUrl: string;
+    /**
+     * Optional the verification URL with the user code appended as query parameter.
+     * @readonly
+     */
+    readonly verificationUrlWithCode?: string | undefined;
+    /**
+     * Await the user to approve the device login request.
+     * @param options Options for the device login flow
+     * @param options.openVerificationUrl Open the verification URL in the default browser; defaults to `true` when not specified
+     * @param cancelToken Cancellation token to cancel the login flow
+     * @returns A promise that resolves when the user has approved the device login request
+     */
+    awaitDeviceApproval(options?: { openVerificationUrl?: boolean }, cancelToken?: CancellationToken): Promise<SalesforceAuthResult>;
+}
 
 /**
  * A shim for accessing SFDX functionality
@@ -52,6 +94,12 @@ export namespace sfdx {
      */
     export const defaultConfigPath = `.sfdx/sfdx-config.json`;
 
+    /**
+     * Default login URL to use when no login URL is specified.
+     * Defaults to the test login URL used to login to sandboxes.
+     */
+    export const defaultLoginUrl = 'https://test.salesforce.com';
+
     export const logger: {
         debug(...args: any[]): any;
         error(...args: any[]): any;
@@ -60,15 +108,37 @@ export namespace sfdx {
     } = console;
 
     /**
-     * Login to Salesforce using the web based OAuth2 flow.
+     * Login to Salesforce using the OAuth2 Device Flow. In this flow a user code is
+     * displayed to the user that must be entered on the device login page.
+     *
+     * If o options are specified the default login URL is used {@link defaultLoginUrl}.
+     *
+     * @param options Options for the login flow
+     * @returns The authentication result
+     * @example
+     * ```typescript
+     * const loginFlow = await sfdx.deviceLogin();
+     * console.log(`Please login to Salesforce using the code ${loginFlow.userCode} on ${loginFlow.verificationUrl}`);
+     * // Ser the openVerificationUrl option to false to prevent opening the verification URL in the browser
+     * const authResult = await loginFlow.awaitDeviceApproval({ openVerificationUrl: true });
+     * console.log(`Successfully logged in to Salesforce as ${authResult.username}`);
+     * ```
+     */
+    export async function deviceLogin(options?: SalesforceOAuth2LoginOptions) : Promise<DeviceLoginFlow> {
+        return SalesforceDeviceLoginFlow.start(options);
+    }
+
+    /**
+     * Login to Salesforce using the OAuth2 Authentication Code flow.
      * @param options Options for the login flow
      * @param cancelToken Cancellation token to cancel the login flow
      * @returns The authentication result
      */
-    export async function webLogin(options: { instanceUrl?: string; alias?: string, loginHint?: string }, cancelToken?: CancellationToken) : Promise<SalesforceAuthResult> {
+    export async function webLogin(options: SalesforceOAuth2LoginOptions & { loginHint?: string }, cancelToken?: CancellationToken) : Promise<SalesforceAuthResult> {
         const oauthServer = await salesforce.WebOAuthServer.create({
             oauthConfig: {
-                loginUrl: options.instanceUrl ?? 'https://test.salesforce.com'
+                loginUrl: options.instanceUrl ?? defaultLoginUrl,
+                clientId: options.clientId,
             }
         });
 
@@ -357,4 +427,65 @@ export namespace sfdx {
         const newConfig = options.replace ? config : { ...currentConfig?.config, ...config };
         await options.fs.writeFile(configPath, Buffer.from(JSON.stringify(newConfig, undefined, 2)));
     }
+
+    /**
+     * Default implementation of the {@link DeviceLoginFlow} that wraps
+     * the {@link salesforce.DeviceOauthService} and {@link salesforce.DeviceCodeResponse}.
+     */
+    class SalesforceDeviceLoginFlow implements DeviceLoginFlow {
+        public get userCode() : string {
+            return this.loginData.user_code;
+        }
+
+        public get verificationUrl() : string {
+            return this.loginData.verification_uri;
+        }
+
+        public get verificationUrlWithCode() : string {
+            const hasQueryParams = this.loginData.verification_uri.includes('?');
+            return `${this.loginData.verification_uri}${hasQueryParams ? '&' : '?'}${
+                `user_code=${encodeURIComponent(this.userCode)}`
+            }`;
+        }
+
+        private constructor(
+            private readonly oauthService: salesforce.DeviceOauthService,
+            private readonly loginData: salesforce.DeviceCodeResponse,
+            private readonly options?: SalesforceOAuth2LoginOptions
+        ) { }
+
+        /**
+         * Start a new device login flow. Requests a new device code from Salesforce and returns a new {@link SalesforceDeviceLoginFlow} instance
+         * which can be used to await the user to approve the login request.
+         * @param options Options for the device login flow
+         * @returns A promise that resolves when the user has approved the device login request
+         */
+        static async start(options?: SalesforceOAuth2LoginOptions): Promise<SalesforceDeviceLoginFlow> {
+            const deviceOauthService = await salesforce.DeviceOauthService.create({
+                loginUrl: options?.instanceUrl ?? sfdx.defaultLoginUrl,
+                clientId: options?.clientId
+            });
+            const loginData = await deviceOauthService.requestDeviceLogin();
+            return new SalesforceDeviceLoginFlow(deviceOauthService, loginData, options);
+        }
+
+        public async awaitDeviceApproval(options?: { openVerificationUrl?: boolean }, token?: CancellationToken): Promise<SalesforceAuthResult> {
+            if (options?.openVerificationUrl ?? true) {
+                await open(this.verificationUrlWithCode, { wait: false });
+            }
+            const approval = await Promise.race([
+                this.oauthService.awaitDeviceApproval(this.loginData),
+                new Promise<undefined>((resolve) => token?.onCancellationRequested(() => resolve(undefined)))
+            ]);
+
+            if (!approval) {
+                throw 'User did not approve the device login request within the specified timeout';
+            }
+
+            const authInfo = await this.oauthService.authorizeAndSave(approval);
+            const authResult = authInfo.getFields(true) as SalesforceAuthResult;
+            await sfdx.saveOrg(authResult, { alias: this.options?.alias });
+            return authResult;
+        }
+    }
 }
diff --git a/packages/vscode-extension/src/commands/selectOrgCommand.ts b/packages/vscode-extension/src/commands/selectOrgCommand.ts
@@ -1,5 +1,5 @@
 import * as vscode from 'vscode';
-import { SalesforceOrgInfo, sfdx } from '@vlocode/util';
+import { SalesforceAuthResult, SalesforceOrgInfo, sfdx } from '@vlocode/util';
 import { CommandBase } from '../lib/commandBase';
 import { vscodeCommand } from '../lib/commandRouter';
 import { VlocodeCommand } from '../constants';
@@ -32,6 +32,17 @@ export default class SelectOrgCommand extends CommandBase {
         description: 'Provide a custom instance URL'
     }];
 
+    private readonly authFlows = [{
+        label: '$(key) Web Login Flow',
+        description: 'default',
+        type: 'web',
+        detail: 'You will be redirected to Salesforce to login and authorize the org.'
+    }, {
+        label: '$(shield) Device Login Flow',
+        type: 'device',
+        detail: 'A device code will be generated that can be used to authorize the org manually from the browser.'
+    }];
+
     private readonly salesforceUrlValidator = (url?: string) : string | undefined => {
         const urlRegex = /(^http(s){0,1}:\/\/[^/]+\.[a-z]+(:[0-9]+|)$)|(^\s*$)/i;
         if (url && !urlRegex.test(url)) {
@@ -52,8 +63,8 @@ export default class SelectOrgCommand extends CommandBase {
     }
 
     public async execute() : Promise<void> {
-        const selectionOptions = [ 
-            this.newOrgOption, 
+        const selectionOptions = [
+            this.newOrgOption,
             this.refreshTokensOption
         ];
 
@@ -85,6 +96,13 @@ export default class SelectOrgCommand extends CommandBase {
     }
 
     protected async authorizeNewOrg() : Promise<SalesforceOrgInfo | undefined> {
+        const flowType = await vscode.window.showQuickPick(this.authFlows,
+            { placeHolder: 'Select the authorization flows you want use' });
+
+        if (!flowType) {
+            return;
+        }
+
         const newOrgType = await vscode.window.showQuickPick(this.salesforceOrgTypes,
             { placeHolder: 'Select the type of org you want to authorize' });
 
@@ -105,31 +123,87 @@ export default class SelectOrgCommand extends CommandBase {
             placeHolder: 'Enter an org alias or use the default alias (Press \'Enter\' to confirm or \'Escape\' to cancel)'
         });
 
-        this.logger.log(`Opening '${instanceUrl}' in a new browser window`);
+        if (flowType.type === 'web') {
+            return this.authorizeWebLogin({ instanceUrl, alias });
+        } else {
+            return this.authorizeDeviceLogin({ instanceUrl, alias });
+        }
+    }
+
+    protected async authorizeDeviceLogin(options: { instanceUrl: string, alias?: string }) : Promise<SalesforceOrgInfo | undefined> {
+        const authInfo = await this.vlocode.withActivity({
+            location: vscode.ProgressLocation.Notification,
+            progressTitle: 'Salesforce Device Login',
+            cancellable: true
+        }, async (progress, token) => {
+            // Request Code and URL 
+            progress.report({ message: 'Requesting device login...' });
+            const deviceLogin = await sfdx.deviceLogin(options);
+            if (token?.isCancellationRequested) {
+                return;
+            }
+
+            // Show user code and verification URL adn ask
+            // the user to open the verification URL or just copy the code
+            this.logger.log(`Enter user code [${deviceLogin.userCode}] (without brackets) to confirm login at: ${deviceLogin.verificationUrl}`);
+            progress.report({ message: `Waiting for approval of user code: ${deviceLogin.userCode}` });
+
+            const action = await vscode.window.showInformationMessage(
+                `Open verrification URL?`,
+                {
+                    modal: true,
+                    detail: `Click 'Open in Browser' to open the verification URL in your browser or 'Copy Code' ` +
+                        `to copy the code to your clipboard and open the URL manually.\n\nVerification URL: ${
+                            deviceLogin.verificationUrl}\nUser Code: ${deviceLogin.userCode}`
+                },
+                { title: 'Open in Browser', code: 'open' },
+                { title: 'Copy Code', code: 'copy' },
+                { title: 'Copy URL', code: 'copy_url' },
+                { title: 'Cancel', code: 'cancel', isCloseAffordance: true }
+            );
+
+            if (token?.isCancellationRequested || !action || action?.code === 'cancel') {
+                return;
+            }
+
+            if (action?.code === 'copy') {
+                vscode.env.clipboard.writeText(deviceLogin.userCode);
+            } else if (action?.code === 'copy_url') {
+                vscode.env.clipboard.writeText(deviceLogin.verificationUrlWithCode || deviceLogin.verificationUrl);
+            }
+
+            progress.report({ message: 'Waiting for device approval...' });
+            return await deviceLogin.awaitDeviceApproval({
+                openVerificationUrl: action.code === 'open'
+            }, token);
+        });
+
+        return this.processAuthinfo(authInfo, options);
+    }
+
+    protected async authorizeWebLogin(options: { instanceUrl: string, alias?: string }) : Promise<SalesforceOrgInfo | undefined> {
+        this.logger.log(`Opening '${options.instanceUrl}' in a new browser window`);
         const authInfo = await this.vlocode.withActivity({
             location: vscode.ProgressLocation.Notification,
             progressTitle: 'Opening browser to authorize new org...',
             cancellable: true
         }, async (_, token) => {
-            const loginResult = await sfdx.webLogin({ instanceUrl }, token);
-            if (loginResult && loginResult.accessToken) {
-                return loginResult;
-            }
+            return await sfdx.webLogin(options, token);
         });
 
-        if (authInfo) {
-            if (alias) {
-                await sfdx.setAlias(authInfo.username, alias);
-            }
-            const successMessage = `Successfully authorized ${authInfo.username}, you can now close the browser`;
-            this.logger.log(successMessage);
-            void vscode.window.showInformationMessage(successMessage);
-            return authInfo;
-        }
+        return this.processAuthinfo(authInfo, options);
+    }
 
-        this.logger.error(`Unable to authorize at '${instanceUrl}'`);
-        void vscode.window.showErrorMessage('Failed to authorize new org, see the log for more details');
-        return;
+    private async processAuthinfo(authInfo: SalesforceAuthResult | undefined, options: { instanceUrl: string, alias?: string }): Promise<SalesforceOrgInfo | undefined>{
+        if (!authInfo || !authInfo.accessToken) {
+            this.logger.error(`Unable to authorize at '${options.instanceUrl}'`);
+            void vscode.window.showErrorMessage('Failed to authorize new org, see the log for more details');
+            return;
+        }
+        const successMessage = `Successfully authorized ${authInfo.username}, you can now close the browser`;
+        this.logger.log(successMessage);
+        void vscode.window.showInformationMessage(successMessage);
+        return authInfo;
     }
 }