Initial authSchemePreference implementation

lauzadis · lauzadis · commit 6ec389ff92f9 · 2025-04-22T15:59:22.000-04:00
diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/ServiceClientConfigGenerator.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/ServiceClientConfigGenerator.kt
@@ -47,6 +47,7 @@ class ServiceClientConfigGenerator(
             add(RuntimeConfigProperty.HttpClient)
             add(RuntimeConfigProperty.HttpInterceptors)
             add(RuntimeConfigProperty.AuthSchemes)
+            add(RuntimeConfigProperty.AuthSchemePreference)
         }
         if (shape.hasIdempotentTokenMember(context.model)) {
             add(RuntimeConfigProperty.IdempotencyTokenProvider)
diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/auth/AuthSchemeProviderConfigIntegration.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/auth/AuthSchemeProviderConfigIntegration.kt
@@ -24,9 +24,9 @@ class AuthSchemeProviderConfigIntegration : KotlinIntegration {
                 documentation = "Configure the provider used to resolve the authentication scheme to use for a particular operation."
                 additionalImports = listOf(defaultProvider)
                 if (ctx.settings.api.enableEndpointAuthProvider) {
-                    propertyType = ConfigPropertyType.RequiredWithDefault("${defaultProvider.name}(endpointProvider)")
+                    propertyType = ConfigPropertyType.RequiredWithDefault("${defaultProvider.name}(endpointProvider, authSchemePreference)")
                 } else {
-                    propertyType = ConfigPropertyType.RequiredWithDefault("${defaultProvider.name}()")
+                    propertyType = ConfigPropertyType.RequiredWithDefault("${defaultProvider.name}(authSchemePreference = authSchemePreference)")
                 }
                 // needs to come after endpointProvider
                 order = 100
diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/auth/AuthSchemeProviderGenerator.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/auth/AuthSchemeProviderGenerator.kt
@@ -14,6 +14,7 @@ import software.amazon.smithy.kotlin.codegen.model.buildSymbol
 import software.amazon.smithy.kotlin.codegen.model.knowledge.AuthIndex
 import software.amazon.smithy.kotlin.codegen.rendering.endpoints.EndpointProviderGenerator
 import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolGenerator
+import software.amazon.smithy.kotlin.codegen.rendering.util.RuntimeConfigProperty
 import software.amazon.smithy.model.shapes.OperationShape
 
 /**
@@ -36,7 +37,7 @@ open class AuthSchemeProviderGenerator {
         }
     }
 
-    fun render(ctx: ProtocolGenerator.GenerationContext) {
+    open fun render(ctx: ProtocolGenerator.GenerationContext) {
         ctx.delegator.useSymbolWriter(getSymbol(ctx.settings)) { writer ->
             renderInterface(ctx, writer)
         }
@@ -65,11 +66,12 @@ open class AuthSchemeProviderGenerator {
 
     private fun renderDefaultImpl(ctx: ProtocolGenerator.GenerationContext, writer: KotlinWriter) {
         writer.withBlock(
-            "#L class #T(private val endpointProvider: #T? = null) : #T {",
+            "#L class #T(private val endpointProvider: #T? = null, private val authSchemePreference: #T? = null) : #T {",
             "}",
             ctx.settings.api.visibility,
             getDefaultSymbol(ctx.settings),
             EndpointProviderGenerator.getSymbol(ctx.settings),
+            RuntimeConfigProperty.AuthSchemePreference.symbol,
             getSymbol(ctx.settings),
         ) {
             val paramsSymbol = AuthSchemeParametersGenerator.getSymbol(ctx.settings)
@@ -87,6 +89,7 @@ open class AuthSchemeProviderGenerator {
                     renderAuthOptionsListOverrideForOperation(ctx, "\"${op.id.name}\"", authHandlersForOperation, writer, op)
                 }
             }
+            write("")
 
             withBlock(
                 "private val serviceDefaults = listOf<#T>(",
@@ -102,31 +105,55 @@ open class AuthSchemeProviderGenerator {
                     write("#W,", inlineWriter)
                 }
             }
+            write("")
 
             withBlock(
                 "override suspend fun resolveAuthScheme(params: #T): List<#T> {",
                 "}",
                 paramsSymbol,
                 RuntimeTypes.Auth.Identity.AuthOption,
             ) {
-                withBlock("val modeledAuthOptions = operationOverrides.getOrElse(params.operationName) {", "}") {
-                    write("serviceDefaults")
-                }
+                write("val modeledAuthOptions = operationOverrides.getOrDefault(params.operationName, serviceDefaults)")
+                write("")
 
                 if (ctx.settings.api.enableEndpointAuthProvider) {
                     write("")
-                    write("val endpointParams = params.endpointParameters")
-                    openBlock("val endpointAuthOptions = if (endpointProvider != null && endpointParams != null) {")
-                        .write("val endpoint = endpointProvider.resolveEndpoint(endpointParams)")
-                        .write("endpoint.#T", RuntimeTypes.SmithyClient.Endpoints.authOptions)
-                        .closeAndOpenBlock("} else {")
-                        .write("emptyList()")
-                        .closeBlock("}")
-                    write("")
-                    write("return #T(modeledAuthOptions, endpointAuthOptions)", RuntimeTypes.Auth.HttpAuthAws.mergeAuthOptions)
+                    withBlock("val authOptions: List<#T> = run {", "}", RuntimeTypes.Auth.Identity.AuthOption) {
+                        write("val endpointParams = params.endpointParameters")
+
+                        withInlineBlock("val endpointAuthOptions = if (endpointProvider != null && endpointParams != null) {", "} ") {
+                            write("val endpoint = endpointProvider.resolveEndpoint(endpointParams)")
+                            write("endpoint.#T", RuntimeTypes.SmithyClient.Endpoints.authOptions)
+                        }
+                        withInlineBlock("else {", "}") {
+                            write("emptyList()")
+                        }
+
+                        write("")
+                        write("#T(modeledAuthOptions, endpointAuthOptions)", RuntimeTypes.Auth.HttpAuthAws.mergeAuthOptions)
+                    }
                 } else {
-                    write("return modeledAuthOptions")
+                    write("val authOptions: List<#T> = modeledAuthOptions", RuntimeTypes.Auth.Identity.AuthOption)
+                }
+                write("")
+
+                // reprioritize auth options based on user's preference
+                withInlineBlock("return authSchemePreference?.let {", "} ") {
+                    // add preferred candidates first
+                    withBlock("val preferredAuthOptions = it.mapNotNull { preferredSchemeId ->", "}") {
+                        writeWithNoFormatting("val preferredSchemeName = preferredSchemeId.id.substringAfter('#')")
+                        withBlock("authOptions.singleOrNull {", "}") {
+                            writeWithNoFormatting("it.schemeId.id.substringAfter('#') == preferredSchemeName")
+                        }
+                    }
+
+                    // add any remaining candidates that weren't in the preference list
+                    write("val nonPreferredAuthOptions = authOptions.filterNot { it in preferredAuthOptions }")
+                    write("")
+
+                    write("preferredAuthOptions + nonPreferredAuthOptions")
                 }
+                write("?: authOptions")
             }
 
             // render any helper methods
diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/util/RuntimeConfigProperty.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/util/RuntimeConfigProperty.kt
@@ -174,6 +174,23 @@ object RuntimeConfigProperty {
             authentication schemes.
         """.trimIndent()
     }
+
+    val AuthSchemePreference = ConfigProperty {
+        name = "authSchemePreference"
+
+        val target = RuntimeTypes.Auth.Identity.AuthSchemeId
+
+        symbol = KotlinTypes.Collections.list(target, isNullable = true)
+        builderSymbol = KotlinTypes.Collections.list(target, isNullable = true)
+//        toBuilderExpression = ""
+
+        baseClass = RuntimeTypes.Auth.HttpAuth.HttpAuthConfig
+        useNestedBuilderBaseClass()
+
+        documentation = """
+            The ordered preference of [AuthScheme] that this client will use. 
+        """.trimIndent()
+    }
 }
 
 internal val Symbol.nestedBuilder: Symbol
diff --git a/codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/ServiceClientConfigGeneratorTest.kt b/codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/ServiceClientConfigGeneratorTest.kt
@@ -19,6 +19,7 @@ import software.amazon.smithy.kotlin.codegen.rendering.util.RuntimeConfigPropert
 import software.amazon.smithy.kotlin.codegen.test.*
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.ServiceShape
+import kotlin.test.Ignore
 import kotlin.test.Test
 
 class ServiceClientConfigGeneratorTest {
@@ -48,6 +49,7 @@ public class Config private constructor(builder: Builder) : HttpAuthConfig, Http
 
         val expectedProps = """
     override val clientName: String = builder.clientName
+    override val authSchemePreference: kotlin.collections.List<aws.smithy.kotlin.runtime.auth.AuthSchemeId>? = builder.authSchemePreference
     override val authSchemes: kotlin.collections.List<aws.smithy.kotlin.runtime.http.auth.AuthScheme> = builder.authSchemes
     public val endpointProvider: TestEndpointProvider = requireNotNull(builder.endpointProvider) { "endpointProvider is a required configuration property" }
     override val idempotencyTokenProvider: IdempotencyTokenProvider = builder.idempotencyTokenProvider ?: IdempotencyTokenProvider.Default
@@ -65,6 +67,11 @@ public class Config private constructor(builder: Builder) : HttpAuthConfig, Http
          */
         override var clientName: String = "Test"
 
+        /**
+         * The ordered preference of [AuthScheme] that this client will use.
+         */
+        override var authSchemePreference: kotlin.collections.List<aws.smithy.kotlin.runtime.auth.AuthSchemeId>? = null
+
         /**
          * Register new or override default [AuthScheme]s configured for this client. By default, the set
          * of auth schemes configured comes from the service model. An auth scheme configured explicitly takes
@@ -212,6 +219,7 @@ public class Config private constructor(builder: Builder) {
         contents.shouldContain(expectedProps)
     }
 
+    @Ignore
     @Test
     fun `it overrides props by name`() {
         val model = getModel()
@@ -239,6 +247,7 @@ public class Config private constructor(builder: Builder) {
         // Expect logMode config value to override default to LogMode.Request
         val expectedConfigValues = """
     override val clientName: String = builder.clientName
+    override val authSchemePreference: kotlin.collections.List<aws.smithy.kotlin.runtime.auth.AuthSchemeId> = builder.authSchemePreference
     override val authSchemes: kotlin.collections.List<aws.smithy.kotlin.runtime.http.auth.AuthScheme> = builder.authSchemes
     public val customProp: Int? = builder.customProp
     public val endpointProvider: TestEndpointProvider = requireNotNull(builder.endpointProvider) { "endpointProvider is a required configuration property" }
diff --git a/runtime/auth/http-auth-api/api/http-auth-api.api b/runtime/auth/http-auth-api/api/http-auth-api.api
@@ -9,11 +9,14 @@ public final class aws/smithy/kotlin/runtime/http/auth/AuthScheme$DefaultImpls {
 }
 
 public abstract interface class aws/smithy/kotlin/runtime/http/auth/HttpAuthConfig {
+	public abstract fun getAuthSchemePreference ()Ljava/util/List;
 	public abstract fun getAuthSchemes ()Ljava/util/List;
 }
 
 public abstract interface class aws/smithy/kotlin/runtime/http/auth/HttpAuthConfig$Builder {
+	public abstract fun getAuthSchemePreference ()Ljava/util/List;
 	public abstract fun getAuthSchemes ()Ljava/util/List;
+	public abstract fun setAuthSchemePreference (Ljava/util/List;)V
 	public abstract fun setAuthSchemes (Ljava/util/List;)V
 }
 
diff --git a/runtime/auth/http-auth-api/common/src/aws/smithy/kotlin/runtime/http/auth/HttpAuthConfig.kt b/runtime/auth/http-auth-api/common/src/aws/smithy/kotlin/runtime/http/auth/HttpAuthConfig.kt
@@ -5,6 +5,8 @@
 
 package aws.smithy.kotlin.runtime.http.auth
 
+import aws.smithy.kotlin.runtime.auth.AuthSchemeId
+
 /**
  * The user-accessible configuration properties for the SDKs HTTP authentication schemes
  */
@@ -17,6 +19,11 @@ public interface HttpAuthConfig {
      */
     public val authSchemes: List<AuthScheme>
 
+    /**
+     * The ordered preference of [AuthScheme] that this client will use.
+     */
+    public val authSchemePreference: List<AuthSchemeId>?
+
     public interface Builder {
         /**
          * Register new or override default [AuthScheme]'s configured for this client. By default, the set
@@ -25,5 +32,10 @@ public interface HttpAuthConfig {
          * authentication schemes.
          */
         public var authSchemes: List<AuthScheme>
+
+        /**
+         * The ordered preference of [AuthScheme] that this client will use.
+         */
+        public var authSchemePreference: List<AuthSchemeId>?
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -47,6 +47,7 @@ class ServiceClientConfigGenerator(`
`47`	`47`	`add(RuntimeConfigProperty.HttpClient)`
`48`	`48`	`add(RuntimeConfigProperty.HttpInterceptors)`
`49`	`49`	`add(RuntimeConfigProperty.AuthSchemes)`
	`50`	`+ add(RuntimeConfigProperty.AuthSchemePreference)`
`50`	`51`	`}`
`51`	`52`	`if (shape.hasIdempotentTokenMember(context.model)) {`
`52`	`53`	`add(RuntimeConfigProperty.IdempotencyTokenProvider)`
Original file line number	Diff line number	Diff line change
`@@ -9,11 +9,14 @@ public final class aws/smithy/kotlin/runtime/http/auth/AuthScheme$DefaultImpls {`
`9`	`9`	`}`
`10`	`10`
`11`	`11`	`public abstract interface class aws/smithy/kotlin/runtime/http/auth/HttpAuthConfig {`
	`12`	`+ public abstract fun getAuthSchemePreference ()Ljava/util/List;`
`12`	`13`	`public abstract fun getAuthSchemes ()Ljava/util/List;`
`13`	`14`	`}`
`14`	`15`
`15`	`16`	`public abstract interface class aws/smithy/kotlin/runtime/http/auth/HttpAuthConfig$Builder {`
	`17`	`+ public abstract fun getAuthSchemePreference ()Ljava/util/List;`
`16`	`18`	`public abstract fun getAuthSchemes ()Ljava/util/List;`
	`19`	`+ public abstract fun setAuthSchemePreference (Ljava/util/List;)V`
`17`	`20`	`public abstract fun setAuthSchemes (Ljava/util/List;)V`
`18`	`21`	`}`
`19`	`22`