Refactor eventstream message signing to accept any custom signer that…

… conforms to ClientRuntime.Signer protocol.
awslabs · Nov 14, 2023 · aa43d72 · aa43d72
1 parent 75fc81a
commit aa43d72
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 25 deletions.
diff --git a/Sources/Core/AWSClientRuntime/EventStream/AWSMessageSigner.swift b/Sources/Core/AWSClientRuntime/EventStream/AWSMessageSigner.swift
@@ -14,6 +14,9 @@ extension AWSEventStream {
         let signer: () async throws -> ClientRuntime.Signer
         let signingConfig: () async throws -> AWSSigningConfig
         let requestSignature: () -> String
+        // Attribute key used to save AWSSigningConfig into signingProperties argument
+        //  for AWSSigV4Signer::signEvent call that conforms to Signer::signEvent.
+        static let signingConfigKey = AttributeKey<AWSSigningConfig>(name: "EventStreamSigningConfig")
 
         private var _previousSignature: String?
 
@@ -52,11 +55,15 @@ extension AWSEventStream {
             // encode to bytes
             let encodedMessage = try encoder.encode(message: message)
             let signingConfig = try await self.signingConfig()
-            let sigv4signer = try await self.getSigV4Signer()
+            // Fetch signer
+            let signer = try await self.signer()
+            // Wrap config into signingProperties: Attributes
+            var configWrapper = Attributes()
+            configWrapper.set(key: AWSMessageSigner.signingConfigKey, value: signingConfig)
             // Sign encoded bytes
-            let signingResult = try await sigv4signer.signEvent(payload: encodedMessage,
-                                                                   previousSignature: previousSignature,
-                                                                   signingConfig: signingConfig)
+            let signingResult = try await signer.signEvent(payload: encodedMessage,
+                                                                previousSignature: previousSignature,
+                                                                signingProperties: configWrapper)
             previousSignature = signingResult.signature
             return signingResult.output
         }
@@ -65,22 +72,16 @@ extension AWSEventStream {
         /// - Returns: Signed `Message` with `:chunk-signature` & `:date` headers
         public func signEmpty() async throws -> ClientRuntime.EventStream.Message {
             let signingConfig = try await self.signingConfig()
-            let sigv4signer = try await self.getSigV4Signer()
-            let signingResult = try await sigv4signer.signEvent(payload: .init(),
-                                                                   previousSignature: previousSignature,
-                                                                   signingConfig: signingConfig)
-            return signingResult.output
-        }
-
-        private func getSigV4Signer() async throws -> AWSSigV4Signer {
-            // Fetch resolved signer from selected auth scheme in middleware context
+            // Fetch signer
             let signer = try await self.signer()
-            // For now, only AWSSigV4Signer is used for eventstream message signing
-            let sigv4signer = signer as? AWSClientRuntime.AWSSigV4Signer
-            guard let sigv4signer else {
-                throw ClientError.authError("Could not sign event stream message using configured signer, \(signer).")
-            }
-            return sigv4signer
+            // Wrap config into signingProperties: Attributes
+            var configWrapper = Attributes()
+            configWrapper.set(key: AWSMessageSigner.signingConfigKey, value: signingConfig)
+            // Sign empty payload
+            let signingResult = try await signer.signEvent(payload: .init(),
+                                                           previousSignature: previousSignature,
+                                                           signingProperties: configWrapper)
+            return signingResult.output
         }
     }
 }
diff --git a/Sources/Core/AWSClientRuntime/Signing/AWSSigV4Signer.swift b/Sources/Core/AWSClientRuntime/Signing/AWSSigV4Signer.swift
@@ -10,7 +10,7 @@ import ClientRuntime
 import Foundation
 
 public class AWSSigV4Signer: ClientRuntime.Signer {
-    public func sign<IdentityT: Identity>(
+    public func signRequest<IdentityT: Identity>(
         requestBuilder: SdkHttpRequestBuilder,
         identity: IdentityT,
         signingProperties: ClientRuntime.Attributes
@@ -94,6 +94,18 @@ public class AWSSigV4Signer: ClientRuntime.Signer {
         )
     }
 
+    public func signEvent(
+        payload: Data,
+        previousSignature: String,
+        signingProperties: Attributes
+    ) async throws -> SigningResult<EventStream.Message> {
+        let signingConfig = signingProperties.get(key: AWSEventStream.AWSMessageSigner.signingConfigKey)
+        guard let signingConfig else {
+            throw ClientError.dataNotFound("Failed to sign event stream message due to missing signing config.")
+        }
+        return try await signEvent(payload: payload, previousSignature: previousSignature, signingConfig: signingConfig)
+    }
+
     static let logger: SwiftLogger = SwiftLogger(label: "AWSSigV4Signer")
 
     public static func sigV4SignedURL(
@@ -183,8 +195,3 @@ public class AWSSigV4Signer: ClientRuntime.Signer {
         }
     }
 }
-
-public struct SigningResult<T> {
-    public let output: T
-    public let signature: String
-}