feat: add support for service account impersonation (#245)

.envrc.example

@@ -2,3 +2,6 @@ export ALLOYDB_DB="some-db"
 export ALLOYDB_USER="some-user"
 export ALLOYDB_PASS="some-password"
 export ALLOYDB_INSTANCE_NAME="projects/<PROJECT>/locations/<REGION>/clusters/<CLUSTER>/instances/<INSTANCE>"
+export ALLOYDB_INSTANCE_IP="some-IP-address"
+export ALLOYDB_IAM_USER="some-user@my-project.iam"
+export ALLOYDB_IMPERSONATED_USER="some-impersonated-IAM-user"

.github/workflows/ci.yaml

@@ -248,6 +248,7 @@ jobs:
             ALLOYDB_CLUSTER_PASS:${{ secrets.GOOGLE_CLOUD_PROJECT }}/ALLOYDB_CLUSTER_PASS
             ALLOYDB_IAM_USER:${{ secrets.GOOGLE_CLOUD_PROJECT }}/ALLOYDB_JAVA_IAM_USER
             ALLOYDB_INSTANCE_IP:${{ secrets.GOOGLE_CLOUD_PROJECT }}/ALLOYDB_INSTANCE_IP
+            ALLOYDB_IMPERSONATED_USER:${{ secrets.GOOGLE_CLOUD_PROJECT }}/IMPERSONATED_USER
 
       - name: Run tests
         env:
@@ -257,6 +258,7 @@ jobs:
           ALLOYDB_PASS: '${{ steps.secrets.outputs.ALLOYDB_CLUSTER_PASS }}'
           ALLOYDB_INSTANCE_NAME: '${{ steps.secrets.outputs.ALLOYDB_INSTANCE_NAME }}'
           ALLOYDB_INSTANCE_IP: '${{ steps.secrets.outputs.ALLOYDB_INSTANCE_IP }}'
+          ALLOYDB_IMPERSONATED_USER: '${{ steps.secrets.outputs.ALLOYDB_IMPERSONATED_USER }}'
           JOB_TYPE: integration
         run: .kokoro/build.sh
         shell: bash

README.md

@@ -157,6 +157,76 @@ performance from a connection pool.
 [e2e]: https://github.com/GoogleCloudPlatform/alloydb-java-connector/blob/main/alloydb-jdbc-connector/src/test/java/com/google/cloud/alloydb/ITSocketFactoryTest.java
 [pool-sizing]: https://github.com/brettwooldridge/HikariCP/wiki/About-Pool-Sizing
 
+### Service Account Impersonation
+
+The Java Connector supports service account impersonation with the
+`alloydbTargetPrincipal` JDBC connection property. When enabled,
+all API requests are made impersonating the supplied service account.
+The IAM principal must have the IAM role "Service Account Token Creator"
+(i.e., `roles/iam.serviceAccounts.serviceAccountTokenCreator`) on the
+service account provided in the `alloydbTargetPrincipal` property.
+
+#### Example
+
+``` java
+import com.zaxxer.hikari.HikariConfig;
+import com.zaxxer.hikari.HikariDataSource;
+
+public class ExampleApplication {
+
+  private HikariDataSource dataSource;
+
+  HikariDataSource getDataSource() {
+    HikariConfig config = new HikariConfig();
+
+    // There is no need to set a host on the JDBC URL
+    // since the Connector will resolve the correct IP address.
+    config.setJdbcUrl("jdbc:postgresql:///postgres");
+    config.setUsername(System.getenv("ALLOYDB_USER"));
+    config.setPassword(System.getenv("ALLOYDB_PASS"));
+
+    // Tell the driver to use the AlloyDB Java Connector's SocketFactory
+    // when connecting to an instance/
+    config.addDataSourceProperty("socketFactory",
+        "com.google.cloud.alloydb.SocketFactory");
+    // Tell the Java Connector which instance to connect to.
+    config.addDataSourceProperty("alloydbInstanceName",
+        System.getenv("ALLOYDB_INSTANCE_NAME"));
+    config.addDataSourceProperty("alloydbTargetPrincipal", 
+        System.getenv("ALLOYDB_IMPERSONATED_USER"));
+
+    this.dataSource = new HikariDataSource(config);
+  }
+
+  // Use DataSource as usual ...
+
+}
+```
+
+#### Delegated Service Account Impersonation
+
+In addition, the `alloydbDelegates` property controls impersonation delegation.
+The value is a comma-separated list of service accounts containing chained
+list of delegates required to grant the final access_token. If set,
+the sequence of identities must have "Service Account Token Creator" capability
+granted to the preceding identity. For example, if set to
+`"serviceAccountB,serviceAccountC"`, the IAM principal must have the
+Token Creator role on serviceAccountB. Then serviceAccountB must have the
+Token Creator role on serviceAccountC. Finally, serviceAccountC must have the
+Token Creator role on `alloydbTargetPrincipal`. If unset, the IAM principal
+must have the Token Creator role on `alloydbTargetPrincipal`.
+
+```java
+config.addDataSourceProperty("alloydbTargetPrincipal",
+    "TARGET_SERVICE_ACCOUNT");
+config.addDataSourceProperty("alloydbDelegates", 
+    "SERVICE_ACCOUNT_1,SERVICE_ACCOUNT_2");
+```
+
+In this example, the IAM principal impersonates SERVICE_ACCOUNT_1 which
+impersonates SERVICE_ACCOUNT_2 which then impersonates the
+TARGET_SERVICE_ACCOUNT.
+
 ## Support policy
 
 ### Major version lifecycle

alloydb-jdbc-connector/pom.xml

@@ -117,6 +117,11 @@
       <scope>provided</scope>
     </dependency>
 
+    <dependency>
+      <groupId>com.google.auth</groupId>
+      <artifactId>google-auth-library-oauth2-http</artifactId>
+    </dependency>
+
     <!-- Logging -->
     <dependency> 
       <groupId>org.slf4j</groupId>
@@ -160,11 +165,5 @@
       <scope>test</scope>
     </dependency>
 
-    <dependency>
-      <groupId>com.google.auth</groupId>
-      <artifactId>google-auth-library-oauth2-http</artifactId>
-      <scope>test</scope>
-    </dependency>
-
   </dependencies>
 </project>

alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/AlloyDBAdminClientFactory.java

@@ -15,6 +15,7 @@
  */
 package com.google.cloud.alloydb;
 
+import com.google.api.gax.core.FixedCredentialsProvider;
 import com.google.api.gax.rpc.FixedHeaderProvider;
 import com.google.cloud.alloydb.v1beta.AlloyDBAdminClient;
 import com.google.cloud.alloydb.v1beta.AlloyDBAdminSettings;
@@ -30,7 +31,8 @@ class AlloyDBAdminClientFactory {
 
   private static final String DEFAULT_ENDPOINT = "alloydb.googleapis.com:443";
 
-  static AlloyDBAdminClient create() throws IOException {
+  static AlloyDBAdminClient create(FixedCredentialsProvider credentialsProvider)
+      throws IOException {
     AlloyDBAdminSettings.Builder settingsBuilder = AlloyDBAdminSettings.newBuilder();
 
     Map<String, String> headers =
@@ -42,6 +44,7 @@ static AlloyDBAdminClient create() throws IOException {
         settingsBuilder
             .setEndpoint(DEFAULT_ENDPOINT)
             .setHeaderProvider(FixedHeaderProvider.create(headers))
+            .setCredentialsProvider(credentialsProvider)
             .build();
 
     return AlloyDBAdminClient.create(alloyDBAdminSettings);

alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/ConnectionConfig.java

@@ -17,39 +17,98 @@
 package com.google.cloud.alloydb;
 
 import com.google.cloud.alloydb.v1beta.InstanceName;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
 import java.util.Properties;
 
 public class ConnectionConfig {
   public static final String ALLOYDB_INSTANCE_NAME = "alloydbInstanceName";
+  public static final String ALLOYDB_TARGET_PRINCIPAL = "alloydbTargetPrincipal";
+  public static final String ALLOYDB_DELEGATES = "alloydbDelegates";
+  public static final String ALLOYDB_NAMED_CONNECTION = "alloydbNamedConnection";
+  public static final String DEFAULT_NAMED_CONNECTION = "default";
   private final InstanceName instanceName;
+  private final String namedConnection;
+  private final String targetPrincipal;
+  private final List<String> delegates;
 
   /** Create a new ConnectionConfig from the well known JDBC Connection properties. */
   public static ConnectionConfig fromConnectionProperties(Properties props) {
     final String instanceNameStr = props.getProperty(ALLOYDB_INSTANCE_NAME, "");
     final InstanceName instanceName = InstanceName.parse(instanceNameStr);
+    final String namedConnection = props.getProperty(ConnectionConfig.ALLOYDB_NAMED_CONNECTION);
+    final String targetPrincipal = props.getProperty(ConnectionConfig.ALLOYDB_TARGET_PRINCIPAL);
+    final String delegatesStr = props.getProperty(ConnectionConfig.ALLOYDB_DELEGATES);
+    final List<String> delegates;
+    if (delegatesStr != null && !delegatesStr.isEmpty()) {
+      delegates = Arrays.asList(delegatesStr.split(","));
+    } else {
+      delegates = Collections.emptyList();
+    }
 
-    return new ConnectionConfig(instanceName);
+    return new ConnectionConfig(instanceName, namedConnection, targetPrincipal, delegates);
   }
 
-  private ConnectionConfig(InstanceName instanceName) {
+  private ConnectionConfig(
+      InstanceName instanceName,
+      String namedConnection,
+      String targetPrincipal,
+      List<String> delegates) {
     this.instanceName = instanceName;
+    this.namedConnection = namedConnection;
+    this.targetPrincipal = targetPrincipal;
+    this.delegates = delegates;
   }
 
   public InstanceName getInstanceName() {
     return instanceName;
   }
 
+  public String getNamedConnection() {
+    if (namedConnection == null || namedConnection.isEmpty()) {
+      return DEFAULT_NAMED_CONNECTION;
+    }
+    return namedConnection;
+  }
+
+  public String getTargetPrincipal() {
+    return targetPrincipal;
+  }
+
+  public List<String> getDelegates() {
+    return delegates;
+  }
+
   /** The builder for the ConnectionConfig. */
   public static class Builder {
     private InstanceName instanceName;
+    private String namedConnection;
+    private String targetPrincipal;
+    private List<String> delegates;
 
     public Builder withInstanceName(InstanceName instanceName) {
       this.instanceName = instanceName;
       return this;
     }
 
+    public Builder withNamedConnection(String namedConnection) {
+      this.namedConnection = namedConnection;
+      return this;
+    }
+
+    public Builder withTargetPrincipal(String targetPrincipal) {
+      this.targetPrincipal = targetPrincipal;
+      return this;
+    }
+
+    public Builder withDelegates(List<String> delegates) {
+      this.delegates = delegates;
+      return this;
+    }
+
     public ConnectionConfig build() {
-      return new ConnectionConfig(instanceName);
+      return new ConnectionConfig(instanceName, namedConnection, targetPrincipal, delegates);
     }
   }
 }

alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/Connector.java

@@ -137,16 +137,16 @@ private static KeyManager[] initializeKeyManager(
     return keyManagerFactory.getKeyManagers();
   }
 
-  Socket connect(ConnectionConfig config) throws IOException {
+  Socket connect(InstanceName instanceName) throws IOException {
     ConnectionInfoCache connectionInfoCache =
         instances.computeIfAbsent(
-            config.getInstanceName(),
+            instanceName,
             k -> {
               DefaultRateLimiter rateLimiter = new DefaultRateLimiter(RATE_LIMIT_PER_SEC);
               return connectionInfoCacheFactory.create(
                   this.executor,
                   this.connectionInfoRepo,
-                  config.getInstanceName(),
+                  instanceName,
                   this.clientConnectorKeyPair,
                   new RefreshCalculator(),
                   rateLimiter);

alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/ConnectorRegistry.java

@@ -15,6 +15,7 @@
  */
 package com.google.cloud.alloydb;
 
+import com.google.api.gax.core.FixedCredentialsProvider;
 import com.google.cloud.alloydb.v1beta.AlloyDBAdminClient;
 import java.io.Closeable;
 import java.io.IOException;
@@ -34,38 +35,41 @@ public enum ConnectorRegistry implements Closeable {
   private final ScheduledExecutorService executor;
 
   @SuppressWarnings("ImmutableEnumChecker")
-  private final AlloyDBAdminClient alloyDBAdminClient;
-
-  @SuppressWarnings("ImmutableEnumChecker")
-  private final Connector connector;
+  private ConcurrentHashMap<String, Connector> registeredConnectors;
 
   ConnectorRegistry() {
     // During refresh, each instance consumes 2 threads from the thread pool. By using 8 threads,
     // there should be enough free threads so that there will not be a deadlock. Most users
     // configure 3 or fewer instances, requiring 6 threads during refresh. By setting
     // this to 8, it's enough threads for most users, plus a safety factor of 2.
     this.executor = Executors.newScheduledThreadPool(8);
-    try {
-      alloyDBAdminClient = AlloyDBAdminClient.create();
-    } catch (IOException e) {
-      throw new RuntimeException(e);
-    }
-    this.connector =
-        new Connector(
-            executor,
-            new DefaultConnectionInfoRepository(executor, alloyDBAdminClient),
-            RsaKeyPairGenerator.generateKeyPair(),
-            new DefaultConnectionInfoCacheFactory(),
-            new ConcurrentHashMap<>());
+    this.registeredConnectors = new ConcurrentHashMap<>();
   }
 
-  Connector getConnector() {
-    return this.connector;
+  Connector getConnector(ConnectionConfig config) {
+    return registeredConnectors.computeIfAbsent(
+        config.getNamedConnection(),
+        k -> {
+          try {
+            FixedCredentialsProvider credentialsProvider =
+                CredentialsProviderFactory.create(config);
+            AlloyDBAdminClient alloyDBAdminClient =
+                AlloyDBAdminClientFactory.create(credentialsProvider);
+
+            return new Connector(
+                executor,
+                new DefaultConnectionInfoRepository(executor, alloyDBAdminClient),
+                RsaKeyPairGenerator.generateKeyPair(),
+                new DefaultConnectionInfoCacheFactory(),
+                new ConcurrentHashMap<>());
+          } catch (IOException e) {
+            throw new RuntimeException(e);
+          }
+        });
   }
 
   @Override
   public void close() {
     this.executor.shutdown();
-    this.alloyDBAdminClient.close();
   }
 }

alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/CredentialsProviderFactory.java

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * 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
+ *
+ *     https://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 com.google.cloud.alloydb;
+
+import com.google.api.gax.core.FixedCredentialsProvider;
+import com.google.auth.oauth2.GoogleCredentials;
+import com.google.auth.oauth2.ImpersonatedCredentials;
+import java.io.IOException;
+import java.util.Arrays;
+
+class CredentialsProviderFactory {
+
+  private static final String CLOUD_PLATFORM = "https://www.googleapis.com/auth/cloud-platform";
+  private static final String ALLOYDB_LOGIN = "https://www.googleapis.com/auth/alloydb.login";
+
+  static FixedCredentialsProvider create(ConnectionConfig config) {
+    GoogleCredentials credentials;
+    try {
+      credentials = GoogleCredentials.getApplicationDefault();
+    } catch (IOException e) {
+      throw new RuntimeException("failed to retrieve OAuth2 access token", e);
+    }
+
+    if (config.getTargetPrincipal() != null && !config.getTargetPrincipal().isEmpty()) {
+      credentials =
+          ImpersonatedCredentials.newBuilder()
+              .setSourceCredentials(credentials)
+              .setTargetPrincipal(config.getTargetPrincipal())
+              .setDelegates(config.getDelegates())
+              .setScopes(Arrays.asList(ALLOYDB_LOGIN, CLOUD_PLATFORM))
+              .build();
+    }
+
+    return FixedCredentialsProvider.create(credentials);
+  }
+}