feat(auth): discover accounts/roles linked to IdC (AWS SSO) #3023

Problem:
Users can add connections to IAM Identity Center but there's no easy way to use
them for the AWS explorer.

Solution:
Automatically discover linked accounts/roles and list them in the connections picker.

Author: JadenSimon

.changes/next-release/Feature-6b974904-ffc3-46c0-87a8-01ae9bf2370d.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "auth: AWS accounts and roles from IAM Identity Center are automatically discovered by the Toolkit when selecting a connection."
+}

src/credentials/auth.ts

@@ -86,6 +86,8 @@ export function credentialsProviderToTelemetryType(o: CredentialsProviderType):
             return 'envVars'
         case 'profile':
             return 'sharedCredentials'
+        case 'sso':
+            return 'iamIdentityCenter'
         default:
             return 'other'
     }

src/credentials/providers/credentials.ts

@@ -87,6 +87,7 @@ export class CredentialsProviderManager {
     }
 
     public addProvider(provider: CredentialsProvider) {
+        this.removeProvider(provider.getCredentialsId())
         this.providers.push(provider)
     }
 
@@ -102,6 +103,13 @@ export class CredentialsProviderManager {
         this.providerFactories.push(...factory)
     }
 
+    public removeProvider(id: CredentialsId) {
+        const idx = this.providers.findIndex(p => isEqual(id, p.getCredentialsId()))
+        if (idx !== -1) {
+            this.providers.splice(idx, 1)
+        }
+    }
+
     public static getInstance(): CredentialsProviderManager {
         if (!CredentialsProviderManager.INSTANCE) {
             CredentialsProviderManager.INSTANCE = new CredentialsProviderManager()

src/credentials/providers/credentialsProviderManager.ts

@@ -0,0 +1,64 @@
+/*!
+ * Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { Credentials } from '@aws-sdk/types'
+import { CredentialType } from '../../shared/telemetry/telemetry.gen'
+import { getStringHash } from '../../shared/utilities/textUtilities'
+import { CredentialsId, CredentialsProvider, CredentialsProviderType } from './credentials'
+import { SsoClient } from '../sso/clients'
+import { SsoAccessTokenProvider } from '../sso/ssoAccessTokenProvider'
+
+export class SsoCredentialsProvider implements CredentialsProvider {
+    public constructor(
+        private readonly id: CredentialsId,
+        private readonly client: SsoClient,
+        private readonly tokenProvider: SsoAccessTokenProvider,
+        private readonly accountId: string,
+        private readonly roleName: string
+    ) {}
+
+    public async isAvailable(): Promise<boolean> {
+        return true
+    }
+
+    public getCredentialsId(): CredentialsId {
+        return this.id
+    }
+
+    public getProviderType(): CredentialsProviderType {
+        return this.id.credentialSource
+    }
+
+    public getTelemetryType(): CredentialType {
+        return 'ssoProfile'
+    }
+
+    public getDefaultRegion(): string | undefined {
+        return this.client.region
+    }
+
+    public getHashCode(): string {
+        return getStringHash(this.accountId + this.roleName)
+    }
+
+    public async canAutoConnect(): Promise<boolean> {
+        return this.hasToken()
+    }
+
+    public async getCredentials(): Promise<Credentials> {
+        if (!(await this.hasToken())) {
+            await this.tokenProvider.createToken()
+        }
+
+        return this.client.getRoleCredentials({
+            accountId: this.accountId,
+            roleName: this.roleName,
+        })
+    }
+
+    private async hasToken() {
+        return (await this.tokenProvider.getToken()) !== undefined
+    }
+}

src/credentials/providers/ssoCredentialsProvider.ts

@@ -95,6 +95,12 @@ type PromisifyClient<T> = {
 }
 
 export class SsoClient {
+    public get region() {
+        const region = this.client.config.region
+
+        return typeof region === 'string' ? (region as string) : undefined
+    }
+
     public constructor(
         private readonly client: PromisifyClient<SSO>,
         private readonly provider: SsoAccessTokenProvider

src/credentials/sso/clients.ts

@@ -31,6 +31,15 @@ export interface AsyncCollection<T> extends AsyncIterable<T> {
     find<U extends T>(predicate: (item: T) => item is U): Promise<U | undefined>
     find<U extends T>(predicate: (item: T) => boolean): Promise<U | undefined>
 
+    /**
+     * Catches all errors from the underlying iterable(s).
+     *
+     * Note that currently the contiuation behavior is highly dependent on the
+     * underlying implementations. For example, a `for await` loop cannot be
+     * continued if any of the resulting values are rejected.
+     */
+    catch<U>(handler: (error: unknown) => Promise<U> | U): AsyncCollection<T | U>
+
     /**
      * Uses only the first 'count' number of values returned by the generator.
      */
@@ -86,6 +95,8 @@ export function toCollection<T>(generator: () => AsyncGenerator<T, T | undefined
         flatten: () => toCollection<SafeUnboxIterable<T>>(() => delegateGenerator(generator(), flatten)),
         filter: <U extends T>(predicate: Predicate<T, U>) =>
             toCollection<U>(() => filterGenerator<T, U>(generator(), predicate)),
+        catch: <U>(fn: (error: unknown) => Promise<U> | U) =>
+            toCollection<T | U>(() => catchGenerator(generator(), fn)),
         map: <U>(fn: (item: T) => Promise<U> | U) => toCollection<U>(() => mapGenerator(generator(), fn)),
         limit: (count: number) => toCollection(() => delegateGenerator(generator(), takeFrom(count))),
         promise: () => promise(iterable),
@@ -244,3 +255,23 @@ async function find<T, U extends T>(iterable: AsyncIterable<T>, predicate: (item
         }
     }
 }
+
+async function* catchGenerator<T, U, R = T>(
+    generator: AsyncGenerator<T, R | undefined | void>,
+    fn: (error: unknown) => Promise<U> | U
+): AsyncGenerator<T | U, R | U | undefined | void> {
+    while (true) {
+        try {
+            const { value, done } = await generator.next()
+            if (done) {
+                return value
+            }
+            yield value
+        } catch (err) {
+            // Catching an error when the generator would normally
+            // report 'done' means that the 'done' value would be
+            // replaced by `undefined`.
+            yield fn(err)
+        }
+    }
+}

src/shared/utilities/asyncCollection.ts

@@ -37,7 +37,7 @@ export function selectFrom<T, K extends keyof T>(obj: T, ...props: K[]): { [P in
     return props.map(p => [p, obj[p]] as const).reduce((a, [k, v]) => ((a[k] = v), a), {} as { [P in K]: T[P] })
 }
 
-export function isNonNullable<T>(obj: T): obj is NonNullable<T> {
+export function isNonNullable<T>(obj: T | void): obj is NonNullable<T> {
     // eslint-disable-next-line no-null/no-null
     return obj !== undefined && obj !== null
 }

src/shared/utilities/tsUtils.ts

@@ -4,12 +4,22 @@
  */
 
 import * as assert from 'assert'
-import { AuthNode, isSsoConnection, promptForConnection, SsoConnection } from '../../credentials/auth'
+import * as sinon from 'sinon'
+import {
+    AuthNode,
+    Connection,
+    isIamConnection,
+    isSsoConnection,
+    promptForConnection,
+    ssoAccountAccessScopes,
+} from '../../credentials/auth'
 import { ToolkitError } from '../../shared/errors'
 import { assertTreeItem } from '../shared/treeview/testUtil'
 import { getTestWindow } from '../shared/vscode/window'
 import { captureEventOnce } from '../testUtil'
-import { createSsoProfile, createTestAuth } from './testUtil'
+import { createBuilderIdProfile, createSsoProfile, createTestAuth } from './testUtil'
+import { toCollection } from '../../shared/utilities/asyncCollection'
+import globals from '../../shared/extensionGlobals'
 
 const ssoProfile = createSsoProfile()
 const scopedSsoProfile = createSsoProfile({ scopes: ['foo'] })
@@ -181,15 +191,15 @@ describe('Auth', function () {
         assert.ok(await conn.getToken())
     })
 
-    describe('SSO Connections', function () {
-        async function runExpiredGetTokenFlow(conn: SsoConnection, selection: string | RegExp) {
-            const token = conn.getToken()
-            const message = await getTestWindow().waitForMessage(expiredConnPattern)
-            message.selectItem(selection)
+    async function runExpiredConnectionFlow(conn: Connection, selection: string | RegExp) {
+        const creds = conn.type === 'sso' ? conn.getToken() : conn.getCredentials()
+        const message = await getTestWindow().waitForMessage(expiredConnPattern)
+        message.selectItem(selection)
 
-            return token
-        }
+        return creds
+    }
 
+    describe('SSO Connections', function () {
         it('creates a new token if one does not exist', async function () {
             const conn = await auth.createConnection(ssoProfile)
             const provider = auth.getTestTokenProvider(conn)
@@ -198,7 +208,7 @@ describe('Auth', function () {
 
         it('prompts the user if the token is invalid or expired', async function () {
             const conn = await auth.createInvalidSsoConnection(ssoProfile)
-            const token = await runExpiredGetTokenFlow(conn, /login/i)
+            const token = await runExpiredConnectionFlow(conn, /login/i)
             assert.notStrictEqual(token, undefined)
         })
 
@@ -207,7 +217,7 @@ describe('Auth', function () {
             await auth.useConnection(conn)
             await auth.invalidateCachedCredentials(conn)
 
-            const token = runExpiredGetTokenFlow(conn, /no/i)
+            const token = runExpiredConnectionFlow(conn, /no/i)
             await assert.rejects(token, ToolkitError)
 
             assert.strictEqual(auth.activeConnection?.state, 'invalid')
@@ -217,12 +227,139 @@ describe('Auth', function () {
             const err1 = new ToolkitError('test', { code: 'test' })
             const conn = await auth.createConnection(ssoProfile)
             auth.getTestTokenProvider(conn)?.getToken.rejects(err1)
-            const err2 = await runExpiredGetTokenFlow(conn, /no/i).catch(e => e)
+            const err2 = await runExpiredConnectionFlow(conn, /no/i).catch(e => e)
             assert.ok(err2 instanceof ToolkitError)
             assert.strictEqual(err2.cause, err1)
         })
     })
 
+    describe('Linked Connections', function () {
+        const linkedSsoProfile = createSsoProfile({ scopes: ssoAccountAccessScopes })
+        const accountRoles = [
+            { accountId: '1245678910', roleName: 'foo' },
+            { accountId: '9876543210', roleName: 'foo' },
+            { accountId: '9876543210', roleName: 'bar' },
+        ]
+
+        beforeEach(function () {
+            auth.ssoClient.listAccounts.returns(
+                toCollection(async function* () {
+                    yield [{ accountId: '1245678910' }, { accountId: '9876543210' }]
+                })
+            )
+
+            auth.ssoClient.listAccountRoles.callsFake(req =>
+                toCollection(async function* () {
+                    yield accountRoles.filter(i => i.accountId === req.accountId)
+                })
+            )
+
+            auth.ssoClient.getRoleCredentials.resolves({
+                accessKeyId: 'xxx',
+                secretAccessKey: 'xxx',
+                expiration: new Date(Date.now() + 1000000),
+            })
+
+            sinon.stub(globals.loginManager, 'validateCredentials').resolves('')
+        })
+
+        afterEach(function () {
+            sinon.restore()
+        })
+
+        it('lists linked conections for SSO connections', async function () {
+            await auth.createConnection(linkedSsoProfile)
+            const connections = await auth.listAndTraverseConnections().promise()
+            assert.deepStrictEqual(
+                connections.map(c => c.type),
+                ['sso', 'iam', 'iam', 'iam']
+            )
+        })
+
+        it('does not gather linked accounts when calling `listConnections`', async function () {
+            await auth.createConnection(linkedSsoProfile)
+            const connections = await auth.listConnections()
+            assert.deepStrictEqual(
+                connections.map(c => c.type),
+                ['sso']
+            )
+        })
+
+        it('caches linked conections when the source connection becomes invalid', async function () {
+            const conn = await auth.createConnection(linkedSsoProfile)
+            await auth.listAndTraverseConnections().promise()
+            await auth.invalidateCachedCredentials(conn)
+
+            const connections = await auth.listConnections()
+            assert.deepStrictEqual(
+                connections.map(c => c.type),
+                ['sso', 'iam', 'iam', 'iam']
+            )
+        })
+
+        it('gracefully handles source connections becoming invalid when discovering linked accounts', async function () {
+            await auth.createConnection(linkedSsoProfile)
+            auth.ssoClient.listAccounts.rejects(new Error('No access'))
+            const connections = await auth.listAndTraverseConnections().promise()
+            assert.deepStrictEqual(
+                connections.map(c => c.type),
+                ['sso']
+            )
+        })
+
+        it('removes linked connections when the source connection is deleted', async function () {
+            const conn = await auth.createConnection(linkedSsoProfile)
+            await auth.listAndTraverseConnections().promise()
+            await auth.deleteConnection(conn)
+
+            assert.deepStrictEqual(await auth.listAndTraverseConnections().promise(), [])
+        })
+
+        it('prompts the user to reauthenticate if the source connection becomes invalid', async function () {
+            const source = await auth.createConnection(linkedSsoProfile)
+            const conn = await auth.listAndTraverseConnections().find(c => isIamConnection(c) && c.id.includes('sso'))
+            assert.ok(conn)
+            await auth.useConnection(conn)
+            await auth.reauthenticate(conn)
+            await auth.invalidateCachedCredentials(conn)
+            await auth.invalidateCachedCredentials(source)
+
+            await runExpiredConnectionFlow(conn, /login/i)
+            assert.strictEqual(auth.getConnectionState(source), 'valid')
+            assert.strictEqual(auth.getConnectionState(conn), 'valid')
+        })
+
+        describe('Multiple Connections', function () {
+            const otherProfile = createBuilderIdProfile({ scopes: ssoAccountAccessScopes })
+
+            // Equivalent profiles from multiple sources is a fairly rare situation right now
+            // Ideally they would be de-duped although the implementation can be tricky
+            it('can handle multiple SSO connection and does not de-dupe', async function () {
+                await auth.createConnection(linkedSsoProfile)
+                await auth.createConnection(otherProfile)
+
+                const connections = await auth.listAndTraverseConnections().promise()
+                assert.deepStrictEqual(
+                    connections.map(c => c.type),
+                    ['sso', 'sso', 'iam', 'iam', 'iam', 'iam', 'iam', 'iam'],
+                    'Expected two SSO connections and 3 IAM connections for each SSO connection'
+                )
+            })
+
+            it('does not stop discovery if one connection fails', async function () {
+                const otherProfile = createBuilderIdProfile({ scopes: ssoAccountAccessScopes })
+                await auth.createConnection(linkedSsoProfile)
+                await auth.createConnection(otherProfile)
+                auth.ssoClient.listAccounts.onFirstCall().rejects(new Error('No access'))
+                const connections = await auth.listAndTraverseConnections().promise()
+                assert.deepStrictEqual(
+                    connections.map(c => c.type),
+                    ['sso', 'sso', 'iam', 'iam', 'iam']
+                )
+            })
+        })
+    })
+
     describe('AuthNode', function () {
         it('shows a message to create a connection if no connections exist', async function () {
             const node = new AuthNode(auth)

src/test/credentials/auth.test.ts

@@ -304,7 +304,7 @@ describe('SharedCredentialsProvider', async function () {
         `
 
         beforeEach(function () {
-            const client = stub(SsoClient)
+            const client = stub(SsoClient, { region: 'foo' })
             client.getRoleCredentials.callsFake(async request => {
                 assert.strictEqual(request.accountId, '012345678910')
                 assert.strictEqual(request.roleName, 'MyRole')