Skip to content

Commit 06eb846

Browse files
authored
feat(auth): add SecondaryAuth (#2991)
## Problem * No "Edit Credentials" button when switching connections * Persisting connections is not cohesive ## Solution * Add "Edit Credentials" button * Add `SecondaryAuth` class (dead code)
1 parent ed54b8e commit 06eb846

File tree

2 files changed

+193
-3
lines changed

2 files changed

+193
-3
lines changed

src/credentials/auth.ts

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -413,6 +413,10 @@ export class Auth implements AuthService, ConnectionManager {
413413
return connections.find(c => c.id === connection.id)
414414
}
415415

416+
public getConnectionState(connection: Pick<Connection, 'id'>): StatefulConnection['state'] | undefined {
417+
return this.store.getProfile(connection.id)?.metadata.connectionState
418+
}
419+
416420
/**
417421
* Attempts to remove all auth state related to the connection.
418422
*
@@ -666,25 +670,32 @@ function toPickerItem(conn: Connection) {
666670

667671
export async function promptForConnection(auth: Auth, type?: 'iam' | 'sso') {
668672
const addNewConnection = {
669-
label: codicon`${getIcon('vscode-plus')} Add new connection`,
673+
label: codicon`${getIcon('vscode-plus')} Add New Connection`,
670674
data: 'addNewConnection' as const,
671675
}
672676

677+
const editCredentials = {
678+
label: codicon`${getIcon('vscode-pencil')} Edit Credentials`,
679+
data: 'editCredentials' as const,
680+
}
681+
673682
const items = (async function () {
674683
// TODO: list linked connections
675684
const connections = await auth.listConnections()
676685
connections.sort((a, b) => (a.type === 'sso' ? -1 : b.type === 'sso' ? 1 : a.label.localeCompare(b.label)))
677686
const filtered = type !== undefined ? connections.filter(c => c.type === type) : connections
687+
const items = [...filtered.map(toPickerItem), addNewConnection]
688+
const canShowEdit = connections.filter(isIamConnection).filter(c => c.label.startsWith('profile')).length > 0
678689

679-
return [...filtered.map(toPickerItem), addNewConnection]
690+
return canShowEdit ? [...items, editCredentials] : items
680691
})()
681692

682693
const placeholder =
683694
type === 'iam'
684695
? localize('aws.auth.promptConnection.iam.placeholder', 'Select an IAM credential')
685696
: localize('aws.auth.promptConnection.all.placeholder', 'Select a connection')
686697

687-
const resp = await showQuickPick<Connection | 'addNewConnection'>(items, {
698+
const resp = await showQuickPick<Connection | 'addNewConnection' | 'editCredentials'>(items, {
688699
placeholder,
689700
title: localize('aws.auth.promptConnection.title', 'Switch Connection'),
690701
buttons: createCommonButtons(),
@@ -698,6 +709,10 @@ export async function promptForConnection(auth: Auth, type?: 'iam' | 'sso') {
698709
return addConnection.execute()
699710
}
700711

712+
if (resp === 'editCredentials') {
713+
return globals.awsContextCommands.onCommandEditCredentials()
714+
}
715+
701716
return resp
702717
}
703718

src/credentials/secondaryAuth.ts

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
/*!
2+
* Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
import globals from '../shared/extensionGlobals'
7+
8+
import * as vscode from 'vscode'
9+
import { getLogger } from '../shared/logger'
10+
import { showQuickPick } from '../shared/ui/pickerPrompter'
11+
import { cast, Optional } from '../shared/utilities/typeConstructors'
12+
import { Auth, Connection } from './auth'
13+
import { once } from '../shared/utilities/functionUtils'
14+
import { UnknownError } from '../shared/errors'
15+
16+
async function promptUseNewConnection(newConn: Connection, oldConn: Connection, tools: string[]) {
17+
// Multi-select picker would be better ?
18+
const saveConnectionItem = {
19+
label: `Yes, keep using ${newConn.label} with ${tools.join(', ')} while using ${
20+
oldConn.label
21+
} with other services.`,
22+
detail: `To remove later, select "Remove Connection from Tool" from the tool's context (right-click) menu.`,
23+
data: 'yes',
24+
} as const
25+
26+
const useConnectionItem = {
27+
label: `No, switch everything to authenticate with ${newConn.label}.`,
28+
detail: 'This will not log you out; you can reconnect at any time by switching connections.',
29+
data: 'no',
30+
} as const
31+
32+
const resp = await showQuickPick([saveConnectionItem, useConnectionItem], {
33+
title: `Some tools you've been using don't work with ${newConn.label}. Keep using ${newConn.label} in the background while using ${oldConn.label}?`,
34+
placeholder: 'Confirm choice',
35+
})
36+
37+
return resp
38+
}
39+
40+
let oldConn: Auth['activeConnection']
41+
const auths = new Map<string, SecondaryAuth>()
42+
const registerAuthListener = once(() => {
43+
Auth.instance.onDidChangeActiveConnection(async conn => {
44+
const potentialConn = oldConn
45+
if (conn !== undefined && potentialConn?.state === 'valid') {
46+
const saveableAuths = Array.from(auths.values()).filter(
47+
a => !a.isUsingSavedConnection && a.isUsable(potentialConn) && !a.isUsable(conn)
48+
)
49+
const toolNames = saveableAuths.map(a => a.toolLabel)
50+
if (saveableAuths.length > 0 && (await promptUseNewConnection(potentialConn, conn, toolNames)) === 'yes') {
51+
await Promise.all(saveableAuths.map(a => a.saveConnection(potentialConn)))
52+
}
53+
}
54+
55+
oldConn = conn
56+
})
57+
})
58+
59+
export function getSecondaryAuth<T extends Connection>(
60+
toolId: string,
61+
toolLabel: string,
62+
isValid: (conn: Connection) => conn is T
63+
): SecondaryAuth<T> {
64+
const auth = new SecondaryAuth(toolId, toolLabel, isValid)
65+
auths.set(toolId, auth)
66+
registerAuthListener()
67+
68+
return auth
69+
}
70+
71+
/**
72+
* Enables a tool to bind to a connection independently from the global {@link Auth} service.
73+
*
74+
* Not all connections are usable by every tool, so callers of this class must provide a function
75+
* that can identify usable connections. Toolkit users are notified whenever a loss of functionality
76+
* would occur after switching connections. Users can then choose to save the usable connection to
77+
* the tool, allowing the global connection to move freely.
78+
*/
79+
export class SecondaryAuth<T extends Connection = Connection> {
80+
#activeConnection: Connection | undefined
81+
#savedConnection: T | undefined
82+
83+
private readonly key = `${this.toolId}.savedConnectionId`
84+
private readonly onDidChangeActiveConnectionEmitter = new vscode.EventEmitter<T | undefined>()
85+
public readonly onDidChangeActiveConnection = this.onDidChangeActiveConnectionEmitter.event
86+
87+
public constructor(
88+
public readonly toolId: string,
89+
public readonly toolLabel: string,
90+
public readonly isUsable: (conn: Connection) => conn is T,
91+
private readonly auth = Auth.instance,
92+
private readonly memento = globals.context.globalState
93+
) {
94+
this.auth.onDidChangeActiveConnection(async conn => {
95+
if (
96+
conn === undefined &&
97+
this.#savedConnection &&
98+
this.#savedConnection.id === this.#activeConnection?.id
99+
) {
100+
await this.removeConnection()
101+
} else {
102+
this.#activeConnection = conn
103+
this.onDidChangeActiveConnectionEmitter.fire(this.activeConnection)
104+
}
105+
})
106+
}
107+
108+
public get activeConnection(): T | undefined {
109+
return (
110+
this.#savedConnection ??
111+
(this.#activeConnection && this.isUsable(this.#activeConnection) ? this.#activeConnection : undefined)
112+
)
113+
}
114+
115+
public get isUsingSavedConnection() {
116+
return this.#savedConnection !== undefined
117+
}
118+
119+
public get isConnectionExpired() {
120+
return !!this.activeConnection && this.auth.getConnectionState(this.activeConnection) === 'invalid'
121+
}
122+
123+
public async saveConnection(conn: T) {
124+
await this.memento.update(this.key, conn.id)
125+
this.#savedConnection = conn
126+
this.onDidChangeActiveConnectionEmitter.fire(this.activeConnection)
127+
}
128+
129+
public async removeConnection() {
130+
await this.memento.update(this.key, undefined)
131+
this.#savedConnection = undefined
132+
this.onDidChangeActiveConnectionEmitter.fire(this.activeConnection)
133+
}
134+
135+
public async useNewConnection(conn: T) {
136+
if (this.auth.activeConnection !== undefined && !this.isUsable(this.auth.activeConnection)) {
137+
if ((await promptUseNewConnection(conn, this.auth.activeConnection, [this.toolLabel])) === 'yes') {
138+
await this.saveConnection(conn)
139+
} else {
140+
await this.auth.useConnection(conn)
141+
}
142+
} else {
143+
await this.auth.useConnection(conn)
144+
}
145+
}
146+
147+
// Used to lazily restore persisted connections.
148+
// Kind of clunky. We need an async module loader layer to make things ergonomic.
149+
public readonly restoreConnection: () => Promise<T | undefined> = once(async () => {
150+
try {
151+
await this.auth.tryAutoConnect()
152+
this.#savedConnection = await this.loadSavedConnection()
153+
this.onDidChangeActiveConnectionEmitter.fire(this.activeConnection)
154+
155+
return this.#savedConnection
156+
} catch (err) {
157+
getLogger().warn(`auth (${this.toolId}): failed to restore connection: ${UnknownError.cast(err).message}`)
158+
}
159+
})
160+
161+
private async loadSavedConnection() {
162+
const id = cast(this.memento.get(this.key), Optional(String))
163+
const conn = id !== undefined ? await this.auth.getConnection({ id }) : undefined
164+
165+
if (conn === undefined) {
166+
getLogger().warn(`auth (${this.toolId}): removing saved connection "${this.key}" as it no longer exists`)
167+
await this.memento.update(this.key, undefined)
168+
} else if (!this.isUsable(conn)) {
169+
getLogger().warn(`auth (${this.toolId}): saved connection "${this.key}" is not valid`)
170+
await this.memento.update(this.key, undefined)
171+
} else {
172+
return conn
173+
}
174+
}
175+
}

0 commit comments

Comments
 (0)