Follow-up feedback for arn-anywhere (#3039)

Author: rli

jetbrains-core/it/software/aws/toolkits/jetbrains/services/federation/AwsConsoleUrlFactoryIntegrationTest.kt

@@ -7,11 +7,14 @@ import com.intellij.testFramework.TestApplicationManager
 import com.intellij.util.io.HttpRequests
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.Assume.assumeFalse
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.Parameterized
 import software.amazon.awssdk.auth.credentials.ProfileCredentialsProvider
+import software.aws.toolkits.core.ConnectionSettings
 import software.aws.toolkits.core.region.AwsRegion
+import software.aws.toolkits.jetbrains.core.credentials.MockCredentialManagerRule
 import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider
 
 @RunWith(Parameterized::class)
@@ -28,6 +31,10 @@ class AwsConsoleUrlFactoryIntegrationTest(@Suppress("UNUSED_PARAMETER") regionId
         }
     }
 
+    @Rule
+    @JvmField
+    val credRule = MockCredentialManagerRule()
+
     /**
      * There is currently no good way to test this in our integration CI fleet, so this test suite only runs locally
      */
@@ -42,8 +49,9 @@ class AwsConsoleUrlFactoryIntegrationTest(@Suppress("UNUSED_PARAMETER") regionId
         }
         assumeFalse("Skipping console sign-in test for $region since a credentials profile was not available", profileName.isNullOrBlank())
 
-        val credProvider = ProfileCredentialsProvider.create(profileName)
-        val signinUrl = AwsConsoleUrlFactory().getSigninUrl(credentials = credProvider.resolveCredentials(), destination = "", region = region)
+        val credProvider = credRule.createCredentialProvider(profileName, ProfileCredentialsProvider.create(profileName).resolveCredentials())
+        val credSettings = ConnectionSettings(credProvider, region)
+        val signinUrl = AwsConsoleUrlFactory.getSigninUrl(credSettings, destination = "")
         val responseCode = HttpRequests.request(signinUrl)
             // don't throw because it'll print the signin token as part of the exception
             .throwStatusCodeException(false)

jetbrains-core/src/software/aws/toolkits/jetbrains/services/federation/AwsConsoleUrlFactory.kt

@@ -11,18 +11,16 @@ import org.apache.http.client.entity.UrlEncodedFormEntity
 import org.apache.http.client.methods.HttpPost
 import org.apache.http.impl.client.HttpClientBuilder
 import org.apache.http.message.BasicNameValuePair
-import software.amazon.awssdk.auth.credentials.AwsCredentials
-import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider
 import software.amazon.awssdk.auth.credentials.AwsSessionCredentials
-import software.amazon.awssdk.regions.Region
 import software.amazon.awssdk.services.sts.StsClient
+import software.aws.toolkits.core.ConnectionSettings
 import software.aws.toolkits.core.region.AwsRegion
 import software.aws.toolkits.jetbrains.core.AwsClientManager
 import java.time.Duration
 
-class AwsConsoleUrlFactory(
-    private val httpClientBuilder: HttpClientBuilder = HttpClientBuilder.create()
-) {
+object AwsConsoleUrlFactory {
+    private val defaultHttpClientBuilder: HttpClientBuilder by lazy { HttpClientBuilder.create() }
+
     fun federationUrl(region: AwsRegion): String {
         // https://docs.aws.amazon.com/general/latest/gr/signin-service.html
         // https://docs.amazonaws.cn/en_us/aws/latest/userguide/endpoints-Beijing.html
@@ -63,10 +61,14 @@ class AwsConsoleUrlFactory(
         return "$consoleHome${fragment ?: "/"}"
     }
 
-    fun getSigninToken(credentials: AwsCredentials, region: AwsRegion): String {
-        val creds = if (credentials !is AwsSessionCredentials) {
-            val stsClient: StsClient = AwsClientManager.getInstance()
-                .createUnmanagedClient(AwsCredentialsProvider { credentials }, Region.of(region.id))
+    fun getSigninUrl(connectionSettings: ConnectionSettings, destination: String?, httpClientBuilder: HttpClientBuilder = defaultHttpClientBuilder): String {
+        return getSigninUrl(getSigninToken(connectionSettings, httpClientBuilder), destination, connectionSettings.region)
+    }
+
+    fun getSigninToken(connectionSettings: ConnectionSettings, httpClientBuilder: HttpClientBuilder = defaultHttpClientBuilder): String {
+        val resolvedCreds = connectionSettings.credentials.resolveCredentials()
+        val sessionCredentials = if (resolvedCreds !is AwsSessionCredentials) {
+            val stsClient = AwsClientManager.getInstance().getClient<StsClient>(connectionSettings)
 
             val tokenResponse = stsClient.use { client ->
                 client.getFederationToken {
@@ -82,14 +84,14 @@ class AwsConsoleUrlFactory(
 
             tokenResponse.credentials().let { AwsSessionCredentials.create(it.accessKeyId(), it.secretAccessKey(), it.sessionToken()) }
         } else {
-            credentials
+            resolvedCreds
         }
 
         val sessionJson = mapper.writeValueAsString(
             GetSigninTokenRequest(
-                sessionId = creds.accessKeyId(),
-                sessionKey = creds.secretAccessKey(),
-                sessionToken = creds.sessionToken()
+                sessionId = sessionCredentials.accessKeyId(),
+                sessionKey = sessionCredentials.secretAccessKey(),
+                sessionToken = sessionCredentials.sessionToken()
             )
         )
 
@@ -99,7 +101,7 @@ class AwsConsoleUrlFactory(
             "Session" to sessionJson
         ).map { BasicNameValuePair(it.key, it.value) }
 
-        val request = HttpPost(federationUrl(region))
+        val request = HttpPost(federationUrl(connectionSettings.region))
             .apply {
                 entity = UrlEncodedFormEntity(params)
             }
@@ -130,10 +132,6 @@ class AwsConsoleUrlFactory(
         return "${federationUrl(region)}?${UrlEncodedFormEntity(params).toUrlEncodedString()}"
     }
 
-    fun getSigninUrl(credentials: AwsCredentials, destination: String?, region: AwsRegion): String {
-        return getSigninUrl(getSigninToken(credentials, region), destination, region)
-    }
-
     private val mapper = jacksonObjectMapper().disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
 }
 

jetbrains-core/src/software/aws/toolkits/jetbrains/services/federation/psireferences/ArnReference.kt

@@ -34,10 +34,9 @@ class ArnReference(element: PsiElement, textRange: TextRange, private val arn: S
                 return
             }
 
-            val (credProvider, region) = connectionSettings
             ApplicationManager.getApplication().executeOnPooledThread {
                 try {
-                    BrowserUtil.browse(AwsConsoleUrlFactory().getSigninUrl(credProvider.resolveCredentials(), "/go/view/$arn", region))
+                    BrowserUtil.browse(AwsConsoleUrlFactory.getSigninUrl(connectionSettings, "/go/view/$arn"))
                 } catch (e: StsException) {
                     val message = message("general.open_in_aws_console.no_permission")
                     notifyError(content = message, project = project)

jetbrains-core/tst/software/aws/toolkits/jetbrains/services/federation/AwsConsoleUrlFactoryTest.kt

@@ -29,30 +29,34 @@ import software.amazon.awssdk.auth.credentials.AwsSessionCredentials
 import software.amazon.awssdk.services.sts.StsClient
 import software.amazon.awssdk.services.sts.model.GetFederationTokenRequest
 import software.amazon.awssdk.services.sts.model.GetFederationTokenResponse
+import software.aws.toolkits.core.ConnectionSettings
 import software.aws.toolkits.core.utils.tryOrNull
 import software.aws.toolkits.jetbrains.core.MockClientManagerRule
-import software.aws.toolkits.jetbrains.core.region.US_EAST_1
+import software.aws.toolkits.jetbrains.core.credentials.MockCredentialManagerRule
+import software.aws.toolkits.jetbrains.core.region.getDefaultRegion
 import java.net.InetAddress
 
 class AwsConsoleUrlFactoryTest {
     val applicationRule = ApplicationRule()
     val mockClientManager = MockClientManagerRule()
+    val mockCredentialManager = MockCredentialManagerRule()
 
     @Rule
     @JvmField
-    val ruleChain = RuleChain(applicationRule, mockClientManager)
+    val ruleChain = RuleChain(applicationRule, mockClientManager, mockCredentialManager)
 
     @Rule
     @JvmField
     val wireMock = WireMockRule(WireMockConfiguration.wireMockConfig().dynamicPort())
 
     private lateinit var stsMock: StsClient
+    private lateinit var httpBuilderMock: HttpClientBuilder
     private lateinit var httpClient: CloseableHttpClient
     private lateinit var sut: AwsConsoleUrlFactory
 
     @Before
     fun setUp() {
-        val httpBuilderMock = mock<HttpClientBuilder>()
+        httpBuilderMock = mock<HttpClientBuilder>()
         httpClient = HttpClientBuilder.create()
             .setRoutePlanner(HttpRoutePlanner { _, _, _ -> HttpRoute(HttpHost(InetAddress.getLoopbackAddress(), wireMock.port(), "http")) })
             .build()
@@ -73,7 +77,7 @@ class AwsConsoleUrlFactoryTest {
             )
         )
 
-        sut = AwsConsoleUrlFactory(httpBuilderMock)
+        sut = AwsConsoleUrlFactory
     }
 
     @After
@@ -84,14 +88,17 @@ class AwsConsoleUrlFactoryTest {
     @Test
     fun `getSigninToken uses temporary credentials as-is`() {
         val tempCreds = AwsSessionCredentials.create("accessKey", "secretKey", "sessionToken")
+        val connectionSettings = ConnectionSettings(mockCredentialManager.createCredentialProvider("tempCreds", tempCreds), getDefaultRegion())
 
         verifyNoMoreInteractions(stsMock)
-        assertThat(sut.getSigninToken(tempCreds, US_EAST_1)).isEqualTo("signinToken")
+        assertThat(sut.getSigninToken(connectionSettings, httpBuilderMock)).isEqualTo("signinToken")
     }
 
     @Test
     fun `getSigninToken requests federation token for long-term credentials`() {
         val longCreds = AwsBasicCredentials.create("basicKeyId", "basicSecretKey")
+        val connectionSettings = ConnectionSettings(mockCredentialManager.createCredentialProvider("longCreds", longCreds), getDefaultRegion())
+
         whenever(stsMock.getFederationToken(any<GetFederationTokenRequest>())).thenReturn(
             GetFederationTokenResponse.builder()
                 .credentials {
@@ -102,14 +109,15 @@ class AwsConsoleUrlFactoryTest {
                 .build()
         )
 
-        assertThat(sut.getSigninToken(longCreds, US_EAST_1)).isEqualTo("signinToken")
+        assertThat(sut.getSigninToken(connectionSettings, httpBuilderMock)).isEqualTo("signinToken")
     }
 
     @Test
     fun `build sign-in url`() {
         val tempCreds = AwsSessionCredentials.create("accessKey", "secretKey", "sessionToken")
+        val connectionSettings = ConnectionSettings(mockCredentialManager.createCredentialProvider("tempCreds", tempCreds), getDefaultRegion())
 
-        assertThat(sut.getSigninUrl(tempCreds, destination = "/something", region = US_EAST_1))
+        assertThat(sut.getSigninUrl(connectionSettings = connectionSettings, destination = "/something", httpClientBuilder = httpBuilderMock))
             .isEqualTo(
                 """
                 https://signin.aws.amazon.com/federation?Action=login&SigninToken=signinToken&Destination=https%3A%2F%2Fus-east-1.console.aws.amazon.com%2Fsomething