Add "invalid xml body root" check exemption for S3's `GetObjectAttrib…

…utes` (#1665)

* add: ability for certain operations to be exempt from XML body root checking
* add: XML body root checking exemption for com.amazonaws.s3#GetObjectAttributesOutput

CHANGELOG.next.toml

@@ -225,4 +225,19 @@ author = "Velfi"
 message = "Bump [MSRV](https://github.com/awslabs/aws-sdk-rust#supported-rust-versions-msrv) from 1.58.1 to 1.61.0 per our policy."
 references = ["smithy-rs#1699"]
 meta = { "breaking" = true, "tada" = false, "bug" = false, "target" = "all"}
+author = "Velfi"
+
+[[aws-sdk-rust]]
+message = "The AWS S3 `GetObjectAttributes` operation will no longer fail with an XML error."
+references = ["aws-sdk-rust#609"]
+meta = { "breaking" = false, "tada" = true, "bug" = true }
+author = "Velfi"
+
+[[smithy-rs]]
+message = """
+It is now possible to exempt specific operations from XML body root checking. To do this, add the `AllowInvalidXmlRoot`
+trait to the output struct of the operation you want to exempt.
+"""
+references = ["aws-sdk-rust#609"]
+meta = { "breaking" = false, "tada" = false, "bug" = false, "target" = "client"}
 author = "Velfi"

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/route53/TrimResourceId.kt

@@ -12,7 +12,7 @@ import software.amazon.smithy.model.traits.AnnotationTrait
 /**
  * Indicates that a member should have their resource ID prefix stripped
  */
-class TrimResourceId() : AnnotationTrait(ID, Node.objectNode()) {
+class TrimResourceId : AnnotationTrait(ID, Node.objectNode()) {
     companion object {
         val ID: ShapeId = ShapeId.from("aws.api.internal#trimResourceId")
     }

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/customize/s3/S3Decorator.kt

@@ -6,8 +6,13 @@
 package software.amazon.smithy.rustsdk.customize.s3
 
 import software.amazon.smithy.aws.traits.protocols.RestXmlTrait
+import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.shapes.ServiceShape
+import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.model.shapes.ShapeId
+import software.amazon.smithy.model.shapes.StructureShape
+import software.amazon.smithy.model.transform.ModelTransformer
 import software.amazon.smithy.rust.codegen.rustlang.CargoDependency
 import software.amazon.smithy.rust.codegen.rustlang.RustModule
 import software.amazon.smithy.rust.codegen.rustlang.Writable
@@ -23,17 +28,24 @@ import software.amazon.smithy.rust.codegen.smithy.customize.RustCodegenDecorator
 import software.amazon.smithy.rust.codegen.smithy.generators.LibRsCustomization
 import software.amazon.smithy.rust.codegen.smithy.generators.LibRsSection
 import software.amazon.smithy.rust.codegen.smithy.letIf
+import software.amazon.smithy.rust.codegen.smithy.protocols.AllowInvalidXmlRoot
 import software.amazon.smithy.rust.codegen.smithy.protocols.ProtocolMap
 import software.amazon.smithy.rust.codegen.smithy.protocols.RestXml
 import software.amazon.smithy.rust.codegen.smithy.protocols.RestXmlFactory
 import software.amazon.smithy.rustsdk.AwsRuntimeType
+import java.util.logging.Logger
 
 /**
  * Top level decorator for S3
  */
 class S3Decorator : RustCodegenDecorator<ClientCodegenContext> {
-    override val name: String = "S3ExtendedError"
+    override val name: String = "S3"
     override val order: Byte = 0
+    private val logger: Logger = Logger.getLogger(javaClass.name)
+    private val invalidXmlRootAllowList = setOf(
+        // API returns GetObjectAttributes_Response_ instead of Output
+        ShapeId.from("com.amazonaws.s3#GetObjectAttributesOutput"),
+    )
 
     private fun applies(serviceId: ShapeId) =
         serviceId == ShapeId.from("com.amazonaws.s3#AmazonS3")
@@ -50,6 +62,17 @@ class S3Decorator : RustCodegenDecorator<ClientCodegenContext> {
             )
         }
 
+    override fun transformModel(service: ServiceShape, model: Model): Model {
+        return model.letIf(applies(service.id)) {
+            ModelTransformer.create().mapShapes(model) { shape ->
+                shape.letIf(isInInvalidXmlRootAllowList(shape)) {
+                    logger.info("Adding AllowInvalidXmlRoot trait to $shape")
+                    (shape as StructureShape).toBuilder().addTrait(AllowInvalidXmlRoot()).build()
+                }
+            }
+        }
+    }
+
     override fun libRsCustomizations(
         codegenContext: ClientCodegenContext,
         baseCustomizations: List<LibRsCustomization>,
@@ -59,6 +82,10 @@ class S3Decorator : RustCodegenDecorator<ClientCodegenContext> {
 
     override fun supportsCodegenContext(clazz: Class<out CoreCodegenContext>): Boolean =
         clazz.isAssignableFrom(ClientCodegenContext::class.java)
+
+    private fun isInInvalidXmlRootAllowList(shape: Shape): Boolean {
+        return shape.isStructureShape && invalidXmlRootAllowList.contains(shape.id)
+    }
 }
 
 class S3(coreCodegenContext: CoreCodegenContext) : RestXml(coreCodegenContext) {

aws/sdk/integration-tests/s3/tests/ignore-invalid-xml-body-root.rs

@@ -0,0 +1,79 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use aws_http::user_agent::AwsUserAgent;
+use aws_sdk_s3::{
+    middleware::DefaultMiddleware, model::ObjectAttributes, operation::GetObjectAttributes,
+    Credentials, Region,
+};
+use aws_smithy_client::{test_connection::TestConnection, Client as CoreClient};
+use aws_smithy_http::body::SdkBody;
+use std::time::{Duration, UNIX_EPOCH};
+
+pub type Client<C> = CoreClient<C, DefaultMiddleware>;
+
+const RESPONSE_BODY_XML: &[u8] = b"<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n<GetObjectAttributesResponse xmlns=\"http://s3.amazonaws.com/doc/2006-03-01/\"><Checksum><ChecksumSHA1>e1AsOh9IyGCa4hLN+2Od7jlnP14=</ChecksumSHA1></Checksum></GetObjectAttributesResponse>";
+
+#[tokio::test]
+async fn ignore_invalid_xml_body_root() {
+    tracing_subscriber::fmt::init();
+
+    let conn = TestConnection::new(vec![
+        (http::Request::builder()
+             .header("x-amz-object-attributes", "Checksum")
+             .header("x-amz-user-agent", "aws-sdk-rust/0.123.test api/test-service/0.123 os/windows/XPSP3 lang/rust/1.50.0")
+             .header("x-amz-date", "20210618T170728Z")
+             .header("authorization", "AWS4-HMAC-SHA256 Credential=ANOTREAL/20210618/us-east-1/s3/aws4_request, SignedHeaders=host;x-amz-content-sha256;x-amz-date;x-amz-object-attributes;x-amz-security-token;x-amz-user-agent, Signature=0e6ec749db5a0af07890a83f553319eda95be0e498d058c64880471a474c5378")
+             .header("x-amz-content-sha256", "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855")
+             .header("x-amz-security-token", "notarealsessiontoken")
+             .uri(http::Uri::from_static("https://s3.us-east-1.amazonaws.com/some-test-bucket/test.txt?attributes"))
+             .body(SdkBody::empty())
+             .unwrap(),
+         http::Response::builder()
+             .header(
+                 "x-amz-id-2",
+                 "rbipIUyF3YKPIcqpz6hrP9x9mzYMSqkHzDEp6TEN/STcKvylDIE/LLN6x9t6EKJRrgctNsdNHWk=",
+             )
+             .header("x-amz-request-id", "K8036R3D4NZNMMVC")
+             .header("date", "Tue, 23 Aug 2022 18:17:23 GMT")
+             .header("last-modified", "Tue, 21 Jun 2022 16:30:01 GMT")
+             .header("server", "AmazonS3")
+             .header("content-length", "224")
+             .status(200)
+             .body(RESPONSE_BODY_XML)
+             .unwrap())
+    ]);
+    let creds = Credentials::new(
+        "ANOTREAL",
+        "notrealrnrELgWzOk3IfjzDKtFBhDby",
+        Some("notarealsessiontoken".to_string()),
+        None,
+        "test",
+    );
+    let conf = aws_sdk_s3::Config::builder()
+        .credentials_provider(creds)
+        .region(Region::new("us-east-1"))
+        .build();
+    let client = Client::new(conn.clone());
+
+    let mut op = GetObjectAttributes::builder()
+        .bucket("some-test-bucket")
+        .key("test.txt")
+        .object_attributes(ObjectAttributes::Checksum)
+        .build()
+        .unwrap()
+        .make_operation(&conf)
+        .await
+        .unwrap();
+    op.properties_mut()
+        .insert(UNIX_EPOCH + Duration::from_secs(1624036048));
+    op.properties_mut().insert(AwsUserAgent::for_tests());
+
+    let res = client.call(op).await.unwrap();
+
+    conn.assert_requests_match(&[]);
+
+    println!("res: {:#?}", res)
+}

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/protocols/RestXml.kt

@@ -6,7 +6,10 @@
 package software.amazon.smithy.rust.codegen.smithy.protocols
 
 import software.amazon.smithy.aws.traits.protocols.RestXmlTrait
+import software.amazon.smithy.model.node.Node
 import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.shapes.ShapeId
+import software.amazon.smithy.model.traits.AnnotationTrait
 import software.amazon.smithy.model.traits.TimestampFormatTrait
 import software.amazon.smithy.rust.codegen.rustlang.CargoDependency
 import software.amazon.smithy.rust.codegen.rustlang.RustModule
@@ -110,3 +113,12 @@ open class RestXml(private val coreCodegenContext: CoreCodegenContext) : Protoco
 
     override fun serverRouterRuntimeConstructor() = "new_rest_xml_router"
 }
+
+/**
+ * Indicates that a service is expected to send XML where the root element name does not match the modeled member name.
+ */
+class AllowInvalidXmlRoot : AnnotationTrait(ID, Node.objectNode()) {
+    companion object {
+        val ID: ShapeId = ShapeId.from("smithy.api.internal#allowInvalidXmlRoot")
+    }
+}

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/protocols/parse/RestXmlParserGenerator.kt

@@ -8,6 +8,11 @@ package software.amazon.smithy.rust.codegen.smithy.protocols.parse
 import software.amazon.smithy.rust.codegen.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.smithy.CoreCodegenContext
 import software.amazon.smithy.rust.codegen.smithy.RuntimeType
+import software.amazon.smithy.rust.codegen.smithy.protocols.AllowInvalidXmlRoot
+import software.amazon.smithy.rust.codegen.smithy.traits.SyntheticOutputTrait
+import software.amazon.smithy.rust.codegen.util.getTrait
+import software.amazon.smithy.rust.codegen.util.hasTrait
+import software.amazon.smithy.rust.codegen.util.orNull
 
 class RestXmlParserGenerator(
     coreCodegenContext: CoreCodegenContext,
@@ -18,14 +23,29 @@ class RestXmlParserGenerator(
             xmlErrors,
         ) { context, inner ->
             val shapeName = context.outputShapeName
-            rustTemplate(
-                """
-                if !(${XmlBindingTraitParserGenerator.XmlName(shapeName).matchExpression("start_el")}) {
-                    return Err(#{XmlError}::custom(format!("invalid root, expected $shapeName got {:?}", start_el)))
+            // Get the non-synthetic version of the outputShape and check to see if it has the `AllowInvalidXmlRoot` trait
+            val allowInvalidRoot = context.model.getShape(context.shape.outputShape).orNull().let { shape ->
+                shape?.getTrait<SyntheticOutputTrait>()?.originalId.let { shapeId ->
+                    context.model.getShape(shapeId).orNull()?.hasTrait<AllowInvalidXmlRoot>() ?: false
                 }
-                """,
-                "XmlError" to context.xmlErrorType,
-            )
+            }
+
+            // If we DON'T allow the XML root to be invalid, insert code to check for and report a mismatch
+            if (!allowInvalidRoot) {
+                rustTemplate(
+                    """
+                    if !${XmlBindingTraitParserGenerator.XmlName(shapeName).matchExpression("start_el")} {
+                        return Err(
+                            #{XmlError}::custom(
+                                format!("encountered invalid XML root: expected $shapeName but got {:?}. This is likely a bug in the SDK.", start_el)
+                            )
+                        )
+                    }
+                    """,
+                    "XmlError" to context.xmlErrorType,
+                )
+            }
+
             inner("decoder")
         },
 ) : StructuredDataParserGenerator by xmlBindingTraitParserGenerator

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/protocols/parse/XmlBindingTraitParserGenerator.kt

@@ -7,6 +7,7 @@ package software.amazon.smithy.rust.codegen.smithy.protocols.parse
 
 import software.amazon.smithy.aws.traits.customizations.S3UnwrappedXmlOutputTrait
 import software.amazon.smithy.codegen.core.CodegenException
+import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.knowledge.HttpBinding
 import software.amazon.smithy.model.knowledge.HttpBindingIndex
 import software.amazon.smithy.model.shapes.BlobShape
@@ -64,6 +65,7 @@ data class OperationWrapperContext(
     val shape: OperationShape,
     val outputShapeName: String,
     val xmlErrorType: RuntimeType,
+    val model: Model,
 )
 
 class XmlBindingTraitParserGenerator(
@@ -194,11 +196,12 @@ class XmlBindingTraitParserGenerator(
 
                     ##[allow(unused_mut)]
                     let mut decoder = doc.root_element()?;
+                    ##[allow(unused_variables)]
                     let start_el = decoder.start_el();
                     """,
                     *codegenScope,
                 )
-                val context = OperationWrapperContext(operationShape, shapeName, xmlError)
+                val context = OperationWrapperContext(operationShape, shapeName, xmlError, model)
                 if (operationShape.hasTrait<S3UnwrappedXmlOutputTrait>()) {
                     unwrappedResponseParser("builder", "decoder", "start_el", outputShape.members())
                 } else {
@@ -264,7 +267,7 @@ class XmlBindingTraitParserGenerator(
                     """,
                     *codegenScope,
                 )
-                val context = OperationWrapperContext(operationShape, shapeName, xmlError)
+                val context = OperationWrapperContext(operationShape, shapeName, xmlError, model)
                 writeOperationWrapper(context) { tagName ->
                     parseStructureInner(members, builder = "builder", Ctx(tag = tagName, accum = null))
                 }

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/traits/SyntheticInputTrait.kt

@@ -12,14 +12,14 @@ import software.amazon.smithy.model.traits.AnnotationTrait
 /**
  * Indicates that a shape is a synthetic input (see `OperationNormalizer.kt`)
  *
- * All operations are normalized to have an input, even when they are defined without on. This is done for backwards compatibility
- * and to produce a consistent API.
+ * All operations are normalized to have an input, even when they are defined without on. This is done for backwards
+ * compatibility and to produce a consistent API.
  */
 class SyntheticInputTrait(
     val operation: ShapeId,
     val originalId: ShapeId?,
 ) : AnnotationTrait(ID, Node.objectNode()) {
     companion object {
-        val ID = ShapeId.from("smithy.api.internal#syntheticInput")
+        val ID: ShapeId = ShapeId.from("smithy.api.internal#syntheticInput")
     }
 }

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/traits/SyntheticOutputTrait.kt

@@ -10,14 +10,14 @@ import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.model.traits.AnnotationTrait
 
 /**
- * Indicates that a shape is a synthetic input (see `OperationNormalizer.kt`)
+ * Indicates that a shape is a synthetic output (see `OperationNormalizer.kt`)
  *
- * All operations are normalized to have an input, even when they are defined without on. This is done for backwards compatibility
- * and to produce a consistent API.
+ * All operations are normalized to have an output, even when they are defined without on. This is done for backwards
+ * compatibility and to produce a consistent API.
  */
 class SyntheticOutputTrait constructor(val operation: ShapeId, val originalId: ShapeId?) :
     AnnotationTrait(ID, Node.objectNode()) {
     companion object {
-        val ID = ShapeId.from("smithy.api.internal#syntheticOutput")
+        val ID: ShapeId = ShapeId.from("smithy.api.internal#syntheticOutput")
     }
 }