fix: application of sigv4a

.changes/517318e5-0ba6-472c-a51f-bb4e758215e2.json

@@ -0,0 +1,5 @@
+{
+    "id": "517318e5-0ba6-472c-a51f-bb4e758215e2",
+    "type": "bugfix",
+    "description": "Fix application of sigv4a authentication scheme for S3, Eventbridge, and CloudFront KeyValueStore"
+}

codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/SigV4AsymmetricTraitCustomization.kt

@@ -9,7 +9,8 @@ import software.amazon.smithy.aws.traits.auth.SigV4ATrait
 import software.amazon.smithy.aws.traits.auth.SigV4Trait
 import software.amazon.smithy.kotlin.codegen.KotlinSettings
 import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
-import software.amazon.smithy.kotlin.codegen.model.expectTrait
+import software.amazon.smithy.kotlin.codegen.model.getTrait
+import software.amazon.smithy.kotlin.codegen.model.hasTrait
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.model.traits.AuthTrait
@@ -24,28 +25,40 @@ class SigV4AsymmetricTraitCustomization : KotlinIntegration {
     // Needs to happen before the `SigV4AsymmetricAuthSchemeIntegration` & `SigV4AuthSchemeIntegration` (-50 & -50)
     override val order: Byte = -60
 
+    // services which support SigV4A but don't model it
+    private val unmodeledSigV4aServices = listOf("s3", "eventbridge")
+
     override fun enabledForService(model: Model, settings: KotlinSettings): Boolean =
-        when (settings.sdkId.lowercase()) {
-            "s3", "eventbridge", "cloudfront keyvaluestore" -> true
-            else -> false
-        }
+        unmodeledSigV4aServices.contains(settings.sdkId.lowercase()) && !model.isTraitApplied(SigV4ATrait::class.java)
 
     override fun preprocessModel(model: Model, settings: KotlinSettings): Model =
         ModelTransformer.create().mapShapes(model) { shape ->
             when (shape.isServiceShape) {
-                true ->
-                    (shape as ServiceShape)
-                        .toBuilder()
-                        .addTraits(
-                            mutableSetOf(
-                                SigV4ATrait
-                                    .builder()
-                                    .name(shape.expectTrait<ServiceTrait>().arnNamespace)
-                                    .build(),
-                                AuthTrait(mutableSetOf(SigV4ATrait.ID, SigV4Trait.ID)),
-                            ),
+                true -> {
+                    val builder = (shape as ServiceShape).toBuilder()
+
+                    if (!shape.hasTrait<SigV4ATrait>()) {
+                        builder.addTrait(
+                            SigV4ATrait.builder()
+                                .name(shape.getTrait<SigV4Trait>()?.name ?: shape.getTrait<ServiceTrait>()?.arnNamespace)
+                                .build(),
                         )
-                        .build()
+                    }
+
+                    // SigV4A is added at the end because these services model SigV4A through endpoint rules instead of the service shape.
+                    // Because of that, SigV4A can apply to any operation, and the safest thing to do is add it at the end
+                    // and let the endpoint rules change priority as needed.
+                    val authTrait = shape.getTrait<AuthTrait>()?.let {
+                        if (it.valueSet.contains(SigV4ATrait.ID)) {
+                            it
+                        } else {
+                            AuthTrait(it.valueSet + mutableSetOf(SigV4ATrait.ID))
+                        }
+                    } ?: AuthTrait(mutableSetOf(SigV4Trait.ID, SigV4ATrait.ID))
+                    builder.addTrait(authTrait)
+
+                    builder.build()
+                }
                 false -> shape
             }
         }

codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/cloudfrontkeyvaluestore/BackfillSigV4ACustomization.kt

@@ -0,0 +1,52 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.codegen.customization.cloudfrontkeyvaluestore
+
+import software.amazon.smithy.aws.traits.ServiceTrait
+import software.amazon.smithy.aws.traits.auth.SigV4ATrait
+import software.amazon.smithy.aws.traits.auth.SigV4Trait
+import software.amazon.smithy.kotlin.codegen.KotlinSettings
+import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+import software.amazon.smithy.kotlin.codegen.model.getTrait
+import software.amazon.smithy.model.Model
+import software.amazon.smithy.model.shapes.ServiceShape
+import software.amazon.smithy.model.traits.AuthTrait
+import software.amazon.smithy.model.transform.ModelTransformer
+
+/**
+ * Modify the [AuthTrait] of the CloudFront KeyValueStore service, placing SigV4A before SigV4 in the
+ * prioritized list.
+ */
+class BackfillSigV4ACustomization : KotlinIntegration {
+    override fun enabledForService(model: Model, settings: KotlinSettings): Boolean = settings.sdkId.lowercase() == "cloudfront keyvaluestore"
+
+    override fun preprocessModel(model: Model, settings: KotlinSettings): Model =
+        ModelTransformer.create().mapShapes(model) { shape ->
+            when (shape.isServiceShape) {
+                true -> {
+                    val builder = (shape as ServiceShape).toBuilder()
+
+                    builder.addTrait(
+                        SigV4ATrait.builder()
+                            .name(shape.getTrait<SigV4Trait>()?.name ?: shape.getTrait<ServiceTrait>()?.arnNamespace)
+                            .build(),
+                    )
+
+                    val authTraitValueSet = shape.getTrait<AuthTrait>()?.valueSet ?: mutableSetOf()
+
+                    if (authTraitValueSet.firstOrNull() != SigV4ATrait.ID) {
+                        // remove existing auth trait
+                        builder.removeTrait(AuthTrait.ID)
+
+                        // add a new auth trait with SigV4A in front of the existing values
+                        builder.addTrait(AuthTrait(mutableSetOf(SigV4ATrait.ID) + authTraitValueSet))
+                    }
+
+                    builder.build()
+                }
+                false -> shape
+            }
+        }
+}

codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/s3/UnsupportedSigningAlgorithmIntegration.kt

@@ -8,12 +8,10 @@ import aws.sdk.kotlin.codegen.AwsRuntimeTypes
 import software.amazon.smithy.kotlin.codegen.KotlinSettings
 import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
 import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
-import software.amazon.smithy.kotlin.codegen.model.expectShape
 import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolGenerator
 import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolMiddleware
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.OperationShape
-import software.amazon.smithy.model.shapes.ServiceShape
 
 // FIXME: Remove this once sigV4a is supported by default AWS signer
 /**
@@ -22,7 +20,10 @@ import software.amazon.smithy.model.shapes.ServiceShape
  */
 class UnsupportedSigningAlgorithmIntegration : KotlinIntegration {
     override fun enabledForService(model: Model, settings: KotlinSettings): Boolean =
-        model.expectShape<ServiceShape>(settings.service).isS3
+        when (settings.sdkId.lowercase()) {
+            "s3", "eventbridge", "cloudfront keyvaluestore" -> true
+            else -> false
+        }
 
     override fun customizeMiddleware(
         ctx: ProtocolGenerator.GenerationContext,

codegen/aws-sdk-codegen/src/main/resources/META-INF/services/software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration

@@ -38,3 +38,4 @@ aws.sdk.kotlin.codegen.customization.ec2.EC2MakePrimitivesOptional
 aws.sdk.kotlin.codegen.customization.RemoveDefaults
 aws.sdk.kotlin.codegen.customization.s3.UnsupportedSigningAlgorithmIntegration
 aws.sdk.kotlin.codegen.customization.SigV4AsymmetricTraitCustomization
+aws.sdk.kotlin.codegen.customization.cloudfrontkeyvaluestore.BackfillSigV4ACustomization