feat: continue header (#802)

Author: ianbotsf

.changes/8f1ff251-7fed-4543-a077-51b2fe0aa684.json

@@ -0,0 +1,8 @@
+{
+    "id": "8f1ff251-7fed-4543-a077-51b2fe0aa684",
+    "type": "feature",
+    "description": "Add an interceptor for adding `Expect: 100-continue` headers to HTTP requests",
+    "issues": [
+        "awslabs/aws-sdk-kotlin#839"
+    ]
+}

codegen/smithy-kotlin-codegen-testutils/src/main/kotlin/software/amazon/smithy/kotlin/codegen/test/ModelTestUtils.kt

@@ -10,6 +10,7 @@ import software.amazon.smithy.kotlin.codegen.KotlinCodegenPlugin
 import software.amazon.smithy.kotlin.codegen.KotlinSettings
 import software.amazon.smithy.kotlin.codegen.core.*
 import software.amazon.smithy.kotlin.codegen.inferService
+import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
 import software.amazon.smithy.kotlin.codegen.model.OperationNormalizer
 import software.amazon.smithy.kotlin.codegen.model.shapes
 import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolGenerator
@@ -116,6 +117,7 @@ fun Model.newTestContext(
     packageName: String = TestModelDefault.NAMESPACE,
     settings: KotlinSettings = this.defaultSettings(serviceName, packageName),
     generator: ProtocolGenerator = MockHttpProtocolGenerator(),
+    integrations: List<KotlinIntegration> = listOf(),
 ): TestContext {
     val manifest = MockManifest()
     val provider: SymbolProvider = KotlinCodegenPlugin.createSymbolProvider(model = this, rootNamespace = packageName, serviceName = serviceName)
@@ -127,7 +129,7 @@ fun Model.newTestContext(
         this,
         service,
         provider,
-        listOf(),
+        integrations,
         generator.protocol,
         delegator,
     )
@@ -238,14 +240,18 @@ fun String.prependNamespaceAndService(
         protocol.annotation to imports + listOf(protocol.import)
     }
 
-    val importExpr = modelImports.map { "use $it" }.joinToString(separator = "\n") { it }
+    val importExpr = modelImports
+        .map { "use $it" }
+        .plus("use aws.api#service")
+        .joinToString(separator = "\n")
 
     return (
         """
         $versionExpr
         namespace $namespace
         $importExpr
         $modelProtocol
+        @service(sdkId: "$serviceName")
         service $serviceName { 
             version: "${TestModelDefault.MODEL_VERSION}",
             operations: $operations

codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/core/RuntimeTypes.kt

@@ -92,6 +92,7 @@ object RuntimeTypes {
         }
 
         object Interceptors : RuntimeTypePackage(KotlinDependency.HTTP, "interceptors") {
+            val ContinueInterceptor = symbol("ContinueInterceptor")
             val HttpInterceptor = symbol("HttpInterceptor")
             val Md5ChecksumInterceptor = symbol("Md5ChecksumInterceptor")
             val FlexibleChecksumsRequestInterceptor = symbol("FlexibleChecksumsRequestInterceptor")

runtime/protocol/http-client/common/src/aws/smithy/kotlin/runtime/http/interceptors/ContinueInterceptor.kt

@@ -0,0 +1,33 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.smithy.kotlin.runtime.http.interceptors
+
+import aws.smithy.kotlin.runtime.InternalApi
+import aws.smithy.kotlin.runtime.client.ProtocolRequestInterceptorContext
+import aws.smithy.kotlin.runtime.http.request.HttpRequest
+import aws.smithy.kotlin.runtime.http.request.header
+import aws.smithy.kotlin.runtime.http.request.toBuilder
+
+/**
+ * An interceptor that adds an HTTP `Expect: 100-continue` header to requests with bodies at a certain length threshold.
+ * Bodies with an unset `contentLength` will get the continue header added regardless of length.
+ * @param thresholdLengthBytes The body length (in bytes) at which a continue header will be set. Bodies under this
+ * length will not get a continue header.
+ */
+@InternalApi
+public class ContinueInterceptor(public val thresholdLengthBytes: Long) : HttpInterceptor {
+    override suspend fun modifyBeforeSigning(context: ProtocolRequestInterceptorContext<Any, HttpRequest>): HttpRequest {
+        val req = context.protocolRequest
+
+        return if ((req.body.contentLength ?: Long.MAX_VALUE) >= thresholdLengthBytes) {
+            req
+                .toBuilder()
+                .apply { header("Expect", "100-continue") }
+                .build()
+        } else {
+            req
+        }
+    }
+}

runtime/protocol/http-client/common/test/aws/smithy/kotlin/runtime/http/interceptors/ContinueInterceptorTest.kt

@@ -0,0 +1,67 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.smithy.kotlin.runtime.http.interceptors
+
+import aws.smithy.kotlin.runtime.client.ProtocolRequestInterceptorContext
+import aws.smithy.kotlin.runtime.http.HttpBody
+import aws.smithy.kotlin.runtime.http.HttpMethod
+import aws.smithy.kotlin.runtime.http.request.HttpRequest
+import aws.smithy.kotlin.runtime.http.request.header
+import aws.smithy.kotlin.runtime.http.request.url
+import aws.smithy.kotlin.runtime.io.SdkSource
+import aws.smithy.kotlin.runtime.net.Url
+import aws.smithy.kotlin.runtime.operation.ExecutionContext
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runTest
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertNull
+import kotlin.test.fail
+
+@OptIn(ExperimentalCoroutinesApi::class)
+class ContinueInterceptorTest {
+    private fun context(request: HttpRequest) = object : ProtocolRequestInterceptorContext<Any, HttpRequest> {
+        override val protocolRequest: HttpRequest = request
+        override val executionContext: ExecutionContext get() = fail("Shouldn't have invoked `executionContext`")
+        override val request: Any get() = fail("Shouldn't have invoked `request`")
+    }
+
+    private fun request(contentLength: Long?) = HttpRequest {
+        method = HttpMethod.POST
+        url(Url.parse("https://localhost"))
+        header("foo", "bar")
+        body = object : HttpBody.SourceContent() {
+            override val contentLength: Long? = contentLength
+            override fun readFrom(): SdkSource = fail("Shouldn't have invoked `readFrom`")
+        }
+    }
+
+    @Test
+    fun testInterceptorSmallBody() = runTest {
+        val input = request(50)
+        val interceptor = ContinueInterceptor(100)
+        val output = interceptor.modifyBeforeSigning(context(input))
+        assertEquals("bar", output.headers["foo"])
+        assertNull(output.headers["Expect"])
+    }
+
+    @Test
+    fun testInterceptorLargeBody() = runTest {
+        val input = request(150)
+        val interceptor = ContinueInterceptor(100)
+        val output = interceptor.modifyBeforeSigning(context(input))
+        assertEquals("bar", output.headers["foo"])
+        assertEquals("100-continue", output.headers["Expect"])
+    }
+
+    @Test
+    fun testInterceptorUnknownLengthBody() = runTest {
+        val input = request(null)
+        val interceptor = ContinueInterceptor(100)
+        val output = interceptor.modifyBeforeSigning(context(input))
+        assertEquals("bar", output.headers["foo"])
+        assertEquals("100-continue", output.headers["Expect"])
+    }
+}