fix(bigquery-jdbc): propagate connection proxy settings to auth library

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java

@@ -24,6 +24,7 @@
 import com.google.api.gax.rpc.HeaderProvider;
 import com.google.api.gax.rpc.TransportChannelProvider;
 import com.google.auth.Credentials;
+import com.google.auth.http.HttpTransportFactory;
 import com.google.cloud.bigquery.BigQuery;
 import com.google.cloud.bigquery.BigQueryException;
 import com.google.cloud.bigquery.BigQueryOptions;
@@ -265,11 +266,34 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
               String.valueOf(ds.getRequestGoogleDriveScope()),
               BigQueryJdbcUrlUtility.REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME);
 
+      Map<String, String> proxyProperties =
+          BigQueryJdbcProxyUtility.parseProxyProperties(ds, this.connectionClassName);
+
+      this.sslTrustStorePath = ds.getSSLTrustStorePath();
+      this.sslTrustStorePassword = ds.getSSLTrustStorePassword();
+      this.httpConnectTimeout = ds.getHttpConnectTimeout();
+      this.httpReadTimeout = ds.getHttpReadTimeout();
+
+      this.httpTransportOptions =
+          BigQueryJdbcProxyUtility.getHttpTransportOptions(
+              proxyProperties,
+              this.sslTrustStorePath,
+              this.sslTrustStorePassword,
+              this.httpConnectTimeout,
+              this.httpReadTimeout,
+              this.connectionClassName);
+
+      HttpTransportFactory httpTransportFactory =
+          this.httpTransportOptions != null
+              ? this.httpTransportOptions.getHttpTransportFactory()
+              : null;
+
       this.credentials =
           BigQueryJdbcOAuthUtility.getCredentials(
               authProperties,
               overrideProperties,
               this.reqGoogleDriveScope,
+              httpTransportFactory,
               this.connectionClassName);
       String defaultDatasetString = ds.getDefaultDataset();
       if (defaultDatasetString == null || defaultDatasetString.trim().isEmpty()) {
@@ -302,22 +326,6 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
       this.destinationDataset = ds.getDestinationDataset();
       this.destinationDatasetExpirationTime = ds.getDestinationDatasetExpirationTime();
       this.kmsKeyName = ds.getKmsKeyName();
-      Map<String, String> proxyProperties =
-          BigQueryJdbcProxyUtility.parseProxyProperties(ds, this.connectionClassName);
-
-      this.sslTrustStorePath = ds.getSSLTrustStorePath();
-      this.sslTrustStorePassword = ds.getSSLTrustStorePassword();
-      this.httpConnectTimeout = ds.getHttpConnectTimeout();
-      this.httpReadTimeout = ds.getHttpReadTimeout();
-
-      this.httpTransportOptions =
-          BigQueryJdbcProxyUtility.getHttpTransportOptions(
-              proxyProperties,
-              this.sslTrustStorePath,
-              this.sslTrustStorePassword,
-              this.httpConnectTimeout,
-              this.httpReadTimeout,
-              this.connectionClassName);
       this.transportChannelProvider =
           BigQueryJdbcProxyUtility.getTransportChannelProvider(
               proxyProperties,

java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtility.java

@@ -21,6 +21,7 @@
 
 import com.google.api.client.util.PemReader;
 import com.google.api.client.util.SecurityUtils;
+import com.google.auth.http.HttpTransportFactory;
 import com.google.auth.oauth2.AccessToken;
 import com.google.auth.oauth2.ClientId;
 import com.google.auth.oauth2.ExternalAccountCredentials;
@@ -51,6 +52,7 @@
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.net.URL;
+import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Paths;
 import java.security.GeneralSecurityException;
@@ -263,6 +265,16 @@ static GoogleCredentials getCredentials(
       Map<String, String> overrideProperties,
       Boolean reqGoogleDriveScopeBool,
       String callerClassName) {
+    return getCredentials(
+        authProperties, overrideProperties, reqGoogleDriveScopeBool, null, callerClassName);
+  }
+
+  static GoogleCredentials getCredentials(
+      Map<String, String> authProperties,
+      Map<String, String> overrideProperties,
+      Boolean reqGoogleDriveScopeBool,
+      HttpTransportFactory httpTransportFactory,
+      String callerClassName) {
     LOG.finer("++enter++\t" + callerClassName);
 
     AuthType authType =
@@ -272,11 +284,13 @@ static GoogleCredentials getCredentials(
     switch (authType) {
       case GOOGLE_SERVICE_ACCOUNT:
         credentials =
-            getGoogleServiceAccountCredentials(authProperties, overrideProperties, callerClassName);
+            getGoogleServiceAccountCredentials(
+                authProperties, overrideProperties, httpTransportFactory, callerClassName);
         break;
       case GOOGLE_USER_ACCOUNT:
         credentials =
-            getGoogleUserAccountCredentials(authProperties, overrideProperties, callerClassName);
+            getGoogleUserAccountCredentials(
+                authProperties, overrideProperties, httpTransportFactory, callerClassName);
         break;
       case PRE_GENERATED_TOKEN:
         credentials =
@@ -286,7 +300,9 @@ static GoogleCredentials getCredentials(
         credentials = getApplicationDefaultCredentials(callerClassName);
         break;
       case EXTERNAL_ACCOUNT_AUTH:
-        credentials = getExternalAccountAuthCredentials(authProperties, callerClassName);
+        credentials =
+            getExternalAccountAuthCredentials(
+                authProperties, httpTransportFactory, callerClassName);
         break;
       default:
         IllegalStateException ex = new IllegalStateException(OAUTH_TYPE_ERROR_MESSAGE);
@@ -295,7 +311,7 @@ static GoogleCredentials getCredentials(
     }
 
     return getServiceAccountImpersonatedCredentials(
-        credentials, reqGoogleDriveScopeBool, authProperties);
+        credentials, reqGoogleDriveScopeBool, authProperties, httpTransportFactory);
   }
 
   private static boolean isFileExists(String filename) {
@@ -326,6 +342,7 @@ private static boolean isJson(byte[] value) {
   private static GoogleCredentials getGoogleServiceAccountCredentials(
       Map<String, String> authProperties,
       Map<String, String> overrideProperties,
+      HttpTransportFactory httpTransportFactory,
       String callerClassName) {
     LOG.finer("++enter++\t" + callerClassName);
 
@@ -370,6 +387,10 @@ private static GoogleCredentials getGoogleServiceAccountCredentials(
         throw new BigQueryJdbcRuntimeException("No valid credentials provided.");
       }
 
+      if (httpTransportFactory != null) {
+        builder.setHttpTransportFactory(httpTransportFactory);
+      }
+
       if (overrideProperties.containsKey(BigQueryJdbcUrlUtility.OAUTH2_TOKEN_URI_PROPERTY_NAME)) {
         builder.setTokenServerUri(
             new URI(overrideProperties.get(BigQueryJdbcUrlUtility.OAUTH2_TOKEN_URI_PROPERTY_NAME)));
@@ -393,6 +414,16 @@ static UserAuthorizer getUserAuthorizer(
       int port,
       String callerClassName)
       throws URISyntaxException {
+    return getUserAuthorizer(authProperties, overrideProperties, port, null, callerClassName);
+  }
+
+  static UserAuthorizer getUserAuthorizer(
+      Map<String, String> authProperties,
+      Map<String, String> overrideProperties,
+      int port,
+      HttpTransportFactory httpTransportFactory,
+      String callerClassName)
+      throws URISyntaxException {
     LOG.finer("++enter++\t" + callerClassName);
 
     List<String> scopes = new ArrayList<>();
@@ -411,6 +442,10 @@ static UserAuthorizer getUserAuthorizer(
             .setScopes(scopes)
             .setCallbackUri(URI.create("http://localhost:" + port));
 
+    if (httpTransportFactory != null) {
+      userAuthorizerBuilder.setHttpTransportFactory(httpTransportFactory);
+    }
+
     if (overrideProperties.containsKey(BigQueryJdbcUrlUtility.OAUTH2_TOKEN_URI_PROPERTY_NAME)) {
       userAuthorizerBuilder.setTokenServerUri(
           new URI(overrideProperties.get(BigQueryJdbcUrlUtility.OAUTH2_TOKEN_URI_PROPERTY_NAME)));
@@ -421,21 +456,36 @@ static UserAuthorizer getUserAuthorizer(
 
   static UserCredentials getCredentialsFromCode(
       UserAuthorizer userAuthorizer, String code, String callerClassName) throws IOException {
+    return getCredentialsFromCode(userAuthorizer, code, null, callerClassName);
+  }
+
+  static UserCredentials getCredentialsFromCode(
+      UserAuthorizer userAuthorizer,
+      String code,
+      HttpTransportFactory httpTransportFactory,
+      String callerClassName)
+      throws IOException {
     LOG.finer("++enter++\t" + callerClassName);
-    return userAuthorizer.getCredentialsFromCode(code, URI.create(""));
+    UserCredentials credentials = userAuthorizer.getCredentialsFromCode(code, URI.create(""));
+    if (httpTransportFactory != null) {
+      credentials = credentials.toBuilder().setHttpTransportFactory(httpTransportFactory).build();
+    }
+    return credentials;
   }
 
   private static GoogleCredentials getGoogleUserAccountCredentials(
       Map<String, String> authProperties,
       Map<String, String> overrideProperties,
+      HttpTransportFactory httpTransportFactory,
       String callerClassName) {
     LOG.finer("++enter++\t" + callerClassName);
     try {
       ServerSocket serverSocket = new ServerSocket(0);
       serverSocket.setSoTimeout(USER_AUTH_TIMEOUT_MS);
       int port = serverSocket.getLocalPort();
       UserAuthorizer userAuthorizer =
-          getUserAuthorizer(authProperties, overrideProperties, port, callerClassName);
+          getUserAuthorizer(
+              authProperties, overrideProperties, port, httpTransportFactory, callerClassName);
 
       URL authURL = userAuthorizer.getAuthorizationUrl("user", "", URI.create(""));
       String code;
@@ -468,7 +518,7 @@ private static GoogleCredentials getGoogleUserAccountCredentials(
         throw new BigQueryJdbcRuntimeException("User auth only supported in desktop environments");
       }
 
-      return getCredentialsFromCode(userAuthorizer, code, callerClassName);
+      return getCredentialsFromCode(userAuthorizer, code, httpTransportFactory, callerClassName);
     } catch (IOException | URISyntaxException ex) {
       throw new BigQueryJdbcRuntimeException(
           "Failed to establish connection using User Account authentication", ex);
@@ -572,7 +622,9 @@ private static GoogleCredentials getApplicationDefaultCredentials(String callerC
   }
 
   private static GoogleCredentials getExternalAccountAuthCredentials(
-      Map<String, String> authProperties, String callerClassName) {
+      Map<String, String> authProperties,
+      HttpTransportFactory httpTransportFactory,
+      String callerClassName) {
     LOG.finer("++enter++\t" + callerClassName);
     try {
       JsonObject jsonObject = null;
@@ -609,18 +661,23 @@ private static GoogleCredentials getExternalAccountAuthCredentials(
         }
       }
 
+      ExternalAccountCredentials credentials;
       if (credentialsPath != null) {
-        return ExternalAccountCredentials.fromStream(
-            Files.newInputStream(Paths.get(credentialsPath)));
+        try (InputStream stream = Files.newInputStream(Paths.get(credentialsPath))) {
+          credentials = ExternalAccountCredentials.fromStream(stream, httpTransportFactory);
+        }
       } else if (jsonObject != null) {
-        return ExternalAccountCredentials.fromStream(
-            new ByteArrayInputStream(jsonObject.toString().getBytes()));
+        credentials =
+            ExternalAccountCredentials.fromStream(
+                new ByteArrayInputStream(jsonObject.toString().getBytes(StandardCharsets.UTF_8)),
+                httpTransportFactory);
       } else {
         IllegalArgumentException ex =
             new IllegalArgumentException("Insufficient info provided for external authentication");
         LOG.severe(ex.getMessage(), ex);
         throw ex;
       }
+      return credentials;
     } catch (IOException e) {
       throw new BigQueryJdbcRuntimeException(
           "IOException during getExternalAccountAuthCredentials", e);
@@ -634,7 +691,8 @@ private static GoogleCredentials getExternalAccountAuthCredentials(
   private static GoogleCredentials getServiceAccountImpersonatedCredentials(
       GoogleCredentials credentials,
       Boolean reqGoogleDriveScopeBool,
-      Map<String, String> authProperties) {
+      Map<String, String> authProperties,
+      HttpTransportFactory httpTransportFactory) {
 
     String impersonationEmail =
         authProperties.get(BigQueryJdbcUrlUtility.OAUTH_SA_IMPERSONATION_EMAIL_PROPERTY_NAME);
@@ -684,6 +742,15 @@ private static GoogleCredentials getServiceAccountImpersonatedCredentials(
       throw ex;
     }
 
+    if (httpTransportFactory != null) {
+      return ImpersonatedCredentials.create(
+          credentials,
+          impersonationEmail,
+          impersonationChain,
+          impersonationScopes,
+          impersonationLifetimeInt,
+          httpTransportFactory);
+    }
     return ImpersonatedCredentials.create(
         credentials,
         impersonationEmail,

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOAuthUtilityTest.java

@@ -22,8 +22,10 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import com.google.auth.http.HttpTransportFactory;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.auth.oauth2.ImpersonatedCredentials;
+import com.google.auth.oauth2.ServiceAccountCredentials;
 import com.google.auth.oauth2.UserAuthorizer;
 import com.google.auth.oauth2.UserCredentials;
 import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException;
@@ -489,4 +491,53 @@ public void testPrivateKeyFromP12Bytes_wrong_password() {
       assertTrue(false);
     }
   }
+
+  @Test
+  public void testGetCredentialsPropagatesHttpTransportFactory() {
+    Map<String, String> authProperties =
+        BigQueryJdbcOAuthUtility.parseOAuthProperties(
+            DataSource.fromUrl(
+                "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;"
+                    + "ProjectId=MyBigQueryProject;OAuthType=0;"
+                    + "OAuthServiceAcctEmail=dummytest@dummytest.iam.gserviceaccount.com;"
+                    + "OAuthPvtKey="
+                    + fake_pkcs8_key
+                    + ";"),
+            null);
+
+    HttpTransportFactory dummyFactory = () -> null;
+
+    GoogleCredentials credentials =
+        BigQueryJdbcOAuthUtility.getCredentials(
+            authProperties, Collections.emptyMap(), false, dummyFactory, null);
+
+    assertThat(credentials).isInstanceOf(ServiceAccountCredentials.class);
+    assertThat(((ServiceAccountCredentials) credentials).toBuilder().getHttpTransportFactory())
+        .isEqualTo(dummyFactory);
+  }
+
+  @Test
+  public void testGetImpersonatedCredentialsPropagatesHttpTransportFactory() {
+    Map<String, String> authProperties =
+        BigQueryJdbcOAuthUtility.parseOAuthProperties(
+            DataSource.fromUrl(
+                "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;"
+                    + "ProjectId=MyBigQueryProject;OAuthType=0;"
+                    + "OAuthServiceAcctEmail=dummytest@dummytest.iam.gserviceaccount.com;"
+                    + "OAuthPvtKey="
+                    + fake_pkcs8_key
+                    + ";"
+                    + "ServiceAccountImpersonationEmail=impersonated@email.com;"),
+            null);
+
+    HttpTransportFactory dummyFactory = () -> null;
+
+    GoogleCredentials credentials =
+        BigQueryJdbcOAuthUtility.getCredentials(
+            authProperties, Collections.emptyMap(), false, dummyFactory, null);
+
+    assertThat(credentials).isInstanceOf(ImpersonatedCredentials.class);
+    assertThat(((ImpersonatedCredentials) credentials).toBuilder().getHttpTransportFactory())
+        .isEqualTo(dummyFactory);
+  }
 }