fix: Make Context thread-safe (#1819)

awslabs · Nov 13, 2024 · 41a69aa · 41a69aa
1 parent 1abd1f1
commit 41a69aa
Show file tree

Hide file tree

Showing 7 changed files with 10 additions and 10 deletions.
diff --git a/...Core/AWSClientRuntime/Sources/AWSClientRuntime/Middlewares/Sha256TreeHashMiddleware.swift b/...Core/AWSClientRuntime/Sources/AWSClientRuntime/Middlewares/Sha256TreeHashMiddleware.swift
@@ -46,7 +46,7 @@ public struct Sha256TreeHashMiddleware<OperationStackInput, OperationStackOutput
             if let treeHash = treeHash, let linearHash = linearHash {
                 builder.withHeader(name: X_AMZ_SHA256_TREE_HASH_HEADER_NAME, value: treeHash)
                 // provide the value but let CRT add the SHA256 header during signing
-                context.attributes.set(key: AttributeKey(name: X_AMZ_CONTENT_SHA256_HEADER_NAME), value: linearHash)
+                context.set(key: AttributeKey(name: X_AMZ_CONTENT_SHA256_HEADER_NAME), value: linearHash)
             }
         case .noStream:
             break

diff --git a/Sources/Core/AWSClientRuntime/Sources/AWSClientRuntime/UserAgent/BusinessMetrics.swift b/Sources/Core/AWSClientRuntime/Sources/AWSClientRuntime/UserAgent/BusinessMetrics.swift
@@ -48,11 +48,11 @@ private extension String {
 
 public extension Context {
     var businessMetrics: Dictionary<String, String> {
-        get { attributes.get(key: businessMetricsKey) ?? [:] }
+        get { get(key: businessMetricsKey) ?? [:] }
         set(newPair) {
             var combined = businessMetrics
             combined.merge(newPair) { (_, new) in new }
-            attributes.set(key: businessMetricsKey, value: combined)
+            set(key: businessMetricsKey, value: combined)
         }
     }
 }

diff --git a/...entRuntime/Tests/AWSClientRuntimeTests/Middlewares/FlexibleChecksumsMiddlewareTests.swift b/...entRuntime/Tests/AWSClientRuntimeTests/Middlewares/FlexibleChecksumsMiddlewareTests.swift
@@ -267,7 +267,7 @@ class FlexibleChecksumsMiddlewareTests: XCTestCase {
             "Stream is not 'chunked eligible'",
             file: file, line: line
         )
-        if let validatedChecksum = self.builtContext.attributes.get(key: AttributeKey<String>(name: "ChecksumHeaderValidated")), validatedChecksum == expectedHeader {
+        if let validatedChecksum = self.builtContext.get(key: AttributeKey<String>(name: "ChecksumHeaderValidated")), validatedChecksum == expectedHeader {
             isChecksumValidated = true
         }
         if expectedChecksum != "" {
@@ -304,7 +304,7 @@ class FlexibleChecksumsMiddlewareTests: XCTestCase {
             }
         }
 
-        if let validatedChecksum = self.builtContext.attributes.get(key: AttributeKey<String>(name: "ChecksumHeaderValidated")), validatedChecksum == expectedHeader {
+        if let validatedChecksum = self.builtContext.get(key: AttributeKey<String>(name: "ChecksumHeaderValidated")), validatedChecksum == expectedHeader {
             isChecksumValidated = true
         }
         if expectedChecksum != "" {
@@ -329,7 +329,7 @@ class FlexibleChecksumsMiddlewareTests: XCTestCase {
             }
         } else {
             try await handleMiddleware()
-            validatedChecksum = self.builtContext.attributes.get(key: AttributeKey<String>(name: "ChecksumHeaderValidated"))
+            validatedChecksum = self.builtContext.get(key: AttributeKey<String>(name: "ChecksumHeaderValidated"))
             if validatedChecksum != nil && validatedChecksum == expectedValidationHeader {
                 isChecksumValidated = true
             }

diff --git a/...re/amazon/smithy/aws/swift/codegen/customization/RulesBasedAuthSchemeResolverGenerator.kt b/...re/amazon/smithy/aws/swift/codegen/customization/RulesBasedAuthSchemeResolverGenerator.kt
@@ -131,7 +131,7 @@ class RulesBasedAuthSchemeResolverGenerator {
                 }
 
                 // Get endpoint param from middleware context
-                openBlock("guard let endpointParam = context.attributes.get(key: \$N<EndpointParams>(name: \"EndpointParams\")) else {", "}", SmithyTypes.AttributeKey) {
+                openBlock("guard let endpointParam = context.get(key: \$N<EndpointParams>(name: \"EndpointParams\")) else {", "}", SmithyTypes.AttributeKey) {
                     write("throw \$N.dataNotFound(\"Endpoint param not configured in middleware context for rules-based auth scheme resolver params construction.\")", SmithyTypes.ClientError)
                 }
 

diff --git a/...oftware/amazon/smithy/aws/swift/codegen/middleware/OperationEndpointResolverMiddleware.kt b/...oftware/amazon/smithy/aws/swift/codegen/middleware/OperationEndpointResolverMiddleware.kt
@@ -52,7 +52,7 @@ class OperationEndpointResolverMiddleware(
 
         // Write code that saves endpoint params to middleware context for use in auth scheme middleware when using rules-based auth scheme resolvers
         if (AuthSchemeResolverGenerator.usesRulesBasedAuthResolver(ctx)) {
-            writer.write("context.attributes.set(key: \$N<EndpointParams>(name: \"EndpointParams\"), value: endpointParams)", SmithyTypes.AttributeKey)
+            writer.write("context.set(key: \$N<EndpointParams>(name: \"EndpointParams\"), value: endpointParams)", SmithyTypes.AttributeKey)
         }
 
         super.renderSpecific(ctx, writer, op, operationStackName, "applyEndpoint")

diff --git a/...degen/src/test/kotlin/software/amazon/smithy/aws/swift/codegen/PresignerGeneratorTests.kt b/...degen/src/test/kotlin/software/amazon/smithy/aws/swift/codegen/PresignerGeneratorTests.kt
@@ -270,7 +270,7 @@ extension PutObjectInput {
         builder.retryErrorInfoProvider(AWSClientRuntime.AWSRetryErrorInfoProvider.errorInfo(for:))
         builder.applySigner(ClientRuntime.SignerMiddleware<PutObjectOutput>())
         let endpointParams = EndpointParams()
-        context.attributes.set(key: Smithy.AttributeKey<EndpointParams>(name: "EndpointParams"), value: endpointParams)
+        context.set(key: Smithy.AttributeKey<EndpointParams>(name: "EndpointParams"), value: endpointParams)
         builder.applyEndpoint(AWSClientRuntime.EndpointResolverMiddleware<PutObjectOutput, EndpointParams>(endpointResolverBlock: { [config] in try config.endpointResolver.resolve(params: ${'$'}0) }, endpointParams: endpointParams))
         builder.interceptors.add(AWSClientRuntime.UserAgentMiddleware<PutObjectInput, PutObjectOutput>(serviceID: serviceName, version: S3Client.version, config: config))
         builder.selectAuthScheme(ClientRuntime.AuthSchemeMiddleware<PutObjectOutput>())

diff --git a/...zon/smithy/aws/swift/codegen/customizations/RulesBasedAuthSchemeResolverGeneratorTests.kt b/...zon/smithy/aws/swift/codegen/customizations/RulesBasedAuthSchemeResolverGeneratorTests.kt
@@ -129,7 +129,7 @@ public struct DefaultS3AuthSchemeResolver: S3AuthSchemeResolver {
         guard let opName = context.getOperation() else {
             throw Smithy.ClientError.dataNotFound("Operation name not configured in middleware context for auth scheme resolver params construction.")
         }
-        guard let endpointParam = context.attributes.get(key: Smithy.AttributeKey<EndpointParams>(name: "EndpointParams")) else {
+        guard let endpointParam = context.get(key: Smithy.AttributeKey<EndpointParams>(name: "EndpointParams")) else {
             throw Smithy.ClientError.dataNotFound("Endpoint param not configured in middleware context for rules-based auth scheme resolver params construction.")
         }
         return S3AuthSchemeResolverParameters(operation: opName)