Skip to content

Commit 9b9297c

Browse files
lauzadislauzadis
authored
feat(rt): use aws-chunked content-encoding (#746)
1 parent afe997e commit 9b9297c

File tree

10 files changed

+242
-90
lines changed

10 files changed

+242
-90
lines changed

‎.changes/0b8435af-23b7-478a-b31a-260a7dfdfaab.json‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
{
2+
"id": "0b8435af-23b7-478a-b31a-260a7dfdfaab",
3+
"type": "feature",
4+
"description": "Use `aws-chunked` content encoding for streaming requests"
5+
}

‎runtime/auth/aws-signing-common/common/src/aws/smithy/kotlin/runtime/auth/awssigning/HashSpecification.kt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,11 @@ public sealed class HashSpecification {
3434
*/
3535
public object StreamingAws4HmacSha256Payload : HashLiteral("STREAMING-AWS4-HMAC-SHA256-PAYLOAD")
3636

37+
/**
38+
* The hash value indicates that the streaming request will have trailers
39+
*/
40+
public object StreamingAws4HmacSha256PayloadWithTrailers : HashLiteral("STREAMING-AWS4-HMAC-SHA256-PAYLOAD-TRAILER")
41+
3742
/**
3843
* The hash value should indicate ???
3944
*/

‎runtime/auth/aws-signing-common/common/src/aws/smithy/kotlin/runtime/auth/awssigning/internal/AwsChunkedReader.kt‎

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,10 @@
55

66
package aws.smithy.kotlin.runtime.auth.awssigning.internal
77

8+
import aws.smithy.kotlin.runtime.auth.awssigning.AwsSignatureType
89
import aws.smithy.kotlin.runtime.auth.awssigning.AwsSigner
910
import aws.smithy.kotlin.runtime.auth.awssigning.AwsSigningConfig
11+
import aws.smithy.kotlin.runtime.auth.awssigning.HashSpecification
1012
import aws.smithy.kotlin.runtime.http.Headers
1113
import aws.smithy.kotlin.runtime.io.SdkBuffer
1214

@@ -133,7 +135,7 @@ internal class AwsChunkedReader(
133135
// signer takes a ByteArray unfortunately...
134136
val chunkBody = bodyBuffer?.readByteArray() ?: return null
135137

136-
val chunkSignature = signer.signChunk(chunkBody, previousSignature, signingConfig).signature
138+
val chunkSignature = signer.signChunk(chunkBody, previousSignature, signingConfig.toChunkSigningConfig()).signature
137139
previousSignature = chunkSignature
138140

139141
val signedChunk = SdkBuffer()
@@ -164,11 +166,31 @@ internal class AwsChunkedReader(
164166
* @return a [SdkBuffer] containing the trailing headers in aws-chunked encoding, ready to send on the wire
165167
*/
166168
private suspend fun getTrailingHeadersChunk(trailingHeaders: Headers): SdkBuffer {
167-
val trailerSignature = signer.signChunkTrailer(trailingHeaders, previousSignature, signingConfig).signature
169+
val trailerSignature = signer.signChunkTrailer(trailingHeaders, previousSignature, signingConfig.toTrailingHeadersSigningConfig()).signature
168170
previousSignature = trailerSignature
169171

170172
val trailerBody = SdkBuffer()
171173
trailerBody.writeTrailers(trailingHeaders, trailerSignature.decodeToString())
172174
return trailerBody
173175
}
176+
177+
/**
178+
* Make a copy of the signing config, changing the signatureType and hashSpecification configuration values
179+
* to specify chunk signing.
180+
* @return an [AwsSigningConfig] which can be used by a signer to sign chunks
181+
*/
182+
private fun AwsSigningConfig.toChunkSigningConfig(): AwsSigningConfig = this.toBuilder().apply {
183+
signatureType = AwsSignatureType.HTTP_REQUEST_CHUNK // signature is for a chunk
184+
hashSpecification = HashSpecification.CalculateFromPayload // calculate the hash from the chunk payload
185+
}.build()
186+
187+
/**
188+
* Make a copy of the signing config, changing the signatureType and hashSpecification configuration values
189+
* to specify trailing headers signing.
190+
* @return an [AwsSigningConfig] which can be used by a signer to sign trailing headers
191+
*/
192+
private fun AwsSigningConfig.toTrailingHeadersSigningConfig(): AwsSigningConfig = this.toBuilder().apply {
193+
signatureType = AwsSignatureType.HTTP_REQUEST_TRAILING_HEADERS // signature is for trailing headers
194+
hashSpecification = HashSpecification.CalculateFromPayload // calculate the hash from the trailing headers payload
195+
}.build()
174196
}

‎runtime/auth/aws-signing-common/common/src/aws/smithy/kotlin/runtime/auth/awssigning/internal/AwsChunkedUtil.kt‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,13 @@
55

66
package aws.smithy.kotlin.runtime.auth.awssigning.internal
77

8+
import aws.smithy.kotlin.runtime.auth.awssigning.AwsSigner
9+
import aws.smithy.kotlin.runtime.auth.awssigning.AwsSigningConfig
10+
import aws.smithy.kotlin.runtime.auth.awssigning.HashSpecification
11+
import aws.smithy.kotlin.runtime.auth.awssigning.middleware.AwsSigningMiddleware
812
import aws.smithy.kotlin.runtime.http.Headers
13+
import aws.smithy.kotlin.runtime.http.HttpBody
14+
import aws.smithy.kotlin.runtime.http.request.HttpRequestBuilder
915
import aws.smithy.kotlin.runtime.io.SdkBuffer
1016

1117
/**
@@ -28,3 +34,33 @@ internal fun SdkBuffer.writeTrailers(
2834
}
2935
writeUtf8("x-amz-trailer-signature:${signature}\r\n")
3036
}
37+
38+
/**
39+
* @return a boolean representing if this HttpBody is eligible to send via aws-chunked content encoding
40+
*/
41+
internal val HttpBody.isEligibleForAwsChunkedStreaming: Boolean
42+
get() = (this is HttpBody.SourceContent || this is HttpBody.ChannelContent) && contentLength != null &&
43+
(isOneShot || contentLength!! > AwsSigningMiddleware.AWS_CHUNKED_THRESHOLD)
44+
45+
/**
46+
* @return a boolean representing if the signing configuration is configured (via [HashSpecification]) for aws-chunked content encoding
47+
*/
48+
internal val AwsSigningConfig.useAwsChunkedEncoding: Boolean
49+
get() = when (hashSpecification) {
50+
is HashSpecification.StreamingAws4HmacSha256Payload, is HashSpecification.StreamingAws4HmacSha256PayloadWithTrailers -> true
51+
else -> false
52+
}
53+
54+
/**
55+
* Set the HTTP headers required for the aws-chunked content encoding
56+
*/
57+
internal fun HttpRequestBuilder.setAwsChunkedHeaders() {
58+
headers.setMissing("Content-Encoding", "aws-chunked")
59+
headers.setMissing("Transfer-Encoding", "chunked")
60+
headers.setMissing("X-Amz-Decoded-Content-Length", body.contentLength!!.toString())
61+
}
62+
63+
/**
64+
* Update the HTTP body to use aws-chunked content encoding
65+
*/
66+
internal expect fun HttpRequestBuilder.setAwsChunkedBody(signer: AwsSigner, signingConfig: AwsSigningConfig, signature: ByteArray)

‎runtime/auth/aws-signing-common/common/src/aws/smithy/kotlin/runtime/auth/awssigning/middleware/AwsSigningMiddleware.kt‎

Lines changed: 23 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -6,14 +6,17 @@ package aws.smithy.kotlin.runtime.auth.awssigning.middleware
66

77
import aws.smithy.kotlin.runtime.auth.awscredentials.CredentialsProvider
88
import aws.smithy.kotlin.runtime.auth.awssigning.*
9+
import aws.smithy.kotlin.runtime.auth.awssigning.internal.*
10+
import aws.smithy.kotlin.runtime.auth.awssigning.internal.isEligibleForAwsChunkedStreaming
11+
import aws.smithy.kotlin.runtime.auth.awssigning.internal.setAwsChunkedBody
12+
import aws.smithy.kotlin.runtime.auth.awssigning.internal.setAwsChunkedHeaders
13+
import aws.smithy.kotlin.runtime.auth.awssigning.internal.useAwsChunkedEncoding
914
import aws.smithy.kotlin.runtime.http.HttpBody
1015
import aws.smithy.kotlin.runtime.http.operation.*
1116
import aws.smithy.kotlin.runtime.http.request.HttpRequest
1217
import aws.smithy.kotlin.runtime.http.request.HttpRequestBuilder
13-
import aws.smithy.kotlin.runtime.tracing.warn
1418
import aws.smithy.kotlin.runtime.util.InternalApi
1519
import aws.smithy.kotlin.runtime.util.get
16-
import kotlin.coroutines.coroutineContext
1720
import kotlin.time.Duration
1821

1922
/**
@@ -29,6 +32,10 @@ public class AwsSigningMiddleware(private val config: Config) : ModifyRequestMid
2932
requireNotNull(config.signer) { "A signer must be specified for the middleware" }
3033
return AwsSigningMiddleware(config)
3134
}
35+
36+
@InternalApi
37+
// The minimum size of a streaming body before the SDK will begin using aws-chunked content encoding.
38+
public const val AWS_CHUNKED_THRESHOLD: Int = CHUNK_SIZE_BYTES * 16
3239
}
3340

3441
public class Config {
@@ -126,38 +133,28 @@ public class AwsSigningMiddleware(private val config: Config) : ModifyRequestMid
126133
signedBodyHeader = contextSignedBodyHeader ?: config.signedBodyHeader
127134

128135
// SDKs are supposed to default to signed payload _always_ when possible (and when `unsignedPayload` trait
129-
// isn't present).
130-
//
131-
// There are a few escape hatches/special cases:
132-
// 1. Customer explicitly disables signed payload (via Config.isUnsignedPayload)
133-
// 2. Customer provides a (potentially) unbounded stream (via HttpBody.Streaming)
134-
//
135-
// When an unbounded stream (2) is given we proceed as follows:
136-
// 2.1. is it replayable?
137-
// (2.1.1) yes -> sign the payload (stream can be consumed more than once)
138-
// (2.1.2) no -> unsigned payload
139-
//
140-
// NOTE: Chunked signing is NOT enabled through this middleware.
141-
// NOTE: 2.1.2 is handled below
142-
143-
// FIXME - see: https://github.com/awslabs/smithy-kotlin/issues/296
144-
// if we know we have a (streaming) body and toSignableRequest() fails to convert it to a CRT equivalent
145-
// then we must decide how to compute the payload hash ourselves (defaults to unsigned payload)
136+
// isn't present). The only exception is when the customer explicitly disables signed payloads (via Config.isUnsignedPayload).
137+
146138
hashSpecification = when {
147139
contextHashSpecification != null -> contextHashSpecification
148140
config.isUnsignedPayload -> HashSpecification.UnsignedPayload
149141
body is HttpBody.Empty -> HashSpecification.EmptyBody
150-
body.isOneShot -> {
151-
coroutineContext.warn<AwsSigningMiddleware> {
152-
"unable to compute hash for unbounded stream; defaulting to unsigned payload"
142+
body.isEligibleForAwsChunkedStreaming -> {
143+
if (req.subject.headers.contains("x-amz-trailer")) {
144+
HashSpecification.StreamingAws4HmacSha256PayloadWithTrailers
145+
} else {
146+
HashSpecification.StreamingAws4HmacSha256Payload
153147
}
154-
HashSpecification.UnsignedPayload
155148
}
156149
// use the payload to compute the hash
157150
else -> HashSpecification.CalculateFromPayload
158151
}
159152
}
160153

154+
if (signingConfig.useAwsChunkedEncoding) {
155+
req.subject.setAwsChunkedHeaders()
156+
}
157+
161158
val signingResult = checkNotNull(config.signer).sign(req.subject.build(), signingConfig)
162159
val signedRequest = signingResult.output
163160

@@ -166,6 +163,9 @@ public class AwsSigningMiddleware(private val config: Config) : ModifyRequestMid
166163

167164
req.subject.update(signedRequest)
168165

166+
if (signingConfig.useAwsChunkedEncoding) {
167+
req.subject.setAwsChunkedBody(checkNotNull(config.signer), signingConfig, signingResult.signature)
168+
}
169169
return req
170170
}
171171
}

‎runtime/auth/aws-signing-common/jvm/src/aws/smithy/kotlin/runtime/auth/awssigning/internal/AwsChunkedUtilJVM.kt‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package aws.smithy.kotlin.runtime.auth.awssigning.internal
6+
7+
import aws.smithy.kotlin.runtime.ClientException
8+
import aws.smithy.kotlin.runtime.auth.awssigning.AwsChunkedByteReadChannel
9+
import aws.smithy.kotlin.runtime.auth.awssigning.AwsChunkedSource
10+
import aws.smithy.kotlin.runtime.auth.awssigning.AwsSigner
11+
import aws.smithy.kotlin.runtime.auth.awssigning.AwsSigningConfig
12+
import aws.smithy.kotlin.runtime.http.HttpBody
13+
import aws.smithy.kotlin.runtime.http.request.HttpRequestBuilder
14+
import aws.smithy.kotlin.runtime.http.toHttpBody
15+
import aws.smithy.kotlin.runtime.http.toSdkByteReadChannel
16+
17+
internal actual fun HttpRequestBuilder.setAwsChunkedBody(signer: AwsSigner, signingConfig: AwsSigningConfig, signature: ByteArray) {
18+
body = when (body) {
19+
is HttpBody.ChannelContent -> AwsChunkedByteReadChannel(checkNotNull(body.toSdkByteReadChannel()), signer, signingConfig, signature).toHttpBody(-1)
20+
is HttpBody.SourceContent -> AwsChunkedSource((body as HttpBody.SourceContent).readFrom(), signer, signingConfig, signature).toHttpBody(-1)
21+
else -> throw ClientException("HttpBody type is not supported")
22+
}
23+
}

‎runtime/auth/aws-signing-tests/common/src/aws/smithy/kotlin/runtime/auth/awssigning/tests/AwsChunkedByteReadChannelTestBase.kt‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,16 +18,16 @@ import kotlin.test.*
1818
import kotlin.time.Duration.Companion.milliseconds
1919

2020
@OptIn(ExperimentalCoroutinesApi::class)
21-
public abstract class AwsChunkedByteReadChannelTestBase : AwsChunkedTestBase(AwsChunkedReaderFactory.Channel) {
21+
abstract class AwsChunkedByteReadChannelTestBase : AwsChunkedTestBase(AwsChunkedReaderFactory.Channel) {
2222
@Test
23-
public fun testSlowProducerMultipleChunksPartialLast(): TestResult = runTest {
23+
fun testSlowProducerMultipleChunksPartialLast(): TestResult = runTest {
2424
val numChunks = 6
2525
val dataLengthBytes = CHUNK_SIZE_BYTES * (numChunks - 1) + CHUNK_SIZE_BYTES / 2 // 5 full chunks, 1 half-full chunk
2626

2727
val data = ByteArray(dataLengthBytes) { Random.Default.nextBytes(1)[0] }
2828
val chan = SdkByteChannel(true)
2929
var previousSignature: ByteArray = byteArrayOf()
30-
val awsChunked = AwsChunkedByteReadChannel(chan, signer, testSigningConfig, previousSignature)
30+
val awsChunked = AwsChunkedByteReadChannel(chan, signer, testChunkSigningConfig, previousSignature)
3131

3232
// launch a coroutine and fill the channel slowly
3333
val writeJob = launch {
@@ -66,7 +66,7 @@ public abstract class AwsChunkedByteReadChannelTestBase : AwsChunkedTestBase(Aws
6666
val expectedChunkSignature = signer.signChunk(
6767
data.slice(CHUNK_SIZE_BYTES * chunk until (CHUNK_SIZE_BYTES * (chunk + 1))).toByteArray(),
6868
previousSignature,
69-
testSigningConfig,
69+
testChunkSigningConfig,
7070
).signature
7171
previousSignature = expectedChunkSignature
7272

@@ -78,14 +78,14 @@ public abstract class AwsChunkedByteReadChannelTestBase : AwsChunkedTestBase(Aws
7878
var expectedChunkSignature = signer.signChunk(
7979
data.slice(CHUNK_SIZE_BYTES * (numChunks - 1) until data.size).toByteArray(),
8080
previousSignature,
81-
testSigningConfig,
81+
testChunkSigningConfig,
8282
).signature
8383
previousSignature = expectedChunkSignature
8484
assertEquals(expectedChunkSignature.decodeToString(), chunkSignatures[chunkSignatures.size - 2])
8585
assertEquals(CHUNK_SIZE_BYTES / 2, chunkSizes[chunkSizes.size - 2])
8686

8787
// validate terminal chunk
88-
expectedChunkSignature = signer.signChunk(byteArrayOf(), previousSignature, testSigningConfig).signature
88+
expectedChunkSignature = signer.signChunk(byteArrayOf(), previousSignature, testChunkSigningConfig).signature
8989
assertEquals(expectedChunkSignature.decodeToString(), chunkSignatures.last())
9090
assertEquals(0, chunkSizes.last())
9191
}

0 commit comments

Comments
 (0)