- move autxhmanager creation to user space

- introduce authenticationeventlisteners - clenaup in connectionpool - add entraidtestcontext - add redisintegrationtests - fix failing tokenbasedauthentication unit&integ tests
redis · Dec 5, 2024 · 4e71535 · 4e71535
1 parent 523fe42
commit 4e71535
Show file tree

Hide file tree

Showing 12 changed files with 734 additions and 145 deletions.
diff --git a/pom.xml b/pom.xml
@@ -156,6 +156,13 @@
 			<scope>test</scope>
 		</dependency>
 
+		<dependency>
+			<groupId>redis.clients.authentication</groupId>
+			<artifactId>redis-authx-entraid</artifactId>
+			<version>0.1.0-SNAPSHOT</version>
+			<scope>test</scope>
+		</dependency>
+
 		<!-- circuit breaker / failover -->
 		<dependency>
 			<groupId>io.github.resilience4j</groupId>

diff --git a/src/main/java/redis/clients/jedis/ConnectionFactory.java b/src/main/java/redis/clients/jedis/ConnectionFactory.java
@@ -11,7 +11,8 @@
 import java.util.function.Supplier;
 
 import redis.clients.jedis.annots.Experimental;
-import redis.clients.jedis.authentication.JedisAuthXManager;
+import redis.clients.jedis.authentication.AuthXManager;
+import redis.clients.jedis.authentication.AuthXEventListener;
 import redis.clients.jedis.csc.Cache;
 import redis.clients.jedis.csc.CacheConnection;
 import redis.clients.jedis.exceptions.JedisException;
@@ -28,41 +29,43 @@ public class ConnectionFactory implements PooledObjectFactory<Connection> {
   private final Cache clientSideCache;
   private final Supplier<Connection> objectMaker;
 
+  private final AuthXEventListener authenticationEventListener;
+
   public ConnectionFactory(final HostAndPort hostAndPort) {
-    this(hostAndPort, DefaultJedisClientConfig.builder().build(), null, null);
+    this(hostAndPort, DefaultJedisClientConfig.builder().build(), null);
   }
 
   public ConnectionFactory(final HostAndPort hostAndPort, final JedisClientConfig clientConfig) {
-    this(hostAndPort, clientConfig, null, null);
+    this(hostAndPort, clientConfig, null);
   }
 
   @Experimental
   public ConnectionFactory(final HostAndPort hostAndPort, final JedisClientConfig clientConfig,
-      Cache csCache, JedisAuthXManager authXManager) {
-    this(new DefaultJedisSocketFactory(hostAndPort, clientConfig), clientConfig, csCache,
-        authXManager);
+      Cache csCache) {
+    this(new DefaultJedisSocketFactory(hostAndPort, clientConfig), clientConfig, csCache);
   }
 
   public ConnectionFactory(final JedisSocketFactory jedisSocketFactory,
       final JedisClientConfig clientConfig) {
-    this(jedisSocketFactory, clientConfig, null, null);
+    this(jedisSocketFactory, clientConfig, null);
   }
 
   private ConnectionFactory(final JedisSocketFactory jedisSocketFactory,
-      final JedisClientConfig clientConfig, Cache csCache, JedisAuthXManager authXManager) {
+      final JedisClientConfig clientConfig, Cache csCache) {
 
     this.jedisSocketFactory = jedisSocketFactory;
     this.clientSideCache = csCache;
+    AuthXManager authXManager = clientConfig.getAuthXManager();
 
     if (authXManager == null) {
       this.clientConfig = clientConfig;
       this.objectMaker = connectionSupplier();
+      this.authenticationEventListener = AuthXEventListener.NOOP_LISTENER;
     } else {
-      this.clientConfig = replaceCredentialsProvider(clientConfig,
-        authXManager);
+      this.clientConfig = replaceCredentialsProvider(clientConfig, authXManager);
       Supplier<Connection> supplier = connectionSupplier();
       this.objectMaker = () -> (Connection) authXManager.addConnection(supplier.get());
-
+      this.authenticationEventListener = authXManager.getListener();
       try {
         authXManager.start(true);
       } catch (InterruptedException | ExecutionException | TimeoutException e) {
@@ -114,15 +117,25 @@ public PooledObject<Connection> makeObject() throws Exception {
   public void passivateObject(PooledObject<Connection> pooledConnection) throws Exception {
     // TODO maybe should select db 0? Not sure right now.
     Connection jedis = pooledConnection.getObject();
-    jedis.reAuth();
+    try {
+      jedis.reAuth();
+    } catch (Exception e) {
+      authenticationEventListener.onConnectionAuthenticationError(e);
+      throw e;
+    }
   }
 
   @Override
   public boolean validateObject(PooledObject<Connection> pooledConnection) {
     final Connection jedis = pooledConnection.getObject();
     try {
       // check HostAndPort ??
-      jedis.reAuth();
+      try {
+        jedis.reAuth();
+      } catch (Exception e) {
+        authenticationEventListener.onConnectionAuthenticationError(e);
+        throw e;
+      }
       return jedis.isConnected() && jedis.ping();
     } catch (final Exception e) {
       logger.warn("Error while validating pooled Connection object.", e);

diff --git a/src/main/java/redis/clients/jedis/ConnectionPool.java b/src/main/java/redis/clients/jedis/ConnectionPool.java
@@ -4,36 +4,25 @@
 import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
 
 import redis.clients.jedis.annots.Experimental;
-import redis.clients.jedis.authentication.JedisAuthXManager;
+import redis.clients.jedis.authentication.AuthXManager;
 import redis.clients.jedis.csc.Cache;
 import redis.clients.jedis.exceptions.JedisException;
 import redis.clients.jedis.util.Pool;
 
 public class ConnectionPool extends Pool<Connection> {
 
-  private JedisAuthXManager authXManager;
+  private AuthXManager authXManager;
 
   public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig) {
-    this(hostAndPort, clientConfig, createAuthXManager(clientConfig));
-  }
-
-  public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig,
-      JedisAuthXManager authXManager) {
-    this(new ConnectionFactory(hostAndPort, clientConfig, null, authXManager));
-    attachAuthenticationListener(authXManager);
+    this(new ConnectionFactory(hostAndPort, clientConfig));
+    attachAuthenticationListener(clientConfig.getAuthXManager());
   }
 
   @Experimental
   public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig,
       Cache clientSideCache) {
-    this(hostAndPort, clientConfig, clientSideCache, createAuthXManager(clientConfig));
-  }
-
-  @Experimental
-  public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig,
-      Cache clientSideCache, JedisAuthXManager authXManager) {
-    this(new ConnectionFactory(hostAndPort, clientConfig, clientSideCache, authXManager));
-    attachAuthenticationListener(authXManager);
+    this(new ConnectionFactory(hostAndPort, clientConfig, clientSideCache));
+    attachAuthenticationListener(clientConfig.getAuthXManager());
   }
 
   public ConnectionPool(PooledObjectFactory<Connection> factory) {
@@ -42,22 +31,15 @@ public ConnectionPool(PooledObjectFactory<Connection> factory) {
 
   public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig,
       GenericObjectPoolConfig<Connection> poolConfig) {
-    this(hostAndPort, clientConfig, null, createAuthXManager(clientConfig), poolConfig);
+    this(new ConnectionFactory(hostAndPort, clientConfig), poolConfig);
+    attachAuthenticationListener(clientConfig.getAuthXManager());
   }
 
   @Experimental
   public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig,
       Cache clientSideCache, GenericObjectPoolConfig<Connection> poolConfig) {
-    this(hostAndPort, clientConfig, clientSideCache, createAuthXManager(clientConfig), poolConfig);
-  }
-
-  @Experimental
-  public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig,
-      Cache clientSideCache, JedisAuthXManager authXManager,
-      GenericObjectPoolConfig<Connection> poolConfig) {
-    this(new ConnectionFactory(hostAndPort, clientConfig, clientSideCache, authXManager),
-        poolConfig);
-        attachAuthenticationListener(authXManager);
+    this(new ConnectionFactory(hostAndPort, clientConfig, clientSideCache), poolConfig);
+    attachAuthenticationListener(clientConfig.getAuthXManager());
   }
 
   public ConnectionPool(PooledObjectFactory<Connection> factory,
@@ -80,17 +62,10 @@ public void close() {
     super.close();
   }
 
-  private static JedisAuthXManager createAuthXManager(JedisClientConfig config) {
-    if (config.getTokenAuthConfig() != null) {
-      return new JedisAuthXManager(config.getTokenAuthConfig());
-    }
-    return null;
-  }
-
-  private void attachAuthenticationListener(JedisAuthXManager authXManager) {
+  private void attachAuthenticationListener(AuthXManager authXManager) {
     this.authXManager = authXManager;
     if (authXManager != null) {
-      authXManager.setListener(token -> {
+      authXManager.addPostAuthenticationHook(token -> {
         try {
           // this is to trigger validations on each connection via ConnectionFactory
           evict();

diff --git a/src/main/java/redis/clients/jedis/DefaultJedisClientConfig.java b/src/main/java/redis/clients/jedis/DefaultJedisClientConfig.java
@@ -5,7 +5,7 @@
 import javax.net.ssl.SSLParameters;
 import javax.net.ssl.SSLSocketFactory;
 
-import  redis.clients.authentication.core.TokenAuthConfig;
+import redis.clients.jedis.authentication.AuthXManager;
 
 public final class DefaultJedisClientConfig implements JedisClientConfig {
 
@@ -30,15 +30,15 @@ public final class DefaultJedisClientConfig implements JedisClientConfig {
 
   private final boolean readOnlyForRedisClusterReplicas;
 
-  private final TokenAuthConfig tokenAuthConfig;
+  private final AuthXManager authXManager;
 
   private DefaultJedisClientConfig(RedisProtocol protocol, int connectionTimeoutMillis,
       int soTimeoutMillis, int blockingSocketTimeoutMillis,
       Supplier<RedisCredentials> credentialsProvider, int database, String clientName, boolean ssl,
       SSLSocketFactory sslSocketFactory, SSLParameters sslParameters,
       HostnameVerifier hostnameVerifier, HostAndPortMapper hostAndPortMapper,
       ClientSetInfoConfig clientSetInfoConfig, boolean readOnlyForRedisClusterReplicas,
-      TokenAuthConfig tokenAuthConfig) {
+      AuthXManager authXManager) {
     this.redisProtocol = protocol;
     this.connectionTimeoutMillis = connectionTimeoutMillis;
     this.socketTimeoutMillis = soTimeoutMillis;
@@ -53,7 +53,8 @@ private DefaultJedisClientConfig(RedisProtocol protocol, int connectionTimeoutMi
     this.hostAndPortMapper = hostAndPortMapper;
     this.clientSetInfoConfig = clientSetInfoConfig;
     this.readOnlyForRedisClusterReplicas = readOnlyForRedisClusterReplicas;
-    this.tokenAuthConfig = tokenAuthConfig;
+    this.authXManager = authXManager;
+
   }
 
   @Override
@@ -93,8 +94,8 @@ public Supplier<RedisCredentials> getCredentialsProvider() {
   }
 
   @Override
-  public TokenAuthConfig getTokenAuthConfig() {
-    return tokenAuthConfig;
+  public AuthXManager getAuthXManager() {
+    return authXManager;
   }
 
   @Override
@@ -171,7 +172,7 @@ public static class Builder {
 
     private boolean readOnlyForRedisClusterReplicas = false;
 
-    private TokenAuthConfig tokenAuthConfig = null;
+    private AuthXManager authXManager;
 
     private Builder() {
     }
@@ -185,7 +186,7 @@ public DefaultJedisClientConfig build() {
       return new DefaultJedisClientConfig(redisProtocol, connectionTimeoutMillis,
           socketTimeoutMillis, blockingSocketTimeoutMillis, credentialsProvider, database,
           clientName, ssl, sslSocketFactory, sslParameters, hostnameVerifier, hostAndPortMapper,
-          clientSetInfoConfig, readOnlyForRedisClusterReplicas, tokenAuthConfig);
+          clientSetInfoConfig, readOnlyForRedisClusterReplicas, authXManager);
     }
 
     /**
@@ -287,8 +288,8 @@ public Builder readOnlyForRedisClusterReplicas() {
       return this;
     }
 
-    public Builder tokenAuthConfig(TokenAuthConfig tokenAuthConfig) {
-      this.tokenAuthConfig = tokenAuthConfig;
+    public Builder authXManager(AuthXManager authXManager) {
+      this.authXManager = authXManager;
       return this;
     }
 
@@ -307,7 +308,7 @@ public Builder from(JedisClientConfig instance) {
       this.hostAndPortMapper = instance.getHostAndPortMapper();
       this.clientSetInfoConfig = instance.getClientSetInfoConfig();
       this.readOnlyForRedisClusterReplicas = instance.isReadOnlyForRedisClusterReplicas();
-      this.tokenAuthConfig = instance.getTokenAuthConfig();
+      this.authXManager = instance.getAuthXManager();
       return this;
     }
   }
@@ -316,12 +317,12 @@ public static DefaultJedisClientConfig create(int connectionTimeoutMillis, int s
       int blockingSocketTimeoutMillis, String user, String password, int database,
       String clientName, boolean ssl, SSLSocketFactory sslSocketFactory,
       SSLParameters sslParameters, HostnameVerifier hostnameVerifier,
-      HostAndPortMapper hostAndPortMapper, TokenAuthConfig tokenAuthConfig) {
+      HostAndPortMapper hostAndPortMapper, AuthXManager authXManager) {
     return new DefaultJedisClientConfig(null, connectionTimeoutMillis, soTimeoutMillis,
         blockingSocketTimeoutMillis,
         new DefaultRedisCredentialsProvider(new DefaultRedisCredentials(user, password)), database,
         clientName, ssl, sslSocketFactory, sslParameters, hostnameVerifier, hostAndPortMapper, null,
-        false, tokenAuthConfig);
+        false, authXManager);
   }
 
   public static DefaultJedisClientConfig copyConfig(JedisClientConfig copy) {
@@ -330,6 +331,6 @@ public static DefaultJedisClientConfig copyConfig(JedisClientConfig copy) {
         copy.getCredentialsProvider(), copy.getDatabase(), copy.getClientName(), copy.isSsl(),
         copy.getSslSocketFactory(), copy.getSslParameters(), copy.getHostnameVerifier(),
         copy.getHostAndPortMapper(), copy.getClientSetInfoConfig(),
-        copy.isReadOnlyForRedisClusterReplicas(), copy.getTokenAuthConfig());
+        copy.isReadOnlyForRedisClusterReplicas(), copy.getAuthXManager());
   }
 }
diff --git a/src/main/java/redis/clients/jedis/JedisClientConfig.java b/src/main/java/redis/clients/jedis/JedisClientConfig.java
@@ -5,7 +5,7 @@
 import javax.net.ssl.SSLParameters;
 import javax.net.ssl.SSLSocketFactory;
 
-import  redis.clients.authentication.core.TokenAuthConfig;
+import redis.clients.jedis.authentication.AuthXManager;
 
 public interface JedisClientConfig {
 
@@ -47,10 +47,11 @@ default String getPassword() {
   }
 
   default Supplier<RedisCredentials> getCredentialsProvider() {
-    return new DefaultRedisCredentialsProvider(new DefaultRedisCredentials(getUser(), getPassword()));
+    return new DefaultRedisCredentialsProvider(
+        new DefaultRedisCredentials(getUser(), getPassword()));
   }
 
-  default TokenAuthConfig getTokenAuthConfig() {
+  default AuthXManager getAuthXManager() {
     return null;
   }
 

diff --git a/src/main/java/redis/clients/jedis/authentication/AuthXEventListener.java b/src/main/java/redis/clients/jedis/authentication/AuthXEventListener.java
@@ -0,0 +1,21 @@
+package redis.clients.jedis.authentication;
+
+public interface AuthXEventListener {
+
+    static AuthXEventListener NOOP_LISTENER = new AuthXEventListener() {
+
+        @Override
+        public void onIdentityProviderError(Exception reason) {
+        }
+
+        @Override
+        public void onConnectionAuthenticationError(Exception reason) {
+        }
+
+    };
+
+    public void onIdentityProviderError(Exception reason);
+
+    public void onConnectionAuthenticationError(Exception reason);
+
+}