Validate recovery scenario and step when redeeming a confirmation code

components/org.wso2.carbon.identity.recovery/src/main/java/org/wso2/carbon/identity/recovery/executor/ConfirmationCodeValidationExecutor.java

@@ -37,6 +37,8 @@
 import org.wso2.carbon.identity.recovery.IdentityRecoveryConstants;
 import org.wso2.carbon.identity.recovery.IdentityRecoveryException;
 import org.wso2.carbon.identity.recovery.IdentityRecoveryServerException;
+import org.wso2.carbon.identity.recovery.RecoveryScenarios;
+import org.wso2.carbon.identity.recovery.RecoverySteps;
 import org.wso2.carbon.identity.recovery.internal.IdentityRecoveryServiceDataHolder;
 import org.wso2.carbon.identity.recovery.model.UserRecoveryData;
 import org.wso2.carbon.identity.recovery.store.JDBCRecoveryDataStore;
@@ -179,9 +181,42 @@ private UserRecoveryData validateConfirmationCode(String code) throws IdentityRe
         if (!StringUtils.equals(contextTenantDomain, userTenantDomain)) {
             throw new IdentityRecoveryClientException("Invalid tenant domain: " + userTenantDomain);
         }
+        validateRecoveryScenario(userRecoveryData, code);
         return userRecoveryData;
     }
 
+    /**
+     * Validates that the confirmation code was issued for a recovery scenario and step this executor serves.
+     *
+     * @param userRecoveryData Recovery data resolved from the confirmation code.
+     * @param code             Confirmation code, included in the error for consistency with other code failures.
+     * @throws IdentityRecoveryClientException If the scenario or the step is not one this executor serves.
+     */
+    private void validateRecoveryScenario(UserRecoveryData userRecoveryData, String code)
+            throws IdentityRecoveryClientException {
+
+        Enum<?> recoveryScenario = userRecoveryData.getRecoveryScenario();
+        Enum<?> recoveryStep = userRecoveryData.getRecoveryStep();
+
+        // Recovery scenarios of the invitations this executor serves.
+        boolean isInvitationScenario = RecoveryScenarios.ASK_PASSWORD.equals(recoveryScenario)
+                || RecoveryScenarios.ASK_PASSWORD_VIA_EMAIL_OTP.equals(recoveryScenario)
+                || RecoveryScenarios.ASK_PASSWORD_VIA_SMS_OTP.equals(recoveryScenario)
+                || RecoveryScenarios.TENANT_ADMIN_ASK_PASSWORD.equals(recoveryScenario)
+                || RecoveryScenarios.ADMIN_INVITE_SET_PASSWORD_OFFLINE.equals(recoveryScenario);
+        // Recovery steps that authorise a password change, as validated in NotificationPasswordRecoveryManager.
+        boolean isPasswordUpdateStep = RecoverySteps.UPDATE_PASSWORD.equals(recoveryStep)
+                || RecoverySteps.SET_PASSWORD.equals(recoveryStep);
+
+        if (!isInvitationScenario || !isPasswordUpdateStep) {
+            if (LOG.isDebugEnabled()) {
+                LOG.debug("Confirmation code issued for recovery scenario: " + recoveryScenario + " and step: "
+                        + recoveryStep + " cannot be redeemed by the invited user registration flow.");
+            }
+            throw Utils.handleClientException(IdentityRecoveryConstants.ErrorMessages.ERROR_CODE_INVALID_CODE, code);
+        }
+    }
+
     @Override
     public List<String> getInitiationData() {
 

components/org.wso2.carbon.identity.recovery/src/test/java/org/wso2/carbon/identity/recovery/executor/ConfirmationCodeValidationExecutorTest.java

@@ -21,6 +21,7 @@
 import org.mockito.MockedStatic;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 import org.wso2.carbon.context.PrivilegedCarbonContext;
 import org.wso2.carbon.identity.application.common.model.User;
@@ -31,6 +32,8 @@
 import org.wso2.carbon.identity.flow.execution.engine.model.ExecutorResponse;
 import org.wso2.carbon.identity.flow.execution.engine.model.FlowExecutionContext;
 import org.wso2.carbon.identity.flow.execution.engine.model.FlowUser;
+import org.wso2.carbon.identity.recovery.RecoveryScenarios;
+import org.wso2.carbon.identity.recovery.RecoverySteps;
 import org.wso2.carbon.identity.recovery.internal.IdentityRecoveryServiceDataHolder;
 import org.wso2.carbon.identity.recovery.model.UserRecoveryData;
 import org.wso2.carbon.identity.recovery.store.JDBCRecoveryDataStore;
@@ -124,6 +127,8 @@ public void testExecuteWithValidConfirmationCode() throws Exception {
 
         UserRecoveryData mockRecoveryData = mock(UserRecoveryData.class);
         when(mockRecoveryData.getUser()).thenReturn(mockUser);
+        when(mockRecoveryData.getRecoveryScenario()).thenReturn(RecoveryScenarios.ASK_PASSWORD);
+        when(mockRecoveryData.getRecoveryStep()).thenReturn(RecoverySteps.UPDATE_PASSWORD);
 
         // Mock UserRecoveryDataStore.
         UserRecoveryDataStore mockStore = mock(UserRecoveryDataStore.class);
@@ -202,4 +207,47 @@ public void testExecuteWithInvalidConfirmationCode() throws Exception {
         assertEquals(response.getResult(), Constants.ExecutorStatus.STATUS_ERROR);
         assertNotNull(response.getErrorMessage());
     }
+
+    @DataProvider(name = "rejectedRecoveryData")
+    public Object[][] rejectedRecoveryData() {
+
+        return new Object[][]{
+                // The code the password recovery page hands out, which this executor must not redeem.
+                {RecoveryScenarios.NOTIFICATION_BASED_PW_RECOVERY, RecoverySteps.SEND_RECOVERY_INFORMATION},
+                // An invitation scenario at a step that does not authorise setting a password.
+                {RecoveryScenarios.ASK_PASSWORD, RecoverySteps.RESEND_CONFIRMATION_CODE}
+        };
+    }
+
+    @Test(dataProvider = "rejectedRecoveryData")
+    public void testExecuteWithConfirmationCodeOfAnotherScenario(Enum<?> scenario, Enum<?> step) throws Exception {
+
+        FlowExecutionContext context = mock(FlowExecutionContext.class);
+        Map<String, String> userInputData = new HashMap<>();
+        userInputData.put(CONFIRMATION_CODE, "code-of-another-scenario");
+        when(context.getUserInputData()).thenReturn(userInputData);
+
+        User mockUser = new User();
+        mockUser.setUserName(USERNAME);
+        mockUser.setTenantDomain(TENANT_DOMAIN);
+        mockUser.setUserStoreDomain(DOMAIN_NAME);
+
+        UserRecoveryData mockRecoveryData = mock(UserRecoveryData.class);
+        when(mockRecoveryData.getUser()).thenReturn(mockUser);
+        when(mockRecoveryData.getRecoveryScenario()).thenReturn(scenario);
+        when(mockRecoveryData.getRecoveryStep()).thenReturn(step);
+
+        UserRecoveryDataStore mockStore = mock(UserRecoveryDataStore.class);
+        mockedJdbcStore.when(JDBCRecoveryDataStore::getInstance).thenReturn(mockStore);
+        when(mockStore.load(anyString())).thenReturn(mockRecoveryData);
+
+        PrivilegedCarbonContext carbonContext = mock(PrivilegedCarbonContext.class);
+        mockedCarbonContext.when(PrivilegedCarbonContext::getThreadLocalCarbonContext).thenReturn(carbonContext);
+        when(carbonContext.getTenantDomain()).thenReturn(TENANT_DOMAIN);
+
+        ExecutorResponse response = executor.execute(context);
+
+        assertEquals(response.getResult(), Constants.ExecutorStatus.STATUS_USER_ERROR);
+        assertNotNull(response.getErrorMessage());
+    }
 }