last one worked ok, but how about this?

BIBSYSDEV · Nov 28, 2024 · ec6f1a4 · ec6f1a4
1 parent 52f026e
commit ec6f1a4
Showing 1 changed file with 15 additions and 7 deletions.
diff --git a/...en-generation-openid/src/main/java/no/unit/nva/cognito/UserSelectionUponLoginHandler.java b/...en-generation-openid/src/main/java/no/unit/nva/cognito/UserSelectionUponLoginHandler.java
@@ -32,6 +32,7 @@
 import com.amazonaws.services.lambda.runtime.events.CognitoUserPoolPreTokenGenerationEventV2.ClaimsAndScopeOverrideDetails;
 import com.amazonaws.services.lambda.runtime.events.CognitoUserPoolPreTokenGenerationEventV2.IdTokenGeneration;
 import com.amazonaws.services.lambda.runtime.events.CognitoUserPoolPreTokenGenerationEventV2.Response;
+import java.io.IOException;
 import java.net.URI;
 import java.time.Instant;
 import java.util.ArrayList;
@@ -494,7 +495,7 @@ private ClaimsAndScopeOverrideDetails buildOverrideClaims(List<String> groupsToO
 
         return ClaimsAndScopeOverrideDetails.builder()
                    //.withGroupOverrideDetails(groups)
-                   .withAccessTokenGeneration(buildAccessTokenGeneration(userAttributes))
+                   .withAccessTokenGeneration(buildAccessTokenGeneration(userAttributes, groupsToOverride))
                    .withIdTokenGeneration(buildIdTokenGeneration())
                    .build();
     }
@@ -505,13 +506,20 @@ private IdTokenGeneration buildIdTokenGeneration() {
     }
 
     @SuppressWarnings("PMD.UnusedFormalParameter")
-    private AccessTokenGeneration buildAccessTokenGeneration(List<AttributeType> userAttributes) {
+    private AccessTokenGeneration buildAccessTokenGeneration(List<AttributeType> userAttributes,
+                                                             List<String> groupsToOverride) {
+        var claims = userAttributes.stream()
+                         .filter(a -> !Arrays.stream(CLAIMS_TO_BE_SUPPRESSED_FROM_PUBLIC)
+                                           .toList()
+                                           .contains(a.name()))
+                         .collect(Collectors.toMap(AttributeType::name, AttributeType::value));
+        try {
+            claims.put("cognito:groups", JsonConfig.writeValueAsString(groupsToOverride));
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
         return AccessTokenGeneration.builder()
-                   .withClaimsToAddOrOverride(userAttributes.stream()
-                          .filter(a -> !Arrays.stream(CLAIMS_TO_BE_SUPPRESSED_FROM_PUBLIC)
-                                            .toList()
-                                            .contains(a.name()))
-                          .collect(Collectors.toMap(AttributeType::name, AttributeType::value)))
+                   .withClaimsToAddOrOverride(claims)
                    .build();
     }
 }