Updating JIT preverified flow (#2611)

spetrescu84 · nilo-ms · web-flow · commit 50420323ad3d · 2025-04-29T11:52:03.000+01:00
* Updating JIT preverified flow

* Update MSAL/src/native_auth/controllers/jit/MSALNativeAuthJITController.swift

Co-authored-by: Danilo Raspa &lt;105228698+nilo-ms@users.noreply.github.com&gt;

---------

Co-authored-by: Danilo Raspa &lt;105228698+nilo-ms@users.noreply.github.com&gt;
diff --git a/MSAL/src/native_auth/controllers/jit/MSALNativeAuthJITController.swift b/MSAL/src/native_auth/controllers/jit/MSALNativeAuthJITController.swift
@@ -277,11 +277,19 @@ final class MSALNativeAuthJITController: MSALNativeAuthBaseController, MSALNativ
                                                          telemetryId: .telemetryApiISignInAfterJIT,
                                                          context: context)
             switch response.result {
-            case .success(let account):
+            case .completed(let account):
                 return .init(.completed(account), correlationId: context.correlationId(), telemetryUpdate: { [weak self] result in
                     self?.stopTelemetryEvent(signInEvent, context: context, delegateDispatcherResult: result)
                 })
-            case .failure(let error):
+            case .jitAuthMethodsSelectionRequired(_, _):
+                return .init(.error(error: .init(type: .generalError,
+                                                 message: "Unexpected result received when trying to signIn: strong authentication method registration required.",
+                                                 correlationId: context.correlationId(),
+                                                 errorCodes: [],
+                                                 errorUri: nil),
+                                    newState: nil),
+                             correlationId: context.correlationId())
+            case .error(let error):
                 return .init(.error(error: .init(type: .generalError,
                                                  message: error.errorDescription,
                                                  correlationId: error.correlationId,
diff --git a/MSAL/src/native_auth/controllers/responses/SignInResults.swift b/MSAL/src/native_auth/controllers/responses/SignInResults.swift
@@ -46,3 +46,9 @@ enum SignInVerifyCodeResult {
     case completed(MSALNativeAuthUserAccountResult)
     case error(error: VerifyCodeError, newState: SignInCodeRequiredState?)
 }
+
+enum SignInAfterPreviousFlowResult {
+    case completed(MSALNativeAuthUserAccountResult)
+    case jitAuthMethodsSelectionRequired(authMethods: [MSALAuthMethod], newState: RegisterStrongAuthState)
+    case error(error: MSALNativeAuthError)
+}
diff --git a/MSAL/src/native_auth/controllers/sign_in/MSALNativeAuthSignInController.swift b/MSAL/src/native_auth/controllers/sign_in/MSALNativeAuthSignInController.swift
@@ -128,7 +128,7 @@ final class MSALNativeAuthSignInController: MSALNativeAuthTokenController, MSALN
                                      format: "SignIn after previous flow not available because continuationToken is nil")
             let error = SignInAfterSignUpError(message: MSALNativeAuthErrorMessage.signInNotAvailable, correlationId: context.correlationId())
             stopTelemetryEvent(telemetryInfo, error: error)
-            return .init(.failure(error), correlationId: context.correlationId())
+            return .init(.error(error: error), correlationId: context.correlationId())
         }
         let scopes = joinScopes(scopes)
         guard let request = createTokenRequest(
@@ -141,7 +141,7 @@ final class MSALNativeAuthSignInController: MSALNativeAuthTokenController, MSALN
         ) else {
             let error = SignInAfterSignUpError(correlationId: context.correlationId())
             stopTelemetryEvent(telemetryInfo, error: error)
-            return .init(.failure(error), correlationId: context.correlationId())
+            return .init(.error(error: error), correlationId: context.correlationId())
         }
         let response = await performAndValidateTokenRequest(request, context: context)
         let result = await handleTokenResponse(response,
@@ -150,56 +150,30 @@ final class MSALNativeAuthSignInController: MSALNativeAuthTokenController, MSALN
                                                telemetryInfo: telemetryInfo)
         switch result {
         case .success(let accountResult):
-        return .init(.success(accountResult), correlationId: context.correlationId(), telemetryUpdate: { [weak self] result in
+            return .init(.completed(accountResult), correlationId: context.correlationId(), telemetryUpdate: { [weak self] result in
                             self?.stopTelemetryEvent(telemetryInfo.event, context: context, delegateDispatcherResult: result)
                         })
         case .awaitingMFA(_):
             let error = SignInAfterSignUpError(correlationId: context.correlationId())
             MSALNativeAuthLogger.log(level: .error, context: context, format: "SignIn: received unexpected MFA required API result")
             self.stopTelemetryEvent(telemetryInfo.event, context: context, error: error)
-            return .init(.failure(error), correlationId: context.correlationId())
+            return .init(.error(error: error), correlationId: context.correlationId())
         case .jitAuthMethodsSelectionRequired(let authMethods, let jitRequiredState):
             MSALNativeAuthLogger.log(level: .info, context: context, format: "JIT required after sing in after previous flow")
-            let jitController = createJITController()
-            guard let authMethod = authMethods.first else {
-                let error = SignInAfterSignUpError(correlationId: context.correlationId())
-                MSALNativeAuthLogger.log(level: .error, context: context, format: "JIT required, did not receive any default methods")
-                self.stopTelemetryEvent(telemetryInfo.event, context: context, error: error)
-                return .init(.failure(error), correlationId: context.correlationId())
-            }
-            let jitChallengeResponse = await jitController.requestJITChallenge(
-                continuationToken: jitRequiredState.continuationToken,
-                authMethod: authMethod,
-                verificationContact: nil,
-                context: context)
-            switch jitChallengeResponse.result {
-            case .completed(let accountResult):
-                return .init(.success(accountResult), correlationId: context.correlationId(), telemetryUpdate: { [weak self] result in
+            return .init(
+                .jitAuthMethodsSelectionRequired(authMethods: authMethods, newState: jitRequiredState),
+                correlationId: context.correlationId(),
+                telemetryUpdate: { [weak self] result in
                     self?.stopTelemetryEvent(telemetryInfo.event, context: context, delegateDispatcherResult: result)
                 })
-            case .verificationRequired(_, _, _, _):
-                let error = SignInAfterSignUpError(correlationId: context.correlationId())
-                MSALNativeAuthLogger.log(level: .error,
-                                  context: context,
-                                  format: "Request JIT challenge, received verification required on SignInAfterPreviousFlow")
-                self.stopTelemetryEvent(telemetryInfo.event, context: context, error: error)
-                return .init(.failure(error), correlationId: context.correlationId())
-            case .error(let apiError, _):
-                let error = SignInAfterSignUpError(correlationId: context.correlationId())
-                MSALNativeAuthLogger.logPII(level: .error,
-                                  context: context,
-                                  format: "Request JIT challenge, received invalid response \(MSALLogMask.maskPII(apiError.errorDescription))")
-                self.stopTelemetryEvent(telemetryInfo.event, context: context, error: error)
-                return .init(.failure(error), correlationId: context.correlationId())
-            }
         case .error(let error):
             let error = SignInAfterSignUpError(
                         message: error.errorDescription,
                         correlationId: error.correlationId,
                         errorCodes: error.errorCodes,
                         errorUri: error.errorUri
                     )
-                    return .init(.failure(error), correlationId: context.correlationId())
+                    return .init(.error(error: error), correlationId: context.correlationId())
         }
     }
 
diff --git a/MSAL/src/native_auth/controllers/sign_in/MSALNativeAuthSignInControlling.swift b/MSAL/src/native_auth/controllers/sign_in/MSALNativeAuthSignInControlling.swift
@@ -28,7 +28,7 @@ protocol MSALNativeAuthSignInControlling {
 
     typealias SignInControllerResponse = MSALNativeAuthControllerTelemetryWrapper<SignInStartResult>
     typealias SignInAfterPreviousFlowControllerResponse =
-        MSALNativeAuthControllerTelemetryWrapper<Result<MSALNativeAuthUserAccountResult, MSALNativeAuthError>>
+        MSALNativeAuthControllerTelemetryWrapper<SignInAfterPreviousFlowResult>
     typealias SignInSubmitCodeControllerResponse = MSALNativeAuthControllerTelemetryWrapper<SignInVerifyCodeResult>
     typealias SignInSubmitPasswordControllerResponse = MSALNativeAuthControllerTelemetryWrapper<SignInPasswordRequiredResult>
     typealias SignInResendCodeControllerResponse = MSALNativeAuthControllerTelemetryWrapper<SignInResendCodeResult>
diff --git a/MSAL/src/native_auth/public/state_machine/delegate_dispatcher/SignInAfterResetPasswordDelegateDispatcher.swift b/MSAL/src/native_auth/public/state_machine/delegate_dispatcher/SignInAfterResetPasswordDelegateDispatcher.swift
@@ -36,4 +36,18 @@ final class SignInAfterResetPasswordDelegateDispatcher: DelegateDispatcher<SignI
             await delegate.onSignInAfterResetPasswordError(error: error)
         }
     }
+
+    func dispatchJITRequired(authMethods: [MSALAuthMethod], newState: RegisterStrongAuthState, correlationId: UUID) async {
+        if let onSignInJITRequired = delegate.onSignInStrongAuthMethodRegistration {
+            telemetryUpdate?(.success(()))
+            await onSignInJITRequired(authMethods, newState)
+        } else {
+            let error = SignInAfterResetPasswordError(
+                message: requiredErrorMessage(for: "onSignInStrongAuthMethodRegistration"),
+                correlationId: correlationId
+            )
+            telemetryUpdate?(.failure(error))
+            await delegate.onSignInAfterResetPasswordError(error: error)
+        }
+    }
 }
diff --git a/MSAL/src/native_auth/public/state_machine/delegate_dispatcher/SignInAfterSignUpDelegateDispatcher.swift b/MSAL/src/native_auth/public/state_machine/delegate_dispatcher/SignInAfterSignUpDelegateDispatcher.swift
@@ -36,4 +36,18 @@ final class SignInAfterSignUpDelegateDispatcher: DelegateDispatcher<SignInAfterS
             await delegate.onSignInAfterSignUpError(error: error)
         }
     }
+
+    func dispatchJITRequired(authMethods: [MSALAuthMethod], newState: RegisterStrongAuthState, correlationId: UUID) async {
+        if let onSignInJITRequired = delegate.onSignInStrongAuthMethodRegistration {
+            telemetryUpdate?(.success(()))
+            await onSignInJITRequired(authMethods, newState)
+        } else {
+            let error = SignInAfterSignUpError(
+                message: requiredErrorMessage(for: "onSignInStrongAuthMethodRegistration"),
+                correlationId: correlationId
+            )
+            telemetryUpdate?(.failure(error))
+            await delegate.onSignInAfterSignUpError(error: error)
+        }
+    }
 }
diff --git a/MSAL/src/native_auth/public/state_machine/state/SignInAfterResetPasswordState.swift b/MSAL/src/native_auth/public/state_machine/state/SignInAfterResetPasswordState.swift
@@ -44,9 +44,13 @@ import Foundation
             )
 
             switch controllerResponse.result {
-            case .success(let accountResult):
+            case .completed(let accountResult):
                 await delegateDispatcher.dispatchSignInCompleted(result: accountResult, correlationId: controllerResponse.correlationId)
-            case .failure(let error):
+            case .jitAuthMethodsSelectionRequired(authMethods: let authMethods, newState: let newState):
+                await delegateDispatcher.dispatchJITRequired(authMethods: authMethods,
+                                                             newState: newState,
+                                                             correlationId: controllerResponse.correlationId)
+            case .error(let error):
                 let error = SignInAfterResetPasswordError(
                     message: error.errorDescription,
                     correlationId: error.correlationId,
diff --git a/MSAL/src/native_auth/public/state_machine/state/SignInAfterSignUpState.swift b/MSAL/src/native_auth/public/state_machine/state/SignInAfterSignUpState.swift
@@ -42,9 +42,13 @@ import Foundation
             let delegateDispatcher = SignInAfterSignUpDelegateDispatcher(delegate: delegate, telemetryUpdate: controllerResponse.telemetryUpdate)
 
             switch controllerResponse.result {
-            case .success(let accountResult):
+            case .completed(let accountResult):
                 await delegateDispatcher.dispatchSignInCompleted(result: accountResult, correlationId: controllerResponse.correlationId)
-            case .failure(let error):
+            case .jitAuthMethodsSelectionRequired(authMethods: let authMethods, newState: let newState):
+                await delegateDispatcher.dispatchJITRequired(authMethods: authMethods,
+                                                             newState: newState,
+                                                             correlationId: controllerResponse.correlationId)
+            case .error(let error):
                 await delegate.onSignInAfterSignUpError(
                     error: SignInAfterSignUpError(message: error.errorDescription, correlationId: error.correlationId, errorCodes: error.errorCodes)
                 )
diff --git a/MSAL/test/unit/native_auth/controllers/MSALNativeAuthJITControllerTests.swift b/MSAL/test/unit/native_auth/controllers/MSALNativeAuthJITControllerTests.swift
@@ -144,7 +144,7 @@ class MSALNativeAuthJITControllerTests: MSALNativeAuthTestCase {
             codeLength: 8
         )
         let userAccountResult = MSALNativeAuthUserAccountResultStub.result
-        signInControllerMock.continuationTokenResult = .init(.success(userAccountResult), correlationId: defaultUUID)
+        signInControllerMock.continuationTokenResult = .init(.completed(userAccountResult), correlationId: defaultUUID)
         jitResponseValidatorMock.challengeValidatedResponse = .preverified(continuationToken: "continuationToken")
         jitResponseValidatorMock.continueValidatedResponse = .success(continuationToken: "continuationToken 2")
         let result = await sut.requestJITChallenge(continuationToken: expectedContinuationToken, authMethod: authMethod, verificationContact: verificationContact, context: expectedContext)
@@ -204,7 +204,7 @@ class MSALNativeAuthJITControllerTests: MSALNativeAuthTestCase {
         jitResponseValidatorMock.continueValidatedResponse = .success(continuationToken: expectedContinuationToken)
         let userAccountResult = MSALNativeAuthUserAccountResultStub.result
 
-        signInControllerMock.continuationTokenResult = .init(.success(userAccountResult), correlationId: defaultUUID)
+        signInControllerMock.continuationTokenResult = .init(.completed(userAccountResult), correlationId: defaultUUID)
         jitResponseValidatorMock.continueValidatedResponse = .success(continuationToken: expectedContinuationToken)
         let result = await sut.submitJITChallenge(challenge: "123456", continuationToken: expectedContinuationToken, grantType: .oobCode, context: expectedContext)
         result.telemetryUpdate?(.success(()))
@@ -218,6 +218,34 @@ class MSALNativeAuthJITControllerTests: MSALNativeAuthTestCase {
         }
     }
 
+    func test_whenRequestJITContinueReturnsJITRequired_ErrorShouldBeReturned() async {
+        let expectedContinuationToken = "continuationToken"
+        let expectedContext = MSALNativeAuthRequestContext(correlationId: defaultUUID)
+        let authMethod = MSALAuthMethod(id: "1", challengeType: "oob", loginHint: "hint", channelTargetType: MSALNativeAuthChannelType(value:"email"))
+        let authMethods = [authMethod]
+        let newState = RegisterStrongAuthState(
+            controller: sut,
+            continuationToken: expectedContinuationToken,
+            correlationId: defaultUUID
+        )
+
+        jitRequestProviderMock.mockContinueRequestFunc(MSALNativeAuthHTTPRequestMock.prepareMockRequest())
+        jitResponseValidatorMock.continueValidatedResponse = .success(continuationToken: expectedContinuationToken)
+
+        signInControllerMock.continuationTokenResult = .init(.jitAuthMethodsSelectionRequired(authMethods: authMethods, newState: newState), correlationId: defaultUUID)
+        jitResponseValidatorMock.continueValidatedResponse = .success(continuationToken: expectedContinuationToken)
+        let result = await sut.submitJITChallenge(challenge: "123456", continuationToken: expectedContinuationToken, grantType: .oobCode, context: expectedContext)
+        result.telemetryUpdate?(.success(()))
+
+        checkTelemetryEventResult(id: .telemetryApiIdJITContinue, isSuccessful: true)
+        if case .error(let error, let newState) = result.result {
+            XCTAssertEqual(error.type, .generalError)
+            XCTAssertNil(newState)
+        } else {
+            XCTFail("Expected verificationRequired result")
+        }
+    }
+
     func test_whenRequestJITContinueRequestFails_ErrorShouldBeReturned() async {
         let expectedContinuationToken = "continuationToken"
         let expectedContext = MSALNativeAuthRequestContext(correlationId: defaultUUID)
diff --git a/MSAL/test/unit/native_auth/controllers/MSALNativeAuthResetPasswordControllerTests.swift b/MSAL/test/unit/native_auth/controllers/MSALNativeAuthResetPasswordControllerTests.swift
@@ -882,8 +882,8 @@ final class MSALNativeAuthResetPasswordControllerTests: MSALNativeAuthTestCase {
 
         let exp2 = expectation(description: "SignInAfterResetPassword expectation")
         signInControllerMock.expectation = exp2
-        signInControllerMock.continuationTokenResult = .init(.failure(SignInAfterResetPasswordError(correlationId: correlationId)), correlationId: correlationId)
-        
+        signInControllerMock.continuationTokenResult = .init(.error(error: SignInAfterResetPasswordError(correlationId: correlationId)), correlationId: correlationId)
+
         let parameters = MSALNativeAuthSignInAfterResetPasswordParameters()
         helper.signInAfterResetPasswordState?.signIn(parameters: parameters, delegate: SignInAfterResetPasswordDelegateStub())
         await fulfillment(of: [exp2], timeout: 1)
@@ -934,7 +934,7 @@ final class MSALNativeAuthResetPasswordControllerTests: MSALNativeAuthTestCase {
 
         let exp2 = expectation(description: "SignInAfterResetPassword expectation")
         signInControllerMock.expectation = exp2
-        signInControllerMock.continuationTokenResult = .init(.failure(SignInAfterResetPasswordError(correlationId: correlationId)), correlationId: correlationId)
+        signInControllerMock.continuationTokenResult = .init(.error(error: SignInAfterResetPasswordError(correlationId: correlationId)), correlationId: correlationId)
 
         let parameters = MSALNativeAuthSignInAfterResetPasswordParameters()
         helper.signInAfterResetPasswordState?.signIn(parameters: parameters, delegate: SignInAfterResetPasswordDelegateStub())
diff --git a/MSAL/test/unit/native_auth/controllers/MSALNativeAuthSignInControllerTests.swift b/MSAL/test/unit/native_auth/controllers/MSALNativeAuthSignInControllerTests.swift
@@ -205,7 +205,7 @@ class MSALNativeAuthSignInControllerTests: MSALNativeAuthTestCase {
 
         let result = await sut.signIn(username: expectedUsername, grantType: nil, continuationToken: continuationToken, scopes: ["openid","profile","offline_access"], claimsRequestJson: expectedClaimsRequestJson, telemetryId: MSALNativeAuthTelemetryApiId.telemetryApiIdSignInAfterSignUp, context: MSALNativeAuthRequestContextMock())
         
-        guard case let .failure(_) = result.result else {
+        guard case let .error(error: _) = result.result else {
             return XCTFail("input should be .error")
         }
     }
diff --git a/MSAL/test/unit/native_auth/controllers/MSALNativeAuthSignUpControllerTests.swift b/MSAL/test/unit/native_auth/controllers/MSALNativeAuthSignUpControllerTests.swift
@@ -1842,7 +1842,7 @@ final class MSALNativeAuthSignUpControllerTests: MSALNativeAuthTestCase {
 
         let exp2 = expectation(description: "SignInAfterSignUp expectation")
         signInControllerMock.expectation = exp2
-        signInControllerMock.continuationTokenResult = .init(.init(.failure(SignInAfterSignUpError(correlationId: correlationId)), correlationId: correlationId))
+        signInControllerMock.continuationTokenResult = .init(.init(.error(error: SignInAfterSignUpError(correlationId: correlationId)), correlationId: correlationId))
         let parameters = MSALNativeAuthSignInAfterSignUpParameters()
         helper.signInAfterSignUpState?.signIn(parameters: parameters, delegate: SignInAfterSignUpDelegateStub())
         await fulfillment(of: [exp2], timeout: 1)
@@ -1883,7 +1883,7 @@ final class MSALNativeAuthSignUpControllerTests: MSALNativeAuthTestCase {
 
         let exp2 = expectation(description: "SignInAfterSignUp expectation")
         signInControllerMock.expectation = exp2
-        signInControllerMock.continuationTokenResult = .init(.init(.failure(SignInAfterSignUpError(correlationId: correlationId)), correlationId: correlationId))
+        signInControllerMock.continuationTokenResult = .init(.init(.error(error: SignInAfterSignUpError(correlationId: correlationId)), correlationId: correlationId))
         let parameters = MSALNativeAuthSignInAfterSignUpParameters()
         helper.signInAfterSignUpState?.signIn(parameters: parameters, delegate: SignInAfterSignUpDelegateStub())
         await fulfillment(of: [exp2], timeout: 1)
diff --git a/MSAL/test/unit/native_auth/mock/delegate/SignInDelegatesSpies.swift b/MSAL/test/unit/native_auth/mock/delegate/SignInDelegatesSpies.swift
diff --git a/MSAL/test/unit/native_auth/public/state_machine/sign_in/SignInAfterResetPasswordTests.swift b/MSAL/test/unit/native_auth/public/state_machine/sign_in/SignInAfterResetPasswordTests.swift
diff --git a/MSAL/test/unit/native_auth/public/state_machine/sign_in/SignInAfterSignUpStateTests.swift b/MSAL/test/unit/native_auth/public/state_machine/sign_in/SignInAfterSignUpStateTests.swift

Original file line number	Diff line number	Diff line change
`@@ -36,4 +36,18 @@ final class SignInAfterResetPasswordDelegateDispatcher: DelegateDispatcher<SignI`
`36`	`36`	`await delegate.onSignInAfterResetPasswordError(error: error)`
`37`	`37`	`}`
`38`	`38`	`}`
	`39`	`+`
	`40`	`+ func dispatchJITRequired(authMethods: [MSALAuthMethod], newState: RegisterStrongAuthState, correlationId: UUID) async {`
	`41`	`+ if let onSignInJITRequired = delegate.onSignInStrongAuthMethodRegistration {`
	`42`	`+ telemetryUpdate?(.success(()))`
	`43`	`+ await onSignInJITRequired(authMethods, newState)`
	`44`	`+ } else {`
	`45`	`+ let error = SignInAfterResetPasswordError(`
	`46`	`+ message: requiredErrorMessage(for: "onSignInStrongAuthMethodRegistration"),`
	`47`	`+ correlationId: correlationId`
	`48`	`+ )`
	`49`	`+ telemetryUpdate?(.failure(error))`
	`50`	`+ await delegate.onSignInAfterResetPasswordError(error: error)`
	`51`	`+ }`
	`52`	`+ }`
`39`	`53`	`}`