From 7e78d29f8dd5a74f02be90d75addee392dd7536e Mon Sep 17 00:00:00 2001
From: Douglas Palmer <dpalmer@redhat.com>
Date: Wed, 22 Nov 2023 11:51:17 -0800
Subject: [PATCH] NPE in User Session Note mapper on Token Exchange Closes
 #24200

Signed-off-by: Douglas Palmer <dpalmer@redhat.com>
---
 .../broker/provider/BrokeredIdentityContext.java         | 9 +++++++++
 .../oidc/mappers/ClaimToUserSessionNoteMapper.java       | 7 ++++++-
 .../protocol/oidc/DefaultTokenExchangeProvider.java      | 2 ++
 3 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/server-spi-private/src/main/java/org/keycloak/broker/provider/BrokeredIdentityContext.java b/server-spi-private/src/main/java/org/keycloak/broker/provider/BrokeredIdentityContext.java
index d882bdd271..b83bb53e30 100755
--- a/server-spi-private/src/main/java/org/keycloak/broker/provider/BrokeredIdentityContext.java
+++ b/server-spi-private/src/main/java/org/keycloak/broker/provider/BrokeredIdentityContext.java
@@ -48,6 +48,7 @@ public class BrokeredIdentityContext {
     private IdentityProviderModel idpConfig;
     private IdentityProvider idp;
     private Map<String, Object> contextData = new HashMap<>();
+    private Map<String, String> claims = new HashMap<>();
     private AuthenticationSessionModel authenticationSession;
 
     public BrokeredIdentityContext(String id) {
@@ -161,6 +162,14 @@ public class BrokeredIdentityContext {
         this.contextData = contextData;
     }
 
+    public Map<String, String> getClaims() {
+        return claims;
+    }
+
+    public void setClaims(Map<String, String> claims) {
+        this.claims = claims;
+    }
+
     // Set the attribute, which will be available on "Update profile" page and in authenticators
     public void setUserAttribute(String attributeName, String attributeValue) {
         List<String> list = new ArrayList<>();
diff --git a/services/src/main/java/org/keycloak/broker/oidc/mappers/ClaimToUserSessionNoteMapper.java b/services/src/main/java/org/keycloak/broker/oidc/mappers/ClaimToUserSessionNoteMapper.java
index d75f012c3b..440994ccd5 100644
--- a/services/src/main/java/org/keycloak/broker/oidc/mappers/ClaimToUserSessionNoteMapper.java
+++ b/services/src/main/java/org/keycloak/broker/oidc/mappers/ClaimToUserSessionNoteMapper.java
@@ -127,7 +127,12 @@ public class ClaimToUserSessionNoteMapper extends AbstractClaimMapper {
                             : valueEquals(value, claimValue);
 
                     if (claimValuesMatch) {
-                        context.getAuthenticationSession().setUserSessionNote(claim.getKey(), claimValue);
+                        if(context.getAuthenticationSession() != null) {
+                            context.getAuthenticationSession().setUserSessionNote(claim.getKey(), claimValue);
+                        }
+                        else {
+                            context.getClaims().put(claim.getKey(), claimValue);
+                        }
                     }
                 }
             }
diff --git a/services/src/main/java/org/keycloak/protocol/oidc/DefaultTokenExchangeProvider.java b/services/src/main/java/org/keycloak/protocol/oidc/DefaultTokenExchangeProvider.java
index 58cc266d91..49977e36f6 100644
--- a/services/src/main/java/org/keycloak/protocol/oidc/DefaultTokenExchangeProvider.java
+++ b/services/src/main/java/org/keycloak/protocol/oidc/DefaultTokenExchangeProvider.java
@@ -515,6 +515,8 @@ public class DefaultTokenExchangeProvider implements TokenExchangeProvider {
         userSession.setNote(IdentityProvider.EXTERNAL_IDENTITY_PROVIDER, externalIdpModel.get().getAlias());
         userSession.setNote(IdentityProvider.FEDERATED_ACCESS_TOKEN, subjectToken);
 
+        context.getClaims().forEach((k, v) -> userSession.setNote(k, v));
+
         return exchangeClientToClient(user, userSession, null, false);
     }