Skip to content

Commit

Permalink
refactor(auth): remove sharing of direct sessions across extensions (a…
Browse files Browse the repository at this point in the history
…ws#5169)

* refactor(auth): remove sharing of direct sessions across extensions

**Separate sessions commit**

Instead of sharing sessions directly (e.g. with tokens, logout logs out of both extensions, etc), we will just share SSO parameters across extensions. The auth can re-use starturl and region but will build its own connection.
Existing connections UI remains for this purpose.

- Remove unused auth and API code. This includes special handling logic that imports direct sessions from the other extension.
- Builder ID will not be displayed as an existing connection because the parameters for this are already known. Users can just log in with buidler ID normally.
- Misc fixes in comments, log statements.

* refactor: rename "existing" to "imported" logins

* feat(auth): re-add checks for existing start urls in same extension
  • Loading branch information
hayemaxi authored Jun 17, 2024
1 parent 0e7b9d4 commit ef74cea
Show file tree
Hide file tree
Showing 14 changed files with 73 additions and 421 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
{
"type": "Removal",
"description": "Auth: No longer share SSO sessions with AWS Toolkit."
}
68 changes: 0 additions & 68 deletions packages/amazonq/src/auth/util.ts

This file was deleted.

12 changes: 6 additions & 6 deletions packages/amazonq/src/extensionCommon.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import {
activate as activateCodeWhisperer,
shutdown as shutdownCodeWhisperer,
amazonQDismissedKey,
AuthUtil,
} from 'aws-core-vscode/codewhisperer'
import {
ExtContext,
Expand All @@ -25,11 +26,10 @@ import {
getLogger,
getMachineId,
} from 'aws-core-vscode/shared'
import { initializeAuth, CredentialsStore, LoginManager, AuthUtils } from 'aws-core-vscode/auth'
import { initializeAuth, CredentialsStore, LoginManager, AuthUtils, SsoConnection } from 'aws-core-vscode/auth'
import { CommonAuthWebview } from 'aws-core-vscode/login'
import { VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils'
import { telemetry, ExtStartUpSources } from 'aws-core-vscode/telemetry'
import { getAuthStatus } from './auth/util'
import { makeEndpointsProvider, registerGenericCommands } from 'aws-core-vscode/common'
import { registerCommands } from './commands'

Expand Down Expand Up @@ -131,11 +131,11 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
telemetry.record({ source: ExtStartUpSources.reload })
}

const { authStatus, authEnabledConnections, authScopes } = await getAuthStatus()
const authState = (await AuthUtil.instance.getChatAuthState()).codewhispererChat
telemetry.record({
authStatus,
authEnabledConnections,
authScopes,
authStatus: authState === 'connected' || authState === 'expired' ? authState : 'notConnected',
authEnabledConnections: AuthUtils.getAuthFormIdsFromConnection(AuthUtil.instance.conn).join(','),
authScopes: ((AuthUtil.instance.conn as SsoConnection)?.scopes ?? []).join(','),
})
})
}
Expand Down
46 changes: 1 addition & 45 deletions packages/core/src/auth/auth.ts
Original file line number Diff line number Diff line change
Expand Up @@ -952,22 +952,7 @@ export class Auth implements AuthService, ConnectionManager {
: `${localizedText.iamIdentityCenter} (${truncatedUrl})`
}

// Used by Amazon Q to re-use connection from AWS Toolkit listConnection API response
public async createConnectionFromApi(connection: AwsConnection) {
getLogger().info(`Reusing connection ${connection.id}`)
const profile = {
type: connection.type,
ssoRegion: connection.ssoRegion,
scopes: connection.scopes,
startUrl: connection.startUrl,
} as SsoProfile
const id = connection.id
const storedProfile = await this.store.addProfile(id, profile)
await this.updateConnectionState(id, connection.state)
return this.getSsoConnection(id, storedProfile)
}

// Used by AWS Toolkit to update connection status & scope when this connection is updated by Amazon Q
// Used by AWS Toolkit to update connection status & scope when this connection is updated by another extension.
// If such connection does not exist, create one with same id.
// Otherwise, update its scope and/or state.
public async setConnectionFromApi(connection: AwsConnection) {
Expand Down Expand Up @@ -999,35 +984,6 @@ export class Auth implements AuthService, ConnectionManager {
}
}

// Used by Amazon Q to update connection status & scope when this connection is updated by AWS Toolkit
// do not create connection in Q for each change event from Toolkit
public async onConnectionUpdate(connection: AwsConnection) {
const conn = await this.getConnection({ id: connection.id })
if (conn) {
const profile = {
type: connection.type,
ssoRegion: connection.ssoRegion,
scopes: connection.scopes,
startUrl: connection.startUrl,
} as SsoProfile
await this.store.updateProfile(connection.id, profile)

await this.store.updateMetadata(connection.id, { connectionState: connection.state })
}
}

// Used by Amazon Q to delete connection status & scope when this deletion is made by AWS Toolkit
// NO event should be emitted from this deletion
// Do not actually perform the delete because toolkit has done the deletion
// Delete the momento states only.
public async onDeleteConnection(id: string) {
const profile = this.store.getProfile(id)
if (profile) {
await this.store.deleteProfile(id)
await this.store.setCurrentProfileId(undefined)
}
}

public declareConnectionFromApi(conn: Pick<AwsConnection, 'startUrl' | 'ssoRegion'>, source: string) {
getLogger().debug(`Declared connection '${conn.startUrl}' from ${source}.`)
this._declaredConnections[conn.startUrl] = {
Expand Down
11 changes: 0 additions & 11 deletions packages/core/src/auth/secondaryAuth.ts
Original file line number Diff line number Diff line change
Expand Up @@ -265,17 +265,6 @@ export class SecondaryAuth<T extends Connection = Connection> {
return conn
}
}

// Used by Amazon Q to delete connection status & scope when this deletion is made by AWS Toolkit
// NO event should be emitted from this deletion to avoid infinite loop
public async onDeleteConnection(id: string) {
await this.auth.onDeleteConnection(id)
if (id === this.activeConnection?.id) {
await this.memento.update(this.key, undefined)
this.#savedConnection = undefined
this.#activeConnection = undefined
}
}
}

/**
Expand Down
16 changes: 0 additions & 16 deletions packages/core/src/codecatalyst/auth.ts
Original file line number Diff line number Diff line change
Expand Up @@ -345,22 +345,6 @@ export class CodeCatalystAuthenticationProvider {
}
}
}
/**
* Try to use an existing connection for CodeCatalyst Login
*/
public async tryUseConnection(conn: SsoConnection) {
const connId = conn.id
if (isValidCodeCatalystConnection(conn)) {
getLogger().info(`auth: re-use connection from existing connection id ${connId}`)
await this.secondaryAuth.useNewConnection(conn)
await this.isConnectionOnboarded(conn, true)
} else {
getLogger().info(`auth: re-use(new scope) to connection from existing connection id ${connId}`)
const newConn = await this.secondaryAuth.addScopes(conn, defaultScopes)
await this.secondaryAuth.useNewConnection(newConn)
await this.isConnectionOnboarded(newConn, true)
}
}

public async reauthenticate(conn: SsoConnection) {
try {
Expand Down
40 changes: 0 additions & 40 deletions packages/core/src/codewhisperer/util/authUtil.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@ import {
scopesFeatureDev,
scopesGumby,
isIdcSsoConnection,
AwsConnection,
hasExactScopes,
} from '../../auth/connection'
import { getLogger } from '../../shared/logger'
Expand Down Expand Up @@ -151,26 +150,6 @@ export class AuthUtil {
)
}

/* Callback used by Amazon Q to delete connection status & scope when this deletion is made by AWS Toolkit
** 1. NO event should be emitted from this deletion
** 2. Should update the context key to update UX
*/
public async onDeleteConnection(id: string) {
await this.secondaryAuth.onDeleteConnection(id)
await this.setVscodeContextProps()
await vscode.commands.executeCommand('aws.amazonq.refreshStatusBar')
}

/* Callback used by Amazon Q to delete connection status & scope when this deletion is made by AWS Toolkit
** 1. NO event should be emitted from this deletion
** 2. Should update the context key to update UX
*/
public async onUpdateConnection(connection: AwsConnection) {
await this.auth.onConnectionUpdate(connection)
await this.setVscodeContextProps()
await vscode.commands.executeCommand('aws.amazonq.refreshStatusBar')
}

public reformatStartUrl(startUrl: string | undefined) {
return !startUrl ? undefined : startUrl.replace(/[\/#]+$/g, '')
}
Expand Down Expand Up @@ -433,25 +412,6 @@ export class AuthUtil {

return state
}

/**
* From the given connections, returns a connection that has some connection to Q.
*
* HACK: There is an edge case where we want to connect to the connection that only has
* the old CW scopes, but not all Q scopes. So this function at the bare minimum returns
* a connection if it has some CW scopes.
*/
findMinimalQConnection(connections: AwsConnection[]): AwsConnection | undefined {
const hasQScopes = (c: AwsConnection) => codeWhispererCoreScopes.every(s => c.scopes?.includes(s))
const score = (c: AwsConnection) => Number(hasQScopes(c)) * 10 + Number(c.state === 'valid')
connections.sort(function (a, b) {
return score(b) - score(a)
})
if (hasQScopes(connections[0])) {
return connections[0]
}
return undefined
}
}

/**
Expand Down
Loading

0 comments on commit ef74cea

Please sign in to comment.