Skip to content

Commit cf855e6

Browse files
rlirli
authored
Add support for OAuth login flow using authorization_grant with PKCE (#4252)
1 parent 937546d commit cf855e6

File tree

29 files changed

+2342
-155
lines changed

29 files changed

+2342
-155
lines changed

‎buildSrc/src/main/kotlin/temp-toolkit-intellij-root-conventions.gradle.kts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,11 +95,17 @@ dependencies {
9595
}
9696

9797
configurations {
98+
all {
99+
// IDE provides netty
100+
exclude("io.netty")
101+
}
102+
98103
// Make sure we exclude stuff we either A) ships with IDE, B) we don't use to cut down on size
99104
runtimeClasspath {
100105
exclude(group = "org.slf4j")
101106
exclude(group = "org.jetbrains.kotlin")
102107
exclude(group = "org.jetbrains.kotlinx")
108+
103109
}
104110
}
105111

‎buildSrc/src/main/kotlin/toolkit-generate-sdks.gradle.kts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ sourceSets {
1919
}
2020

2121
java {
22-
setSrcDirs(listOf(sdkGenerator.srcDir()))
22+
setSrcDirs(listOf(sdkGenerator.srcDir(), "src"))
2323
}
2424
}
2525

‎buildSrc/src/main/kotlin/toolkit-intellij-subplugin.gradle.kts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,9 @@ configurations {
5353
}
5454

5555
all {
56+
// IDE provides netty
57+
exclude("io.netty")
58+
5659
if (name.startsWith("detekt")) {
5760
return@all
5861
}

‎plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/AmazonQTestBase.kt

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ import software.aws.toolkits.core.credentials.ToolkitBearerTokenProvider
2020
import software.aws.toolkits.core.utils.test.aString
2121
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
2222
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
23-
import software.aws.toolkits.jetbrains.core.credentials.sso.AccessToken
23+
import software.aws.toolkits.jetbrains.core.credentials.sso.DeviceAuthorizationGrantToken
2424
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider
2525
import software.aws.toolkits.jetbrains.services.amazonq.clients.AmazonQStreamingClient
2626
import software.aws.toolkits.jetbrains.utils.rules.CodeInsightTestFixtureRule
@@ -47,7 +47,7 @@ open class AmazonQTestBase(
4747
project = projectRule.project
4848
toolkitConnectionManager = spy(ToolkitConnectionManager.getInstance(project))
4949

50-
val accessToken = AccessToken(aString(), aString(), aString(), aString(), Instant.MAX, Instant.now())
50+
val accessToken = DeviceAuthorizationGrantToken(aString(), aString(), aString(), aString(), Instant.MAX, Instant.now())
5151

5252
val provider = mock<BearerTokenProvider> {
5353
doReturn(accessToken).whenever(it).refresh()

‎plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevTestBase.kt

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ import software.aws.toolkits.core.credentials.ToolkitBearerTokenProvider
2929
import software.aws.toolkits.core.utils.test.aString
3030
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
3131
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
32-
import software.aws.toolkits.jetbrains.core.credentials.sso.AccessToken
32+
import software.aws.toolkits.jetbrains.core.credentials.sso.DeviceAuthorizationGrantToken
3333
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider
3434
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.clients.FeatureDevClient
3535
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.clients.GenerateTaskAssistPlanResult
@@ -91,7 +91,7 @@ open class FeatureDevTestBase(
9191
open fun setup() {
9292
project = projectRule.project
9393
toolkitConnectionManager = spy(ToolkitConnectionManager.getInstance(project))
94-
val accessToken = AccessToken(aString(), aString(), aString(), aString(), Instant.MAX, Instant.now())
94+
val accessToken = DeviceAuthorizationGrantToken(aString(), aString(), aString(), aString(), Instant.MAX, Instant.now())
9595
val provider = mock<BearerTokenProvider> {
9696
doReturn(accessToken).whenever(it).refresh()
9797
}

‎plugins/amazonq/codetransform/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codemodernizer/CodeWhispererCodeModernizerTestBase.kt

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ import software.aws.toolkits.core.credentials.ToolkitBearerTokenProvider
4545
import software.aws.toolkits.core.utils.test.aString
4646
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
4747
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
48-
import software.aws.toolkits.jetbrains.core.credentials.sso.AccessToken
48+
import software.aws.toolkits.jetbrains.core.credentials.sso.DeviceAuthorizationGrantToken
4949
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider
5050
import software.aws.toolkits.jetbrains.services.codemodernizer.client.GumbyClient
5151
import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeModernizerArtifact
@@ -218,7 +218,7 @@ open class CodeWhispererCodeModernizerTestBase(
218218
project = projectRule.project
219219
toolkitConnectionManager = spy(ToolkitConnectionManager.getInstance(project))
220220

221-
val accessToken = AccessToken(aString(), aString(), aString(), aString(), Instant.MAX, Instant.now())
221+
val accessToken = DeviceAuthorizationGrantToken(aString(), aString(), aString(), aString(), Instant.MAX, Instant.now())
222222
val provider = mock<BearerTokenProvider> {
223223
doReturn(accessToken).whenever(it).refresh()
224224
}

‎plugins/core/jetbrains-community/resources/META-INF/aws.toolkit.core.xml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
</extensionPoints>
1212

1313
<extensions defaultExtensionNs="com.intellij">
14+
<httpRequestHandler implementation="software.aws.toolkits.jetbrains.core.credentials.sso.pkce.ToolkitOAuthCallbackHandler"/>
15+
1416
<applicationService serviceInterface="software.aws.toolkits.jetbrains.settings.AwsSettings"
1517
serviceImplementation="software.aws.toolkits.jetbrains.settings.DefaultAwsSettings"
1618
testServiceImplementation="software.aws.toolkits.jetbrains.settings.MockAwsSettings"/>
@@ -52,6 +54,8 @@
5254
<projectService serviceInterface="software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager"
5355
serviceImplementation="software.aws.toolkits.jetbrains.core.credentials.DefaultToolkitConnectionManager"/>
5456

57+
<registryKey key="aws.dev.pkceAuth" description="True if new authorization requests should be using the PKCE grant flow"
58+
defaultValue="false" restartRequired="false"/>
5559
<registryKey key="aws.telemetry.endpoint" description="Endpoint to use for publishing AWS client-side telemetry"
5660
defaultValue="https://client-telemetry.us-east-1.amazonaws.com" restartRequired="true"/>
5761
<registryKey key="aws.telemetry.identityPool" description="Cognito identity pool to use for publishing AWS client-side telemetry"

‎plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/AccessToken.kt

Lines changed: 63 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,11 @@
33

44
package software.aws.toolkits.jetbrains.core.credentials.sso
55

6+
import com.fasterxml.jackson.annotation.JsonIgnore
67
import com.fasterxml.jackson.annotation.JsonInclude
8+
import com.fasterxml.jackson.annotation.JsonSubTypes
9+
import com.fasterxml.jackson.annotation.JsonTypeInfo
10+
import com.intellij.collaboration.auth.credentials.Credentials
711
import software.amazon.awssdk.auth.token.credentials.SdkToken
812
import software.amazon.awssdk.services.sso.SsoClient
913
import software.amazon.awssdk.services.ssooidc.SsoOidcClient
@@ -15,29 +19,76 @@ import java.util.Optional
1519
/**
1620
* Access token returned from [SsoOidcClient.createToken] used to retrieve AWS Credentials from [SsoClient.getRoleCredentials].
1721
*/
18-
data class AccessToken(
19-
val startUrl: String,
20-
val region: String,
22+
@JsonTypeInfo(use = JsonTypeInfo.Id.DEDUCTION)
23+
@JsonSubTypes(value = [JsonSubTypes.Type(DeviceAuthorizationGrantToken::class), JsonSubTypes.Type(PKCEAuthorizationGrantToken::class) ])
24+
sealed interface AccessToken : SdkToken, Credentials {
25+
val region: String
26+
2127
@SensitiveField
22-
val accessToken: String,
28+
override val accessToken: String
29+
2330
@SensitiveField
24-
@JsonInclude(JsonInclude.Include.NON_NULL)
25-
val refreshToken: String? = null,
26-
val expiresAt: Instant,
27-
val createdAt: Instant = Instant.EPOCH
28-
) : SdkToken {
31+
@get:JsonInclude(JsonInclude.Include.NON_NULL)
32+
val refreshToken: String?
33+
34+
val expiresAt: Instant
35+
val createdAt: Instant
36+
2937
override fun token() = accessToken
3038

3139
override fun expirationTime() = Optional.of(expiresAt)
3240

41+
@get:JsonIgnore
42+
val ssoUrl: String
43+
}
44+
45+
data class DeviceAuthorizationGrantToken(
46+
val startUrl: String,
47+
override val region: String,
48+
override val accessToken: String,
49+
override val refreshToken: String? = null,
50+
override val expiresAt: Instant,
51+
override val createdAt: Instant = Instant.EPOCH
52+
) : AccessToken {
53+
override val ssoUrl: String
54+
get() = startUrl
55+
3356
override fun toString() = redactedString(this)
3457
}
3558

59+
data class PKCEAuthorizationGrantToken(
60+
val issuerUrl: String,
61+
override val region: String,
62+
override val accessToken: String,
63+
override val refreshToken: String,
64+
override val expiresAt: Instant,
65+
override val createdAt: Instant
66+
) : AccessToken {
67+
override val ssoUrl: String
68+
get() = issuerUrl
69+
70+
override fun toString() = redactedString(this)
71+
}
72+
73+
// we really don't need to differentitate since they refresh the same way, but to save some mental cycles,
74+
// treat them as independent so we don't need to worry about intermingling the token/registration combos
75+
@JsonTypeInfo(use = JsonTypeInfo.Id.DEDUCTION)
76+
@JsonSubTypes(value = [JsonSubTypes.Type(DeviceGrantAccessTokenCacheKey::class), JsonSubTypes.Type(PKCEAccessTokenCacheKey::class) ])
77+
sealed interface AccessTokenCacheKey {
78+
val scopes: List<String>
79+
}
80+
3681
// diverging from SDK/CLI impl here since they do: sha1sum(sessionName ?: startUrl)
3782
// which isn't good enough for us
3883
// only used in scoped case
39-
data class AccessTokenCacheKey(
84+
data class DeviceGrantAccessTokenCacheKey(
4085
val connectionId: String,
4186
val startUrl: String,
42-
val scopes: List<String>
43-
)
87+
override val scopes: List<String>
88+
) : AccessTokenCacheKey
89+
90+
data class PKCEAccessTokenCacheKey(
91+
val issuerUrl: String,
92+
val region: String,
93+
override val scopes: List<String>
94+
) : AccessTokenCacheKey

‎plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/Authorization.kt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import java.time.Instant
1111
/**
1212
* Returned by [SsoOidcClient.startDeviceAuthorization] that contains the required data to construct the user visible SSO login flow.
1313
*/
14+
@Deprecated("Device authorization grant flow is deprecated")
1415
data class Authorization(
1516
@SensitiveField
1617
val deviceCode: String,

‎plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/ClientRegistration.kt

Lines changed: 50 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
package software.aws.toolkits.jetbrains.core.credentials.sso
55

66
import com.fasterxml.jackson.annotation.JsonInclude
7+
import com.fasterxml.jackson.annotation.JsonSubTypes
8+
import com.fasterxml.jackson.annotation.JsonTypeInfo
79
import software.amazon.awssdk.services.ssooidc.SsoOidcClient
810
import software.aws.toolkits.core.utils.SensitiveField
911
import software.aws.toolkits.core.utils.redactedString
@@ -14,22 +16,60 @@ import java.time.Instant
1416
*
1517
* It should be persisted for reuse through many authentication requests.
1618
*/
17-
data class ClientRegistration(
18-
@SensitiveField
19-
val clientId: String,
19+
@JsonTypeInfo(use = JsonTypeInfo.Id.DEDUCTION, defaultImpl = DeviceAuthorizationClientRegistration::class)
20+
@JsonSubTypes(value = [JsonSubTypes.Type(DeviceAuthorizationClientRegistration::class), JsonSubTypes.Type(PKCEClientRegistration::class) ])
21+
sealed interface ClientRegistration {
22+
val clientId: String
23+
2024
@SensitiveField
21-
val clientSecret: String,
22-
val expiresAt: Instant,
23-
@JsonInclude(JsonInclude.Include.NON_EMPTY)
24-
val scopes: List<String> = emptyList()
25-
) {
25+
val clientSecret: String
26+
27+
val expiresAt: Instant
28+
29+
@get:JsonInclude(JsonInclude.Include.NON_EMPTY)
30+
val scopes: List<String>
31+
}
32+
33+
data class DeviceAuthorizationClientRegistration(
34+
override val clientId: String,
35+
override val clientSecret: String,
36+
override val expiresAt: Instant,
37+
override val scopes: List<String> = emptyList(),
38+
) : ClientRegistration {
39+
override fun toString(): String = redactedString(this)
40+
}
41+
42+
data class PKCEClientRegistration(
43+
override val clientId: String,
44+
override val clientSecret: String,
45+
override val expiresAt: Instant,
46+
override val scopes: List<String>,
47+
// fields below are implied from the key, but trying reverse the key is annoying
48+
val issuerUrl: String,
49+
val region: String,
50+
val clientType: String,
51+
val grantTypes: List<String>,
52+
val redirectUris: List<String>,
53+
) : ClientRegistration {
2654
override fun toString(): String = redactedString(this)
2755
}
2856

57+
sealed interface ClientRegistrationCacheKey
58+
2959
// only applicable in scoped registration path
3060
// based on internal development branch @da780a4,L2574-2586
31-
data class ClientRegistrationCacheKey(
61+
data class DeviceAuthorizationClientRegistrationCacheKey(
3262
val startUrl: String,
3363
val scopes: List<String>,
3464
val region: String,
35-
)
65+
) : ClientRegistrationCacheKey
66+
67+
data class PKCEClientRegistrationCacheKey(
68+
val issuerUrl: String,
69+
val region: String,
70+
val scopes: List<String>,
71+
// assume clientType, grantTypes, redirectUris are static, but throw them in just in case
72+
val clientType: String,
73+
val grantTypes: List<String>,
74+
val redirectUris: List<String>
75+
) : ClientRegistrationCacheKey

0 commit comments

Comments
 (0)