Add EC2 Query protocol support

codegen-test/build.gradle.kts

@@ -33,6 +33,7 @@ val CodegenTests = listOf(
     CodegenTest("aws.protocoltests.restjson#RestJsonExtras", "rest_json_extas"),
     CodegenTest("aws.protocoltests.restxml#RestXml", "rest_xml"),
     CodegenTest("aws.protocoltests.query#AwsQuery", "aws_query"),
+    CodegenTest("aws.protocoltests.ec2#AwsEc2", "ec2_query"),
     CodegenTest(
         "aws.protocoltests.restxml.xmlns#RestXmlWithNamespace",
         "rest_xml_namespace"

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/rustlang/CargoDependency.kt

@@ -107,6 +107,9 @@ class InlineDependency(
             CargoDependency.SmithyHttp(runtimeConfig)
         )
 
+        fun ec2QueryErrors(runtimeConfig: RuntimeConfig): InlineDependency =
+            forRustFile("ec2_query_errors", CargoDependency.smithyXml(runtimeConfig))
+
         fun wrappedXmlErrors(runtimeConfig: RuntimeConfig): InlineDependency =
             forRustFile("rest_xml_wrapped_errors", CargoDependency.smithyXml(runtimeConfig))
 

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/RuntimeTypes.kt

@@ -225,6 +225,9 @@ data class RuntimeType(val name: String?, val dependency: RustDependency?, val n
             namespace = "smithy_http::response"
         )
 
+        fun ec2QueryErrors(runtimeConfig: RuntimeConfig) =
+            forInlineDependency(InlineDependency.ec2QueryErrors(runtimeConfig))
+
         fun wrappedXmlErrors(runtimeConfig: RuntimeConfig) =
             forInlineDependency(InlineDependency.wrappedXmlErrors(runtimeConfig))
 

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

@@ -0,0 +1,86 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+package software.amazon.smithy.rust.codegen.smithy.protocols
+
+import software.amazon.smithy.model.Model
+import software.amazon.smithy.model.pattern.UriPattern
+import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.traits.HttpTrait
+import software.amazon.smithy.model.traits.TimestampFormatTrait
+import software.amazon.smithy.rust.codegen.rustlang.CargoDependency
+import software.amazon.smithy.rust.codegen.rustlang.asType
+import software.amazon.smithy.rust.codegen.rustlang.rust
+import software.amazon.smithy.rust.codegen.rustlang.rustBlockTemplate
+import software.amazon.smithy.rust.codegen.smithy.RuntimeType
+import software.amazon.smithy.rust.codegen.smithy.generators.ProtocolConfig
+import software.amazon.smithy.rust.codegen.smithy.generators.ProtocolGeneratorFactory
+import software.amazon.smithy.rust.codegen.smithy.generators.ProtocolSupport
+import software.amazon.smithy.rust.codegen.smithy.protocols.parse.Ec2QueryParserGenerator
+import software.amazon.smithy.rust.codegen.smithy.protocols.parse.StructuredDataParserGenerator
+import software.amazon.smithy.rust.codegen.smithy.protocols.serialize.Ec2QuerySerializerGenerator
+import software.amazon.smithy.rust.codegen.smithy.protocols.serialize.StructuredDataSerializerGenerator
+import software.amazon.smithy.rust.codegen.smithy.transformers.OperationNormalizer
+import software.amazon.smithy.rust.codegen.smithy.transformers.RemoveEventStreamOperations
+
+class Ec2QueryFactory : ProtocolGeneratorFactory<HttpBoundProtocolGenerator> {
+    override fun buildProtocolGenerator(protocolConfig: ProtocolConfig): HttpBoundProtocolGenerator =
+        HttpBoundProtocolGenerator(protocolConfig, Ec2QueryProtocol(protocolConfig))
+
+    override fun transformModel(model: Model): Model {
+        return OperationNormalizer(model).transformModel(
+            inputBodyFactory = OperationNormalizer.NoBody,
+            outputBodyFactory = OperationNormalizer.NoBody
+        ).let(RemoveEventStreamOperations::transform)
+    }
+
+    override fun support(): ProtocolSupport {
+        return ProtocolSupport(
+            requestSerialization = true,
+            requestBodySerialization = true,
+            responseDeserialization = true,
+            errorDeserialization = true,
+        )
+    }
+}
+
+class Ec2QueryProtocol(private val protocolConfig: ProtocolConfig) : Protocol {
+    private val runtimeConfig = protocolConfig.runtimeConfig
+    private val ec2QueryErrors: RuntimeType = RuntimeType.ec2QueryErrors(runtimeConfig)
+    override val httpBindingResolver: HttpBindingResolver = StaticHttpBindingResolver(
+        protocolConfig.model,
+        HttpTrait.builder()
+            .code(200)
+            .method("POST")
+            .uri(UriPattern.parse("/"))
+            .build(),
+        "application/x-www-form-urlencoded"
+    )
+
+    override val defaultTimestampFormat: TimestampFormatTrait.Format = TimestampFormatTrait.Format.DATE_TIME
+
+    override fun structuredDataParser(operationShape: OperationShape): StructuredDataParserGenerator =
+        Ec2QueryParserGenerator(protocolConfig, ec2QueryErrors)
+
+    override fun structuredDataSerializer(operationShape: OperationShape): StructuredDataSerializerGenerator =
+        Ec2QuerySerializerGenerator(protocolConfig)
+
+    override fun parseGenericError(operationShape: OperationShape): RuntimeType {
+        /**
+         fn parse_generic(response: &Response<Bytes>) -> Result<smithy_types::error::Generic, T: Error>
+         **/
+        return RuntimeType.forInlineFun("parse_generic_error", "xml_deser") {
+            it.rustBlockTemplate(
+                "pub fn parse_generic_error(response: &#{Response}<#{Bytes}>) -> Result<#{Error}, #{XmlError}>",
+                "Response" to RuntimeType.http.member("Response"),
+                "Bytes" to RuntimeType.Bytes,
+                "Error" to RuntimeType.GenericError(runtimeConfig),
+                "XmlError" to CargoDependency.smithyXml(runtimeConfig).asType().member("decode::XmlError")
+            ) {
+                rust("#T::parse_generic_error(response.body().as_ref())", ec2QueryErrors)
+            }
+        }
+    }
+}

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

@@ -8,6 +8,7 @@ package software.amazon.smithy.rust.codegen.smithy.protocols
 import software.amazon.smithy.aws.traits.protocols.AwsJson1_0Trait
 import software.amazon.smithy.aws.traits.protocols.AwsJson1_1Trait
 import software.amazon.smithy.aws.traits.protocols.AwsQueryTrait
+import software.amazon.smithy.aws.traits.protocols.Ec2QueryTrait
 import software.amazon.smithy.aws.traits.protocols.RestJson1Trait
 import software.amazon.smithy.aws.traits.protocols.RestXmlTrait
 import software.amazon.smithy.codegen.core.CodegenException
@@ -40,6 +41,7 @@ class ProtocolLoader(private val supportedProtocols: ProtocolMap) {
             AwsJson1_0Trait.ID to BasicAwsJsonFactory(AwsJsonVersion.Json10),
             AwsJson1_1Trait.ID to BasicAwsJsonFactory(AwsJsonVersion.Json11),
             AwsQueryTrait.ID to AwsQueryFactory(),
+            Ec2QueryTrait.ID to Ec2QueryFactory(),
             RestJson1Trait.ID to RestJsonFactory(),
             RestXmlTrait.ID to RestXmlFactory(),
         )

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

@@ -0,0 +1,40 @@
+package software.amazon.smithy.rust.codegen.smithy.protocols.parse
+
+import software.amazon.smithy.rust.codegen.rustlang.rustTemplate
+import software.amazon.smithy.rust.codegen.smithy.RuntimeType
+import software.amazon.smithy.rust.codegen.smithy.generators.ProtocolConfig
+
+/**
+ * The EC2 query protocol's responses are identical to REST XML's, except that they are wrapped
+ * in a Response tag:
+ *
+ * ```
+ * <SomeOperationResponse>
+ *    <ActualData /> <!-- This part is the same as REST XML -->
+ * </SomeOperationResponse>
+ * ```
+ *
+ * This class wraps [XmlBindingTraitParserGenerator] and uses it to render the vast majority
+ * of the response parsing, but it overrides [operationParser] to add the protocol differences.
+ */
+class Ec2QueryParserGenerator(
+    protocolConfig: ProtocolConfig,
+    xmlErrors: RuntimeType,
+    private val xmlBindingTraitParserGenerator: XmlBindingTraitParserGenerator =
+        XmlBindingTraitParserGenerator(
+            protocolConfig,
+            xmlErrors
+        ) { context, inner ->
+            val operationName = protocolConfig.symbolProvider.toSymbol(context.shape).name
+            val responseWrapperName = operationName + "Response"
+            rustTemplate(
+                """
+                if !(${XmlBindingTraitParserGenerator.XmlName(responseWrapperName).matchExpression("start_el")}) {
+                    return Err(#{XmlError}::custom(format!("invalid root, expected $responseWrapperName got {:?}", start_el)))
+                }
+                """,
+                "XmlError" to context.xmlErrorType
+            )
+            inner("decoder")
+        }
+) : StructuredDataParserGenerator by xmlBindingTraitParserGenerator