Merge pull request #1535 from AzureAD/jarias/jit-troubleshooting-legacy-flow

JIT troubleshooting in legacy auth flow

Author: juan-arias

IdentityCore/src/MSIDConstants.h

@@ -213,6 +213,7 @@ extern NSString * _Nonnull const MSID_BROWSER_RESPONSE_SWITCH_BROWSER_RESUME;
 
 extern NSString * _Nonnull const MSID_FLIGHT_USE_V2_WEB_RESPONSE_FACTORY;
 extern NSString * _Nonnull const MSID_FLIGHT_SUPPORT_DUNA_CBA;
+extern NSString * _Nonnull const MSID_FLIGHT_DISABLE_JIT_TROUBLESHOOTING_LEGACY_AUTH;
 extern NSString * _Nonnull const MSID_FLIGHT_CLIENT_SFRT_STATUS;
 extern NSString * _Nonnull const MSID_FLIGHT_DISABLE_PREFERRED_IDENTITY_CBA;
 

IdentityCore/src/MSIDConstants.m

@@ -85,6 +85,7 @@
 
 NSString *const MSID_FLIGHT_USE_V2_WEB_RESPONSE_FACTORY = @"use_v2_web_response_factory";
 NSString *const MSID_FLIGHT_SUPPORT_DUNA_CBA = @"support_duna_cba_v2";
+NSString *const MSID_FLIGHT_DISABLE_JIT_TROUBLESHOOTING_LEGACY_AUTH = @"disable_jit_remediation_legacy_auth";
 NSString *const MSID_FLIGHT_CLIENT_SFRT_STATUS = @"sfrt_v2";
 NSString *const MSID_FLIGHT_DISABLE_PREFERRED_IDENTITY_CBA = @"dis_pre_iden_cba";
 

IdentityCore/src/webview/embeddedWebview/MSIDAADOAuthEmbeddedWebviewController.m

@@ -29,6 +29,9 @@
 #import "MSIDWorkPlaceJoinConstants.h"
 #import "MSIDPKeyAuthHandler.h"
 #import "MSIDWorkPlaceJoinUtil.h"
+#import "MSIDWebAuthNUtil.h"
+#import "MSIDFlightManager.h"
+#import "MSIDConstants.h"
 
 #if !MSID_EXCLUDE_WEBKIT
 
@@ -67,6 +70,28 @@ - (BOOL)decidePolicyAADForNavigationAction:(WKNavigationAction *)navigationActio
     BOOL isBrokerUrl = [@"msauth" caseInsensitiveCompare:requestURL.scheme] == NSOrderedSame;
     BOOL isBrowserUrl = [@"browser" caseInsensitiveCompare:requestURL.scheme] == NSOrderedSame;
 
+    if (![MSIDFlightManager.sharedInstance boolForKey:MSID_FLIGHT_DISABLE_JIT_TROUBLESHOOTING_LEGACY_AUTH])
+    {
+        // When not running in SSO extension, the CA block page will return with "https" scheme instead of "browser"
+        if (requestURL && ![MSIDWebAuthNUtil amIRunningInExtension] &&
+            self.externalDecidePolicyForBrowserAction &&
+            [@"https" caseInsensitiveCompare:requestURL.scheme] == NSOrderedSame)
+        {
+            // Create new URL replacing 'https' scheme with 'browser' scheme
+            NSURL *legacyFlowUrl = [NSURL URLWithString:[NSString stringWithFormat:@"browser%@", [requestURL.absoluteString substringFromIndex:5]]];
+            NSURLRequest *challengeResponse = self.externalDecidePolicyForBrowserAction(self, legacyFlowUrl);
+
+            if (challengeResponse)
+            {
+                MSID_LOG_WITH_CTX(MSIDLogLevelInfo, self.context, @"Found AAD policy for navigation using https url and externalDecidePolicyForBrowserAction in legacy auth flow.");
+                decisionHandler(WKNavigationActionPolicyCancel);
+                [self loadRequest:challengeResponse];
+
+                return YES;
+            }
+        }
+    }
+    
     if (isBrokerUrl || isBrowserUrl)
     {
         // Let external code decide if browser url is allowed to continue

IdentityCore/tests/MSIDAADOAuthEmbeddedWebviewControllerTests.m

@@ -26,6 +26,7 @@
 #import <XCTest/XCTest.h>
 #import "MSIDAADOAuthEmbeddedWebviewController.h"
 #import "MSIDWKNavigationActionMock.h"
+#import "MSIDWebAuthNUtil.h"
 
 #if !MSID_EXCLUDE_WEBKIT
 
@@ -195,6 +196,67 @@ - (void)testDecidePolicyForNavigationAction_whenExternalDecidePolicyForBrowserAc
     XCTAssertTrue(result);
 }
 
+- (void)testDecidePolicyForNavigationAction_whenExternalDecidePolicyForBrowserActionLegacyFlow_shouldCancelActionAndReturnYesAndCallExternalMethod
+{
+    [MSIDWebAuthNUtil setAmIRunningInExtension:NO];
+    
+    MSIDAADOAuthEmbeddedWebviewController *webVC = [[MSIDAADOAuthEmbeddedWebviewController alloc]
+            initWithStartURL:[NSURL URLWithString:@"https://contoso.com/oauth/authorize"]
+                      endURL:[NSURL URLWithString:@"endurl://host"]
+                     webview:nil
+               customHeaders:nil
+              platfromParams:nil
+                     context:nil];
+
+    NSURLRequest *request = [[NSURLRequest alloc] initWithURL:[[NSURL alloc] initWithString:@"https://www.web-cp.com/check"]];
+    MSIDWKNavigationActionMock *action = [[MSIDWKNavigationActionMock alloc] initWithRequest:request];
+
+    XCTestExpectation *expectationExternalDecisionHandler = [self expectationWithDescription:@"external decision handler"];
+    XCTestExpectation *expectationDecisionHandler = [self expectationWithDescription:@"decision handler"];
+
+    webVC.externalDecidePolicyForBrowserAction = ^NSURLRequest *(MSIDOAuth2EmbeddedWebviewController *webView, NSURL *url) {
+
+        XCTAssertNotNil(webView);
+        XCTAssertEqualObjects([url absoluteString], @"browser://www.web-cp.com/check");
+        [expectationExternalDecisionHandler fulfill];
+
+        return [[NSURLRequest alloc] initWithURL:url];
+    };
+
+
+    BOOL result = [webVC decidePolicyAADForNavigationAction:action decisionHandler:^(WKNavigationActionPolicy decision) {
+
+        XCTAssertEqual(decision, WKNavigationActionPolicyCancel);
+        [expectationDecisionHandler fulfill];
+    }];
+
+    [self waitForExpectationsWithTimeout:1.0 handler:nil];
+
+    XCTAssertTrue(result);
+}
+
+- (void)testDecidePolicyForNavigationAction_whenExternalDecidePolicyForBrowserActionLegacyFlowNonHttps_shouldCancelActionAndReturnNoAndCallExternalMethod
+{
+    [MSIDWebAuthNUtil setAmIRunningInExtension:NO];
+    
+    MSIDAADOAuthEmbeddedWebviewController *webVC = [[MSIDAADOAuthEmbeddedWebviewController alloc]
+            initWithStartURL:[NSURL URLWithString:@"https://contoso.com/oauth/authorize"]
+                      endURL:[NSURL URLWithString:@"endurl://host"]
+                     webview:nil
+               customHeaders:nil
+              platfromParams:nil
+                     context:nil];
+
+    NSURLRequest *request = [[NSURLRequest alloc] initWithURL:[[NSURL alloc] initWithString:@"http://www.web-cp.com/check"]];
+    MSIDWKNavigationActionMock *action = [[MSIDWKNavigationActionMock alloc] initWithRequest:request];
+
+    BOOL result = [webVC decidePolicyAADForNavigationAction:action decisionHandler:^(WKNavigationActionPolicy decision) {
+        XCTAssertEqual(decision, WKNavigationActionPolicyCancel);
+    }];
+
+    XCTAssertFalse(result);
+}
+
 @end
 
 #endif