updated RefreshUserPoolTokens to use GetTokensFromRefreshToken API to enable refresh token rotation, also updated test mock clients and added unit tests

Author: ekjotmultani

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Actions/RefreshAuthorizationSession/UserPool/RefreshUserPoolTokens.swift

@@ -28,32 +28,34 @@ struct RefreshUserPoolTokens: Action {
                 return
             }
 
-            let authEnv = try environment.authEnvironment()
             let config = environment.userPoolConfiguration
             let client = try? environment.cognitoUserPoolFactory()
             let existingTokens = existingSignedIndata.cognitoUserPoolTokens
 
             let deviceMetadata = await DeviceMetadataHelper.getDeviceMetadata(
                 for: existingSignedIndata.username,
                 with: environment)
-
-            let asfDeviceId = try await CognitoUserPoolASF.asfDeviceID(
-                for: existingSignedIndata.username,
-                credentialStoreClient: authEnv.credentialsClient)
-
-            let input = await InitiateAuthInput.refreshAuthInput(
-                username: existingSignedIndata.username,
-                refreshToken: existingTokens.refreshToken,
+            
+            let deviceKey: String? = {
+                if case .metadata(let data) = deviceMetadata {
+                    return data.deviceKey
+                }
+                return nil
+            }()
+
+            let input = GetTokensFromRefreshTokenInput(
+                clientId: config.clientId,
                 clientMetadata: [:],
-                asfDeviceId: asfDeviceId,
-                deviceMetadata: deviceMetadata,
-                environment: environment)
+                clientSecret: config.clientSecret,
+                deviceKey: deviceKey,
+                refreshToken: existingTokens.refreshToken
+            )
 
-            logVerbose("\(#fileID) Starting initiate auth refresh token", environment: environment)
+            logVerbose("\(#fileID) Starting get tokens from refresh token", environment: environment)
 
-            let response = try await client?.initiateAuth(input: input)
+            let response = try await client?.getTokensFromRefreshToken(input: input)
 
-            logVerbose("\(#fileID) Initiate auth response received", environment: environment)
+            logVerbose("\(#fileID) Get tokens from refresh token response received", environment: environment)
 
             guard let authenticationResult = response?.authenticationResult,
                   let idToken = authenticationResult.idToken,
@@ -69,8 +71,7 @@ struct RefreshUserPoolTokens: Action {
             let userPoolTokens = AWSCognitoUserPoolTokens(
                 idToken: idToken,
                 accessToken: accessToken,
-                refreshToken: existingTokens.refreshToken,
-                expiresIn: authenticationResult.expiresIn
+                refreshToken: authenticationResult.refreshToken ?? existingTokens.refreshToken
             )
             let signedInData = SignedInData(
                 signedInDate: existingSignedIndata.signedInDate,
@@ -96,7 +97,7 @@ struct RefreshUserPoolTokens: Action {
             await dispatcher.send(event)
         }
 
-        logVerbose("\(#fileID) Initiate auth complete", environment: environment)
+        logVerbose("\(#fileID) Get tokens from refresh token complete", environment: environment)
     }
 }
 

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Service/CognitoUserPoolBehavior.swift

@@ -29,6 +29,9 @@ protocol CognitoUserPoolBehavior {
     /// Throws RevokeTokenOutputError
     func revokeToken(input: RevokeTokenInput) async throws -> RevokeTokenOutput
 
+    /// Throws GetTokensFromRefreshTokenOutputError
+    func getTokensFromRefreshToken(input: GetTokensFromRefreshTokenInput) async throws -> GetTokensFromRefreshTokenOutput
+
     // MARK: - User Attribute API's
 
     /// Throws GetUserAttributeVerificationCodeOutputError

AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/ActionTests/FetchAuthSession/FetchUserPoolTokens/RefreshUserPoolTokensTests.swift

@@ -45,8 +45,8 @@ class RefreshUserPoolTokensTests: XCTestCase {
         let expectation = expectation(description: "refreshUserPoolTokens")
         let identityProviderFactory: BasicSRPAuthEnvironment.CognitoUserPoolFactory = {
             MockIdentityProvider(
-                mockInitiateAuthResponse: { _ in
-                    return InitiateAuthOutput()
+                mockGetTokensFromRefreshTokenResponse: { _ in
+                    return GetTokensFromRefreshTokenOutput()
                 }
             )
         }
@@ -77,8 +77,8 @@ class RefreshUserPoolTokensTests: XCTestCase {
         let expectation = expectation(description: "refreshUserPoolTokens")
         let identityProviderFactory: BasicSRPAuthEnvironment.CognitoUserPoolFactory = {
             MockIdentityProvider(
-                mockInitiateAuthResponse: { _ in
-                    return InitiateAuthOutput(
+                mockGetTokensFromRefreshTokenResponse: { _ in
+                    return GetTokensFromRefreshTokenOutput(
                         authenticationResult: .init(
                             accessToken: "accessTokenNew",
                             expiresIn: 100,
@@ -114,7 +114,7 @@ class RefreshUserPoolTokensTests: XCTestCase {
 
         let identityProviderFactory: BasicSRPAuthEnvironment.CognitoUserPoolFactory = {
             MockIdentityProvider(
-                mockInitiateAuthResponse: { _ in
+                mockGetTokensFromRefreshTokenResponse: { _ in
                     throw testError
                 }
             )
@@ -144,4 +144,74 @@ class RefreshUserPoolTokensTests: XCTestCase {
         )
     }
 
+    func testRefreshTokenRotation() async {
+
+        let expectation = expectation(description: "refreshTokenRotation")
+        let identityProviderFactory: BasicSRPAuthEnvironment.CognitoUserPoolFactory = {
+            MockIdentityProvider(
+                mockGetTokensFromRefreshTokenResponse: { _ in
+                    return GetTokensFromRefreshTokenOutput(
+                        authenticationResult: .init(
+                            accessToken: "accessTokenNew",
+                            expiresIn: 100,
+                            idToken: "idTokenNew",
+                            refreshToken: "refreshTokenRotated"))
+                }
+            )
+        }
+
+        let action = RefreshUserPoolTokens(existingSignedIndata: .testData)
+
+        await action.execute(withDispatcher: MockDispatcher { event in
+
+            if let userPoolEvent = event as? RefreshSessionEvent,
+               case let .refreshIdentityInfo(signedInData, _) = userPoolEvent.eventType {
+                XCTAssertEqual(signedInData.cognitoUserPoolTokens.refreshToken, "refreshTokenRotated")
+                expectation.fulfill()
+            }
+        }, environment: Defaults.makeDefaultAuthEnvironment(
+            userPoolFactory: identityProviderFactory)
+        )
+
+        await fulfillment(
+            of: [expectation],
+            timeout: 0.1
+        )
+    }
+
+    func testRefreshTokenNoRotation() async {
+
+        let expectation = expectation(description: "refreshTokenNoRotation")
+        let identityProviderFactory: BasicSRPAuthEnvironment.CognitoUserPoolFactory = {
+            MockIdentityProvider(
+                mockGetTokensFromRefreshTokenResponse: { _ in
+                    return GetTokensFromRefreshTokenOutput(
+                        authenticationResult: .init(
+                            accessToken: "accessTokenNew",
+                            expiresIn: 100,
+                            idToken: "idTokenNew",
+                            refreshToken: nil))
+                }
+            )
+        }
+
+        let action = RefreshUserPoolTokens(existingSignedIndata: .testData)
+
+        await action.execute(withDispatcher: MockDispatcher { event in
+
+            if let userPoolEvent = event as? RefreshSessionEvent,
+               case let .refreshIdentityInfo(signedInData, _) = userPoolEvent.eventType {
+                XCTAssertEqual(signedInData.cognitoUserPoolTokens.refreshToken, "refreshToken")
+                expectation.fulfill()
+            }
+        }, environment: Defaults.makeDefaultAuthEnvironment(
+            userPoolFactory: identityProviderFactory)
+        )
+
+        await fulfillment(
+            of: [expectation],
+            timeout: 0.1
+        )
+    }
+
 }

AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/Support/MockIdentityProvider.swift

@@ -20,6 +20,9 @@ struct MockIdentityProvider: CognitoUserPoolBehavior {
     typealias MockInitiateAuthResponse = (InitiateAuthInput) async throws
     -> InitiateAuthOutput
 
+    typealias MockGetTokensFromRefreshTokenResponse = (GetTokensFromRefreshTokenInput) async throws
+    -> GetTokensFromRefreshTokenOutput
+
     typealias MockConfirmSignUpResponse = (ConfirmSignUpInput) async throws
     -> ConfirmSignUpOutput
 
@@ -88,6 +91,7 @@ struct MockIdentityProvider: CognitoUserPoolBehavior {
     let mockSignUpResponse: MockSignUpResponse?
     let mockRevokeTokenResponse: MockRevokeTokenResponse?
     let mockInitiateAuthResponse: MockInitiateAuthResponse?
+    let mockGetTokensFromRefreshTokenResponse: MockGetTokensFromRefreshTokenResponse?
     let mockGlobalSignOutResponse: MockGlobalSignOutResponse?
     let mockConfirmSignUpResponse: MockConfirmSignUpResponse?
     let mockRespondToAuthChallengeResponse: MockRespondToAuthChallengeResponse?
@@ -116,6 +120,7 @@ struct MockIdentityProvider: CognitoUserPoolBehavior {
         mockSignUpResponse: MockSignUpResponse? = nil,
         mockRevokeTokenResponse: MockRevokeTokenResponse? = nil,
         mockInitiateAuthResponse: MockInitiateAuthResponse? = nil,
+        mockGetTokensFromRefreshTokenResponse: MockGetTokensFromRefreshTokenResponse? = nil,
         mockGlobalSignOutResponse: MockGlobalSignOutResponse? = nil,
         mockConfirmSignUpResponse: MockConfirmSignUpResponse? = nil,
         mockRespondToAuthChallengeResponse: MockRespondToAuthChallengeResponse? = nil,
@@ -139,6 +144,7 @@ struct MockIdentityProvider: CognitoUserPoolBehavior {
         self.mockSignUpResponse = mockSignUpResponse
         self.mockRevokeTokenResponse = mockRevokeTokenResponse
         self.mockInitiateAuthResponse = mockInitiateAuthResponse
+        self.mockGetTokensFromRefreshTokenResponse = mockGetTokensFromRefreshTokenResponse
         self.mockGlobalSignOutResponse = mockGlobalSignOutResponse
         self.mockConfirmSignUpResponse = mockConfirmSignUpResponse
         self.mockRespondToAuthChallengeResponse = mockRespondToAuthChallengeResponse
@@ -192,6 +198,11 @@ struct MockIdentityProvider: CognitoUserPoolBehavior {
         return try await mockRevokeTokenResponse!(input)
     }
 
+    /// Throws GetTokensFromRefreshTokenOutputError
+    func getTokensFromRefreshToken(input: GetTokensFromRefreshTokenInput) async throws -> GetTokensFromRefreshTokenOutput {
+        return try await mockGetTokensFromRefreshTokenResponse!(input)
+    }
+
     func getUserAttributeVerificationCode(input: GetUserAttributeVerificationCodeInput) async throws -> GetUserAttributeVerificationCodeOutput {
         return try await mockGetUserAttributeVerificationCodeOutput!(input)
     }

Package.resolved

@@ -9,7 +9,7 @@ let platforms: [SupportedPlatform] = [
     .watchOS(.v9)
 ]
 let dependencies: [Package.Dependency] = [
-    .package(url: "https://github.com/awslabs/aws-sdk-swift", exact: "1.5.14"),
+    .package(url: "https://github.com/awslabs/aws-sdk-swift", exact: "1.5.18"),
     .package(url: "https://github.com/stephencelis/SQLite.swift.git", exact: "0.15.3"),
     .package(url: "https://github.com/mattgallagher/CwlPreconditionTesting.git", from: "2.1.0"),
     .package(url: "https://github.com/aws-amplify/amplify-swift-utils-notifications.git", from: "1.1.0")