Enable modification of sub-attributes in complex OIDC claims & fix org switch default resolution in pre-issue id token action

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/openidconnect/action/preissueidtoken/execution/PreIssueIDTokenRequestBuilder.java

@@ -63,7 +63,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.stream.Collectors;
 
 import javax.servlet.http.HttpServletRequestWrapper;
 
@@ -207,24 +206,45 @@ private List<AllowedOperation> getAllowedOperations(Map<String, Object> claimsTo
 
     private List<String> getRemoveOrReplacePaths(Map<String, Object> oidcClaims) {
 
-        List<String> removeOrReplacePaths =  oidcClaims.entrySet().stream()
-                .filter(entry -> entry.getValue() instanceof String ||
-                        entry.getValue() instanceof Number ||
-                        entry.getValue() instanceof Boolean || entry.getValue() instanceof List ||
-                        entry.getValue() instanceof String[])
-                .map(this::generatePathForClaim)
-                .collect(Collectors.toList());
+        List<String> removeOrReplacePaths = new ArrayList<>();
+        for (Map.Entry<String, Object> entry : oidcClaims.entrySet()) {
+
+            String basePath = CLAIMS_PATH_PREFIX + entry.getKey();
+            Object value = entry.getValue();
+
+            removeOrReplacePaths.add(basePath);
+            if (value instanceof String || value instanceof Number || value instanceof Boolean) {
+                continue;
+            }
+            if (value instanceof List || value instanceof String[]) {
+                removeOrReplacePaths.add(basePath + "/");
+                continue;
+            }
+
+            // Handle nested objects
+            if (value instanceof Map) {
+                collectNestedClaimPaths(basePath, value, removeOrReplacePaths);
+            }
+        }
         removeOrReplacePaths.add(CLAIMS_PATH_PREFIX + AccessToken.ClaimNames.AUD.getName() + "/");
         return removeOrReplacePaths;
     }
 
-    private String generatePathForClaim(Map.Entry<String, Object> entry) {
+    private void collectNestedClaimPaths(String basePath, Object value, List<String> paths) {
+
+        if (value instanceof Map) {
+            Map<?, ?> map = (Map<?, ?>) value;
 
-        String basePath = CLAIMS_PATH_PREFIX + entry.getKey();
-        if (entry.getValue() instanceof List || entry.getValue() instanceof String[]) {
-            basePath += "/";
+            for (Map.Entry<?, ?> entry : map.entrySet()) {
+                if (!(entry.getKey() instanceof String)) {
+                    continue;
+                }
+                String childPath = basePath + "/" + entry.getKey();
+                paths.add(childPath);
+
+                collectNestedClaimPaths(childPath, entry.getValue(), paths);
+            }
         }
-        return basePath;
     }
 
     private AllowedOperation createAllowedOperation(Operation op, List<String> paths) {
@@ -298,8 +318,17 @@ private User resolveLocalUser(AuthenticatedUser authenticatedUser, String grantT
         userBuilder.organization(resolveUserAuthenticatedOrganization(authenticatedUser));
 
         if (OAuthConstants.GrantTypes.ORGANIZATION_SWITCH.equals(grantType)) {
-            userBuilder.accessingOrganization(buildOrganization(authenticatedUser.getAccessingOrganization(),
-                    authenticatedUser.getTenantDomain()));
+            Organization accessingOrg;
+            if (authenticatedUser.getAccessingOrganization() != null) {
+                accessingOrg = buildOrganization(authenticatedUser.getAccessingOrganization(),
+                        authenticatedUser.getTenantDomain());
+                // In case of org switch, if accessing org is not set, it means user is switching to root org.
+            } else {
+                accessingOrg = buildOrganization(
+                        resolveOrganizationId(authenticatedUser.getTenantDomain()),
+                        authenticatedUser.getTenantDomain());
+            }
+            userBuilder.accessingOrganization(accessingOrg);
         }
         return userBuilder.build();
     }

components/org.wso2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/openidconnect/action/preissueidtoken/execution/PreIssueIDTokenResponseProcessor.java

@@ -54,6 +54,7 @@
 import org.wso2.carbon.utils.DiagnosticLog;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.regex.Matcher;
@@ -351,7 +352,9 @@ private OperationExecutionResult addToOtherClaims(PerformableOperation operation
             }
 
             Object claimValue = claim.getValue();
-            if (isValidPrimitiveValue(claimValue) || isValidListValue(claimValue)) {
+            if (isValidPrimitiveValue(claimValue)
+                    || isValidListValue(claimValue)
+                    || isValidMapValue(claimValue)) {
                 responseIDTokenDTO.getCustomOIDCClaims().put(claim.getName(), claim.getValue());
                 return new OperationExecutionResult(operation, OperationExecutionResult.Status.SUCCESS, "Claim added.");
             } else {
@@ -402,6 +405,15 @@ private boolean isValidListValue(Object value) {
         return list.stream().allMatch(item -> item instanceof String);
     }
 
+    private boolean isValidMapValue(Object value) {
+
+        if (!(value instanceof Map<?, ?>)) {
+            return false;
+        }
+        Map<?, ?> map = (Map<?, ?>) value;
+        return true;
+    }
+
     private OperationExecutionResult handleRemoveOperation(PerformableOperation operation, IDToken requestedIDToken,
                                                           IDTokenDTO responseIDTokenDTO) {
 
@@ -455,9 +467,15 @@ private OperationExecutionResult removeOtherClaims(PerformableOperation operatio
                                                        IDToken requestIDToken,
                                                        IDTokenDTO responseIDTokenDTO) {
 
-       ClaimPathInfo claimPathInfo = parseOperationPath(operation.getPath());
-       IDToken.Claim claim = requestIDToken.getClaim(claimPathInfo.getClaimName());
+        List<String> pathSegments = extractNestedClaimPath(operation.getPath());
 
+        // Nested removal
+        if (pathSegments.size() > 1 && !isArrayIndexPath(pathSegments)) {
+            return removeNestedClaim(pathSegments, requestIDToken, responseIDTokenDTO, operation);
+        }
+
+        ClaimPathInfo claimPathInfo = parseOperationPath(operation.getPath());
+        IDToken.Claim claim = requestIDToken.getClaim(claimPathInfo.getClaimName());
         if (claim == null) {
             return new OperationExecutionResult(operation, OperationExecutionResult.Status.FAILURE,
                     "Claim not found.");
@@ -471,6 +489,57 @@ private OperationExecutionResult removeOtherClaims(PerformableOperation operatio
         }
     }
 
+    private OperationExecutionResult removeNestedClaim(List<String> pathSegments, IDToken requestIDToken,
+                                                       IDTokenDTO responseIDTokenDTO,
+                                                       PerformableOperation operation) {
+
+        String rootClaimName = pathSegments.get(0);
+        List<String> nestedPath = pathSegments.subList(1, pathSegments.size());
+
+        IDToken.Claim rootClaim = requestIDToken.getClaim(rootClaimName);
+        if (rootClaim == null || !(rootClaim.getValue() instanceof Map)) {
+            return new OperationExecutionResult(operation, OperationExecutionResult.Status.FAILURE,
+                    "Root claim is not a complex object.");
+        }
+
+        Map<String, Object> rootValue = new HashMap<>((Map<String, Object>) rootClaim.getValue());
+        boolean removed = removeFromNestedMap(rootValue, nestedPath, 0);
+        if (!removed) {
+            return new OperationExecutionResult(operation, OperationExecutionResult.Status.FAILURE,
+                    "Nested claim not found.");
+        }
+
+        if (rootValue.isEmpty()) {
+            responseIDTokenDTO.getCustomOIDCClaims().remove(rootClaimName);
+        } else {
+            responseIDTokenDTO.getCustomOIDCClaims().put(rootClaimName, rootValue);
+        }
+
+        return new OperationExecutionResult(operation, OperationExecutionResult.Status.SUCCESS,
+                "Nested claim removed.");
+    }
+
+    private boolean removeFromNestedMap(Map<String, Object> current, List<String> path, int index) {
+
+        String key = path.get(index);
+        if (index == path.size() - 1) {
+            return current.remove(key) != null;
+        }
+
+        Object next = current.get(key);
+        if (!(next instanceof Map)) {
+            return false;
+        }
+
+        return removeFromNestedMap((Map<String, Object>) next, path, index + 1);
+    }
+
+    private List<String> extractNestedClaimPath(String operationPath) {
+
+        String relativePath = operationPath.substring(CLAIMS_PATH_PREFIX.length());
+        return List.of(relativePath.split("/"));
+    }
+
     private OperationExecutionResult removeClaimValueAtIndexFromArrayTypeClaim(PerformableOperation operation,
                                                                                ClaimPathInfo claimPathInfo,
                                                                                IDToken.Claim claim,
@@ -570,6 +639,13 @@ private OperationExecutionResult replaceOtherClaims(PerformableOperation operati
                                                         IDToken requestIDToken,
                                                         IDTokenDTO responseIDTokenDTO) {
 
+        List<String> pathSegments = extractNestedClaimPath(operation.getPath());
+
+        // Nested replace
+        if (pathSegments.size() > 1 && !isArrayIndexPath(pathSegments)) {
+            return replaceNestedClaim(pathSegments, requestIDToken, responseIDTokenDTO, operation);
+        }
+
         ClaimPathInfo claimPathInfo = parseOperationPath(operation.getPath());
         IDToken.Claim claim = requestIDToken.getClaim(claimPathInfo.getClaimName());
 
@@ -585,6 +661,63 @@ private OperationExecutionResult replaceOtherClaims(PerformableOperation operati
         }
     }
 
+    private OperationExecutionResult replaceNestedClaim(List<String> pathSegments, IDToken requestIDToken,
+                                                        IDTokenDTO responseIDTokenDTO,
+                                                        PerformableOperation operation) {
+
+        String rootClaimName = pathSegments.get(0);
+        List<String> nestedPath = pathSegments.subList(1, pathSegments.size());
+
+        IDToken.Claim rootClaim = requestIDToken.getClaim(rootClaimName);
+        if (rootClaim == null || !(rootClaim.getValue() instanceof Map)) {
+            return new OperationExecutionResult(operation, OperationExecutionResult.Status.FAILURE,
+                    "Root claim is not a complex object.");
+        }
+
+        Map<String, Object> rootValue = new HashMap<>((Map<String, Object>) rootClaim.getValue());
+
+        boolean replaced = replaceInNestedMap(rootValue, nestedPath, 0, operation.getValue());
+        if (!replaced) {
+            return new OperationExecutionResult(operation, OperationExecutionResult.Status.FAILURE,
+                    "Nested claim not found.");
+        }
+
+        if (rootValue.isEmpty()) {
+            responseIDTokenDTO.getCustomOIDCClaims().remove(rootClaimName);
+        } else {
+            responseIDTokenDTO.getCustomOIDCClaims().put(rootClaimName, rootValue);
+        }
+
+        return new OperationExecutionResult(operation, OperationExecutionResult.Status.SUCCESS,
+                "Nested claim replaced.");
+    }
+
+    private boolean replaceInNestedMap(Map<String, Object> current, List<String> path, int index, Object newValue) {
+
+        String key = path.get(index);
+
+        if (index == path.size() - 1) {
+            if (newValue == null) {
+                return current.remove(key) != null;
+            }
+            current.put(key, newValue);
+            return true;
+        }
+
+        Object next = current.get(key);
+        if (!(next instanceof Map)) {
+            return false;
+        }
+        boolean updated = replaceInNestedMap((Map<String, Object>) next, path, index + 1, newValue);
+
+        Map<String, Object> nextMap = (Map<String, Object>) next;
+        if (updated && nextMap.isEmpty()) {
+            current.remove(key);
+        }
+
+        return updated;
+    }
+
     private OperationExecutionResult replaceClaimValueAtIndexFromArrayTypeClaim(PerformableOperation operation,
                                                                                 ClaimPathInfo claimPathInfo,
                                                                                 IDToken.Claim claim,
@@ -680,6 +813,25 @@ private boolean isValidStringOrURI(String input) {
         return matcher.matches();
     }
 
+    private boolean isArrayIndexPath(List<String> pathSegments) {
+
+        if (pathSegments.size() != 2) {
+            return false;
+        }
+
+        String lastSegment = pathSegments.get(1);
+        if (LAST_ELEMENT_CHARACTER.equals(lastSegment)) {
+            return true;
+        }
+
+        try {
+            Integer.parseInt(lastSegment);
+            return true;
+        } catch (NumberFormatException e) {
+            return false;
+        }
+    }
+
     private ClaimPathInfo parseOperationPath(String operationPath) {
 
         String[] pathSegments = operationPath.split("/");

components/org.wso2.carbon.identity.oauth/src/test/java/org/wso2/carbon/identity/openidconnect/action/preissueidtoken/execution/PreIssueIDTokenRequestBuilderTest.java

@@ -1197,6 +1197,19 @@ private IDTokenDTO getMockIDTokenDTO() {
         return idTokenDTO;
     }
 
+    private void assertComplexPaths(List<String> paths) {
+
+        Assert.assertTrue(paths.contains("/idToken/claims/simple"));
+        Assert.assertTrue(paths.contains("/idToken/claims/nested"));
+        Assert.assertTrue(paths.contains("/idToken/claims/nested/intermediate"));
+        Assert.assertTrue(paths.contains("/idToken/claims/nested/intermediate/leaf"));
+
+        Assert.assertTrue(paths.contains("/idToken/claims/list_claim"));
+        Assert.assertTrue(paths.contains("/idToken/claims/list_claim/"));
+
+        Assert.assertTrue(paths.contains("/idToken/claims/aud/"));
+    }
+
     @Test
     public void buildActionExecutionRequestWithMultipleHeadersForTokenFlow()
             throws ActionExecutionRequestBuilderException {
@@ -1613,4 +1626,67 @@ public void buildActionExecutionRequestWithHeadersAndParametersForAuthzFlow()
         Assert.assertNotNull(request.getAdditionalParams());
         Assert.assertEquals(request.getAdditionalParams().size(), 2);
     }
+
+    @Test
+    public void testGetRemoveOrReplacePathsWithDeeplyNestedClaims() throws ActionExecutionRequestBuilderException {
+
+        Map<String, Object> customClaims = new HashMap<>();
+        customClaims.put("simple", "value");
+        Map<String, Object> level1 = new HashMap<>();
+        Map<String, Object> level2 = new HashMap<>();
+        level2.put("leaf", "value");
+        level1.put("intermediate", level2);
+        customClaims.put("nested", level1);
+        customClaims.put("list_claim", Arrays.asList("a", "b"));
+
+        IDTokenDTO idTokenDTO = getMockIDTokenDTO();
+        idTokenDTO.setCustomOIDCClaims(customClaims);
+
+        FlowContext flowContext = FlowContext.create()
+                .add(TOKEN_REQUEST_MESSAGE_CONTEXT, getMockTokenMessageContext())
+                .add(ID_TOKEN_DTO, idTokenDTO)
+                .add(REQUEST_TYPE, REQUEST_TYPE_TOKEN);
+        ActionExecutionRequest request = preIssueIDTokenRequestBuilder.buildActionExecutionRequest(
+                flowContext, null);
+
+        List<String> removePaths = request.getAllowedOperations().stream()
+                .filter(op -> op.getOp() == Operation.REMOVE)
+                .findFirst()
+                .orElseThrow(() -> new AssertionError("REMOVE operation not found"))
+                .getPaths();
+        assertComplexPaths(removePaths);
+
+        List<String> replacePaths = request.getAllowedOperations().stream()
+                .filter(op -> op.getOp() == Operation.REPLACE)
+                .findFirst()
+                .orElseThrow(() -> new AssertionError("REPLACE operation not found"))
+                .getPaths();
+        assertComplexPaths(replacePaths);
+    }
+
+    @Test
+    public void testGetRemoveAndReplacePathsWithArrays() throws ActionExecutionRequestBuilderException {
+
+        Map<String, Object> customClaims = new HashMap<>();
+        customClaims.put("array_claim", new String[]{"val1", "val2"});
+
+        IDTokenDTO idTokenDTO = getMockIDTokenDTO();
+        idTokenDTO.setCustomOIDCClaims(customClaims);
+
+        FlowContext flowContext = FlowContext.create()
+                .add(TOKEN_REQUEST_MESSAGE_CONTEXT, getMockTokenMessageContext())
+                .add(ID_TOKEN_DTO, idTokenDTO)
+                .add(REQUEST_TYPE, REQUEST_TYPE_TOKEN);
+
+        ActionExecutionRequest request = preIssueIDTokenRequestBuilder.buildActionExecutionRequest(
+                flowContext, null);
+        List<String> removePaths = request.getAllowedOperations().get(1).getPaths();
+        List<String> replacePaths = request.getAllowedOperations().get(2).getPaths();
+
+        Assert.assertTrue(removePaths.contains("/idToken/claims/array_claim"));
+        Assert.assertTrue(removePaths.contains("/idToken/claims/array_claim/"));
+
+        Assert.assertTrue(replacePaths.contains("/idToken/claims/array_claim"));
+        Assert.assertTrue(replacePaths.contains("/idToken/claims/array_claim/"));
+    }
 }