Skip to content

Commit 9fb3bb0

Browse files
kneekey23kneekey23
authored
feat: Glacier Checksums Customization (#421)
1 parent d887aff commit 9fb3bb0

File tree

8 files changed

+226
-0
lines changed

8 files changed

+226
-0
lines changed

‎AWSClientRuntime/Sources/Middlewares/Sha256TreeHashMiddleware.swift‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0.
3+
4+
import AwsCommonRuntimeKit
5+
import ClientRuntime
6+
import AwsCCal
7+
8+
public struct Sha256TreeHashMiddleware<OperationStackOutput: HttpResponseBinding,
9+
OperationStackError: HttpResponseBinding>: Middleware {
10+
public let id: String = "Sha256TreeHash"
11+
12+
private let X_AMZ_SHA256_TREE_HASH_HEADER_NAME = "X-Amz-Sha256-Tree-Hash"
13+
14+
private let X_AMZ_CONTENT_SHA256_HEADER_NAME = "X-Amz-Content-Sha256"
15+
16+
public init() {}
17+
18+
public func handle<H>(context: Context,
19+
input: MInput,
20+
next: H) -> Result<MOutput, MError>
21+
where H: Handler,
22+
Self.MInput == H.Input,
23+
Self.MOutput == H.Output,
24+
Self.Context == H.Context,
25+
Self.MError == H.MiddlewareError {
26+
let request = input.build()
27+
28+
switch request.body {
29+
case .data(let data):
30+
guard let data = data else {
31+
return next.handle(context: context, input: input)
32+
}
33+
if !request.headers.exists(name: X_AMZ_CONTENT_SHA256_HEADER_NAME) {
34+
let sha256 = ByteBuffer(data: data).sha256().encodeToHexString()
35+
input.withHeader(name: X_AMZ_CONTENT_SHA256_HEADER_NAME, value: sha256)
36+
}
37+
case .stream(let stream):
38+
let streamBytes = stream.toBytes()
39+
guard streamBytes.length > 0 else {
40+
return next.handle(context: context, input: input)
41+
}
42+
let (linearHash, treeHash) = computeHashes(bytes: streamBytes)
43+
if let treeHash = treeHash, let linearHash = linearHash {
44+
input.withHeader(name: X_AMZ_SHA256_TREE_HASH_HEADER_NAME, value: treeHash)
45+
input.withHeader(name: X_AMZ_CONTENT_SHA256_HEADER_NAME, value: linearHash)
46+
}
47+
case .empty, .none:
48+
break
49+
}
50+
51+
return next.handle(context: context, input: input)
52+
}
53+
54+
/// Computes the tree-hash and linear hash of a `ByteBuffer`.
55+
/// See http://docs.aws.amazon.com/amazonglacier/latest/dev/checksum-calculations.html for more information.
56+
private func computeHashes(bytes: ByteBuffer) -> (String?, String?) {
57+
let bufferSize = 1024 * 1024
58+
var hashes = [[UInt8]]()
59+
60+
while true {
61+
var oneMbTempBuffer = ByteBuffer(size: bufferSize)
62+
let bytesRead = bytes.readIntoBuffer(buffer: &oneMbTempBuffer)
63+
if bytesRead == 0 {
64+
break
65+
}
66+
let hash = oneMbTempBuffer.sha256()
67+
hashes.append(hash.toByteArray())
68+
}
69+
70+
return (bytes.sha256().encodeToHexString(), computeTreeHash(hashes: hashes))
71+
}
72+
73+
/// Builds a tree hash root node given a slice of hashes. Glacier tree hash to be derived from SHA256 hashes of 1MB chunks of the data.
74+
/// See http://docs.aws.amazon.com/amazonglacier/latest/dev/checksum-calculations.html for more information.
75+
private func computeTreeHash(hashes: [[UInt8]]) -> String? {
76+
guard !hashes.isEmpty else {
77+
return nil
78+
}
79+
var previousLevelHashes = hashes
80+
while previousLevelHashes.count > 1 {
81+
var currentLevelHashes = [[UInt8]]()
82+
for index in stride(from: 0, to: previousLevelHashes.count, by: 2) {
83+
if previousLevelHashes.count - index > 1 {
84+
var concatenatedLevelHash = [UInt8]()
85+
concatenatedLevelHash.append(contentsOf: previousLevelHashes[index])
86+
concatenatedLevelHash.append(contentsOf: previousLevelHashes[index + 1])
87+
let concatenatedLevelHashByteBuffer = ByteBuffer(bytes: concatenatedLevelHash)
88+
89+
let md = concatenatedLevelHashByteBuffer.sha256()
90+
currentLevelHashes.append(md.toByteArray())
91+
92+
} else {
93+
currentLevelHashes.append(previousLevelHashes[index])
94+
}
95+
}
96+
previousLevelHashes = currentLevelHashes
97+
}
98+
99+
let byteBuf = ByteBuffer(bytes: previousLevelHashes[0])
100+
return byteBuf.encodeToHexString()
101+
}
102+
103+
public typealias MInput = SdkHttpRequestBuilder
104+
public typealias MOutput = OperationOutput<OperationStackOutput>
105+
public typealias Context = HttpContext
106+
public typealias MError = SdkError<OperationStackError>
107+
}

‎AWSClientRuntime/Tests/Middlewares/Sha256TreeHashMiddlewareTests.swift‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
//
2+
// Copyright Amazon.com Inc. or its affiliates.
3+
// All Rights Reserved.
4+
//
5+
// SPDX-License-Identifier: Apache-2.0
6+
//
7+
8+
import XCTest
9+
import ClientRuntime
10+
import AwsCommonRuntimeKit
11+
import SmithyTestUtil
12+
@testable import AWSClientRuntime
13+
14+
class Sha256TreeHashMiddlewareTests: XCTestCase {
15+
override func setUp() {
16+
#if os(Linux)
17+
AwsCommonRuntimeKit.initialize()
18+
#endif
19+
}
20+
21+
override func tearDown() {
22+
#if os(Linux)
23+
AwsCommonRuntimeKit.cleanUp()
24+
#endif
25+
}
26+
27+
func testTreeHashAllZeroes() throws {
28+
let context = HttpContextBuilder().build()
29+
let expectation = XCTestExpectation(description: "closure was run")
30+
let bytesIn5_5MB: Int = Int(1024 * 1024 * 5.5)
31+
let byteArray: [UInt8] = Array(repeating: 0, count: bytesIn5_5MB)
32+
let byteBuffer = ByteBuffer(bytes: byteArray)
33+
let streamInput = MockStreamInput(body: ByteStream.buffer(byteBuffer))
34+
var stack = OperationStack<MockStreamInput, MockOutput, MockMiddlewareError>(id: "TreeHashMiddlewareTestStack")
35+
stack.serializeStep.intercept(position: .before, middleware: MockSerializeStreamMiddleware())
36+
let mockHttpResponse = HttpResponse(body: HttpBody.none, statusCode: .accepted)
37+
let mockOutput = try MockOutput(httpResponse: mockHttpResponse, decoder: nil)
38+
let output = OperationOutput<MockOutput>(httpResponse: mockHttpResponse, output: mockOutput)
39+
stack.finalizeStep.intercept(position: .after, middleware: Sha256TreeHashMiddleware())
40+
_ = stack.handleMiddleware(context: context, input: streamInput, next: MockHandler(handleCallback: { context, input in
41+
let linear = input.headers.value(for: "X-Amz-Content-Sha256")
42+
XCTAssertEqual(linear, "733cf513448ce6b20ad1bc5e50eb27c06aefae0c320713a5dd99f4e51bc1ca60")
43+
let treeHash = input.headers.value(for: "X-Amz-Sha256-Tree-Hash")
44+
XCTAssertEqual(treeHash, "a3a82dbe3644dd6046be472f2e3ec1f8ef47f8f3adb86d0de4de7a254f255455")
45+
expectation.fulfill()
46+
return .success(output)
47+
}))
48+
49+
wait(for: [expectation], timeout: 3.0)
50+
}
51+
}

‎codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/AWSClientRuntimeTypes.kt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ object AWSClientRuntimeTypes {
4949
val UnknownAWSHttpServiceError = runtimeSymbol("UnknownAWSHttpServiceError")
5050
val AWSHttpServiceError = runtimeSymbol("AWSHttpServiceError")
5151
val RegionResolver = runtimeSymbol("RegionResolver")
52+
val Sha256TreeHashMiddleware = runtimeSymbol("Sha256TreeHashMiddleware")
5253
}
5354
}
5455

‎codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/AWSHttpBindingProtocolGenerator.kt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,10 +44,12 @@ abstract class AWSHttpBindingProtocolGenerator : HttpBindingProtocolGenerator()
4444
override val shouldRenderCodingKeysForEncodable = true
4545

4646
override fun generateProtocolUnitTests(ctx: ProtocolGenerator.GenerationContext): Int {
47+
// TODO: enable these tests once this PR is merged/released into Smithy: https://github.com/awslabs/smithy/pull/930
4748
val ignoredTests = setOf(
4849
"GlacierChecksums", // aws-sdk-swift#208
4950
"GlacierMultipartChecksums", // aws-sdk-swift#208
5051
)
52+
val imports = listOf(AWSSwiftDependency.AWS_CLIENT_RUNTIME.target)
5153
return HttpProtocolTestGenerator(
5254
ctx,
5355
requestTestBuilder,
@@ -56,6 +58,7 @@ abstract class AWSHttpBindingProtocolGenerator : HttpBindingProtocolGenerator()
5658
httpProtocolCustomizable,
5759
operationMiddleware,
5860
serdeContext,
61+
imports,
5962
ignoredTests
6063
).generateProtocolTests()
6164
}

‎codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/customization/glacier/GlacierChecksum.kt‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
package software.amazon.smithy.aws.swift.codegen.customization.glacier
2+
3+
import software.amazon.smithy.aws.swift.codegen.middleware.AWSSigningMiddleware
4+
import software.amazon.smithy.aws.swift.codegen.middleware.Sha256TreeHashMiddleware
5+
import software.amazon.smithy.aws.swift.codegen.sdkId
6+
import software.amazon.smithy.aws.traits.auth.UnsignedPayloadTrait
7+
import software.amazon.smithy.model.Model
8+
import software.amazon.smithy.model.shapes.OperationShape
9+
import software.amazon.smithy.model.shapes.ServiceShape
10+
import software.amazon.smithy.swift.codegen.SwiftSettings
11+
import software.amazon.smithy.swift.codegen.integration.ProtocolGenerator
12+
import software.amazon.smithy.swift.codegen.integration.SwiftIntegration
13+
import software.amazon.smithy.swift.codegen.middleware.MiddlewareStep
14+
import software.amazon.smithy.swift.codegen.middleware.OperationMiddleware
15+
import software.amazon.smithy.swift.codegen.model.expectShape
16+
import software.amazon.smithy.swift.codegen.model.hasTrait
17+
import java.util.Locale
18+
19+
/**
20+
* Adds a middleware for Glacier to add checksum headers needed for payloads
21+
* See: https://github.com/awslabs/aws-sdk-swift/issues/208
22+
* See also: https://docs.aws.amazon.com/amazonglacier/latest/dev/checksum-calculations.html
23+
*/
24+
class GlacierChecksum : SwiftIntegration {
25+
override fun enabledForService(model: Model, settings: SwiftSettings) =
26+
model.expectShape<ServiceShape>(settings.service).sdkId.lowercase(Locale.getDefault()) == "glacier"
27+
28+
override fun customizeMiddleware(
29+
ctx: ProtocolGenerator.GenerationContext,
30+
operationShape: OperationShape,
31+
operationMiddleware: OperationMiddleware
32+
) {
33+
operationMiddleware.removeMiddleware(operationShape, MiddlewareStep.FINALIZESTEP, "AWSSigningMiddleware")
34+
operationMiddleware.appendMiddleware(operationShape, AWSSigningMiddleware(::middlewareParamsString))
35+
operationMiddleware.appendMiddleware(operationShape, Sha256TreeHashMiddleware())
36+
}
37+
38+
private fun middlewareParamsString(op: OperationShape): String {
39+
val hasUnsignedPayload = op.hasTrait<UnsignedPayloadTrait>()
40+
return "signedBodyHeader: .contentSha256, unsignedBody: $hasUnsignedPayload"
41+
}
42+
}

‎codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/middleware/Sha256TreeHashMiddleware.kt‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
package software.amazon.smithy.aws.swift.codegen.middleware
2+
3+
import software.amazon.smithy.aws.swift.codegen.AWSClientRuntimeTypes
4+
import software.amazon.smithy.model.shapes.OperationShape
5+
import software.amazon.smithy.swift.codegen.SwiftWriter
6+
import software.amazon.smithy.swift.codegen.middleware.MiddlewarePosition
7+
import software.amazon.smithy.swift.codegen.middleware.MiddlewareRenderable
8+
import software.amazon.smithy.swift.codegen.middleware.MiddlewareStep
9+
10+
class Sha256TreeHashMiddleware : MiddlewareRenderable {
11+
override val name = "Sha256TreeHashMiddleware"
12+
13+
override val middlewareStep = MiddlewareStep.FINALIZESTEP
14+
15+
override val position = MiddlewarePosition.AFTER
16+
17+
override fun render(writer: SwiftWriter, op: OperationShape, operationStackName: String) {
18+
writer.write("$operationStackName.${middlewareStep.stringValue()}.intercept(position: ${position.stringValue()}, middleware: \$N())", AWSClientRuntimeTypes.Core.Sha256TreeHashMiddleware)
19+
}
20+
}

‎codegen/smithy-aws-swift-codegen/src/main/kotlin/software/amazon/smithy/aws/swift/codegen/restxml/RestXmlProtocolGenerator.kt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ class RestXmlProtocolGenerator : AWSHttpBindingProtocolGenerator() {
8181
httpProtocolCustomizable,
8282
operationMiddleware,
8383
serdeContext,
84+
listOf(),
8485
testsToIgnore
8586
).generateProtocolTests()
8687
}

‎codegen/smithy-aws-swift-codegen/src/main/resources/META-INF/services/software.amazon.smithy.swift.codegen.integration.SwiftIntegration‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ software.amazon.smithy.aws.swift.codegen.customization.s3.S3ErrorIntegration
44
software.amazon.smithy.aws.swift.codegen.customization.apigateway.ApiGatewayAddAcceptHeader
55
software.amazon.smithy.aws.swift.codegen.customization.glacier.GlacierAddVersionHeader
66
software.amazon.smithy.aws.swift.codegen.customization.glacier.GlacierAccountIdDefault
7+
software.amazon.smithy.aws.swift.codegen.customization.glacier.GlacierChecksum
78
software.amazon.smithy.aws.swift.codegen.customization.BoxServices
89
software.amazon.smithy.aws.swift.codegen.customization.PresignableModelIntegration
910
software.amazon.smithy.aws.swift.codegen.customization.polly.PollyGetPresignerIntegration

0 commit comments

Comments
 (0)