Skip to content

Commit bbefe6b

Browse files
nkomonen-amazonnkomonen-amazon
authored
feat: Enable Authorization Code + PKCE flow (#4974)
* remove PKCE feature flag - Now we run Device Code if we are in an ssh workspace - Otherwise we run the Authorization code flow Signed-off-by: Nikolas Komonen <[email protected]> * use new oidc client, remove our temp one Since the new OIDC client with Authorization Code flow was not released we had to use a pre-build version. Now that the new client is released we can actually use it. Solution: This commit removes the old temp OIDC client and simply uses our existing one with minor modifications to get things working. Also we were required to update some transitive dependencies which is why there are 'smithy' related dependency changes Signed-off-by: Nikolas Komonen <[email protected]> * get tests working + changelog items Signed-off-by: Nikolas Komonen <[email protected]> --------- Signed-off-by: Nikolas Komonen <[email protected]>
1 parent fe588c5 commit bbefe6b

File tree

11 files changed

+1805
-2583
lines changed

11 files changed

+1805
-2583
lines changed

‎package-lock.json‎

Lines changed: 1740 additions & 2474 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/amazonq/.changes/next-release/Feature-1fd5f451-8d50-4e44-b17e-d0cb77ca0c5e.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type": "Feature",
3+
"description": "New SSO Authorization Code flow for faster logins"
4+
}

‎packages/core/package.json‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4042,7 +4042,7 @@
40424042
"@aws-sdk/client-cognito-identity": "3.490.0",
40434043
"@aws-sdk/client-lambda": "3.385.0",
40444044
"@aws-sdk/client-sso": "^3.342.0",
4045-
"@aws-sdk/client-sso-oidc": "^3.181.0",
4045+
"@aws-sdk/client-sso-oidc": "^3.574.0",
40464046
"@aws-sdk/credential-provider-ini": "3.46.0",
40474047
"@aws-sdk/credential-provider-process": "3.37.0",
40484048
"@aws-sdk/credential-provider-sso": "^3.345.0",
@@ -4052,7 +4052,11 @@
40524052
"@aws/mynah-ui": "^4.8.0",
40534053
"@gerhobbelt/gitignore-parser": "^0.2.0-9",
40544054
"@iarna/toml": "^2.2.5",
4055+
"@smithy/middleware-retry": "^2.3.1",
4056+
"@smithy/protocol-http": "^3.3.0",
4057+
"@smithy/service-error-classification": "^2.1.5",
40554058
"@smithy/shared-ini-file-loader": "^2.2.8",
4059+
"@smithy/util-retry": "^2.2.0",
40564060
"@vscode/debugprotocol": "^1.57.0",
40574061
"adm-zip": "^0.5.10",
40584062
"amazon-s3-uri": "^0.1.1",

‎packages/core/scripts/build/generateServiceClient.ts‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -245,10 +245,6 @@ void (async () => {
245245
serviceJsonPath: 'src/amazonqFeatureDev/client/codewhispererruntime-2022-11-11.json',
246246
serviceName: 'FeatureDevProxyClient',
247247
},
248-
{
249-
serviceJsonPath: 'src/auth/sso/service-2.json',
250-
serviceName: 'OidcClientPKCE',
251-
},
252248
]
253249
await generateServiceClients(serviceClientDefinitions)
254250
})()

‎packages/core/src/auth/sso/clients.ts‎

Lines changed: 22 additions & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -31,19 +31,15 @@ import { SsoAccessTokenProvider } from './ssoAccessTokenProvider'
3131
import { isClientFault } from '../../shared/errors'
3232
import { DevSettings } from '../../shared/settings'
3333
import { SdkError } from '@aws-sdk/types'
34-
import { HttpRequest, HttpResponse } from '@aws-sdk/protocol-http'
35-
import { StandardRetryStrategy, defaultRetryDecider } from '@aws-sdk/middleware-retry'
36-
import OidcClientPKCE from './oidcclientpkce'
34+
import { HttpRequest, HttpResponse } from '@smithy/protocol-http'
35+
import { StandardRetryStrategy, defaultRetryDecider } from '@smithy/middleware-retry'
36+
import { AuthenticationFlow } from './model'
3737
import { toSnakeCase } from '../../shared/utilities/textUtilities'
38-
import { Credentials, Service } from 'aws-sdk'
39-
import apiConfig = require('./service-2.json')
40-
import { ServiceOptions } from '../../shared/awsClientBuilder'
41-
import { ClientRegistration } from './model'
4238

4339
export class OidcClient {
4440
public constructor(private readonly client: SSOOIDC, private readonly clock: { Date: typeof Date }) {}
4541

46-
public async registerClient(request: RegisterClientRequest) {
42+
public async registerClient(request: RegisterClientRequest, flow?: AuthenticationFlow) {
4743
const response = await this.client.registerClient(request)
4844
assertHasProps(response, 'clientId', 'clientSecret', 'clientSecretExpiresAt')
4945

@@ -52,6 +48,7 @@ export class OidcClient {
5248
clientId: response.clientId,
5349
clientSecret: response.clientSecret,
5450
expiresAt: new this.clock.Date(response.clientSecretExpiresAt * 1000),
51+
...(flow ? { flow } : {}),
5552
}
5653
}
5754

@@ -66,6 +63,23 @@ export class OidcClient {
6663
}
6764
}
6865

66+
public async authorize(request: {
67+
responseType: string
68+
clientId: string
69+
redirectUri: string
70+
scopes: string[]
71+
state: string
72+
codeChallenge: string
73+
codeChallengeMethod: string
74+
}) {
75+
// aws sdk doesn't convert to url params until right before you make the request, so we have to do
76+
// it manually ahead of time
77+
const params = toSnakeCase(request)
78+
const searchParams = new URLSearchParams(params).toString()
79+
const region = await this.client.config.region()
80+
return `https://oidc.${region}.amazonaws.com/authorize?${searchParams}`
81+
}
82+
6983
public async createToken(request: CreateTokenRequest) {
7084
const response = await this.client.createToken(request as CreateTokenRequest)
7185
assertHasProps(response, 'accessToken', 'expiresIn')
@@ -100,79 +114,6 @@ export class OidcClient {
100114
}
101115
}
102116

103-
export class OidcClientV2 {
104-
public constructor(private readonly region: string, private readonly clock: { Date: typeof Date }) {}
105-
106-
/**
107-
* TODO remove this when the real client gets created.
108-
*
109-
* Creating a new client is required because the old sdk seems to drop unknown parameters from requests
110-
*/
111-
private async createNewClient() {
112-
return (await globals.sdkClientBuilder.createAwsService(
113-
Service,
114-
{
115-
apiConfig,
116-
region: this.region,
117-
credentials: new Credentials({ accessKeyId: 'xxx', secretAccessKey: 'xxx' }),
118-
} as ServiceOptions,
119-
undefined
120-
)) as OidcClientPKCE
121-
}
122-
123-
public async registerClient(request: OidcClientPKCE.RegisterClientRequest): Promise<ClientRegistration> {
124-
const client = await this.createNewClient()
125-
const response = await client.makeUnauthenticatedRequest('registerClient', request).promise()
126-
assertHasProps(response, 'clientId', 'clientSecret', 'clientSecretExpiresAt')
127-
128-
return {
129-
scopes: request.scopes,
130-
clientId: response.clientId,
131-
clientSecret: response.clientSecret,
132-
expiresAt: new this.clock.Date(response.clientSecretExpiresAt * 1000),
133-
flow: 'auth code',
134-
}
135-
}
136-
137-
public async authorize(request: OidcClientPKCE.AuthorizeRequest) {
138-
// aws sdk doesn't convert to url params until right before you make the request, so we have to do
139-
// it manually ahead of time
140-
const params = toSnakeCase(request)
141-
const searchParams = new URLSearchParams(params).toString()
142-
return `https://oidc.${this.region}.amazonaws.com/authorize?${searchParams}`
143-
}
144-
145-
public async startDeviceAuthorization(request: StartDeviceAuthorizationRequest): Promise<{
146-
expiresAt: Date
147-
interval: number | undefined
148-
deviceCode: string
149-
userCode: string
150-
verificationUri: string
151-
}> {
152-
throw new Error('OidcClientV2 does not support device authorization')
153-
}
154-
155-
public async createToken(request: OidcClientPKCE.CreateTokenRequest): Promise<{
156-
accessToken: string
157-
expiresAt: Date
158-
tokenType?: string | undefined
159-
refreshToken?: string | undefined
160-
}> {
161-
const client = await this.createNewClient()
162-
const response = await client.makeUnauthenticatedRequest('createToken', request).promise()
163-
assertHasProps(response, 'accessToken', 'expiresIn')
164-
165-
return {
166-
...selectFrom(response, 'accessToken', 'refreshToken', 'tokenType'),
167-
expiresAt: new this.clock.Date(response.expiresIn * 1000 + this.clock.Date.now()),
168-
}
169-
}
170-
171-
public static create(region: string) {
172-
return new this(region, globals.clock)
173-
}
174-
}
175-
176117
type OmittedProps = 'accessToken' | 'nextToken'
177118
type ExtractOverload<T, U> = T extends {
178119
(...args: infer P1): infer R1

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

Lines changed: 27 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,12 @@
55

66
import * as vscode from 'vscode'
77
import globals from '../../shared/extensionGlobals'
8-
import { AuthorizationPendingException, SSOOIDCServiceException, SlowDownException } from '@aws-sdk/client-sso-oidc'
8+
import {
9+
AuthorizationPendingException,
10+
CreateTokenRequest,
11+
SSOOIDCServiceException,
12+
SlowDownException,
13+
} from '@aws-sdk/client-sso-oidc'
914
import {
1015
SsoToken,
1116
ClientRegistration,
@@ -18,7 +23,7 @@ import {
1823
} from './model'
1924
import { getCache } from './cache'
2025
import { hasProps, hasStringProps, RequiredProps, selectFrom } from '../../shared/utilities/tsUtils'
21-
import { OidcClient, OidcClientV2 } from './clients'
26+
import { OidcClient } from './clients'
2227
import { loadOr } from '../../shared/utilities/cacheUtils'
2328
import {
2429
ToolkitError,
@@ -33,13 +38,12 @@ import { AwsLoginWithBrowser, AwsRefreshCredentials, Metric, telemetry } from '.
3338
import { indent } from '../../shared/utilities/textUtilities'
3439
import { AuthSSOServer } from './server'
3540
import { CancellationError, sleep } from '../../shared/utilities/timeoutUtils'
36-
import OidcClientPKCE from './oidcclientpkce'
3741
import { getIdeProperties, isCloud9 } from '../../shared/extensionUtilities'
3842
import { randomBytes, createHash } from 'crypto'
3943
import { UriHandler } from '../../shared/vscode/uriHandler'
40-
import { DevSettings } from '../../shared/settings'
4144
import { localize } from '../../shared/utilities/vsCodeUtils'
4245
import { randomUUID } from '../../common/crypto'
46+
import { isRemoteWorkspace } from '../../shared/vscode/env'
4347

4448
export const authenticationPath = 'sso/authenticated'
4549

@@ -98,7 +102,7 @@ export abstract class SsoAccessTokenProvider {
98102
public constructor(
99103
protected readonly profile: Pick<SsoProfile, 'startUrl' | 'region' | 'scopes' | 'identifier'>,
100104
protected readonly cache = getCache(),
101-
protected readonly oidc: OidcClient | OidcClientV2 = OidcClient.create(profile.region)
105+
protected readonly oidc: OidcClient = OidcClient.create(profile.region)
102106
) {}
103107

104108
public async invalidate(): Promise<void> {
@@ -249,12 +253,13 @@ export abstract class SsoAccessTokenProvider {
249253
public static create(
250254
profile: Pick<SsoProfile, 'startUrl' | 'region' | 'scopes' | 'identifier'>,
251255
cache = getCache(),
252-
oidc: OidcClient = OidcClient.create(profile.region)
256+
oidc: OidcClient = OidcClient.create(profile.region),
257+
useDeviceFlow: () => boolean = isRemoteWorkspace
253258
) {
254-
if (!DevSettings.instance.get('pkceAuth', false)) {
259+
if (useDeviceFlow()) {
255260
return new DeviceFlowAuthorization(profile, cache, oidc)
256261
}
257-
return new AuthFlowAuthorization(profile, cache, OidcClientV2.create(profile.region))
262+
return new AuthFlowAuthorization(profile, cache, oidc)
258263
}
259264
}
260265

@@ -403,22 +408,25 @@ class AuthFlowAuthorization extends SsoAccessTokenProvider {
403408
constructor(
404409
profile: Pick<SsoProfile, 'startUrl' | 'region' | 'scopes' | 'identifier'>,
405410
cache = getCache(),
406-
protected override readonly oidc: OidcClientV2
411+
oidc: OidcClient
407412
) {
408413
super(profile, cache, oidc)
409414
}
410415

411416
override async registerClient(): Promise<ClientRegistration> {
412417
const companyName = getIdeProperties().company
413-
return this.oidc.registerClient({
414-
// All AWS extensions (Q, Toolkit) for a given IDE use the same client name.
415-
clientName: isCloud9() ? `${companyName} Cloud9` : `${companyName} IDE Extensions for VSCode`,
416-
clientType: clientRegistrationType,
417-
scopes: this.profile.scopes,
418-
grantTypes: [authorizationGrantType, refreshGrantType],
419-
redirectUris: ['http://127.0.0.1/oauth/callback'],
420-
issuerUrl: this.profile.startUrl,
421-
})
418+
return this.oidc.registerClient(
419+
{
420+
// All AWS extensions (Q, Toolkit) for a given IDE use the same client name.
421+
clientName: isCloud9() ? `${companyName} Cloud9` : `${companyName} IDE Extensions for VSCode`,
422+
clientType: clientRegistrationType,
423+
scopes: this.profile.scopes,
424+
grantTypes: [authorizationGrantType, refreshGrantType],
425+
redirectUris: ['http://127.0.0.1/oauth/callback'],
426+
issuerUrl: this.profile.startUrl,
427+
},
428+
'auth code'
429+
)
422430
}
423431

424432
override async authorize(
@@ -455,7 +463,7 @@ class AuthFlowAuthorization extends SsoAccessTokenProvider {
455463
throw authorizationCode.err()
456464
}
457465

458-
const tokenRequest: OidcClientPKCE.CreateTokenRequest = {
466+
const tokenRequest: CreateTokenRequest = {
459467
clientId: registration.clientId,
460468
clientSecret: registration.clientSecret,
461469
grantType: authorizationGrantType,

‎packages/core/src/shared/clients/codewhispererChatClient.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55
import { CodeWhispererStreaming } from '@amzn/codewhisperer-streaming'
6-
import { ConfiguredRetryStrategy } from '@aws-sdk/util-retry'
6+
import { ConfiguredRetryStrategy } from '@smithy/util-retry'
77
import { getCodewhispererConfig } from '../../codewhisperer/client/codewhisperer'
88
import { AuthUtil } from '../../codewhisperer/util/authUtil'
99

‎packages/core/src/shared/errors.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
import * as vscode from 'vscode'
77
import { AWSError } from 'aws-sdk'
88
import { ServiceException } from '@aws-sdk/smithy-client'
9-
import { isThrottlingError, isTransientError } from '@aws-sdk/service-error-classification'
9+
import { isThrottlingError, isTransientError } from '@smithy/service-error-classification'
1010
import { Result } from './telemetry/telemetry'
1111
import { CancellationError } from './utilities/timeoutUtils'
1212
import { isNonNullable } from './utilities/tsUtils'

‎packages/core/src/shared/settings.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -728,7 +728,6 @@ const devSettings = {
728728
codewhispererService: Record(String, String),
729729
ssoCacheDirectory: String,
730730
enableIamPolicyChecksFeature: Boolean,
731-
pkceAuth: Boolean,
732731
autofillStartUrl: String,
733732
}
734733
type ResolvedDevSettings = FromDescriptor<typeof devSettings>

‎packages/core/src/test/credentials/sso/ssoAccessTokenProvider.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ describe('SsoAccessTokenProvider', function () {
8787
oidcClient = stub(OidcClient)
8888
tempDir = await makeTemporaryTokenCacheFolder()
8989
cache = getCache(tempDir)
90-
sut = SsoAccessTokenProvider.create({ region, startUrl }, cache, oidcClient)
90+
sut = SsoAccessTokenProvider.create({ region, startUrl }, cache, oidcClient, () => true)
9191
})
9292

9393
afterEach(async function () {

0 commit comments

Comments
 (0)