From 68fc5aa44b699f88ff8f3b907e5797d65ce51411 Mon Sep 17 00:00:00 2001
From: Pedro Igor <pigor.craveiro@gmail.com>
Date: Mon, 28 Apr 2025 15:13:30 -0300
Subject: [PATCH] Make sure LDAP connections are released when closing sessions

Closes #38660

Signed-off-by: Pedro Igor <pigor.craveiro@gmail.com>
---
 .../idm/store/ldap/LDAPContextManager.java    |  3 +-
 .../ldap/SessionBoundInitialLdapContext.java  | 47 +++++++++++++++++++
 2 files changed, 48 insertions(+), 2 deletions(-)
 create mode 100644 federation/ldap/src/main/java/org/keycloak/storage/ldap/idm/store/ldap/SessionBoundInitialLdapContext.java

diff --git a/federation/ldap/src/main/java/org/keycloak/storage/ldap/idm/store/ldap/LDAPContextManager.java b/federation/ldap/src/main/java/org/keycloak/storage/ldap/idm/store/ldap/LDAPContextManager.java
index 93f2d6c9c86..f409a9ec026 100644
--- a/federation/ldap/src/main/java/org/keycloak/storage/ldap/idm/store/ldap/LDAPContextManager.java
+++ b/federation/ldap/src/main/java/org/keycloak/storage/ldap/idm/store/ldap/LDAPContextManager.java
@@ -11,7 +11,6 @@ import org.keycloak.vault.VaultStringSecret;
 import javax.naming.AuthenticationException;
 import javax.naming.Context;
 import javax.naming.NamingException;
-import javax.naming.ldap.InitialLdapContext;
 import javax.naming.ldap.LdapContext;
 import javax.naming.ldap.StartTlsRequest;
 import javax.naming.ldap.StartTlsResponse;
@@ -78,7 +77,7 @@ public final class LDAPContextManager implements AutoCloseable {
                 connProp.put(LDAPConstants.CONNECTION_TRACE_BER, System.err);
             }
 
-            ldapContext = new InitialLdapContext(connProp, null);
+            ldapContext = new SessionBoundInitialLdapContext(session, connProp, null);
             if (ldapConfig.isStartTls()) {
                 SSLSocketFactory sslSocketFactory = null;
                 if (LDAPUtil.shouldUseTruststoreSpi(ldapConfig)) {
diff --git a/federation/ldap/src/main/java/org/keycloak/storage/ldap/idm/store/ldap/SessionBoundInitialLdapContext.java b/federation/ldap/src/main/java/org/keycloak/storage/ldap/idm/store/ldap/SessionBoundInitialLdapContext.java
new file mode 100644
index 00000000000..61025f3c84b
--- /dev/null
+++ b/federation/ldap/src/main/java/org/keycloak/storage/ldap/idm/store/ldap/SessionBoundInitialLdapContext.java
@@ -0,0 +1,47 @@
+/*
+ * Copyright 2025 Red Hat, Inc. and/or its affiliates
+ * and other contributors as indicated by the @author tags.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.keycloak.storage.ldap.idm.store.ldap;
+
+import org.keycloak.models.KeycloakSession;
+
+import javax.naming.NamingException;
+import javax.naming.ldap.Control;
+import javax.naming.ldap.InitialLdapContext;
+import java.util.Hashtable;
+
+/**
+ * A {@link InitialLdapContext} that binds instances of this class with the {@link KeycloakSession} so that any resource
+ * acquired during the session lifetime is closed when the session is closed.
+ */
+public final class SessionBoundInitialLdapContext extends InitialLdapContext {
+
+    public SessionBoundInitialLdapContext(KeycloakSession session, Hashtable<?, ?> environment, Control[] connCtls) throws NamingException {
+        super(environment, connCtls);
+        session.enlistForClose(() -> {
+                try {
+                    close();
+                } catch (NamingException e) {
+                    failedToCloseLdapContext(e);
+                }
+        });
+    }
+
+    private void failedToCloseLdapContext(NamingException e) {
+        throw new RuntimeException("Failed to close LDAP context", e);
+    }
+}