fix: handle rejected promise in server tests (#4705)

* fix: handle rejected promise in server tests

Problem:
- We get:
rejected promise not handled within 1 second: Error: AuthSSOServer: foo: foo
rejected promise not handled within 1 second: Error: AuthSSOServer: invalid state
rejected promise not handled within 1 second: Error: AuthSSOServer: missing code
rejected promise not handled within 1 second: Error: AuthSSOServer: missing state

Solution:
- Catch and test the authorization promise

Author: jpinkney-aws

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

@@ -8,6 +8,7 @@ import { getLogger } from '../../shared/logger'
 import { ToolkitError } from '../../shared/errors'
 import { Socket } from 'net'
 import globals from '../../shared/extensionGlobals'
+import { Result } from '../../shared/utilities/result'
 
 export class MissingPortError extends ToolkitError {
     constructor() {
@@ -49,14 +50,14 @@ export class AuthSSOServer {
     private authenticationFlowTimeoutInMs = 600000
     private authenticationWarningTimeoutInMs = 60000
 
-    private readonly authenticationPromise: Promise<string>
-    private deferred: { resolve: (result: string) => void; reject: (reason: any) => void } | undefined
+    private readonly authenticationPromise: Promise<Result<string>>
+    private deferred: { resolve: (result: Result<string>) => void } | undefined
     private server: http.Server
     private connections: Socket[]
 
     constructor(private readonly state: string, private readonly vscodeUriPath: string) {
-        this.authenticationPromise = new Promise<string>((resolve, reject) => {
-            this.deferred = { resolve, reject }
+        this.authenticationPromise = new Promise<Result<string>>(resolve => {
+            this.deferred = { resolve }
         })
 
         this.connections = []
@@ -174,7 +175,7 @@ export class AuthSSOServer {
             return
         }
 
-        this.deferred?.resolve(code)
+        this.deferred?.resolve(Result.ok(code))
         res.setHeader('Content-Type', 'text/html')
         res.writeHead(200)
         res.end(`
@@ -194,13 +195,13 @@ export class AuthSSOServer {
         res.end(error.message)
 
         // Send the response back to the editor
-        this.deferred?.reject(error)
+        this.deferred?.resolve(Result.err(error))
     }
 
-    public waitForAuthorization(): Promise<string> {
+    public waitForAuthorization(): Promise<Result<string>> {
         return Promise.race([
             this.authenticationPromise,
-            new Promise<string>((_, reject) => {
+            new Promise<Result<string>>((_, reject) => {
                 globals.clock.setTimeout(() => {
                     reject(
                         new ToolkitError('Timed-out waiting for browser login flow to complete', {

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

@@ -427,14 +427,17 @@ class AuthFlowAuthorization extends SsoAccessTokenProvider {
                 }
 
                 const authorizationCode = await authServer.waitForAuthorization()
+                if (authorizationCode.isErr()) {
+                    throw authorizationCode.err()
+                }
 
                 const tokenRequest: OidcClientPKCE.CreateTokenRequest = {
                     clientId: registration.clientId,
                     clientSecret: registration.clientSecret,
                     grantType: authorizationGrantType,
                     redirectUri,
                     codeVerifier,
-                    code: authorizationCode,
+                    code: authorizationCode.unwrap(),
                 }
 
                 return this.oidc.createToken(tokenRequest)

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

@@ -13,6 +13,7 @@ import {
 } from '../../../auth/sso/server'
 import request, { RequestError } from '../../../common/request'
 import { URLSearchParams } from 'url'
+import { ToolkitError } from '../../../shared/errors'
 
 describe('AuthSSOServer', function () {
     const code = 'zfhgaiufgsbdfigsdfg'
@@ -39,13 +40,25 @@ describe('AuthSSOServer', function () {
 
     async function assertRequestError(params: Record<string, string>, expectedErrorMsg: string) {
         const url = createURL(server.redirectUri, params)
+        const authorizationPromise = server.waitForAuthorization()
         try {
             const response = await request.fetch('GET', url).response
             assert.fail(`Expected error but found ${response.body}`)
         } catch (err: unknown) {
-            if (err instanceof RequestError) {
-                assert.strictEqual(err.code, 400)
-                assert.deepStrictEqual(err.body, expectedErrorMsg)
+            if (!(err instanceof RequestError)) {
+                assert.fail('Unknown error')
+            }
+
+            const e = err as RequestError
+            assert.strictEqual(e.code, 400)
+            assert.deepStrictEqual(e.body, expectedErrorMsg)
+        }
+
+        try {
+            await authorizationPromise
+        } catch (err: unknown) {
+            if (err instanceof ToolkitError) {
+                assert.deepStrictEqual(err.message, expectedErrorMsg)
                 return
             }
             assert.fail('Unknown error')
@@ -103,7 +116,7 @@ describe('AuthSSOServer', function () {
         assert.deepStrictEqual(response.status, 200)
 
         const token = await server.waitForAuthorization()
-        assert.deepStrictEqual(code, token)
+        assert.deepStrictEqual(code, token.unwrap())
     })
 
     it('address is bound to localhost', function () {