fix(auth): allow flow to continue if users copy url #5007

Problem: Clicking anything other than 'open' in the VSC redirect url modal for the auth page would result in cancelling the auth attempt.

Solution: Do not expect a specific response from the url modal. Allow the auth flow to continue even if the user hits 'Cancel' or 'Copy'. To cancel the auth flow, the user should click 'Cancel' in the login ui OR leave the login ui.

Author: hayemaxi

packages/core/src/auth/sso/server.ts

@@ -11,6 +11,7 @@ import { Socket } from 'net'
 import globals from '../../shared/extensionGlobals'
 import { Result } from '../../shared/utilities/result'
 import { FileSystemCommon } from '../../srcShared/fs'
+import { CancellationError } from '../../shared/utilities/timeoutUtils'
 
 export class MissingPortError extends ToolkitError {
     constructor() {
@@ -47,6 +48,12 @@ export class AuthError extends ToolkitError {
  * back to VSCode
  */
 export class AuthSSOServer {
+    // Last initialized instance of the Auth Server
+    static #lastInstance: AuthSSOServer | undefined
+    public static get lastInstance() {
+        return AuthSSOServer.#lastInstance
+    }
+
     private baseUrl = `http://127.0.0.1`
     private oauthCallback = '/oauth/callback'
     private authenticationFlowTimeoutInMs = 600000
@@ -57,7 +64,9 @@ export class AuthSSOServer {
     private server: http.Server
     private connections: Socket[]
 
-    constructor(private readonly state: string) {
+    private _closed: boolean = false
+
+    private constructor(private readonly state: string) {
         this.authenticationPromise = new Promise<Result<string>>(resolve => {
             this.deferred = { resolve }
         })
@@ -99,6 +108,22 @@ export class AuthSSOServer {
         })
     }
 
+    public static init(state: string) {
+        const lastInstance = AuthSSOServer.#lastInstance
+        if (lastInstance !== undefined && !lastInstance.closed) {
+            lastInstance
+                .close()!
+                .catch(err =>
+                    getLogger().error('Failed to close already existing auth sever in AuthSSOServer.init(): %s', err)
+                )
+        }
+
+        getLogger().debug('AuthSSOServer: Initialized new auth server.')
+        const instance = new AuthSSOServer(state)
+        AuthSSOServer.#lastInstance = instance
+        return instance
+    }
+
     start() {
         if (this.server.listening) {
             throw new ToolkitError('AuthSSOServer: Server already started')
@@ -127,10 +152,17 @@ export class AuthSSOServer {
 
     close() {
         return new Promise<void>((resolve, reject) => {
+            if (this._closed) {
+                resolve()
+                return
+            }
+
             if (!this.server.listening) {
                 reject(new ToolkitError('AuthSSOServer: Server not started'))
             }
 
+            getLogger().debug('AuthSSOServer: Attempting to close server.')
+
             this.connections.forEach(connection => {
                 connection.destroy()
             })
@@ -139,6 +171,8 @@ export class AuthSSOServer {
                 if (err) {
                     reject(err)
                 }
+                this._closed = true
+                getLogger().debug('AuthSSOServer: Server closed successfully.')
                 resolve()
             })
         })
@@ -152,6 +186,10 @@ export class AuthSSOServer {
         return `${this.baseUrl}:${this.getPort()}`
     }
 
+    public get closed() {
+        return this._closed
+    }
+
     public getAddress() {
         return this.server.address()
     }
@@ -228,6 +266,11 @@ export class AuthSSOServer {
         this.deferred?.resolve(Result.err(error))
     }
 
+    public cancelCurrentFlow() {
+        getLogger().debug('AuthSSOServer: Cancelling current login flow')
+        this.deferred?.resolve(Result.err(new CancellationError('user')))
+    }
+
     public waitForAuthorization(): Promise<Result<string>> {
         return Promise.race([
             this.authenticationPromise,
@@ -244,7 +287,7 @@ export class AuthSSOServer {
                     getLogger().warn('AuthSSOServer: Authentication is taking a long time')
                 }, this.authenticationWarningTimeoutInMs)
 
-                void this.authenticationPromise.then(() => {
+                void this.authenticationPromise.finally(() => {
                     clearTimeout(warningTimeout)
                 })
             }),

packages/core/src/auth/sso/ssoAccessTokenProvider.ts

@@ -19,7 +19,6 @@ import {
     builderIdStartUrl,
     openSsoPortalLink,
     isDeprecatedAuth,
-    openSsoUrl,
 } from './model'
 import { getCache } from './cache'
 import { hasProps, hasStringProps, RequiredProps, selectFrom } from '../../shared/utilities/tsUtils'
@@ -436,7 +435,7 @@ class AuthFlowAuthorization extends SsoAccessTokenProvider {
         registration: ClientRegistration
     ): Promise<{ token: SsoToken; registration: ClientRegistration; region: string; startUrl: string }> {
         const state = randomUUID()
-        const authServer = new AuthSSOServer(state)
+        const authServer = AuthSSOServer.init(state)
 
         try {
             await authServer.start()
@@ -457,9 +456,7 @@ class AuthFlowAuthorization extends SsoAccessTokenProvider {
                     codeChallengeMethod: 'S256',
                 })
 
-                if (!(await openSsoUrl(vscode.Uri.parse(location)))) {
-                    throw new CancellationError('user')
-                }
+                await vscode.env.openExternal(vscode.Uri.parse(location))
 
                 const authorizationCode = await authServer.waitForAuthorization()
                 if (authorizationCode.isErr()) {

packages/core/src/login/webview/commonAuthViewProvider.ts

@@ -104,6 +104,7 @@ export class CommonAuthViewProvider implements WebviewViewProvider {
                     this.webView!.server.storeMetricMetadata({ isReAuth: false })
                 }
                 this.webView!.server.emitAuthMetric()
+                this.webView!.server.cancelAuthFlow()
 
                 // Set after emitting. If users use side bar to return to login, this source is correct
                 // for the next iteration. Otherwise, other sources will be set accordingly by whatever

packages/core/src/login/webview/vue/backend.ts

@@ -31,6 +31,7 @@ import { AuthSources } from '../util'
 import { AuthEnabledFeatures, AuthError, AuthFlowState, AuthUiClick, TelemetryMetadata, userCancelled } from './types'
 import { AuthUtil } from '../../../codewhisperer/util/authUtil'
 import { DevSettings } from '../../../shared/settings'
+import { AuthSSOServer } from '../../../auth/sso/server'
 
 export abstract class CommonAuthWebview extends VueWebview {
     private metricMetadata: TelemetryMetadata = {}
@@ -296,4 +297,8 @@ export abstract class CommonAuthWebview extends VueWebview {
     getDefaultStartUrl() {
         return DevSettings.instance.get('autofillStartUrl', '')
     }
+
+    cancelAuthFlow() {
+        AuthSSOServer.lastInstance?.cancelCurrentFlow()
+    }
 }

packages/core/src/login/webview/vue/login.vue

@@ -514,6 +514,8 @@ export default defineComponent({
             void client.emitUiClick('auth_regionSelection')
         },
         async handleCancelButton() {
+            void client.cancelAuthFlow()
+
             await client.storeMetricMetadata({ isReAuth: false, result: 'Cancelled' })
             void client.emitAuthMetric()
             void client.emitUiClick('auth_cancelButton')

packages/core/src/test/credentials/sso/server.test.ts

@@ -13,7 +13,8 @@ import {
 } from '../../../auth/sso/server'
 import request from '../../../common/request'
 import { URLSearchParams } from 'url'
-import { ToolkitError } from '../../../shared/errors'
+import { isUserCancelledError, ToolkitError } from '../../../shared/errors'
+import { sleep } from '../../../shared/utilities/timeoutUtils'
 
 describe('AuthSSOServer', function () {
     const code = 'zfhgaiufgsbdfigsdfg'
@@ -24,7 +25,7 @@ describe('AuthSSOServer', function () {
     let server: AuthSSOServer
 
     beforeEach(async function () {
-        server = new AuthSSOServer(state)
+        server = AuthSSOServer.init(state)
         await server.start()
     })
 
@@ -122,4 +123,21 @@ describe('AuthSSOServer', function () {
         }
         assert.fail('Expected address 127.0.0.1')
     })
+
+    it('can be cancelled while waiting for auth', async function () {
+        const promise = server.waitForAuthorization().catch(e => {
+            return e
+        })
+        server.cancelCurrentFlow()
+
+        const err = await promise
+        assert.ok(isUserCancelledError(err.inner), 'CancellationError not thrown.')
+    })
+
+    it('initializes and closes instances', async function () {
+        const newServer = AuthSSOServer.init('1234')
+        assert.equal(AuthSSOServer.lastInstance, newServer)
+        await sleep(100)
+        assert.ok(server.closed)
+    })
 })