add support to get and set db credentials in an atomic operation (#2189)

brettwooldridge · Sep 23, 2024 · d43c272 · d43c272
1 parent 1d17351
commit d43c272
Show file tree

Hide file tree

Showing 7 changed files with 159 additions and 17 deletions.
diff --git a/src/main/java/com/zaxxer/hikari/HikariConfig.java b/src/main/java/com/zaxxer/hikari/HikariConfig.java
@@ -18,6 +18,7 @@
 
 import com.codahale.metrics.health.HealthCheckRegistry;
 import com.zaxxer.hikari.metrics.MetricsTrackerFactory;
+import com.zaxxer.hikari.util.Credentials;
 import com.zaxxer.hikari.util.PropertyElf;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -36,6 +37,7 @@
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.atomic.AtomicReference;
 
 import static com.zaxxer.hikari.util.UtilityElf.getNullIfEmpty;
 import static com.zaxxer.hikari.util.UtilityElf.safeIsAssignableFrom;
@@ -68,8 +70,7 @@ public class HikariConfig implements HikariConfigMXBean
    private volatile long maxLifetime;
    private volatile int maxPoolSize;
    private volatile int minIdle;
-   private volatile String username;
-   private volatile String password;
+   private final AtomicReference<Credentials> credentials = new AtomicReference<>(Credentials.of(null, null));
 
    // Properties NOT changeable at runtime
    //
@@ -283,7 +284,7 @@ public void setMinimumIdle(int minIdle)
     */
    public String getPassword()
    {
-      return password;
+      return credentials.get().getPassword();
    }
 
    /**
@@ -293,7 +294,7 @@ public String getPassword()
    @Override
    public void setPassword(String password)
    {
-      this.password = password;
+      credentials.updateAndGet(current -> Credentials.of(current.getUsername(), password));
    }
 
    /**
@@ -303,7 +304,7 @@ public void setPassword(String password)
     */
    public String getUsername()
    {
-      return username;
+      return credentials.get().getUsername();
    }
 
    /**
@@ -314,7 +315,28 @@ public String getUsername()
    @Override
    public void setUsername(String username)
    {
-      this.username = username;
+      credentials.updateAndGet(current -> Credentials.of(username, current.getPassword()));
+   }
+
+   /**
+    * Atomically set the default username and password to use for DataSource.getConnection(username, password) calls.
+    *
+    * @param credentials the username and password pair
+    */
+   @Override
+   public void setCredentials(final Credentials credentials)
+   {
+      this.credentials.set(credentials);
+   }
+
+   /**
+    * Atomically get the default username and password to use for DataSource.getConnection(username, password) calls.
+    *
+    * @return the username and password pair
+    */
+   public Credentials getCredentials()
+   {
+      return credentials.get();
    }
 
    /** {@inheritDoc} */
@@ -945,17 +967,20 @@ void seal()
     *
     * @param other Other {@link HikariConfig} to copy the state to.
     */
+   @SuppressWarnings({"rawtypes", "unchecked"})
    public void copyStateTo(HikariConfig other)
    {
       for (var field : HikariConfig.class.getDeclaredFields()) {
-         if (!Modifier.isFinal(field.getModifiers())) {
-            field.setAccessible(true);
-            try {
+         try {
+            if (!Modifier.isFinal(field.getModifiers())) {
+               field.setAccessible(true);
                field.set(other, field.get(this));
+            } else if (field.getType().isAssignableFrom(AtomicReference.class)) {
+               ((AtomicReference) field.get(other)).set(((AtomicReference) field.get(this)).get());
             }
-            catch (Exception e) {
-               throw new RuntimeException("Failed to copy HikariConfig state: " + e.getMessage(), e);
-            }
+         }
+         catch (Exception e) {
+            throw new RuntimeException("Failed to copy HikariConfig state: " + e.getMessage(), e);
          }
       }
 

diff --git a/src/main/java/com/zaxxer/hikari/HikariConfigMXBean.java b/src/main/java/com/zaxxer/hikari/HikariConfigMXBean.java
@@ -16,6 +16,8 @@
 
 package com.zaxxer.hikari;
 
+import com.zaxxer.hikari.util.Credentials;
+
 /**
  * The javax.management MBean for a Hikari pool configuration.
  *
@@ -167,6 +169,14 @@ public interface HikariConfigMXBean
     */
    void setUsername(String username);
 
+   /**
+    * Set the username and password used for authentication. Changing this at runtime will apply to new
+    * connections only. Altering this at runtime only works for DataSource-based connections, not Driver-class
+    * or JDBC URL-based connections.
+    *
+    * @param credentials the database username and password pair
+    */
+   void setCredentials(Credentials credentials);
 
    /**
     * The name of the connection pool.

diff --git a/src/main/java/com/zaxxer/hikari/pool/PoolBase.java b/src/main/java/com/zaxxer/hikari/pool/PoolBase.java
@@ -311,8 +311,7 @@ else if (mBeanServer.isRegistered(beanConfigName)) {
    private void initializeDataSource()
    {
       final var jdbcUrl = config.getJdbcUrl();
-      final var username = config.getUsername();
-      final var password = config.getPassword();
+      final var credentials = config.getCredentials();
       final var dsClassName = config.getDataSourceClassName();
       final var driverClassName = config.getDriverClassName();
       final var dataSourceJNDI = config.getDataSourceJNDI();
@@ -324,7 +323,7 @@ private void initializeDataSource()
          PropertyElf.setTargetFromProperties(ds, dataSourceProperties);
       }
       else if (jdbcUrl != null && ds == null) {
-         ds = new DriverDataSource(jdbcUrl, driverClassName, dataSourceProperties, username, password);
+         ds = new DriverDataSource(jdbcUrl, driverClassName, dataSourceProperties, credentials.getUsername(), credentials.getPassword());
       }
       else if (dataSourceJNDI != null && ds == null) {
          try {
@@ -354,8 +353,9 @@ private Connection newConnection() throws Exception
 
       Connection connection = null;
       try {
-         var username = config.getUsername();
-         var password = config.getPassword();
+         final var credentials = config.getCredentials();
+         final var username = credentials.getUsername();
+         final var password = credentials.getPassword();
 
          connection = (username == null) ? dataSource.getConnection() : dataSource.getConnection(username, password);
          if (connection == null) {

diff --git a/src/main/java/com/zaxxer/hikari/util/Credentials.java b/src/main/java/com/zaxxer/hikari/util/Credentials.java
@@ -0,0 +1,57 @@
+package com.zaxxer.hikari.util;
+
+import javax.management.ConstructorParameters;
+
+/**
+ * A simple class to hold connection credentials and is designed to be immutable.
+ */
+public final class Credentials
+{
+
+   private final String username;
+   private final String password;
+
+   /**
+    * Construct an immutable Credentials object with the supplied username and password.
+    *
+    * @param username the username
+    * @param password the password
+    * @return a new Credentials object
+    */
+   public static Credentials of(final String username, final String password) {
+      return new Credentials(username, password);
+   }
+
+   /**
+    * Construct an immutable Credentials object with the supplied username and password.
+    *
+    * @param username the username
+    * @param password the password
+    */
+   @ConstructorParameters({ "username", "password" })
+   public Credentials(final String username, final String password)
+   {
+      this.username = username;
+      this.password = password;
+   }
+
+   /**
+    * Get the username.
+    *
+    * @return the username
+    */
+   public String getUsername()
+   {
+      return username;
+   }
+
+   /**
+    * Get the password.
+    *
+    * @return the password
+    */
+   public String getPassword()
+   {
+      return password;
+   }
+}
diff --git a/src/test/java/com/zaxxer/hikari/datasource/TestSealedConfig.java b/src/test/java/com/zaxxer/hikari/datasource/TestSealedConfig.java
@@ -2,6 +2,7 @@
 
 import com.zaxxer.hikari.HikariConfig;
 import com.zaxxer.hikari.HikariDataSource;
+import com.zaxxer.hikari.util.Credentials;
 import org.junit.Test;
 
 import java.sql.Connection;
@@ -68,6 +69,7 @@ public void testSealedAccessibleMethods() throws SQLException
          ds.setMaximumPoolSize(8);
          ds.setPassword("password");
          ds.setUsername("username");
+         ds.setCredentials(Credentials.of("anothername", "anotherpassword"));
       }
    }
 }
diff --git a/src/test/java/com/zaxxer/hikari/pool/PostgresTest.java b/src/test/java/com/zaxxer/hikari/pool/PostgresTest.java
@@ -29,6 +29,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import com.zaxxer.hikari.util.Credentials;
 import org.junit.After;
 import org.junit.Before;
 
@@ -105,6 +106,24 @@ public void testCase4() throws Exception
       exerciseConfig(config, 3);
    }
 
+   @Test
+   public void testCredentialRotation()
+   {
+      HikariConfig config = createConfig(postgres);
+      config.setMinimumIdle(3);
+      config.setMaximumPoolSize(10);
+      config.setConnectionTimeout(1000);
+      config.setIdleTimeout(SECONDS.toMillis(20));
+
+      exerciseConfig(config, 3);
+
+      updatePostgresCredentials("newuser", "newpassword");
+      config.setJdbcUrl(postgres.getJdbcUrl());
+      config.setCredentials(Credentials.of("newuser", "newpassword"));
+
+      exerciseConfig(config, 3);
+   }
+
    static private void exerciseConfig(HikariConfig config, int numThreads) {
       try (final HikariDataSource ds = new HikariDataSource(config)) {
          assertTrue(ds.isRunning());
@@ -193,4 +212,12 @@ static private HikariConfig createConfig(PostgreSQLContainer<?> postgres) {
       config.setDriverClassName(postgres.getDriverClassName());
       return config;
    }
+
+   private void updatePostgresCredentials(String username, String password) {
+      postgres.stop();
+      postgres = new PostgreSQLContainer<>(IMAGE_NAME)
+         .withUsername(username)
+         .withPassword(password);
+      postgres.start();
+   }
 }
diff --git a/src/test/java/com/zaxxer/hikari/pool/TestMBean.java b/src/test/java/com/zaxxer/hikari/pool/TestMBean.java
@@ -20,6 +20,7 @@
 import com.zaxxer.hikari.HikariDataSource;
 import com.zaxxer.hikari.HikariPoolMXBean;
 import com.zaxxer.hikari.mocks.StubDataSource;
+import com.zaxxer.hikari.util.Credentials;
 import org.junit.Test;
 
 import javax.management.JMX;
@@ -170,4 +171,24 @@ public void testMBeanConnectionTimeoutChange() throws SQLException {
          System.clearProperty("com.zaxxer.hikari.housekeeping.periodMs");
       }
    }
+
+   @Test
+   public void testMBeanCredentialRotation() {
+      HikariConfig config = newHikariConfig();
+      config.setMinimumIdle(3);
+      config.setMaximumPoolSize(5);
+      config.setRegisterMbeans(true);
+      config.setConnectionTimeout(2800);
+      config.setConnectionTestQuery("VALUES 1");
+      config.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");
+      config.setCredentials(Credentials.of("foo", "bar"));
+
+      try (HikariDataSource ds = new HikariDataSource(config)) {
+         HikariConfigMXBean hikariConfigMXBean = ds.getHikariConfigMXBean();
+         hikariConfigMXBean.setCredentials(Credentials.of("newFoo", "newBar"));
+
+         assertEquals("newFoo", ds.getUsername());
+         assertEquals("newBar", ds.getPassword());
+      }
+   }
 }