Token based authentication integration with core extension

pom.xml

@@ -9,7 +9,7 @@
 	<packaging>jar</packaging>
 	<groupId>redis.clients</groupId>
 	<artifactId>jedis</artifactId>
-	<version>5.3.0-SNAPSHOT</version>
+	<version>5.3.1-SNAPSHOT</version>
 	<name>Jedis</name>
 	<description>Jedis is a blazingly small and sane Redis java client.</description>
 	<url>https://github.com/redis/jedis</url>
@@ -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 -->

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

@@ -14,6 +14,7 @@
 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;
@@ -44,6 +45,8 @@ public class Connection implements Closeable {
   private String strVal;
   protected String server;
   protected String version;
+  protected AtomicReference<RedisCredentials> currentCredentials = new AtomicReference<RedisCredentials>(
+      null);
 
   public Connection() {
     this(Protocol.DEFAULT_HOST, Protocol.DEFAULT_PORT);
@@ -93,8 +96,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 +441,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;
@@ -469,7 +472,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 +529,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 +547,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);
+  }
+
+  public void authenticate(RedisCredentials credentials) {
     if (credentials == null || credentials.getPassword() == null) {
       return;
     }
@@ -559,6 +568,13 @@ private void auth(RedisCredentials credentials) {
     getStatusCodeReply();
   }
 
+  public void reAuth() {
+    RedisCredentials temp = currentCredentials.getAndSet(null);
+    if (temp != null) {
+      authenticate(temp);
+    }
+  }
+
   protected Map<String, Object> hello(byte[]... args) {
     sendCommand(Command.HELLO, args);
     return BuilderFactory.ENCODED_OBJECT_MAP.build(getOne());

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

@@ -6,7 +6,12 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeoutException;
+import java.util.function.Supplier;
+
 import redis.clients.jedis.annots.Experimental;
+import redis.clients.jedis.authentication.JedisAuthXManager;
 import redis.clients.jedis.csc.Cache;
 import redis.clients.jedis.csc.CacheConnection;
 import redis.clients.jedis.exceptions.JedisException;
@@ -20,28 +25,61 @@ 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;
 
   public ConnectionFactory(final HostAndPort hostAndPort) {
-    this.clientConfig = DefaultJedisClientConfig.builder().build();
-    this.jedisSocketFactory = new DefaultJedisSocketFactory(hostAndPort);
+    this(hostAndPort, DefaultJedisClientConfig.builder().build(), null, null);
   }
 
   public ConnectionFactory(final HostAndPort hostAndPort, final JedisClientConfig clientConfig) {
-    this.clientConfig = clientConfig;
-    this.jedisSocketFactory = new DefaultJedisSocketFactory(hostAndPort, this.clientConfig);
+    this(hostAndPort, clientConfig, null, 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, JedisAuthXManager authXManager) {
+    this(new DefaultJedisSocketFactory(hostAndPort, clientConfig), clientConfig, csCache,
+        authXManager);
+  }
+
+  public ConnectionFactory(final JedisSocketFactory jedisSocketFactory,
+      final JedisClientConfig clientConfig) {
+    this(jedisSocketFactory, clientConfig, null, null);
   }
 
-  public ConnectionFactory(final JedisSocketFactory jedisSocketFactory, final JedisClientConfig clientConfig) {
-    this.clientConfig = clientConfig;
+  private ConnectionFactory(final JedisSocketFactory jedisSocketFactory,
+      final JedisClientConfig clientConfig, Cache csCache, JedisAuthXManager authXManager) {
+
     this.jedisSocketFactory = jedisSocketFactory;
+    this.clientSideCache = csCache;
+
+    if (authXManager == null) {
+      this.clientConfig = clientConfig;
+      this.objectMaker = connectionSupplier();
+    } else {
+      this.clientConfig = replaceCredentialsProvider(clientConfig,
+        authXManager);
+      Supplier<Connection> supplier = connectionSupplier();
+      this.objectMaker = () -> (Connection) authXManager.addConnection(supplier.get());
+
+      try {
+        authXManager.start(true);
+      } catch (InterruptedException | ExecutionException | TimeoutException e) {
+        throw new JedisException("AuthXManager failed to start!", e);
+      }
+    }
+  }
+
+  private JedisClientConfig replaceCredentialsProvider(JedisClientConfig origin,
+      Supplier<RedisCredentials> newCredentialsProvider) {
+    return DefaultJedisClientConfig.builder().from(origin)
+        .credentialsProvider(newCredentialsProvider).build();
+  }
+
+  private Supplier<Connection> connectionSupplier() {
+    return clientSideCache == null ? () -> new Connection(jedisSocketFactory, clientConfig)
+        : () -> new CacheConnection(jedisSocketFactory, clientConfig, clientSideCache);
   }
 
   @Override
@@ -64,8 +102,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,13 +113,16 @@ 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();
+    jedis.reAuth();
   }
 
   @Override
   public boolean validateObject(PooledObject<Connection> pooledConnection) {
     final Connection jedis = pooledConnection.getObject();
     try {
       // check HostAndPort ??
+      jedis.reAuth();
       return jedis.isConnected() && jedis.ping();
     } catch (final Exception e) {
       logger.warn("Error while validating pooled Connection object.", e);

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

@@ -2,19 +2,38 @@
 
 import org.apache.commons.pool2.PooledObjectFactory;
 import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
+
 import redis.clients.jedis.annots.Experimental;
+import redis.clients.jedis.authentication.JedisAuthXManager;
 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;
+
   public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig) {
-    this(new ConnectionFactory(hostAndPort, clientConfig));
+    this(hostAndPort, clientConfig, createAuthXManager(clientConfig));
+  }
+
+  public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig,
+      JedisAuthXManager authXManager) {
+    this(new ConnectionFactory(hostAndPort, clientConfig, null, authXManager));
+    attachAuthenticationListener(authXManager);
   }
 
   @Experimental
-  public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig, Cache clientSideCache) {
-    this(new ConnectionFactory(hostAndPort, clientConfig, clientSideCache));
+  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);
   }
 
   public ConnectionPool(PooledObjectFactory<Connection> factory) {
@@ -23,13 +42,22 @@ public ConnectionPool(PooledObjectFactory<Connection> factory) {
 
   public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig,
       GenericObjectPoolConfig<Connection> poolConfig) {
-    this(new ConnectionFactory(hostAndPort, clientConfig), poolConfig);
+    this(hostAndPort, clientConfig, null, createAuthXManager(clientConfig), poolConfig);
+  }
+
+  @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,
+  public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig,
+      Cache clientSideCache, JedisAuthXManager authXManager,
       GenericObjectPoolConfig<Connection> poolConfig) {
-    this(new ConnectionFactory(hostAndPort, clientConfig, clientSideCache), poolConfig);
+    this(new ConnectionFactory(hostAndPort, clientConfig, clientSideCache, authXManager),
+        poolConfig);
+        attachAuthenticationListener(authXManager);
   }
 
   public ConnectionPool(PooledObjectFactory<Connection> factory,
@@ -43,4 +71,33 @@ public Connection getResource() {
     conn.setHandlingPool(this);
     return conn;
   }
+
+  @Override
+  public void close() {
+    if (authXManager != null) {
+      authXManager.stop();
+    }
+    super.close();
+  }
+
+  private static JedisAuthXManager createAuthXManager(JedisClientConfig config) {
+    if (config.getTokenAuthConfig() != null) {
+      return new JedisAuthXManager(config.getTokenAuthConfig());
+    }
+    return null;
+  }
+
+  private void attachAuthenticationListener(JedisAuthXManager authXManager) {
+    this.authXManager = authXManager;
+    if (authXManager != null) {
+      authXManager.setListener(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);
+        }
+      });
+    }
+  }
 }