Caching - Add IAM authentication support for ElastiCache Valkey (PR #4)

- Add ElastiCacheIamTokenUtility for token generation
- Extend DataRemoteCachePlugin with IAM auth detection
- Support serverless and regular ElastiCache endpoints
- Implement 15-minute token refresh and 12-hour re-auth cycles
- Add cacheIamAuthEnabled configuration property

Author: Frank-Gu-81

examples/AWSDriverExample/src/main/java/software/amazon/DatabaseConnectionWithCacheExample.java

@@ -12,6 +12,13 @@ public class DatabaseConnectionWithCacheExample {
   private static final String DB_CONNECTION_STRING = env.get("DB_CONNECTION_STRING");
   private static final String CACHE_RW_SERVER_ADDR = env.get("CACHE_RW_SERVER_ADDR");
   private static final String CACHE_RO_SERVER_ADDR = env.get("CACHE_RO_SERVER_ADDR");
+  // If the cache server is authenticated with IAM
+  private static final String CACHE_NAME = env.get("CACHE_NAME");
+  // Both IAM and traditional auth uses the same CACHE_USERNAME
+  private static final String CACHE_USERNAME = env.get("CACHE_USERNAME"); // e.g., "iam-user-01" / "username"
+  private static final String CACHE_IAM_REGION = env.get("CACHE_IAM_REGION"); // e.g., "us-west-2"
+  // If the cache server is authenticated with traditional username password
+  // private static final String CACHE_PASSWORD = env.get("CACHE_PASSWORD");
   private static final String USERNAME = env.get("DB_USERNAME");
   private static final String PASSWORD = env.get("DB_PASSWORD");
   private static final String USE_SSL = env.get("USE_SSL");
@@ -30,6 +37,12 @@ public static void main(String[] args) throws SQLException {
     properties.setProperty("wrapperPlugins", "dataRemoteCache");
     properties.setProperty("cacheEndpointAddrRw", CACHE_RW_SERVER_ADDR);
     properties.setProperty("cacheEndpointAddrRo", CACHE_RO_SERVER_ADDR);
+    // If the cache server is authenticated with IAM
+    properties.setProperty("cacheName", CACHE_NAME);
+    properties.setProperty("cacheUsername", CACHE_USERNAME);
+    properties.setProperty("cacheIamRegion", CACHE_IAM_REGION);
+    // If the cache server is authenticated with traditional username password
+    // properties.setProperty("cachePassword", PASSWORD);
     properties.setProperty("cacheUseSSL", USE_SSL); // "true" or "false"
     properties.setProperty("wrapperLogUnclosedConnections", "true");
     String queryStr = "/*+ CACHE_PARAM(ttl=300s) */ select * from cinemas";

wrapper/src/main/java/software/amazon/jdbc/plugin/cache/CacheConnection.java

@@ -1,27 +1,52 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * 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 software.amazon.jdbc.plugin.cache;
 
 import io.lettuce.core.RedisClient;
-import io.lettuce.core.RedisCommandExecutionException;
+import io.lettuce.core.RedisCredentials;
+import io.lettuce.core.RedisCredentialsProvider;
 import io.lettuce.core.RedisURI;
+import io.lettuce.core.RedisCommandExecutionException;
 import io.lettuce.core.SetArgs;
 import io.lettuce.core.api.StatefulRedisConnection;
 import io.lettuce.core.api.async.RedisAsyncCommands;
 import io.lettuce.core.codec.ByteArrayCodec;
 import io.lettuce.core.resource.ClientResources;
-import software.amazon.jdbc.AwsWrapperProperty;
+import reactor.core.publisher.Mono;
+import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
+import software.amazon.awssdk.regions.Region;
 import java.nio.charset.StandardCharsets;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.time.Duration;
 import java.util.Properties;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Supplier;
 import java.util.logging.Logger;
 import org.apache.commons.pool2.BasePooledObjectFactory;
 import org.apache.commons.pool2.impl.GenericObjectPool;
 import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
 import org.apache.commons.pool2.impl.DefaultPooledObject;
 import org.apache.commons.pool2.PooledObject;
+import software.amazon.jdbc.AwsWrapperProperty;
 import software.amazon.jdbc.PropertyDefinition;
+import software.amazon.jdbc.authentication.AwsCredentialsManager;
+import software.amazon.jdbc.plugin.iam.ElastiCacheIamTokenUtility;
 import software.amazon.jdbc.util.StringUtils;
 
 // Abstraction layer on top of a connection to a remote cache server
@@ -31,18 +56,21 @@ public class CacheConnection {
   private static volatile GenericObjectPool<StatefulRedisConnection<byte[], byte[]>> readConnectionPool;
   private static volatile GenericObjectPool<StatefulRedisConnection<byte[], byte[]>> writeConnectionPool;
   private static final GenericObjectPoolConfig<StatefulRedisConnection<byte[], byte[]>> poolConfig = createPoolConfig();
-  private final String cacheRwServerAddr; // read-write cache server
-  private final String cacheRoServerAddr; // read-only cache server
-  private MessageDigest msgHashDigest = null;
 
   private static final int DEFAULT_POOL_SIZE  = 20;
   private static final int DEFAULT_POOL_MAX_IDLE = 20;
   private static final int DEFAULT_POOL_MIN_IDLE = 0;
   private static final long DEFAULT_MAX_BORROW_WAIT_MS = 100;
+  private static final long TOKEN_CACHE_DURATION = 15 * 60 - 30;
 
   private static final ReentrantLock READ_LOCK = new ReentrantLock();
   private static final ReentrantLock WRITE_LOCK = new ReentrantLock();
 
+  private final String cacheRwServerAddr; // read-write cache server
+  private final String cacheRoServerAddr; // read-only cache server
+  private final String[] defaultCacheServerHostAndPort;
+  private MessageDigest msgHashDigest = null;
+
   protected static final AwsWrapperProperty CACHE_RW_ENDPOINT_ADDR =
       new AwsWrapperProperty(
           "cacheEndpointAddrRw",
@@ -61,7 +89,38 @@ public class CacheConnection {
           "true",
           "Whether to use SSL for cache connections.");
 
+  protected static final AwsWrapperProperty CACHE_IAM_REGION =
+      new AwsWrapperProperty(
+          "cacheIamRegion",
+          null,
+          "AWS region for ElastiCache IAM authentication.");
+
+  protected static final AwsWrapperProperty CACHE_USERNAME =
+      new AwsWrapperProperty(
+          "cacheUsername",
+          null,
+          "Username for ElastiCache regular authentication.");
+
+  protected static final AwsWrapperProperty CACHE_PASSWORD =
+      new AwsWrapperProperty(
+          "cachePassword",
+          null,
+          "Password for ElastiCache regular authentication.");
+
+  protected static final AwsWrapperProperty CACHE_NAME =
+      new AwsWrapperProperty(
+          "cacheName",
+          null,
+          "Explicit cache name for ElastiCache IAM authentication. ");
+
   private final boolean useSSL;
+  private final boolean iamAuthEnabled;
+  private final String cacheIamRegion;
+  private final String cacheUsername;
+  private final String cacheName;
+  private final String cachePassword;
+  private final Properties awsProfileProperties;
+  private final AwsCredentialsProvider credentialsProvider;
 
   static {
     PropertyDefinition.registerPluginProperties(CacheConnection.class);
@@ -71,6 +130,44 @@ public CacheConnection(final Properties properties) {
     this.cacheRwServerAddr = CACHE_RW_ENDPOINT_ADDR.getString(properties);
     this.cacheRoServerAddr = CACHE_RO_ENDPOINT_ADDR.getString(properties);
     this.useSSL = Boolean.parseBoolean(CACHE_USE_SSL.getString(properties));
+    this.cacheName = CACHE_NAME.getString(properties);
+    this.cacheIamRegion = CACHE_IAM_REGION.getString(properties);
+    this.cacheUsername = CACHE_USERNAME.getString(properties);
+    this.cachePassword = CACHE_PASSWORD.getString(properties);
+    this.iamAuthEnabled = !StringUtils.isNullOrEmpty(this.cacheIamRegion);
+    boolean hasTraditionalAuth = !StringUtils.isNullOrEmpty(this.cachePassword);
+    // Validate authentication configuration
+    if (this.iamAuthEnabled && hasTraditionalAuth) {
+      throw new IllegalArgumentException(
+          "Cannot specify both IAM authentication (cacheIamRegion) and traditional authentication (cachePassword). Choose one authentication method.");
+    }
+    if (this.cacheRwServerAddr == null) {
+      throw new IllegalArgumentException("Cache endpoint address is required");
+    }
+    this.defaultCacheServerHostAndPort = getHostnameAndPort(this.cacheRwServerAddr);
+    if (this.iamAuthEnabled) {
+      if (this.cacheUsername == null || this.defaultCacheServerHostAndPort[0] == null || this.cacheName == null) {
+        throw new IllegalArgumentException("IAM authentication requires cache name, username, region, and hostname");
+      }
+    }
+    if (PropertyDefinition.AWS_PROFILE.getString(properties) != null) {
+      this.awsProfileProperties = new Properties();
+      this.awsProfileProperties.setProperty(
+          PropertyDefinition.AWS_PROFILE.name,
+          PropertyDefinition.AWS_PROFILE.getString(properties)
+      );
+    } else {
+      this.awsProfileProperties = null;
+    }
+    if (this.iamAuthEnabled) {
+      // Handle null case
+      Properties propsToPass = (this.awsProfileProperties != null)
+          ? this.awsProfileProperties
+          : new Properties();
+      this.credentialsProvider = AwsCredentialsManager.getProvider(null, propsToPass);
+    } else {
+      this.credentialsProvider = null;
+    }
   }
 
   /* Here we check if we need to initialise connection pool for read or write to cache.
@@ -113,15 +210,14 @@ private void createConnectionPool(boolean isRead) {
       if (isRead && !StringUtils.isNullOrEmpty(cacheRoServerAddr)) {
         serverAddr = cacheRoServerAddr;
       }
-      String[] hostnameAndPort = serverAddr.split(":");
-      RedisURI redisUriCluster = RedisURI.Builder.redis(hostnameAndPort[0])
-          .withPort(Integer.parseInt(hostnameAndPort[1]))
-          .withSsl(useSSL).withVerifyPeer(false).withLibraryName("aws-sql-jdbc-lettuce").build();
+      String[] hostnameAndPort = getHostnameAndPort(serverAddr);
+      RedisURI redisUriCluster = buildRedisURI(hostnameAndPort[0], Integer.parseInt(hostnameAndPort[1]));
 
       RedisClient client = RedisClient.create(resources, redisUriCluster);
       GenericObjectPool<StatefulRedisConnection<byte[], byte[]>> pool = new GenericObjectPool<>(
           new BasePooledObjectFactory<StatefulRedisConnection<byte[], byte[]>>() {
             public StatefulRedisConnection<byte[], byte[]> create() {
+
               StatefulRedisConnection<byte[], byte[]> connection = client.connect(new ByteArrayCodec());
               // In cluster mode, we need to send READONLY command to the server for reading from replica.
               // Note: we gracefully ignore ERR reply to support non cluster mode.
@@ -148,7 +244,6 @@ public PooledObject<StatefulRedisConnection<byte[], byte[]>> wrap(StatefulRedisC
       } else {
         writeConnectionPool = pool;
       }
-
     } catch (Exception e) {
       String poolType = isRead ? "read" : "write";
       String errorMsg = String.format("Failed to create Cache %s connection pool", poolType);
@@ -247,13 +342,13 @@ public void writeToCache(String key, byte[] value, int expiry) {
   private void returnConnectionBackToPool(StatefulRedisConnection <byte[], byte[]> connection, boolean isConnectionBroken, boolean isRead) {
     GenericObjectPool<StatefulRedisConnection<byte[], byte[]>> pool = isRead ? readConnectionPool : writeConnectionPool;
     if (isConnectionBroken) {
-        try {
-          pool.invalidateObject(connection);
-        } catch (Exception e) {
-          throw new RuntimeException("Could not invalidate connection for the pool", e);
-        }
+      try {
+        pool.invalidateObject(connection);
+      } catch (Exception e) {
+        throw new RuntimeException("Could not invalidate connection for the pool", e);
+      }
     } else {
-        pool.returnObject(connection);
+      pool.returnObject(connection);
     }
   }
 
@@ -263,4 +358,43 @@ protected void setConnectionPools(GenericObjectPool<StatefulRedisConnection<byte
     readConnectionPool = readPool;
     writeConnectionPool = writePool;
   }
+
+  protected RedisURI buildRedisURI(String hostname, int port) {
+    RedisURI.Builder uriBuilder = RedisURI.Builder.redis(hostname)
+        .withPort(port)
+        .withSsl(useSSL)
+        .withVerifyPeer(false)
+        .withLibraryName("aws-sql-jdbc-lettuce");
+
+    if (this.iamAuthEnabled) {
+      // Create a credentials provider that Lettuce will call whenever authentication is needed
+      RedisCredentialsProvider credentialsProvider = () -> {
+        // Create a cached token supplier that automatically refreshes tokens every 14.5 minutes
+        Supplier<String> tokenSupplier = CachedSupplier.memoizeWithExpiration(
+            () -> {
+              ElastiCacheIamTokenUtility tokenUtility = new ElastiCacheIamTokenUtility(this.cacheName);
+              return tokenUtility.generateAuthenticationToken(
+                  this.credentialsProvider,
+                  Region.of(this.cacheIamRegion),
+                  this.defaultCacheServerHostAndPort[0],
+                  Integer.parseInt(this.defaultCacheServerHostAndPort[1]),
+                  this.cacheUsername
+              );
+            },
+            TOKEN_CACHE_DURATION,
+            TimeUnit.SECONDS
+        );
+        // Package the username and token (from cache or freshly generated) into Redis credentials
+        return Mono.just(RedisCredentials.just(this.cacheUsername, tokenSupplier.get()));
+      };
+      uriBuilder.withAuthentication(credentialsProvider);
+    } else if (!StringUtils.isNullOrEmpty(this.cachePassword)) {
+      uriBuilder.withAuthentication(this.cacheUsername, this.cachePassword);
+    }
+    return uriBuilder.build();
+  }
+
+  private String[] getHostnameAndPort(String serverAddr) {
+    return serverAddr.split(":");
+  }
 }

wrapper/src/main/java/software/amazon/jdbc/plugin/cache/CachedSupplier.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * 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 software.amazon.jdbc.plugin.cache;
+
+import java.util.Objects;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Supplier;
+
+public final class CachedSupplier {
+
+  private CachedSupplier() {
+    throw new UnsupportedOperationException("Utility class should not be instantiated");
+  }
+
+  public static <T> Supplier<T> memoizeWithExpiration(
+      Supplier<T> delegate, long duration, TimeUnit unit) {
+
+    Objects.requireNonNull(delegate, "delegate Supplier must not be null");
+    Objects.requireNonNull(unit, "TimeUnit must not be null");
+    if (duration <= 0) {
+      throw new IllegalArgumentException("duration must be > 0");
+    }
+
+    return new ExpiringMemoizingSupplier<>(delegate, duration, unit);
+  }
+
+  private static final class ExpiringMemoizingSupplier<T> implements Supplier<T> {
+
+    private final Supplier<T> delegate;
+    private final long durationNanos;
+    private final ReentrantLock lock = new ReentrantLock();
+
+    private volatile T value;
+    private volatile long expirationNanos; // 0 means not yet initialized
+
+    ExpiringMemoizingSupplier(Supplier<T> delegate, long duration, TimeUnit unit) {
+      this.delegate = delegate;
+      this.durationNanos = unit.toNanos(duration);
+    }
+
+    @Override
+    public T get() {
+      long now = System.nanoTime();
+
+      // Check if value is expired or uninitialized
+      if (expirationNanos == 0 || now - expirationNanos >= 0) {
+        lock.lock();
+        try {
+          if (expirationNanos == 0 || now - expirationNanos >= 0) {
+            value = delegate.get();
+            long next = now + durationNanos;
+            expirationNanos = (next == 0) ? 1 : next; // avoid 0 sentinel
+          }
+        } finally {
+          lock.unlock();
+        }
+      }
+      return value;
+    }
+  }
+}