SNOW-1647589: Fix NullPointerException when MFA is enabled in Okta and native Okta authentication is used (#1887)

Co-authored-by: Rafal Kowalski <[email protected]>

Author: sfc-gh-wfateem

src/main/java/net/snowflake/client/core/SessionUtil.java

@@ -1555,17 +1555,23 @@ private static String federatedFlowStep3(
               loginInput.getHttpClientSettingsKey(),
               null);
 
-      logger.debug("User is authenticated against {}.", loginInput.getAuthenticator());
-
       // session token is in the data field of the returned json response
       final JsonNode jsonNode = mapper.readTree(idpResponse);
+      boolean isMfaEnabledInOkta = "MFA_REQUIRED".equals(jsonNode.get("status").asText());
+      if (isMfaEnabledInOkta) {
+        throw new SnowflakeSQLLoggedException(
+            null,
+            ErrorCode.OKTA_MFA_NOT_SUPPORTED.getMessageCode(),
+            SqlState.FEATURE_NOT_SUPPORTED);
+      }
       oneTimeToken =
           jsonNode.get("sessionToken") != null
               ? jsonNode.get("sessionToken").asText()
               : jsonNode.get("cookieToken").asText();
     } catch (IOException | URISyntaxException ex) {
       handleFederatedFlowError(loginInput, ex);
     }
+    logger.debug("User is authenticated against {}.", loginInput.getAuthenticator());
     return oneTimeToken;
   }
 

src/main/java/net/snowflake/client/jdbc/ErrorCode.java

@@ -84,6 +84,7 @@ public enum ErrorCode {
   OAUTH_CLIENT_CREDENTIALS_FLOW_ERROR(200069, SqlState.CONNECTION_EXCEPTION),
   OAUTH_REFRESH_TOKEN_FLOW_ERROR(200070, SqlState.CONNECTION_EXCEPTION),
   WORKLOAD_IDENTITY_FLOW_ERROR(200071, SqlState.CONNECTION_EXCEPTION),
+  OKTA_MFA_NOT_SUPPORTED(200072, SqlState.FEATURE_NOT_SUPPORTED),
   FILE_TRANSFER_ERROR(253000, SqlState.SYSTEM_ERROR),
   DOWNLOAD_ERROR(253002, SqlState.SYSTEM_ERROR),
   UPLOAD_ERROR(253003, SqlState.SYSTEM_ERROR),

src/main/resources/net/snowflake/client/jdbc/jdbc_error_messages.properties

@@ -90,6 +90,8 @@ Error message={3}, Extended error info={4}
 200069=Error during OAuth Client Credentials authentication: {0}
 200070=Error during obtaining OAuth access token using refresh token: {0}
 200071=Error during Workload Identity authentication: {0}
+200072=MFA enabled in Okta is not supported with this authenticator type. \
+  Please use 'externalbrowser' instead or a different authentication method.
 253000=Error during file transfer: {0}
 253003=Error during file upload to stage: {0}
 253002=Error during file download to stage: {0}

src/test/java/net/snowflake/client/core/SessionUtilWiremockIT.java

@@ -4,6 +4,7 @@
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.fail;
 
 import java.util.Comparator;
@@ -15,6 +16,7 @@
 import net.snowflake.client.jdbc.BaseWiremockTest;
 import net.snowflake.client.jdbc.ErrorCode;
 import net.snowflake.client.jdbc.SnowflakeSQLException;
+import net.snowflake.client.jdbc.SnowflakeSQLLoggedException;
 import org.junit.jupiter.api.Tag;
 import org.junit.jupiter.api.Test;
 
@@ -32,6 +34,8 @@ public class SessionUtilWiremockIT extends BaseWiremockTest {
       "/wiremock/mappings/session/session-util-wiremock-it-multiple-429-in-federated-step-3.json";
   private static final String MULTIPLE_429_IN_FEDERATED_STEP_4 =
       "/wiremock/mappings/session/session-util-wiremock-it-multiple-429-in-federated-step-4.json";
+  private static final String UNSUPPORTED_MFA_IN_FEDERATED_STEP_3 =
+      "/wiremock/mappings/session/session-util-wiremock-it-unsupported-mfa-in-federated-step-3.json";
 
   /**
    * Minimum spacing we expect between consecutive requests, in milliseconds - associated with
@@ -231,6 +235,31 @@ public void testOktaRetriesUntilTimeoutThenRaisesAuthTimeoutExceptionWhen429InFe
         ALLOWED_DIFFERENCE_BETWEEN_LOGIN_TIMEOUT_AND_ACTUAL_DURATION_IN_MS);
   }
 
+  @Test
+  public void testErrorHandlingWhenOktaReturnsUnsupportedMfaInFederatedStep3() throws Throwable {
+    // GIVEN
+    Map<String, Object> placeholders = new HashMap<>();
+    placeholders.put("{{WIREMOCK_HOST_WITH_HTTPS_AND_PORT}}", WIREMOCK_HOST_WITH_HTTPS_AND_PORT);
+    String wireMockMapping =
+        getWireMockMappingFromFile(UNSUPPORTED_MFA_IN_FEDERATED_STEP_3, placeholders);
+    importMapping(wireMockMapping);
+
+    setCustomTrustStorePropertyPath();
+
+    SFLoginInput loginInput = createOktaLoginInputBase();
+    Map<SFSessionProperty, Object> connectionPropertiesMap = initConnectionPropertiesMap();
+
+    SnowflakeSQLLoggedException thrown =
+        assertThrows(
+            SnowflakeSQLLoggedException.class,
+            () -> SessionUtil.openSession(loginInput, connectionPropertiesMap, "ALL"));
+    assertThat(thrown.getErrorCode(), equalTo(ErrorCode.OKTA_MFA_NOT_SUPPORTED.getMessageCode()));
+    assertThat(
+        thrown.getMessage(),
+        equalTo(
+            "MFA enabled in Okta is not supported with this authenticator type. Please use 'externalbrowser' instead or a different authentication method."));
+  }
+
   private void assertThatTotalLoginTimeoutIsKeptWhenRetrying(
       List<MinimalServeEvent> requestEvents, long loginTimeout, long allowedDifferenceInMs) {
     final int SECONDS_TO_MS_FACTOR = 1000;

src/test/resources/wiremock/mappings/session/session-util-wiremock-it-unsupported-mfa-in-federated-step-3.json

@@ -0,0 +1,54 @@
+{
+  "mappings": [
+    {
+      "scenarioName": "Unsupported MFA in federated step 3",
+      "request": {
+        "method": "POST",
+        "urlPath": "/session/authenticator-request",
+        "queryParameters": {
+          "request_guid": {
+            "matches": ".*"
+          }
+        }
+      },
+      "response": {
+        "status": 200,
+        "headers": {
+          "Content-Type": "application/json"
+        },
+        "jsonBody": {
+          "data": {
+            "tokenUrl": "{{WIREMOCK_HOST_WITH_HTTPS_AND_PORT}}/okta-stub/vanity-url/api/v1/authn",
+            "ssoUrl": "{{WIREMOCK_HOST_WITH_HTTPS_AND_PORT}}/okta-stub/vanity-url/app/snowflake/tokenlikepartofurl/sso/saml",
+            "proofKey": null
+          },
+          "code": null,
+          "message": null,
+          "success": true
+        }
+      }
+    },
+    {
+      "scenarioName": "Unsupported MFA in federated step 3",
+      "request": {
+        "method": "POST",
+        "urlPath": "/okta-stub/vanity-url/api/v1/authn"
+      },
+      "response": {
+        "status": 200,
+        "headers": {
+          "Content-Type": "application/json"
+        },
+        "jsonBody": {
+          "status": "MFA_REQUIRED",
+          "sessionToken": null,
+          "cookieToken": null
+        }
+      }
+    }
+  ],
+  "importOptions": {
+    "duplicatePolicy": "IGNORE",
+    "deleteAllNotInImport": true
+  }
+}