|
| 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