feat(api): Add async protocol implementation for token request interceptors (#2870)

Author: royjit

Amplify/Categories/API/AuthProvider/APIAuthProviderFactory.swift

@@ -25,8 +25,20 @@ open class APIAuthProviderFactory {
 }
 
 public protocol AmplifyAuthTokenProvider {
+
     typealias AuthToken = String
+
     func getLatestAuthToken() -> Result<AuthToken, Error>
+
+    func getLatestAuthToken(completion: @escaping (Result<AuthToken, Error>) -> Void)
+}
+
+public extension AmplifyAuthTokenProvider {
+
+    func getLatestAuthToken(completion: @escaping (Result<AuthToken, Error>) -> Void) {
+        let result = getLatestAuthToken()
+        completion(result)
+    }
 }
 
 /// Amplify OIDC Auth Provider

AmplifyPlugins.podspec

@@ -36,7 +36,7 @@ Pod::Spec.new do |s|
 
   s.subspec 'AWSAPIPlugin' do |ss|
     ss.source_files = 'AmplifyPlugins/API/AWSAPICategoryPlugin/**/*.swift'
-    ss.dependency 'AppSyncRealTimeClient', "~> 3.0"
+    ss.dependency 'AppSyncRealTimeClient', "~> 3.1"
   end
 
   s.subspec 'AWSCognitoAuthPlugin' do |ss|

AmplifyPlugins/API/AWSAPICategoryPlugin/Interceptor/RequestInterceptor/AuthTokenProviderWrapper.swift

@@ -29,4 +29,19 @@ class AuthTokenProviderWrapper: AuthTokenProvider {
         }
     }
 
+    func getToken(completion: @escaping (Result<String, AuthError>) -> Void) {
+        wrappedAuthTokenProvider.getLatestAuthToken { result in
+            switch result {
+            case .success(let token):
+                completion(.success(token))
+                return
+            case .failure(let error):
+                completion(.failure(AuthError.service("Unable to get latest auth token",
+                                                      "",
+                                                      error)))
+                return
+            }
+        }
+    }
+
 }

AmplifyPlugins/API/AWSAPICategoryPlugin/Interceptor/SubscriptionInterceptor/AuthenticationTokenAuthInterceptor.swift

@@ -19,35 +19,50 @@ class AuthenticationTokenAuthInterceptor: AuthInterceptor {
         self.authTokenProvider = authTokenProvider
     }
 
-    func interceptMessage(_ message: AppSyncMessage, for endpoint: URL) -> AppSyncMessage {
-        let host = endpoint.host!
-        guard let authToken = getAuthToken() else {
-            log.warn("Missing authentication token for subscription")
-            return message
+    func interceptMessage(
+        _ message: AppSyncMessage,
+        for endpoint: URL,
+        completion: @escaping (AppSyncMessage) -> Void) {
+            let host = endpoint.host!
+
+            guard case .subscribe = message.messageType else {
+                completion(message)
+                return
+            }
+
+            authTokenProvider.getLatestAuthToken { result in
+                let signedMessage = self.signedMessage(
+                    message,
+                    host: host,
+                    tokenResult: result)
+                completion(signedMessage)
+                return
+            }
         }
 
-        guard case .subscribe = message.messageType else {
-            return message
+    func interceptConnection(
+        _ request: AppSyncConnectionRequest,
+        for endpoint: URL,
+        completion: @escaping (AppSyncConnectionRequest) -> Void) {
+            let host = endpoint.host!
+            authTokenProvider.getLatestAuthToken { result in
+                let signedRequest = self.signedRequest(
+                    request,
+                    for: endpoint,
+                    host: host,
+                    tokenResult: result)
+                completion(signedRequest)
+                return
+            }
         }
 
-        let authHeader = TokenAuthHeader(token: authToken, host: host)
-        var payload = message.payload ?? AppSyncMessage.Payload()
-        payload.authHeader = authHeader
-
-        let signedMessage = AppSyncMessage(
-            id: message.id,
-            payload: payload,
-            type: message.messageType
-        )
-        return signedMessage
-    }
-
-    func interceptConnection(
+    private func signedRequest(
         _ request: AppSyncConnectionRequest,
-        for endpoint: URL
+        for endpoint: URL,
+        host: String,
+        tokenResult: Result<String, Error>
     ) -> AppSyncConnectionRequest {
-        let host = endpoint.host!
-        guard let authToken = getAuthToken() else {
+        guard let authToken = try? tokenResult.get() else {
             log.warn("Missing authentication token for subscription request")
             return request
         }
@@ -71,9 +86,54 @@ class AuthenticationTokenAuthInterceptor: AuthInterceptor {
         return signedRequest
     }
 
-    private func getAuthToken() -> AmplifyAuthTokenProvider.AuthToken? {
-        try? authTokenProvider.getLatestAuthToken().get()
+    private func signedMessage(
+        _ message: AppSyncMessage,
+        host: String,
+        tokenResult: Result<String, Error>) -> AppSyncMessage {
+
+            guard let authToken = try? tokenResult.get() else {
+                log.warn("Missing authentication token for subscription")
+                return message
+            }
+
+            let authHeader = TokenAuthHeader(token: authToken, host: host)
+            var payload = message.payload ?? AppSyncMessage.Payload()
+            payload.authHeader = authHeader
+
+            let signedMessage = AppSyncMessage(
+                id: message.id,
+                payload: payload,
+                type: message.messageType
+            )
+            return signedMessage
+        }
+
+    func interceptMessage(_ message: AppSyncMessage, for endpoint: URL) -> AppSyncMessage {
+        let host = endpoint.host!
+
+        guard case .subscribe = message.messageType else {
+            return message
+        }
+        let signedMessage = self.signedMessage(
+            message,
+            host: host,
+            tokenResult: authTokenProvider.getLatestAuthToken())
+        return signedMessage
+    }
+
+    func interceptConnection(
+        _ request: AppSyncConnectionRequest,
+        for endpoint: URL
+    ) -> AppSyncConnectionRequest {
+        let host = endpoint.host!
+        let signedRequest = signedRequest(
+            request,
+            for: endpoint,
+            host: host,
+            tokenResult: authTokenProvider.getLatestAuthToken())
+        return signedRequest
     }
+
 }
 
 // MARK: AuthorizationTokenAuthInterceptor + DefaultLogger

AmplifyPlugins/API/AWSAPICategoryPlugin/Interceptor/SubscriptionInterceptor/OIDCAuthProviderWrapper.swift

@@ -20,4 +20,8 @@ class OIDCAuthProviderWrapper: OIDCAuthProvider {
     func getLatestAuthToken() -> Result<String, Error> {
         return authTokenProvider.getLatestAuthToken()
     }
+
+    func getLatestAuthToken(completion: @escaping (Result<String, Error>) -> Void ) {
+        authTokenProvider.getLatestAuthToken(completion: completion)
+    }
 }

AmplifyPlugins/API/AWSAPICategoryPluginTests/Configuration/AWSAPICategoryPluginConfigurationTests.swift

@@ -126,6 +126,10 @@ class AWSAPICategoryPluginConfigurationTests: XCTestCase {
         func getToken() -> Result<String, AuthError> {
             .success("token")
         }
+
+        func getToken(completion: @escaping (Result<String, AuthError>) -> Void) {
+            completion(.success("token"))
+        }
     }
 
 }

AmplifyPlugins/API/AWSAPICategoryPluginTests/Interceptor/AuthTokenURLRequestInterceptorTests.swift

@@ -37,5 +37,9 @@ extension AuthTokenURLRequestInterceptorTests {
         func getToken() -> Result<String, AuthError> {
             .success(authorizationToken)
         }
+
+        func getToken(completion: @escaping (Result<String, AuthError>) -> Void) {
+            completion(.success(authorizationToken))
+        }
     }
 }

AmplifyPlugins/API/Podfile

@@ -8,7 +8,7 @@ include_build_tools!
 target 'AWSAPICategoryPlugin' do
 	pod 'Amplify', :path => '../../'
   pod 'AWSPluginsCore', :path => '../../'
-  pod "AppSyncRealTimeClient", "~> 3.0"
+  pod "AppSyncRealTimeClient", "~> 3.1"
 
   target "AWSAPICategoryPluginTests" do
     inherit! :complete

AmplifyPlugins/API/Podfile.lock

@@ -54,7 +54,7 @@ DEPENDENCIES:
   - Amplify (from `../../`)
   - AmplifyPlugins/AWSCognitoAuthPlugin (from `../../`)
   - AmplifyTestCommon (from `../../`)
-  - AppSyncRealTimeClient (~> 3.0)
+  - AppSyncRealTimeClient (~> 3.1)
   - AWSPluginsCore (from `../../`)
   - CwlCatchException (= 2.1.1)
   - CwlPreconditionTesting (from `https://github.com/mattgallagher/CwlPreconditionTesting.git`, tag `2.1.0`)
@@ -115,6 +115,6 @@ SPEC CHECKSUMS:
   SwiftFormat: 3b5caa6389b2b9adbc00e133b3ccc8c6e687a6a4
   SwiftLint: 32ee33ded0636d0905ef6911b2b67bbaeeedafa5
 
-PODFILE CHECKSUM: 5170578806036f2ba018abb8868d56e448fb0ada
+PODFILE CHECKSUM: f542a7a3bf7a7e882e6ea4799c9722d5b6495208
 
-COCOAPODS: 1.11.3
+COCOAPODS: 1.12.0

AmplifyPlugins/Core/AWSPluginsCore/Auth/Provider/AuthTokenProvider.swift

@@ -9,7 +9,18 @@ import Foundation
 import Amplify
 
 public protocol AuthTokenProvider {
+
     func getToken() -> Result<String, AuthError>
+
+    func getToken(completion: @escaping (Result<String, AuthError>) -> Void)
+}
+
+public extension AuthTokenProvider {
+
+    func getToken(completion: @escaping (Result<String, AuthError>) -> Void) {
+        let result = getToken()
+        completion(result)
+    }
 }
 
 public struct BasicUserPoolTokenProvider: AuthTokenProvider {
@@ -23,4 +34,8 @@ public struct BasicUserPoolTokenProvider: AuthTokenProvider {
     public func getToken() -> Result<String, AuthError> {
         return authService.getToken()
     }
+
+    public func getToken(completion: @escaping (Result<String, AuthError>) -> Void) {
+        authService.getUserPoolAccessToken(completion: completion)
+    }
 }