feat: add glacier request body checksum (#379)

Author: ianbotsf

codegen/protocol-tests/build.gradle.kts

@@ -106,7 +106,7 @@ fun generateSmithyBuild(tests: List<ProtocolTest>): String {
             "kotlin-codegen": {
               "service": "${test.serviceShapeId}",
               "package": {
-                "name": "aws.sdk.kotlin.protocoltest.${test.packageName}",
+                "name": "aws.sdk.kotlin.services.${test.packageName}",
                 "version": "1.0"
               },
               "build": {
@@ -145,12 +145,18 @@ open class ProtocolTestTask : DefaultTask() {
     @get:Input
     var plugin: String = ""
 
+    /**
+     * The build directory for the task
+     */
+    val generatedBuildDir: File
+        @OutputDirectory
+        get() = project.buildDir.resolve("smithyprojections/${project.name}/$protocol/$plugin")
+
     @TaskAction
     fun runTests() {
         require(protocol.isNotEmpty()) { "protocol name must be specified" }
         require(plugin.isNotEmpty()) { "plugin name must be specified" }
 
-        val generatedBuildDir = project.file("${project.buildDir}/smithyprojections/${project.name}/$protocol/$plugin")
         println("[$protocol] buildDir: $generatedBuildDir")
         if (!generatedBuildDir.exists()) {
             throw GradleException("$generatedBuildDir does not exist")
@@ -169,14 +175,22 @@ open class ProtocolTestTask : DefaultTask() {
     }
 }
 
-
-
 enabledProtocols.forEach {
-    tasks.register<ProtocolTestTask>("testProtocol-${it.projectionName}") {
+    val protocolName = it.projectionName
+
+    val protocolTestTask = tasks.register<ProtocolTestTask>("testProtocol-$protocolName") {
         dependsOn(tasks["generateSdk"])
         group = "Verification"
-        protocol = it.projectionName
+        protocol = protocolName
         plugin = "kotlin-codegen"
+    }.get()
+
+    // FIXME This is a hack to work around how protocol tests aren't in the actual service model and thus codegen
+    // separately from service customizations.
+    tasks.create<Copy>("copyStaticFiles-$protocolName") {
+        from(rootProject.projectDir.resolve("services/$protocolName/common/src"))
+        into(protocolTestTask.generatedBuildDir.resolve("src/main/kotlin/"))
+        tasks["generateSdk"].finalizedBy(this)
     }
 }
 

codegen/smithy-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/glacier/GlacierBodyChecksum.kt

@@ -0,0 +1,57 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+package aws.sdk.kotlin.codegen.customization.glacier
+
+import aws.sdk.kotlin.codegen.sdkId
+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.buildSymbol
+import software.amazon.smithy.kotlin.codegen.model.expectShape
+import software.amazon.smithy.kotlin.codegen.model.isStreaming
+import software.amazon.smithy.kotlin.codegen.rendering.protocol.HttpFeatureMiddleware
+import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolGenerator
+import software.amazon.smithy.kotlin.codegen.rendering.protocol.ProtocolMiddleware
+import software.amazon.smithy.kotlin.codegen.utils.getOrNull
+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.StructureShape
+
+public class GlacierBodyChecksum : KotlinIntegration {
+    override fun enabledForService(model: Model, settings: KotlinSettings): Boolean =
+        model.expectShape<ServiceShape>(settings.service).sdkId.equals("Glacier", ignoreCase = true)
+
+    override fun customizeMiddleware(
+        ctx: ProtocolGenerator.GenerationContext,
+        resolved: List<ProtocolMiddleware>,
+    ): List<ProtocolMiddleware> = resolved + glacierBodyChecksumMiddleware
+
+    private val glacierBodyChecksumMiddleware = object : HttpFeatureMiddleware() {
+        override val order: Byte = 127 // Must come after AwsSignatureVersion4
+        override val name: String = "GlacierBodyChecksum"
+
+        override fun isEnabledFor(ctx: ProtocolGenerator.GenerationContext, op: OperationShape): Boolean {
+            val input = op.input.getOrNull()?.let { ctx.model.expectShape<StructureShape>(it) }
+            return input?.members()?.any { it.isStreaming || ctx.model.expectShape(it.target).isStreaming } == true
+        }
+
+        override fun renderConfigure(writer: KotlinWriter) {
+            writer.addImport(RuntimeTypes.Utils.Sha256)
+            writer.addImport(glacierSymbol("GlacierBodyChecksum"))
+            writer.addImport(glacierSymbol("TreeHasherImpl"))
+
+            writer.write("val chunkSizeBytes = 1024 * 1024 // 1MB")
+            writer.write("treeHasher = TreeHasherImpl(chunkSizeBytes) { Sha256() }")
+        }
+
+        private fun glacierSymbol(name: String) = buildSymbol {
+            this.name = name
+            namespace = "aws.sdk.kotlin.services.glacier.internal"
+        }
+    }
+}

codegen/smithy-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/protocols/core/AwsHttpBindingProtocolGenerator.kt

@@ -78,10 +78,6 @@ abstract class AwsHttpBindingProtocolGenerator : HttpBindingProtocolGenerator()
                 // FIXME - document type not fully supported yet, see https://github.com/awslabs/smithy-kotlin/issues/123
                 "PutAndGetInlineDocumentsInput",
 
-                // Glacier customizations
-                "GlacierChecksums", // smithy-kotlin#164
-                "GlacierMultipartChecksums", // smithy-kotlin#164
-
                 // aws-sdk-kotlin#390
                 "RestJsonHttpWithHeaderMemberNoModeledBody",
                 "RestJsonHttpWithNoModeledBody",

codegen/smithy-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/protocols/middleware/AwsSignatureVersion4.kt

@@ -26,7 +26,7 @@ import software.amazon.smithy.model.traits.OptionalAuthTrait
  */
 open class AwsSignatureVersion4(private val signingServiceName: String) : HttpFeatureMiddleware() {
     override val name: String = AwsRuntimeTypes.Signing.AwsSigV4SigningMiddleware.name
-    override val order: Byte = 127
+    override val order: Byte = 126 // Must come before GlacierBodyChecksum
 
     init {
         require(signingServiceName.isNotEmpty()) { "signingServiceName must be specified" }

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

@@ -14,3 +14,4 @@ aws.sdk.kotlin.codegen.customization.glacier.GlacierAddVersionHeader
 aws.sdk.kotlin.codegen.customization.glacier.GlacierAccountIdDefault
 aws.sdk.kotlin.codegen.customization.polly.PollyPresigner
 aws.sdk.kotlin.codegen.customization.BoxServices
+aws.sdk.kotlin.codegen.customization.glacier.GlacierBodyChecksum

services/build.gradle.kts

@@ -38,6 +38,11 @@ subprojects {
         sourceSets.getByName("test") {
             kotlin.srcDir("common/test")
             kotlin.srcDir("generated-src/test")
+
+            dependencies {
+                implementation(kotlin("test-junit5"))
+                implementation(project(":aws-runtime:testing"))
+            }
         }
     }
 

services/glacier/common/src/aws/sdk/kotlin/services/glacier/internal/GlacierBodyChecksum.kt

@@ -0,0 +1,53 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+package aws.sdk.kotlin.services.glacier.internal
+
+import aws.sdk.kotlin.services.glacier.model.GlacierException
+import aws.smithy.kotlin.runtime.client.operationName
+import aws.smithy.kotlin.runtime.http.Feature
+import aws.smithy.kotlin.runtime.http.FeatureKey
+import aws.smithy.kotlin.runtime.http.HttpBody
+import aws.smithy.kotlin.runtime.http.HttpClientFeatureFactory
+import aws.smithy.kotlin.runtime.http.operation.SdkHttpOperation
+import aws.smithy.kotlin.runtime.http.request.headers
+import aws.smithy.kotlin.runtime.util.Sha256
+import aws.smithy.kotlin.runtime.util.encodeToHex
+
+private const val defaultChunkSizeBytes = 1024 * 1024 // 1MB
+
+internal class GlacierBodyChecksum(config: Config) : Feature {
+    private val treeHasher = config.treeHasher
+
+    internal class Config {
+        internal var treeHasher: TreeHasher = TreeHasherImpl(defaultChunkSizeBytes) { Sha256() }
+    }
+
+    internal companion object Feature : HttpClientFeatureFactory<Config, GlacierBodyChecksum> {
+        override val key: FeatureKey<GlacierBodyChecksum> = FeatureKey("GlacierBodyChecksum")
+
+        override fun create(block: Config.() -> Unit): GlacierBodyChecksum {
+            val config = Config().apply(block)
+            return GlacierBodyChecksum(config)
+        }
+    }
+
+    override fun <I, O> install(operation: SdkHttpOperation<I, O>) {
+        operation.execution.finalize.intercept { req, next ->
+            val body = req.subject.body
+            if (body is HttpBody.Streaming && !body.isReplayable) {
+                val opName = req.context.operationName ?: "This operation"
+                throw GlacierException("$opName requires a byte array or replayable stream")
+            }
+            val hashes = treeHasher.calculateHashes(body)
+            req.subject.headers {
+                set("X-Amz-Content-Sha256", hashes.fullHash.encodeToHex())
+                set("X-Amz-Sha256-Tree-Hash", hashes.treeHash.encodeToHex())
+            }
+
+            next.call(req)
+        }
+    }
+}

services/glacier/common/src/aws/sdk/kotlin/services/glacier/internal/TreeHasher.kt

@@ -0,0 +1,94 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+package aws.sdk.kotlin.services.glacier.internal
+
+import aws.smithy.kotlin.runtime.http.HttpBody
+import aws.smithy.kotlin.runtime.util.HashFunction
+import kotlinx.coroutines.flow.*
+import kotlin.math.min
+
+internal typealias HashSupplier = () -> HashFunction
+
+/**
+ * The result of a [TreeHasher] calculation.
+ * @param fullHash A full hash of the entire byte array (taken at once)
+ * @param treeHash A composite hash of the byte array, taken in chunks.
+ */
+internal class Hashes(public val fullHash: ByteArray, public val treeHash: ByteArray)
+
+/**
+ * A hash calculator that returns [Hashes] derived using a tree. See
+ * [Computing Checksums](https://docs.aws.amazon.com/amazonglacier/latest/dev/checksum-calculations.html) in the Glacier
+ * service guide for more details.
+ */
+internal interface TreeHasher {
+    /**
+     * Perform the hash calculation.
+     * @param body The [HttpBody] over which to calculate hashes.
+     * @return A [Hashes] containing the results of the calculation.
+     */
+    suspend fun calculateHashes(body: HttpBody): Hashes
+}
+
+/**
+ * The default implementation of a [TreeHasher].
+ */
+internal class TreeHasherImpl(private val chunkSizeBytes: Int, private val hashSupplier: HashSupplier) : TreeHasher {
+    override suspend fun calculateHashes(body: HttpBody): Hashes {
+        val full = hashSupplier()
+        val hashTree = ArrayDeque<ByteArray>()
+
+        body.chunks().collect { chunk ->
+            full.update(chunk)
+            hashTree.addLast(chunk.hash())
+        }
+
+        if (hashTree.isEmpty()) {
+            // Edge case for empty bodies
+            hashTree.add(byteArrayOf().hash())
+        }
+
+        while (hashTree.size > 1) {
+            val nextRow = mutableListOf<ByteArray>()
+
+            while (hashTree.isNotEmpty()) {
+                if (hashTree.size == 1) {
+                    nextRow.add(hashTree.removeFirst())
+                } else {
+                    val hash = hashSupplier()
+                    hashTree.removeFirst().let(hash::update)
+                    hashTree.removeFirst().let(hash::update)
+                    nextRow.add(hash.digest())
+                }
+            }
+
+            hashTree.addAll(nextRow)
+        }
+
+        return Hashes(full.digest(), hashTree.first())
+    }
+
+    private fun ByteArray.hash(): ByteArray = hashSupplier().apply { update(this@hash) }.digest()
+
+    private suspend fun HttpBody.chunks(): Flow<ByteArray> = when (this) {
+        is HttpBody.Empty -> flowOf()
+
+        is HttpBody.Bytes -> {
+            val size = bytes().size
+            val chunkStarts = 0 until size step chunkSizeBytes
+            val chunkRanges = chunkStarts.map { it until min(it + chunkSizeBytes, size) }
+            chunkRanges.asFlow().map(bytes()::sliceArray)
+        }
+
+        is HttpBody.Streaming -> flow {
+            val channel = readFrom()
+            while (!channel.isClosedForRead) {
+                emit(channel.readRemaining(chunkSizeBytes))
+            }
+            reset()
+        }
+    }
+}

services/glacier/common/test/aws/sdk/kotlin/services/glacier/internal/TreeHasherTest.kt

@@ -0,0 +1,87 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+package aws.sdk.kotlin.services.glacier.internal
+
+import aws.sdk.kotlin.runtime.testing.runSuspendTest
+import aws.smithy.kotlin.runtime.content.ByteStream
+import aws.smithy.kotlin.runtime.http.content.ByteArrayContent
+import aws.smithy.kotlin.runtime.http.toHttpBody
+import aws.smithy.kotlin.runtime.io.SdkByteChannel
+import aws.smithy.kotlin.runtime.io.SdkByteReadChannel
+import aws.smithy.kotlin.runtime.util.HashFunction
+import aws.smithy.kotlin.runtime.util.Sha256
+import aws.smithy.kotlin.runtime.util.encodeToHex
+import kotlinx.coroutines.async
+import kotlinx.coroutines.withTimeout
+import kotlin.test.Test
+import kotlin.test.assertContentEquals
+import kotlin.test.assertEquals
+import kotlin.test.fail
+
+private const val megabyte = 1024 * 1024
+
+class TreeHasherTest {
+    @Test
+    fun testCalculateHashes() = runSuspendTest {
+        val chunk1 = byteArrayOf(1, 2, 3)
+        val chunk2 = byteArrayOf(4, 5, 6)
+        val payload = chunk1 + chunk2
+
+        val fullHash = byteArrayOf(7, 9, 11) // Each element added once (thus, ∑(c[n] + 1))
+        val treeHash = byteArrayOf(9, 11, 13) // Elements added twice (thus, ∑(c[n] + 2))
+        val chunkSize = 3
+        val hasher = TreeHasherImpl(chunkSize) { RollingSumHashFunction(chunkSize) }
+
+        val body = ByteArrayContent(payload)
+        val hashes = hasher.calculateHashes(body)
+
+        assertContentEquals(fullHash, hashes.fullHash)
+        assertContentEquals(treeHash, hashes.treeHash)
+    }
+
+    @Test
+    fun integrationTestCalculateHashes() = runSuspendTest {
+        val byteStream = object : ByteStream.ReplayableStream() {
+            override fun newReader(): SdkByteReadChannel {
+                val byteChannel = SdkByteChannel()
+                val payloadBytes = "abcdefghijklmnopqrstuvwxyz".encodeToByteArray() // 26 bytes
+                async {
+                    withTimeout(10_000) { // For sanity, bail out after 10s
+                        (0 until megabyte).forEach { // This will yield len(payloadBytes) megabytes of content
+                            byteChannel.writeFully(payloadBytes)
+                        }
+                    }
+                    byteChannel.close()
+                }
+                return byteChannel
+            }
+        }
+
+        val hasher = TreeHasherImpl(megabyte) { Sha256() }
+        val hashes = hasher.calculateHashes(byteStream.toHttpBody())
+
+        assertEquals("74df7872289a84fa31b6ae4cfdbac34ef911cfe9357e842c600a060da6a899ae", hashes.fullHash.encodeToHex())
+        assertEquals("a1c6d421d75f727ce97e6998ab79e6a1cc08ee9502f541f9c5748b462c4dc83f", hashes.treeHash.encodeToHex())
+    }
+}
+
+/**
+ * Calculates a rolling sum for a hash. In this algorithm:
+ * * hash size = chunk size
+ * * every [update] call adds the positional input bytes to a rolling hash plus 1 (differentiates full from tree hashes)
+ */
+class RollingSumHashFunction(private val chunkSize: Int) : HashFunction {
+    private val rollingHash = ByteArray(chunkSize)
+
+    override fun digest(): ByteArray = rollingHash
+    override fun reset() = fail("reset should not have been called")
+    override fun update(input: ByteArray) {
+        assertEquals(chunkSize, input.size, "Chunk size must be exactly $chunkSize")
+        for (i in input.indices) {
+            rollingHash[i] = (rollingHash[i] + input[i] + 1).toByte()
+        }
+    }
+}