Add user session notes back on importNewUser

Closes: #58
Co-authored-by: mikaelkaron <[email protected]>

src/main/java/org/keycloak/broker/bankid/BankidIdentityProvider.java

@@ -2,18 +2,21 @@
 
 import java.net.URI;
 import java.net.URISyntaxException;
+import java.util.Map;
 
 import jakarta.ws.rs.core.Response;
 
 import org.apache.http.client.HttpClient;
 import org.keycloak.broker.provider.AbstractIdentityProvider;
 import org.keycloak.broker.provider.AuthenticationRequest;
+import org.keycloak.broker.provider.BrokeredIdentityContext;
 import org.keycloak.connections.httpclient.HttpClientBuilder;
 import org.keycloak.connections.httpclient.ProxyMappings;
 import org.keycloak.events.EventBuilder;
 import org.keycloak.models.FederatedIdentityModel;
 import org.keycloak.models.KeycloakSession;
 import org.keycloak.models.RealmModel;
+import org.keycloak.models.UserModel;
 
 public class BankidIdentityProvider extends AbstractIdentityProvider<BankidIdentityProviderConfig> {
 
@@ -46,14 +49,38 @@ public Response retrieveToken(KeycloakSession session, FederatedIdentityModel id
 		return Response.ok(identity.getToken()).build();
 	}
 
-	public ProxyMappings generateProxyMapping(){
+	@Override
+	public void preprocessFederatedIdentity(KeycloakSession session, RealmModel realm,
+			BrokeredIdentityContext context) {
+		context.getContextData().putAll(context.getAuthenticationSession().getUserSessionNotes());
+	}
+
+	@Override
+	public void importNewUser(KeycloakSession session, RealmModel realm, UserModel user,
+			BrokeredIdentityContext context) {
+		// Here context.getAuthenticationSession().getUserSessionNotes() is empty
+		// use context data to retrieve information stored in {@link BankidIdentProvider#preprocessFederatedIdentity()}
+		Map<String, Object> contextData = context.getContextData();
+		// Itterate over the context data to extract the required information:
+		for (Map.Entry<String, Object> entry : contextData.entrySet()) {
+			String key = entry.getKey();
+			// Add the value to the user session notes if key starts with provider config
+			// alias since it means it was added by the BankidEndpoint
+			if (key.startsWith(this.getConfig().getAlias())) {
+				Object value = entry.getValue();
+				context.getAuthenticationSession().setUserSessionNote(key, value.toString());
+			}
+		}
+	}
+
+	public ProxyMappings generateProxyMapping() {
 		String httpsProxy = System.getenv("HTTPS_PROXY");
-		if(httpsProxy == null){
+		if (httpsProxy == null) {
 			httpsProxy = System.getenv("https_proxy");
 		}
 
 		String noProxy = System.getenv("NO_PROXY");
-		if(noProxy == null){
+		if (noProxy == null) {
 			noProxy = System.getenv("no_proxy");
 		}