diff --git a/src/main/java/com/scality/keycloak/truststore/JpaCertificateTruststoreProvider.java b/src/main/java/com/scality/keycloak/truststore/JpaCertificateTruststoreProvider.java index 101bee3..521b937 100644 --- a/src/main/java/com/scality/keycloak/truststore/JpaCertificateTruststoreProvider.java +++ b/src/main/java/com/scality/keycloak/truststore/JpaCertificateTruststoreProvider.java @@ -12,6 +12,8 @@ import java.security.cert.X509Certificate; import java.util.Base64; import java.util.List; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; import org.bouncycastle.asn1.x500.RDN; import org.bouncycastle.asn1.x500.X500Name; @@ -19,10 +21,8 @@ import org.bouncycastle.asn1.x500.style.IETFUtils; import org.bouncycastle.cert.jcajce.JcaX509CertificateHolder; import org.hibernate.CacheMode; -import org.hibernate.exception.GenericJDBCException; import org.hibernate.jpa.AvailableHints; import org.jboss.logging.Logger; -import org.keycloak.common.util.Retry; import org.keycloak.connections.jpa.JpaConnectionProvider; import org.keycloak.models.KeycloakSession; import org.keycloak.models.utils.KeycloakModelUtils; @@ -32,7 +32,6 @@ import jakarta.persistence.NoResultException; import jakarta.ws.rs.NotFoundException; import jakarta.ws.rs.core.Response.Status; -import java.sql.SQLException; public class JpaCertificateTruststoreProvider implements CertificateTruststoreProvider { @@ -43,52 +42,25 @@ public JpaCertificateTruststoreProvider(KeycloakSession session) { this.session = session; } - /*** - * - * @return EntityManager - */ private EntityManager getEntityManager() { return session.getProvider(JpaConnectionProvider.class).getEntityManager(); } /** - * Checks if the exception is related to a closed connection or statement. - * These errors can occur when the connection pool closes connections due to timeouts - * or when transactions are committed/rolled back prematurely. + * Executes a database query in a completely new Keycloak session with a fresh JDBC connection. + * Used on retry when the current session's connection is closed — em.clear() on the same + * EntityManager does NOT get a new connection, so we must open a new session entirely. */ - private boolean isConnectionClosedError(Throwable e) { - if (e == null) { - return false; - } - - String message = e.getMessage(); - if (message != null) { - String lowerMessage = message.toLowerCase(); - if (lowerMessage.contains("connection is closed") || - lowerMessage.contains("this statement has been closed") || - lowerMessage.contains("statement has been closed") || - lowerMessage.contains("connection closed")) { - return true; - } - } - - // Check for SQLException with specific SQL states - if (e instanceof SQLException) { - SQLException sqlEx = (SQLException) e; - String sqlState = sqlEx.getSQLState(); - // SQLState 55000 is a generic PostgreSQL error that can indicate connection issues - if ("55000".equals(sqlState) || sqlState == null) { - return true; + private T executeInNewSession(Function query) { + AtomicReference resultRef = new AtomicReference<>(); + KeycloakModelUtils.runJobInTransaction( + session.getKeycloakSessionFactory(), + newSession -> { + EntityManager newEm = newSession.getProvider(JpaConnectionProvider.class).getEntityManager(); + resultRef.set(query.apply(newEm)); } - } - - // Check for GenericJDBCException which wraps SQL exceptions - if (e instanceof GenericJDBCException) { - return isConnectionClosedError(e.getCause()); - } - - // Recursively check cause - return isConnectionClosedError(e.getCause()); + ); + return resultRef.get(); } @Override @@ -124,32 +96,13 @@ private CertificateRepresentation toCertificateRepresentation(TruststoreEntity e @Override public CertificateRepresentation getCertificate(String alias) { try { - return Retry.call((iteration) -> { - try { - TruststoreEntity certificate = getEntityManager() - .createNamedQuery("findByAlias", TruststoreEntity.class) - .setParameter("alias", alias) - .getSingleResult(); - return toCertificateRepresentation(certificate); - } catch (NoResultException e) { - throw new NotFoundException("Certificate not found"); - } catch (RuntimeException e) { - // Only retry on connection closed errors - if (isConnectionClosedError(e) && iteration < 2) { - logger.debugf("Connection closed error on getCertificate, retrying (iteration %d)", iteration); - getEntityManager().clear(); - throw e; - } - throw e; - } - }, 3, 50); // 3 attempts with 50ms delay - } catch (NotFoundException e) { - throw e; - } catch (Exception e) { - if (e.getCause() instanceof NotFoundException) { - throw (NotFoundException) e.getCause(); - } - throw new RuntimeException("Failed to get certificate: " + alias, e); + TruststoreEntity certificate = getEntityManager() + .createNamedQuery("findByAlias", TruststoreEntity.class) + .setParameter("alias", alias) + .getSingleResult(); + return toCertificateRepresentation(certificate); + } catch (NoResultException e) { + throw new NotFoundException("Certificate not found"); } } @@ -246,65 +199,29 @@ public void removeCertificate(String alias) { @Override public CertificateRepresentation[] getCertificates() { - try { - return Retry.call((iteration) -> { - try { - // Removed getEntityManager().clear() as it's unnecessary for read operations - // and can cause issues with connection state - @SuppressWarnings("unchecked") - List list = (List) getEntityManager() - .createNativeQuery("select t.id, t.alias, t.certificate, t.is_root_ca from truststore t", - TruststoreEntity.class) - .setHint(AvailableHints.HINT_CACHEABLE, false) - .setHint(AvailableHints.HINT_CACHE_MODE, CacheMode.IGNORE) - .getResultList(); - return list.stream() - .map(this::toCertificateRepresentation) - .toArray(CertificateRepresentation[]::new); - } catch (RuntimeException e) { - // Only retry on connection closed errors - if (isConnectionClosedError(e) && iteration < 2) { - logger.debugf("Connection closed error on getCertificates, retrying (iteration %d)", iteration); - // Clear the entity manager to force a new connection on retry - getEntityManager().clear(); - throw e; - } - throw e; - } - }, 3, 50); // 3 attempts with 50ms delay - } catch (Exception e) { - logger.error("Failed to get certificates after retries", e); - throw new RuntimeException("Failed to get certificates", e); - } + return executeInNewSession(newEm -> { + @SuppressWarnings("unchecked") + List list = (List) newEm + .createNativeQuery("select t.id, t.alias, t.certificate, t.is_root_ca from truststore t", + TruststoreEntity.class) + .setHint(AvailableHints.HINT_CACHEABLE, false) + .setHint(AvailableHints.HINT_CACHE_MODE, CacheMode.IGNORE) + .getResultList(); + return list.stream() + .map(this::toCertificateRepresentation) + .toArray(CertificateRepresentation[]::new); + }); } @Override public CertificateRepresentation[] getCertificates(boolean isRootCA) { - try { - return Retry.call((iteration) -> { - try { - return getEntityManager() - .createNamedQuery("findByIsRootCA", TruststoreEntity.class) - .setParameter("isRootCA", isRootCA) - .getResultList() - .stream() - .map(this::toCertificateRepresentation) - .toArray(CertificateRepresentation[]::new); - } catch (RuntimeException e) { - // Only retry on connection closed errors - if (isConnectionClosedError(e) && iteration < 2) { - logger.debugf("Connection closed error on getCertificates(isRootCA=%s), retrying (iteration %d)", isRootCA, iteration); - // Clear the entity manager to force a new connection on retry - getEntityManager().clear(); - throw e; - } - throw e; - } - }, 3, 50); // 3 attempts with 50ms delay - } catch (Exception e) { - logger.errorf("Failed to get certificates (isRootCA=%s) after retries", isRootCA, e); - throw new RuntimeException("Failed to get certificates", e); - } + return executeInNewSession(newEm -> newEm + .createNamedQuery("findByIsRootCA", TruststoreEntity.class) + .setParameter("isRootCA", isRootCA) + .getResultList() + .stream() + .map(this::toCertificateRepresentation) + .toArray(CertificateRepresentation[]::new)); } @Override