Skip to content

Commit 8536147

Browse files
harsh62harsh62
authored
fix(auth): caching session during error transitions while auto signin (#4094)
* fix(auth): caching session during error transitions while auto signing in * update test * add unit tests * fix swift format errors * Update AWSAuthSignUpAPITests.swift
1 parent 2fe2727 commit 8536147

File tree

11 files changed

+426
-24
lines changed

11 files changed

+426
-24
lines changed

‎AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Actions/SignUp/ConfirmSignUp.swift‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,15 +49,15 @@ struct ConfirmSignUp: Action {
4949
await dispatcher.send(SignUpEvent(eventType: .signedUp(dataToSend, .init(.done))))
5050
}
5151
} catch let error as SignUpError {
52-
let errorEvent = SignUpEvent(eventType: .throwAuthError(error))
52+
let errorEvent = SignUpEvent(eventType: .throwAuthError(error, data))
5353
logVerbose(
5454
"\(#fileID) Sending event \(errorEvent)",
5555
environment: environment
5656
)
5757
await dispatcher.send(errorEvent)
5858
} catch {
5959
let error = SignUpError.service(error: error)
60-
let errorEvent = SignUpEvent(eventType: .throwAuthError(error))
60+
let errorEvent = SignUpEvent(eventType: .throwAuthError(error, data))
6161
logVerbose(
6262
"\(#fileID) Sending event \(errorEvent)",
6363
environment: environment

‎AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Actions/SignUp/InitiateSignUp.swift‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,15 +63,15 @@ struct InitiateSignUp: Action {
6363
}
6464
await dispatcher.send(event)
6565
} catch let error as SignUpError {
66-
let errorEvent = SignUpEvent(eventType: .throwAuthError(error))
66+
let errorEvent = SignUpEvent(eventType: .throwAuthError(error, data))
6767
logVerbose(
6868
"\(#fileID) Sending event \(errorEvent)",
6969
environment: environment
7070
)
7171
await dispatcher.send(errorEvent)
7272
} catch {
7373
let error = SignUpError.service(error: error)
74-
let errorEvent = SignUpEvent(eventType: .throwAuthError(error))
74+
let errorEvent = SignUpEvent(eventType: .throwAuthError(error, data))
7575
logVerbose(
7676
"\(#fileID) Sending event \(errorEvent)",
7777
environment: environment

‎AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/Events/SignUpEvent.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ struct SignUpEvent: StateMachineEvent {
2121
case initiateSignUpComplete(SignUpEventData, AuthSignUpResult)
2222
case confirmSignUp(SignUpEventData, ConfirmationCode, ForceAliasCreation?)
2323
case signedUp(SignUpEventData, AuthSignUpResult)
24-
case throwAuthError(SignUpError)
24+
case throwAuthError(SignUpError, SignUpEventData)
2525
}
2626

2727
init(

‎AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/States/SignUpState.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ enum SignUpState: State {
1414
case awaitingUserConfirmation(SignUpEventData, AuthSignUpResult)
1515
case confirmingSignUp(SignUpEventData)
1616
case signedUp(SignUpEventData, AuthSignUpResult)
17-
case error(SignUpError)
17+
case error(SignUpError, SignUpEventData)
1818
}
1919

2020
extension SignUpState {

‎AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/Resolvers/SignUp/SignUpState+Resolver.swift‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -46,8 +46,8 @@ extension SignUpState {
4646
case .confirmSignUp(let data, let code, let forceAliasCreation):
4747
let action = ConfirmSignUp(data: data, confirmationCode: code, forceAliasCreation: forceAliasCreation)
4848
return .init(newState: .confirmingSignUp(data), actions: [action])
49-
case .throwAuthError(let error):
50-
return .init(newState: .error(error))
49+
case .throwAuthError(let error, let signUpData):
50+
return .init(newState: .error(error, signUpData))
5151
default:
5252
return .from(oldState)
5353
}
@@ -84,8 +84,8 @@ extension SignUpState {
8484
return .init(newState: .confirmingSignUp(data), actions: [action])
8585
case .signedUp(let data, let result):
8686
return .init(newState: .signedUp(data, result))
87-
case .throwAuthError(let error):
88-
return .init(newState: .error(error))
87+
case .throwAuthError(let error, let signUpData):
88+
return .init(newState: .error(error, signUpData))
8989
}
9090
}
9191

@@ -100,8 +100,8 @@ extension SignUpState {
100100
case .confirmSignUp(let data, let code, let forceAliasCreation):
101101
let action = ConfirmSignUp(data: data, confirmationCode: code, forceAliasCreation: forceAliasCreation)
102102
return .init(newState: .confirmingSignUp(data), actions: [action])
103-
case .throwAuthError(let error):
104-
return .init(newState: .error(error))
103+
case .throwAuthError(let error, let signUpData):
104+
return .init(newState: .error(error, signUpData))
105105
default:
106106
return .from(oldState)
107107
}
@@ -120,8 +120,8 @@ extension SignUpState {
120120
return .init(newState: .confirmingSignUp(data), actions: [action])
121121
case .signedUp(let data, let result):
122122
return .init(newState: .signedUp(data, result))
123-
case .throwAuthError(let error):
124-
return .init(newState: .error(error))
123+
case .throwAuthError(let error, let signUpData):
124+
return .init(newState: .error(error, signUpData))
125125
default:
126126
return .from(oldState)
127127
}
@@ -138,8 +138,8 @@ extension SignUpState {
138138
case .confirmSignUp(let data, let code, let forceAliasCreation):
139139
let action = ConfirmSignUp(data: data, confirmationCode: code, forceAliasCreation: forceAliasCreation)
140140
return .init(newState: .confirmingSignUp(data), actions: [action])
141-
case .throwAuthError(let error):
142-
return .init(newState: .error(error))
141+
case .throwAuthError(let error, let signUpData):
142+
return .init(newState: .error(error, signUpData))
143143
default:
144144
return .from(oldState)
145145
}

‎AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Task/AWSAuthConfirmSignUpTask.swift‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ class AWSAuthConfirmSignUpTask: AuthConfirmSignUpTask, DefaultLogger {
5858
switch signUpState {
5959
case .signedUp(_, let result):
6060
return result
61-
case .error(let signUpError):
61+
case .error(let signUpError, _):
6262
throw signUpError.authError
6363
default:
6464
continue
@@ -81,6 +81,11 @@ class AWSAuthConfirmSignUpTask: AuthConfirmSignUpTask, DefaultLogger {
8181
// only include session if the cached username matches
8282
// the username in confirmSignUp() call
8383
session = data.session
84+
} else if case .error(_, let data) = signUpState,
85+
request.username == data.username {
86+
// only include session if the cached username matches
87+
// the username in confirmSignUp() call
88+
session = data.session
8489
}
8590

8691
let pluginOptions = request.options.pluginOptions as? AWSAuthConfirmSignUpOptions

‎AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Task/AWSAuthSignUpTask.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ class AWSAuthSignUpTask: AuthSignUpTask, DefaultLogger {
6262
return result
6363
case .signedUp(_, let result):
6464
return result
65-
case .error(let signUpError):
65+
case .error(let signUpError, _):
6666
throw signUpError.authError
6767
default:
6868
continue

‎AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/TaskTests/ClientBehaviorTests/SignIn/AWSAuthAutoSignInTests.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -454,7 +454,7 @@ class AWSAuthAutoSignInTests: BasePluginTest {
454454
let initialStateError = AuthState.configured(
455455
.signedOut(.init(lastKnownUserName: nil)),
456456
.configured,
457-
.error(.service(error: AuthError.service("Unknown error", "Unknown error")))
457+
.error(.service(error: AuthError.service("Unknown error", "Unknown error")), .init(username: "username", session: "sessio"))
458458
)
459459

460460
let authPluginError = configureCustomPluginWith(

‎AmplifyPlugins/Auth/Tests/AWSCognitoAuthPluginUnitTests/TaskTests/ClientBehaviorTests/SignUp/AWSAuthConfirmSignUpAPITests.swift‎

Lines changed: 231 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,10 @@ class AWSAuthConfirmSignUpAPITests: BasePluginTest {
140140
let initialStateError = AuthState.configured(
141141
.signedOut(.init(lastKnownUserName: nil)),
142142
.configured,
143-
.error(.service(error: AuthError.service("Unknown error", "Unknown error")))
143+
.error(
144+
.service(error: AuthError.service("Unknown error", "Unknown error")),
145+
.init(username: "username", session: "sessio")
146+
)
144147
)
145148

146149
let authPluginError = configureCustomPluginWith(
@@ -280,7 +283,10 @@ class AWSAuthConfirmSignUpAPITests: BasePluginTest {
280283
let initialStateError = AuthState.configured(
281284
.signedOut(.init(lastKnownUserName: nil)),
282285
.configured,
283-
.error(.service(error: AuthError.service("Unknown error", "Unknown error")))
286+
.error(
287+
.service(error: AuthError.service("Unknown error", "Unknown error")),
288+
.init(username: "username", session: "sessio")
289+
)
284290
)
285291

286292
let authPluginError = configureCustomPluginWith(
@@ -595,4 +601,227 @@ class AWSAuthConfirmSignUpAPITests: BasePluginTest {
595601
XCTAssertEqual(awsCognitoAuthError, expectedCognitoError)
596602
}
597603
}
604+
605+
// MARK: - Session Caching Tests
606+
607+
/// Given: Configured auth machine in `.error` sign up state with a cached session for matching username
608+
/// When: `Auth.confirmSignUp(for:confirmationCode:options:)` is invoked with the same username
609+
/// Then: The cached session is used and confirm sign up completes with `.completeAutoSignIn`
610+
func testConfirmSignUpUsesSessionFromErrorStateWithMatchingUsername() async throws {
611+
let cachedSession = "cached-session-from-error"
612+
let mockIdentityProvider = MockIdentityProvider(
613+
mockConfirmSignUpResponse: { request in
614+
// Verify the cached session is passed through
615+
XCTAssertEqual(request.session, cachedSession)
616+
return .init(session: "new-session")
617+
}
618+
)
619+
620+
let initialStateError = AuthState.configured(
621+
.signedOut(.init(lastKnownUserName: nil)),
622+
.configured,
623+
.error(
624+
.service(error: AuthError.service("Previous error", "Recovery")),
625+
.init(username: "jeffb", session: cachedSession)
626+
)
627+
)
628+
629+
let authPluginError = configureCustomPluginWith(
630+
userPool: { mockIdentityProvider },
631+
initialState: initialStateError
632+
)
633+
634+
let result = try await authPluginError.confirmSignUp(
635+
for: "jeffb",
636+
confirmationCode: "123456",
637+
options: options
638+
)
639+
640+
guard case .completeAutoSignIn(let session) = result.nextStep else {
641+
XCTFail("Result should be .completeAutoSignIn for next step")
642+
return
643+
}
644+
XCTAssertTrue(result.isSignUpComplete, "Sign up result should be complete")
645+
XCTAssertEqual(session, "new-session")
646+
}
647+
648+
/// Given: Configured auth machine in `.error` sign up state with a cached session for different username
649+
/// When: `Auth.confirmSignUp(for:confirmationCode:options:)` is invoked with a different username
650+
/// Then: The cached session is NOT used
651+
func testConfirmSignUpDoesNotUseSessionFromErrorStateWithDifferentUsername() async throws {
652+
let cachedSession = "cached-session-from-error"
653+
let mockIdentityProvider = MockIdentityProvider(
654+
mockConfirmSignUpResponse: { request in
655+
// Verify the cached session is NOT passed through
656+
XCTAssertNil(request.session)
657+
return .init(session: "new-session")
658+
}
659+
)
660+
661+
let initialStateError = AuthState.configured(
662+
.signedOut(.init(lastKnownUserName: nil)),
663+
.configured,
664+
.error(
665+
.service(error: AuthError.service("Previous error", "Recovery")),
666+
.init(username: "different-user", session: cachedSession)
667+
)
668+
)
669+
670+
let authPluginError = configureCustomPluginWith(
671+
userPool: { mockIdentityProvider },
672+
initialState: initialStateError
673+
)
674+
675+
let result = try await authPluginError.confirmSignUp(
676+
for: "jeffb",
677+
confirmationCode: "123456",
678+
options: options
679+
)
680+
681+
guard case .completeAutoSignIn(let session) = result.nextStep else {
682+
XCTFail("Result should be .completeAutoSignIn for next step")
683+
return
684+
}
685+
XCTAssertTrue(result.isSignUpComplete, "Sign up result should be complete")
686+
XCTAssertEqual(session, "new-session")
687+
}
688+
689+
/// Given: Configured auth machine in `.error` sign up state with nil session
690+
/// When: `Auth.confirmSignUp(for:confirmationCode:options:)` is invoked
691+
/// Then: No session is passed and confirm sign up completes successfully
692+
func testConfirmSignUpFromErrorStateWithNilSession() async throws {
693+
let mockIdentityProvider = MockIdentityProvider(
694+
mockConfirmSignUpResponse: { request in
695+
XCTAssertNil(request.session)
696+
return .init()
697+
}
698+
)
699+
700+
let initialStateError = AuthState.configured(
701+
.signedOut(.init(lastKnownUserName: nil)),
702+
.configured,
703+
.error(
704+
.service(error: AuthError.service("Previous error", "Recovery")),
705+
.init(username: "jeffb", session: nil)
706+
)
707+
)
708+
709+
let authPluginError = configureCustomPluginWith(
710+
userPool: { mockIdentityProvider },
711+
initialState: initialStateError
712+
)
713+
714+
let result = try await authPluginError.confirmSignUp(
715+
for: "jeffb",
716+
confirmationCode: "123456",
717+
options: options
718+
)
719+
720+
guard case .done = result.nextStep else {
721+
XCTFail("Result should be .done for next step")
722+
return
723+
}
724+
XCTAssertTrue(result.isSignUpComplete, "Sign up result should be complete")
725+
}
726+
727+
/// Given: Configured auth machine in `.awaitingUserConfirmation` with a session
728+
/// When: `Auth.confirmSignUp(for:confirmationCode:options:)` is invoked with matching username
729+
/// Then: The cached session is used for auto sign-in
730+
func testConfirmSignUpUsesSessionFromAwaitingUserConfirmationState() async throws {
731+
let cachedSession = "cached-session-awaiting"
732+
let mockIdentityProvider = MockIdentityProvider(
733+
mockConfirmSignUpResponse: { request in
734+
XCTAssertEqual(request.session, cachedSession)
735+
return .init(session: "new-session")
736+
}
737+
)
738+
739+
let initialState = AuthState.configured(
740+
.signedOut(.init(lastKnownUserName: nil)),
741+
.configured,
742+
.awaitingUserConfirmation(
743+
SignUpEventData(username: "jeffb", session: cachedSession),
744+
.init(.confirmUser())
745+
)
746+
)
747+
748+
let authPlugin = configureCustomPluginWith(
749+
userPool: { mockIdentityProvider },
750+
initialState: initialState
751+
)
752+
753+
let result = try await authPlugin.confirmSignUp(
754+
for: "jeffb",
755+
confirmationCode: "123456",
756+
options: options
757+
)
758+
759+
guard case .completeAutoSignIn(let session) = result.nextStep else {
760+
XCTFail("Result should be .completeAutoSignIn for next step")
761+
return
762+
}
763+
XCTAssertTrue(result.isSignUpComplete, "Sign up result should be complete")
764+
XCTAssertEqual(session, "new-session")
765+
}
766+
767+
/// Given: Configured auth machine transitions to error state during confirm sign up
768+
/// When: The error state contains session data
769+
/// Then: The session is preserved in the error state for potential retry
770+
func testSessionPreservedInErrorStateDuringConfirmSignUp() async throws {
771+
let initialSession = "initial-session"
772+
let mockIdentityProvider = MockIdentityProvider(
773+
mockConfirmSignUpResponse: { _ in
774+
throw AWSCognitoIdentityProvider.CodeMismatchException()
775+
}
776+
)
777+
778+
let initialState = AuthState.configured(
779+
.signedOut(.init(lastKnownUserName: nil)),
780+
.configured,
781+
.awaitingUserConfirmation(
782+
SignUpEventData(username: "jeffb", session: initialSession),
783+
.init(.confirmUser())
784+
)
785+
)
786+
787+
let authPlugin = configureCustomPluginWith(
788+
userPool: { mockIdentityProvider },
789+
initialState: initialState
790+
)
791+
792+
do {
793+
_ = try await authPlugin.confirmSignUp(
794+
for: "jeffb",
795+
confirmationCode: "wrong-code",
796+
options: options
797+
)
798+
XCTFail("Should throw error")
799+
} catch {
800+
guard let authError = error as? AuthError else {
801+
XCTFail("Should throw Auth error")
802+
return
803+
}
804+
805+
guard case .service = authError else {
806+
XCTFail("Auth error should be of type service error")
807+
return
808+
}
809+
810+
// Verify the state machine is now in error state
811+
let currentState = await authPlugin.authStateMachine.currentState
812+
guard case .configured(_, _, let signUpState) = currentState else {
813+
XCTFail("Should be in configured state")
814+
return
815+
}
816+
817+
guard case .error(_, let data) = signUpState else {
818+
XCTFail("Should be in error state")
819+
return
820+
}
821+
822+
// Verify session is preserved in error state
823+
XCTAssertEqual(data.session, initialSession)
824+
XCTAssertEqual(data.username, "jeffb")
825+
}
826+
}
598827
}

0 commit comments

Comments
 (0)