feat: validate returned content length on S3 GetObject responses (#991)

Author: lauzadis

.changes/ee74d8c0-697e-4ca2-95c9-d0285876c71a.json

@@ -0,0 +1,8 @@
+{
+  "id": "4e8e536e-a39c-41a8-acfb-3877528a321f",
+  "type": "feature",
+  "description": "Validate returned content length on S3 `GetObject` responses.",
+  "issues": [
+    "awslabs/aws-sdk-kotlin#745"
+  ]
+}

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

@@ -0,0 +1,35 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.codegen.customization.s3
+
+import software.amazon.smithy.kotlin.codegen.KotlinSettings
+import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
+import software.amazon.smithy.kotlin.codegen.core.RuntimeTypes
+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
+
+class GetObjectResponseLengthValidationIntegration : KotlinIntegration {
+    override fun enabledForService(model: Model, settings: KotlinSettings) =
+        model.expectShape<ServiceShape>(settings.service).isS3
+
+    override fun customizeMiddleware(ctx: ProtocolGenerator.GenerationContext, resolved: List<ProtocolMiddleware>) =
+        resolved + responseLengthValidationMiddleware
+}
+
+internal val responseLengthValidationMiddleware = object : ProtocolMiddleware {
+    override val name: String = "ResponseLengthValidationMiddleware"
+
+    override fun isEnabledFor(ctx: ProtocolGenerator.GenerationContext, op: OperationShape) = op.id.name == "GetObject"
+
+    override fun render(ctx: ProtocolGenerator.GenerationContext, op: OperationShape, writer: KotlinWriter) {
+        val interceptorSymbol = RuntimeTypes.HttpClient.Interceptors.ResponseLengthValidationInterceptor
+        writer.write("op.interceptors.add(#T())", interceptorSymbol)
+    }
+}

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

@@ -24,6 +24,7 @@ aws.sdk.kotlin.codegen.customization.route53.TrimResourcePrefix
 aws.sdk.kotlin.codegen.customization.route53.ChangeResourceRecordSetsUnmarshallingIntegration
 aws.sdk.kotlin.codegen.customization.s3.ClientConfigIntegration
 aws.sdk.kotlin.codegen.customization.s3.ContinueIntegration
+aws.sdk.kotlin.codegen.customization.s3.GetObjectResponseLengthValidationIntegration
 aws.sdk.kotlin.codegen.customization.s3.HttpPathFilter
 aws.sdk.kotlin.codegen.customization.s3.TruncatablePaginationIntegration
 aws.sdk.kotlin.codegen.customization.s3control.HostPrefixFilter

codegen/smithy-aws-kotlin-codegen/src/test/kotlin/aws/sdk/kotlin/codegen/customization/s3/ContinueIntegrationTest.kt

@@ -5,17 +5,13 @@
 package aws.sdk.kotlin.codegen.customization.s3
 
 import aws.sdk.kotlin.codegen.testutil.lines
+import aws.sdk.kotlin.codegen.testutil.model
 import org.junit.jupiter.api.Test
-import software.amazon.smithy.codegen.core.SymbolProvider
-import software.amazon.smithy.kotlin.codegen.KotlinSettings
-import software.amazon.smithy.kotlin.codegen.core.CodegenContext
 import software.amazon.smithy.kotlin.codegen.core.KotlinWriter
-import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
 import software.amazon.smithy.kotlin.codegen.rendering.ServiceClientConfigGenerator
 import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolGenerator
 import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolMiddleware
 import software.amazon.smithy.kotlin.codegen.test.*
-import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.OperationShape
 import kotlin.test.*
 
@@ -101,25 +97,6 @@ class ContinueIntegrationTest {
     }
 }
 
-private fun Model.codegenContext() = object : CodegenContext {
-    override val model: Model = this@codegenContext
-    override val symbolProvider: SymbolProvider get() = fail("Unexpected call to `symbolProvider`")
-    override val settings: KotlinSettings get() = fail("Unexpected call to `settings`")
-    override val protocolGenerator: ProtocolGenerator? = null
-    override val integrations: List<KotlinIntegration> = listOf()
-}
-
-private fun model(serviceName: String): Model =
-    """
-        @http(method: "PUT", uri: "/foo")
-        operation Foo { }
-        
-        @http(method: "POST", uri: "/bar")
-        operation Bar { }
-    """
-        .prependNamespaceAndService(operations = listOf("Foo", "Bar"), serviceName = serviceName)
-        .toSmithyModel()
-
 object FooMiddleware : ProtocolMiddleware {
     override val name: String = "FooMiddleware"
     override fun render(ctx: ProtocolGenerator.GenerationContext, op: OperationShape, writer: KotlinWriter) =

codegen/smithy-aws-kotlin-codegen/src/test/kotlin/aws/sdk/kotlin/codegen/customization/s3/GetObjectResponseLengthValidationIntegrationTest.kt

@@ -0,0 +1,39 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.codegen.customization.s3
+
+import aws.sdk.kotlin.codegen.testutil.model
+import org.junit.jupiter.api.Test
+import software.amazon.smithy.kotlin.codegen.test.defaultSettings
+import software.amazon.smithy.kotlin.codegen.test.newTestContext
+import kotlin.test.assertEquals
+import kotlin.test.assertFalse
+import kotlin.test.assertTrue
+
+class GetObjectResponseLengthValidationIntegrationTest {
+    @Test
+    fun testNotExpectedForNonS3Model() {
+        val model = model("NotS3")
+        val actual = GetObjectResponseLengthValidationIntegration().enabledForService(model, model.defaultSettings())
+        assertFalse(actual)
+    }
+
+    @Test
+    fun testExpectedForS3Model() {
+        val model = model("S3")
+        val actual = GetObjectResponseLengthValidationIntegration().enabledForService(model, model.defaultSettings())
+        assertTrue(actual)
+    }
+
+    @Test
+    fun testMiddlewareAddition() {
+        val model = model("S3")
+        val preexistingMiddleware = listOf(FooMiddleware)
+        val ctx = model.newTestContext("S3")
+        val actual = GetObjectResponseLengthValidationIntegration().customizeMiddleware(ctx.generationCtx, preexistingMiddleware)
+
+        assertEquals(listOf(FooMiddleware, responseLengthValidationMiddleware), actual)
+    }
+}

codegen/smithy-aws-kotlin-codegen/src/test/kotlin/aws/sdk/kotlin/codegen/testutil/ModelUtil.kt

@@ -0,0 +1,16 @@
+package aws.sdk.kotlin.codegen.testutil
+
+import software.amazon.smithy.kotlin.codegen.test.prependNamespaceAndService
+import software.amazon.smithy.kotlin.codegen.test.toSmithyModel
+import software.amazon.smithy.model.Model
+
+internal fun model(serviceName: String): Model =
+    """
+        @http(method: "PUT", uri: "/foo")
+        operation Foo { }
+        
+        @http(method: "POST", uri: "/bar")
+        operation Bar { }
+    """
+        .prependNamespaceAndService(operations = listOf("Foo", "Bar"), serviceName = serviceName)
+        .toSmithyModel()