fix: support for different types of credentials (#1097)

Co-authored-by: Shubha Rajan <[email protected]>

Author: dirkluijk

core/pom.xml

@@ -200,6 +200,12 @@
       <version>1.13.0</version>
     </dependency>
 
+    <dependency>
+      <groupId>com.google.oauth-client</groupId>
+      <artifactId>google-oauth-client</artifactId>
+      <version>1.34.1</version>
+    </dependency>
+
     <!-- com.google.cloud.sql.nativeimage.CloudSqlFeature needs the GraalVM
       dependencies for compilation. The provided-scope dependencies do not add
       additional dependencies to library users. -->

core/src/main/java/com/google/cloud/sql/core/CloudSqlInstance.java

@@ -18,13 +18,16 @@
 
 import static com.google.common.base.Preconditions.checkArgument;
 
+import com.google.api.client.auth.oauth2.Credential;
 import com.google.api.client.googleapis.json.GoogleJsonResponseException;
+import com.google.api.client.http.HttpRequestInitializer;
 import com.google.api.services.sqladmin.SQLAdmin;
 import com.google.api.services.sqladmin.model.ConnectSettings;
 import com.google.api.services.sqladmin.model.GenerateEphemeralCertRequest;
 import com.google.api.services.sqladmin.model.GenerateEphemeralCertResponse;
 import com.google.api.services.sqladmin.model.IpMapping;
 import com.google.auth.http.HttpCredentialsAdapter;
+import com.google.auth.oauth2.AccessToken;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.auth.oauth2.OAuth2Credentials;
 import com.google.cloud.sql.CredentialFactory;
@@ -154,9 +157,9 @@ class CloudSqlInstance {
     this.keyPair = keyPair;
 
     if (enableIamAuth) {
-      HttpCredentialsAdapter credentialsAdapter = (HttpCredentialsAdapter) tokenSourceFactory
-          .create();
-      this.credentials = Optional.of((OAuth2Credentials) credentialsAdapter.getCredentials());
+      HttpRequestInitializer source = tokenSourceFactory.create();
+
+      this.credentials = Optional.of(parseCredentials(source));
       this.credentials.get().refresh();
     } else {
       this.credentials = Optional.empty();
@@ -169,6 +172,36 @@ class CloudSqlInstance {
     }
   }
 
+  private OAuth2Credentials parseCredentials(HttpRequestInitializer source) {
+    if (source instanceof HttpCredentialsAdapter) {
+      HttpCredentialsAdapter adapter = (HttpCredentialsAdapter) source;
+      return (OAuth2Credentials) adapter.getCredentials();
+    }
+
+    if (source instanceof Credential) {
+      Credential credential = (Credential) source;
+      AccessToken accessToken = new AccessToken(
+          credential.getAccessToken(),
+          new Date(credential.getExpirationTimeMilliseconds())
+      );
+      GoogleCredentials googleCredentials = new GoogleCredentials(accessToken) {
+
+        @Override
+        public AccessToken refreshAccessToken() throws IOException {
+          credential.refreshToken();
+
+          return new AccessToken(
+              credential.getAccessToken(),
+              new Date(credential.getExpirationTimeMilliseconds()));
+        }
+      };
+
+      return googleCredentials;
+    }
+
+    throw new RuntimeException("Not supporting credentials of type " + source.getClass().getName());
+  }
+
   /**
    * Generates public key certificate for which the instance has the matching private key.
    *
@@ -326,7 +359,7 @@ boolean forceRefresh() {
    * would expire.
    */
   private ListenableFuture<InstanceData> performRefresh() {
-    // To avoid unreasonable SQL Admin API usage, use a rate limit to throttle our usage. 
+    // To avoid unreasonable SQL Admin API usage, use a rate limit to throttle our usage.
     forcedRenewRateLimiter.acquire(1);
     // Use the Cloud SQL Admin API to return the Metadata and Certificate
     ListenableFuture<Metadata> metadataFuture = executor.submit(this::fetchMetadata);
@@ -492,7 +525,7 @@ private Metadata fetchMetadata() {
                     + "instance.",
                 connectionName));
       }
-      
+
       checkDatabaseCompatibility(instanceMetadata, enableIamAuth, connectionName);
 
 

core/src/test/java/com/google/cloud/sql/core/CoreSocketFactoryTest.java

@@ -18,24 +18,18 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.isA;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.verifyNoMoreInteractions;
-import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.*;
 
 import com.google.api.client.googleapis.auth.oauth2.GoogleCredential;
 import com.google.api.client.googleapis.json.GoogleJsonError;
 import com.google.api.client.googleapis.json.GoogleJsonError.ErrorInfo;
 import com.google.api.client.googleapis.json.GoogleJsonResponseException;
-import com.google.api.client.http.HttpRequest;
-import com.google.api.client.http.HttpResponse;
-import com.google.api.client.http.HttpStatusCodes;
-import com.google.api.client.http.HttpTransport;
-import com.google.api.client.http.LowLevelHttpRequest;
+import com.google.api.client.http.*;
 import com.google.api.client.json.GenericJson;
 import com.google.api.client.json.Json;
 import com.google.api.client.json.JsonFactory;
@@ -122,8 +116,6 @@ public class CoreSocketFactoryTest {
   @Mock
   private CredentialFactory credentialFactory;
   @Mock
-  private GoogleCredential credential;
-  @Mock
   private SQLAdmin adminApi;
   @Mock
   private SQLAdmin.Connect adminApiConnect;
@@ -191,6 +183,7 @@ public void setup()
                         new IpMapping().setIpAddress(PUBLIC_IP).setType("PRIMARY"),
                         new IpMapping().setIpAddress(PRIVATE_IP).setType("PRIVATE")))
                 .setServerCaCert(new SslCert().setCert(TestKeys.SERVER_CA_CERT))
+                .setDatabaseVersion("POSTGRES14")
                 .setRegion("myRegion"));
     when(adminApiConnectGenerateEphemeralCert.execute())
         .thenReturn(generateEphemeralCertResponse);
@@ -253,10 +246,7 @@ public void create_successfulPrivateConnection()
         .generateEphemeralCert(
             eq("myProject"), eq("myRegion~myInstance"), isA(GenerateEphemeralCertRequest.class));
 
-    BufferedReader bufferedReader =
-        new BufferedReader(new InputStreamReader(socket.getInputStream(), UTF_8));
-    String line = bufferedReader.readLine();
-    assertThat(line).isEqualTo(SERVER_MESSAGE);
+    assertThat(readLine(socket)).isEqualTo(SERVER_MESSAGE);
   }
 
   @Test
@@ -275,10 +265,7 @@ public void create_successfulConnection() throws IOException, InterruptedExcepti
         .generateEphemeralCert(
             eq("myProject"), eq("myRegion~myInstance"), isA(GenerateEphemeralCertRequest.class));
 
-    BufferedReader bufferedReader =
-        new BufferedReader(new InputStreamReader(socket.getInputStream(), UTF_8));
-    String line = bufferedReader.readLine();
-    assertThat(line).isEqualTo(SERVER_MESSAGE);
+    assertThat(readLine(socket)).isEqualTo(SERVER_MESSAGE);
   }
 
   @Test
@@ -298,10 +285,7 @@ public void create_successfulDomainScopedConnection() throws IOException, Interr
             eq("example.com:myProject"), eq("myRegion~myInstance"),
             isA(GenerateEphemeralCertRequest.class));
 
-    BufferedReader bufferedReader =
-        new BufferedReader(new InputStreamReader(socket.getInputStream(), UTF_8));
-    String line = bufferedReader.readLine();
-    assertThat(line).isEqualTo(SERVER_MESSAGE);
+    assertThat(readLine(socket)).isEqualTo(SERVER_MESSAGE);
   }
 
   @Test
@@ -332,10 +316,7 @@ public void create_expiredCertificateOnFirstConnection_certificateRenewed()
         .generateEphemeralCert(
             eq("myProject"), eq("myRegion~myInstance"), isA(GenerateEphemeralCertRequest.class));
 
-    BufferedReader bufferedReader =
-        new BufferedReader(new InputStreamReader(socket.getInputStream(), UTF_8));
-    String line = bufferedReader.readLine();
-    assertThat(line).isEqualTo(SERVER_MESSAGE);
+    assertThat(readLine(socket)).isEqualTo(SERVER_MESSAGE);
   }
 
   @Test
@@ -396,6 +377,47 @@ public void create_notAuthorized() throws IOException {
     }
   }
 
+  @Test
+  public void supportsCustomCredentialFactoryWithIAM() throws InterruptedException, IOException {
+    GoogleCredential customCredential = mock(GoogleCredential.class);
+    when(credentialFactory.create()).thenReturn(customCredential);
+
+    when(customCredential.getAccessToken()).thenReturn("foo");
+    when(customCredential.getExpirationTimeMilliseconds()).thenReturn(new Date().getTime());
+
+    FakeSslServer sslServer = new FakeSslServer();
+    int port = sslServer.start();
+
+    CoreSocketFactory coreSocketFactory =
+        new CoreSocketFactory(clientKeyPair, adminApi, credentialFactory, port, defaultExecutor);
+    Socket socket =
+        coreSocketFactory.createSslSocket(
+            "myProject:myRegion:myInstance", Arrays.asList("PRIMARY"), true);
+
+    assertThat(readLine(socket)).isEqualTo(SERVER_MESSAGE);
+  }
+
+  @Test
+  public void doesNotSupportNonGoogleCredentialWithIAM() throws InterruptedException, IOException {
+    BasicAuthentication nonGoogleCredential = mock(BasicAuthentication.class);
+    when(credentialFactory.create()).thenReturn(nonGoogleCredential);
+
+    FakeSslServer sslServer = new FakeSslServer();
+    int port = sslServer.start();
+
+    CoreSocketFactory coreSocketFactory =
+        new CoreSocketFactory(clientKeyPair, adminApi, credentialFactory, port, defaultExecutor);
+    assertThrows(RuntimeException.class, () -> {
+      coreSocketFactory.createSslSocket(
+          "myProject:myRegion:myInstance", Arrays.asList("PRIMARY"), true);
+    });
+  }
+
+  private String readLine(Socket socket) throws IOException {
+    BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(socket.getInputStream(), UTF_8));
+    return bufferedReader.readLine();
+  }
+
   // Creates a fake "accessNotConfigured" exception that can be used for testing.
   private static GoogleJsonResponseException fakeNotConfiguredException() throws IOException {
     return fakeGoogleJsonResponseException(