Token based authentication integration with core extension

pom.xml

@@ -75,6 +75,12 @@
 			<version>2.11.0</version>
 		</dependency>
 
+		<dependency>
+			<groupId>redis.clients.authentication</groupId>
+			<artifactId>redis-authx-core</artifactId>
+			<version>0.1.0-SNAPSHOT</version>
+		</dependency>
+
 		<!-- Optional dependencies -->
 
 		<!-- UNIX socket connection support -->
@@ -150,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>

src/main/java/redis/clients/jedis/Connection.java

@@ -14,12 +14,14 @@
 import java.util.List;
 import java.util.Map;
 import java.util.function.Supplier;
+import java.util.concurrent.atomic.AtomicReference;
 
 import redis.clients.jedis.Protocol.Command;
 import redis.clients.jedis.Protocol.Keyword;
 import redis.clients.jedis.annots.Experimental;
 import redis.clients.jedis.args.ClientAttributeOption;
 import redis.clients.jedis.args.Rawable;
+import redis.clients.jedis.authentication.AuthXManager;
 import redis.clients.jedis.commands.ProtocolCommand;
 import redis.clients.jedis.exceptions.JedisConnectionException;
 import redis.clients.jedis.exceptions.JedisDataException;
@@ -44,6 +46,9 @@ public class Connection implements Closeable {
   private String strVal;
   protected String server;
   protected String version;
+  protected AtomicReference<RedisCredentials> currentCredentials = new AtomicReference<RedisCredentials>(
+      null);
+  private boolean isTokenBasedAuthenticationEnabled = false;
 
   public Connection() {
     this(Protocol.DEFAULT_HOST, Protocol.DEFAULT_PORT);
@@ -93,8 +98,8 @@ public String toIdentityString() {
     SocketAddress remoteAddr = socket.getRemoteSocketAddress();
     SocketAddress localAddr = socket.getLocalSocketAddress();
     if (remoteAddr != null) {
-      strVal = String.format("%s{id: 0x%X, L:%s %c R:%s}", className, id,
-          localAddr, (broken ? '!' : '-'), remoteAddr);
+      strVal = String.format("%s{id: 0x%X, L:%s %c R:%s}", className, id, localAddr,
+        (broken ? '!' : '-'), remoteAddr);
     } else if (localAddr != null) {
       strVal = String.format("%s{id: 0x%X, L:%s}", className, id, localAddr);
     } else {
@@ -438,8 +443,8 @@ private static boolean validateClientInfo(String info) {
     for (int i = 0; i < info.length(); i++) {
       char c = info.charAt(i);
       if (c < '!' || c > '~') {
-        throw new JedisValidationException("client info cannot contain spaces, "
-            + "newlines or special characters.");
+        throw new JedisValidationException(
+            "client info cannot contain spaces, " + "newlines or special characters.");
       }
     }
     return true;
@@ -451,7 +456,14 @@ protected void initializeFromClientConfig(final JedisClientConfig config) {
 
       protocol = config.getRedisProtocol();
 
-      final Supplier<RedisCredentials> credentialsProvider = config.getCredentialsProvider();
+      Supplier<RedisCredentials> credentialsProvider = config.getCredentialsProvider();
+
+      AuthXManager authXManager = config.getAuthXManager();
+      if (authXManager != null) {
+        isTokenBasedAuthenticationEnabled = true;
+        credentialsProvider = authXManager;
+      }
+
       if (credentialsProvider instanceof RedisCredentialsProvider) {
         final RedisCredentialsProvider redisCredentialsProvider = (RedisCredentialsProvider) credentialsProvider;
         try {
@@ -469,7 +481,8 @@ protected void initializeFromClientConfig(final JedisClientConfig config) {
 
       String clientName = config.getClientName();
       if (clientName != null && validateClientInfo(clientName)) {
-        fireAndForgetMsg.add(new CommandArguments(Command.CLIENT).add(Keyword.SETNAME).add(clientName));
+        fireAndForgetMsg
+            .add(new CommandArguments(Command.CLIENT).add(Keyword.SETNAME).add(clientName));
       }
 
       ClientSetInfoConfig setInfoConfig = config.getClientSetInfoConfig();
@@ -525,12 +538,13 @@ private void helloAndAuth(final RedisProtocol protocol, final RedisCredentials c
     if (protocol != null && credentials != null && credentials.getUser() != null) {
       byte[] rawPass = encodeToBytes(credentials.getPassword());
       try {
-        helloResult = hello(encode(protocol.version()), Keyword.AUTH.getRaw(), encode(credentials.getUser()), rawPass);
+        helloResult = hello(encode(protocol.version()), Keyword.AUTH.getRaw(),
+          encode(credentials.getUser()), rawPass);
       } finally {
         Arrays.fill(rawPass, (byte) 0); // clear sensitive data
       }
     } else {
-      auth(credentials);
+      authenticate(credentials);
       helloResult = protocol == null ? null : hello(encode(protocol.version()));
     }
     if (helloResult != null) {
@@ -542,7 +556,11 @@ private void helloAndAuth(final RedisProtocol protocol, final RedisCredentials c
     // handled in RedisCredentialsProvider.cleanUp()
   }
 
-  private void auth(RedisCredentials credentials) {
+  public void setCredentials(RedisCredentials credentials) {
+    currentCredentials.set(credentials);
+  }
+
+  private void authenticate(RedisCredentials credentials) {
     if (credentials == null || credentials.getPassword() == null) {
       return;
     }
@@ -559,6 +577,10 @@ private void auth(RedisCredentials credentials) {
     getStatusCodeReply();
   }
 
+  public void reAuthenticate() {
+    authenticate(currentCredentials.getAndSet(null));
+  }
+
   protected Map<String, Object> hello(byte[]... args) {
     sendCommand(Command.HELLO, args);
     return BuilderFactory.ENCODED_OBJECT_MAP.build(getOne());
@@ -585,4 +607,8 @@ public boolean ping() {
     }
     return true;
   }
+
+  public boolean isTokenBasedAuthenticationEnabled() {
+    return isTokenBasedAuthenticationEnabled;
+  }
 }

src/main/java/redis/clients/jedis/ConnectionFactory.java

@@ -6,7 +6,11 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.function.Supplier;
+
 import redis.clients.jedis.annots.Experimental;
+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;
@@ -20,28 +24,52 @@ public class ConnectionFactory implements PooledObjectFactory<Connection> {
 
   private final JedisSocketFactory jedisSocketFactory;
   private final JedisClientConfig clientConfig;
-  private Cache clientSideCache = null;
+  private final Cache clientSideCache;
+  private final Supplier<Connection> objectMaker;
+
+  private final AuthXEventListener authXEventListener;
 
   public ConnectionFactory(final HostAndPort hostAndPort) {
-    this.clientConfig = DefaultJedisClientConfig.builder().build();
-    this.jedisSocketFactory = new DefaultJedisSocketFactory(hostAndPort);
+    this(hostAndPort, DefaultJedisClientConfig.builder().build(), null);
   }
 
   public ConnectionFactory(final HostAndPort hostAndPort, final JedisClientConfig clientConfig) {
-    this.clientConfig = clientConfig;
-    this.jedisSocketFactory = new DefaultJedisSocketFactory(hostAndPort, this.clientConfig);
+    this(hostAndPort, clientConfig, null);
   }
 
   @Experimental
-  public ConnectionFactory(final HostAndPort hostAndPort, final JedisClientConfig clientConfig, Cache csCache) {
-    this.clientConfig = clientConfig;
-    this.jedisSocketFactory = new DefaultJedisSocketFactory(hostAndPort, this.clientConfig);
-    this.clientSideCache = csCache;
+  public ConnectionFactory(final HostAndPort hostAndPort, final JedisClientConfig clientConfig,
+      Cache csCache) {
+    this(new DefaultJedisSocketFactory(hostAndPort, clientConfig), clientConfig, csCache);
   }
 
-  public ConnectionFactory(final JedisSocketFactory jedisSocketFactory, final JedisClientConfig clientConfig) {
-    this.clientConfig = clientConfig;
+  public ConnectionFactory(final JedisSocketFactory jedisSocketFactory,
+      final JedisClientConfig clientConfig) {
+    this(jedisSocketFactory, clientConfig, null);
+  }
+
+  private ConnectionFactory(final JedisSocketFactory jedisSocketFactory,
+      final JedisClientConfig clientConfig, Cache csCache) {
+
     this.jedisSocketFactory = jedisSocketFactory;
+    this.clientSideCache = csCache;
+    this.clientConfig = clientConfig;
+
+    AuthXManager authXManager = clientConfig.getAuthXManager();
+    if (authXManager == null) {
+      this.objectMaker = connectionSupplier();
+      this.authXEventListener = AuthXEventListener.NOOP_LISTENER;
+    } else {
+      Supplier<Connection> supplier = connectionSupplier();
+      this.objectMaker = () -> (Connection) authXManager.addConnection(supplier.get());
+      this.authXEventListener = authXManager.getListener();
+      authXManager.start();
+    }
+  }
+
+  private Supplier<Connection> connectionSupplier() {
+    return clientSideCache == null ? () -> new Connection(jedisSocketFactory, clientConfig)
+        : () -> new CacheConnection(jedisSocketFactory, clientConfig, clientSideCache);
   }
 
   @Override
@@ -64,8 +92,7 @@ public void destroyObject(PooledObject<Connection> pooledConnection) throws Exce
   @Override
   public PooledObject<Connection> makeObject() throws Exception {
     try {
-      Connection jedis = clientSideCache == null ? new Connection(jedisSocketFactory, clientConfig)
-          : new CacheConnection(jedisSocketFactory, clientConfig, clientSideCache);
+      Connection jedis = objectMaker.get();
       return new DefaultPooledObject<>(jedis);
     } catch (JedisException je) {
       logger.debug("Error while makeObject", je);
@@ -76,17 +103,32 @@ public PooledObject<Connection> makeObject() throws Exception {
   @Override
   public void passivateObject(PooledObject<Connection> pooledConnection) throws Exception {
     // TODO maybe should select db 0? Not sure right now.
+    Connection jedis = pooledConnection.getObject();
+    reAuthenticate(jedis);
   }
 
   @Override
   public boolean validateObject(PooledObject<Connection> pooledConnection) {
     final Connection jedis = pooledConnection.getObject();
     try {
       // check HostAndPort ??
-      return jedis.isConnected() && jedis.ping();
+      if (!jedis.isConnected()) {
+        return false;
+      }
+      reAuthenticate(jedis);
+      return jedis.ping();
     } catch (final Exception e) {
       logger.warn("Error while validating pooled Connection object.", e);
       return false;
     }
   }
+
+  private void reAuthenticate(Connection jedis) throws Exception {
+    try {
+      jedis.reAuthenticate();
+    } catch (Exception e) {
+      authXEventListener.onConnectionAuthenticationError(e);
+      throw e;
+    }
+  }
 }

src/main/java/redis/clients/jedis/ConnectionPool.java

@@ -2,19 +2,27 @@
 
 import org.apache.commons.pool2.PooledObjectFactory;
 import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
+
 import redis.clients.jedis.annots.Experimental;
+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 AuthXManager authXManager;
+
   public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig) {
     this(new ConnectionFactory(hostAndPort, clientConfig));
+    attachAuthenticationListener(clientConfig.getAuthXManager());
   }
 
   @Experimental
-  public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig, Cache clientSideCache) {
+  public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig,
+      Cache clientSideCache) {
     this(new ConnectionFactory(hostAndPort, clientConfig, clientSideCache));
+    attachAuthenticationListener(clientConfig.getAuthXManager());
   }
 
   public ConnectionPool(PooledObjectFactory<Connection> factory) {
@@ -24,12 +32,14 @@ public ConnectionPool(PooledObjectFactory<Connection> factory) {
   public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig,
       GenericObjectPoolConfig<Connection> poolConfig) {
     this(new ConnectionFactory(hostAndPort, clientConfig), poolConfig);
+    attachAuthenticationListener(clientConfig.getAuthXManager());
   }
 
   @Experimental
-  public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig, Cache clientSideCache,
-      GenericObjectPoolConfig<Connection> poolConfig) {
+  public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig,
+      Cache clientSideCache, GenericObjectPoolConfig<Connection> poolConfig) {
     this(new ConnectionFactory(hostAndPort, clientConfig, clientSideCache), poolConfig);
+    attachAuthenticationListener(clientConfig.getAuthXManager());
   }
 
   public ConnectionPool(PooledObjectFactory<Connection> factory,
@@ -43,4 +53,29 @@ public Connection getResource() {
     conn.setHandlingPool(this);
     return conn;
   }
+
+  @Override
+  public void close() {
+    try {
+      if (authXManager != null) {
+        authXManager.stop();
+      }
+    } finally {
+      super.close();
+    }
+  }
+
+  private void attachAuthenticationListener(AuthXManager authXManager) {
+    this.authXManager = authXManager;
+    if (authXManager != null) {
+      authXManager.addPostAuthenticationHook(token -> {
+        try {
+          // this is to trigger validations on each connection via ConnectionFactory
+          evict();
+        } catch (Exception e) {
+          throw new JedisException("Failed to evict connections from pool", e);
+        }
+      });
+    }
+  }
 }