SNOW-1868467: Add custom headers

parent-pom.xml

@@ -13,6 +13,7 @@
     <module>FIPS</module>
   </modules>
 
+
   <properties>
     <apache.commons.compress.version>1.21</apache.commons.compress.version>
     <apache.commons.lang3.version>3.12.0</apache.commons.lang3.version>

src/main/java/net/snowflake/client/core/HttpClientSettingsKey.java

@@ -24,6 +24,7 @@ public class HttpClientSettingsKey implements Serializable {
   // Adds a suffix to the user agent header in the http requests made by the jdbc driver.
   // More details in SNOW-717606
   private String userAgentSuffix = "";
+  private String proxyCrtFile = "";
 
   private Boolean gzipDisabled = false;
 
@@ -49,6 +50,32 @@ public HttpClientSettingsKey(
     this.userAgentSuffix = !Strings.isNullOrEmpty(userAgentSuffix) ? userAgentSuffix.trim() : "";
   }
 
+
+  public HttpClientSettingsKey(
+          OCSPMode mode,
+          String host,
+          int port,
+          String nonProxyHosts,
+          String user,
+          String password,
+          String scheme,
+          String userAgentSuffix,
+          String proxyCrtFile,
+          Boolean gzipDisabled) {
+    this.useProxy = true;
+    this.ocspMode = mode != null ? mode : OCSPMode.FAIL_OPEN;
+    this.proxyHost = !Strings.isNullOrEmpty(host) ? host.trim() : "";
+    this.proxyPort = port;
+    this.proxyCrtFile = proxyCrtFile;
+    this.nonProxyHosts = !Strings.isNullOrEmpty(nonProxyHosts) ? nonProxyHosts.trim() : "";
+    this.proxyUser = !Strings.isNullOrEmpty(user) ? user.trim() : "";
+    this.proxyPassword = !Strings.isNullOrEmpty(password) ? password.trim() : "";
+    this.proxyProtocol = !Strings.isNullOrEmpty(scheme) ? scheme.trim() : "http";
+    this.gzipDisabled = gzipDisabled;
+    this.userAgentSuffix = !Strings.isNullOrEmpty(userAgentSuffix) ? userAgentSuffix.trim() : "";
+  }
+
+
   public HttpClientSettingsKey(OCSPMode mode) {
     this.useProxy = false;
     this.ocspMode = mode != null ? mode : OCSPMode.FAIL_OPEN;
@@ -155,6 +182,10 @@ public Boolean getGzipDisabled() {
     return gzipDisabled;
   }
 
+  public String getProxyCrtFile() {
+    return proxyCrtFile;
+  }
+
   @Override
   public String toString() {
     return "HttpClientSettingsKey["

src/main/java/net/snowflake/client/core/HttpUtil.java

@@ -12,24 +12,30 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Strings;
 import com.microsoft.azure.storage.OperationContext;
-import java.io.File;
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.PrintWriter;
-import java.io.StringWriter;
+
+import java.io.*;
 import java.net.InetSocketAddress;
 import java.net.Proxy;
 import java.net.Socket;
 import java.security.KeyManagementException;
+import java.security.KeyStore;
 import java.security.NoSuchAlgorithmException;
+import java.security.PrivateKey;
+import java.security.cert.Certificate;
+import java.security.cert.CertificateFactory;
 import java.time.Duration;
-import java.util.Map;
-import java.util.Properties;
+import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import javax.annotation.Nullable;
+import javax.net.ssl.SSLContext;
 import javax.net.ssl.TrustManager;
+import javax.net.ssl.X509ExtendedKeyManager;
+import javax.net.ssl.X509ExtendedTrustManager;
+
 import net.snowflake.client.jdbc.ErrorCode;
 import net.snowflake.client.jdbc.RestRequest;
 import net.snowflake.client.jdbc.SnowflakeDriver;
@@ -43,6 +49,9 @@
 import net.snowflake.client.util.SecretDetector;
 import net.snowflake.client.util.Stopwatch;
 import net.snowflake.common.core.SqlState;
+import nl.altindag.ssl.SSLFactory;
+import nl.altindag.ssl.pem.util.PemUtils;
+import nl.altindag.ssl.util.CertificateUtils;
 import org.apache.commons.io.IOUtils;
 import org.apache.http.HttpHost;
 import org.apache.http.auth.AuthScope;
@@ -52,10 +61,9 @@
 import org.apache.http.client.config.RequestConfig;
 import org.apache.http.client.methods.CloseableHttpResponse;
 import org.apache.http.client.methods.HttpRequestBase;
-import org.apache.http.config.Registry;
-import org.apache.http.config.RegistryBuilder;
 import org.apache.http.conn.socket.ConnectionSocketFactory;
 import org.apache.http.conn.socket.PlainConnectionSocketFactory;
+import org.apache.http.conn.ssl.SSLConnectionSocketFactory;
 import org.apache.http.impl.client.BasicCredentialsProvider;
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.apache.http.impl.client.DefaultRedirectStrategy;
@@ -373,22 +381,22 @@ public static CloseableHttpClient buildHttpClient(
       logger.debug("Omitting trust manager instantiation as configuration is not provided");
     }
     try {
-      logger.debug(
-          "Registering https connection socket factory with socks proxy disabled: {} and http "
-              + "connection socket factory",
-          socksProxyDisabled);
-
-      Registry<ConnectionSocketFactory> registry =
-          RegistryBuilder.<ConnectionSocketFactory>create()
-              .register(
-                  "https", new SFSSLConnectionSocketFactory(trustManagers, socksProxyDisabled))
-              .register("http", new SFConnectionSocketFactory())
-              .build();
-
-      // Build a connection manager with enough connections
-      connectionManager =
-          new PoolingHttpClientConnectionManager(
-              registry, null, null, null, timeToLive, TimeUnit.SECONDS);
+//      logger.debug(
+//          "Registering https connection socket factory with socks proxy disabled: {} and http "
+//              + "connection socket factory",
+//          socksProxyDisabled);
+//
+//      Registry<ConnectionSocketFactory> registry =
+//          RegistryBuilder.<ConnectionSocketFactory>create()
+//              .register(
+//                  "https", new SFSSLConnectionSocketFactory(trustManagers, socksProxyDisabled))
+//              .register("http", new SFSSLConnectionSocketFactory(trustManagers, socksProxyDisabled))
+//              .build();
+
+////       Build a connection manager with enough connections
+//      connectionManager =
+//          new PoolingHttpClientConnectionManager(
+//              registry, null, null, null, timeToLive, TimeUnit.SECONDS);
       int maxConnections =
           SystemUtil.convertSystemPropertyToIntValue(
               JDBC_MAX_CONNECTIONS_PROPERTY, DEFAULT_MAX_CONNECTIONS);
@@ -399,14 +407,14 @@ public static CloseableHttpClient buildHttpClient(
           "Max connections total in connection pooling manager: {}; max connections per route: {}",
           maxConnections,
           maxConnectionsPerRoute);
-      connectionManager.setMaxTotal(maxConnections);
-      connectionManager.setDefaultMaxPerRoute(maxConnectionsPerRoute);
+//      connectionManager.setMaxTotal(maxConnections);
+//      connectionManager.setDefaultMaxPerRoute(maxConnections);
 
       logger.debug("Disabling cookie management for http client");
       String userAgentSuffix = key != null ? key.getUserAgentSuffix() : "";
       HttpClientBuilder httpClientBuilder =
           HttpClientBuilder.create()
-              .setConnectionManager(connectionManager)
+//              .setConnectionManager(connectionManager)
               // Support JVM proxy settings
               .useSystemProperties()
               .setRedirectStrategy(new DefaultRedirectStrategy())
@@ -442,18 +450,39 @@ public static CloseableHttpClient buildHttpClient(
           credentialsProvider.setCredentials(authScope, credentials);
           httpClientBuilder = httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider);
         }
+        if (key.getProxyCrtFile() != null && !key.getProxyCrtFile().isEmpty()) {
+          httpClientBuilder.setSSLContext(buildSslContext(key.getProxyCrtFile()));
+        }
       }
+      httpClientBuilder.setConnectionTimeToLive(timeToLive,TimeUnit.SECONDS);
+      httpClientBuilder.setMaxConnPerRoute(maxConnectionsPerRoute);
+      httpClientBuilder.setMaxConnTotal(maxConnections);
       httpClientBuilder.setDefaultRequestConfig(DefaultRequestConfig);
       if (downloadUnCompressed) {
         logger.debug("Disabling content compression for http client");
         httpClientBuilder = httpClientBuilder.disableContentCompression();
       }
+
       return httpClientBuilder.build();
-    } catch (NoSuchAlgorithmException | KeyManagementException ex) {
+    } catch (Exception ex) {
       throw new SSLInitializationException(ex.getMessage(), ex);
     }
   }
 
+  private static SSLContext buildSslContext(String crtFile) {
+    List<Certificate> certificates = CertificateUtils.loadCertificate(crtFile);
+
+    SSLFactory sslFactory = SSLFactory.builder()
+            .withTrustMaterial(certificates)
+            .build();
+
+
+    return sslFactory.getSslContext();
+  }
+
+
+
+
   public static void updateRoutePlanner(HttpClientSettingsKey key) {
     if (httpClientRoutePlanner.containsKey(key)
         && !httpClientRoutePlanner

src/main/java/net/snowflake/client/core/SFBaseSession.java

@@ -533,6 +533,7 @@ public HttpClientSettingsKey getHttpClientKey() throws SnowflakeSQLException {
       String proxyHost = (String) connectionPropertiesMap.get(SFSessionProperty.PROXY_HOST);
       String proxyUser = (String) connectionPropertiesMap.get(SFSessionProperty.PROXY_USER);
       String proxyPassword = (String) connectionPropertiesMap.get(SFSessionProperty.PROXY_PASSWORD);
+      String proxyCrtFile = (String) connectionPropertiesMap.get(SFSessionProperty.PROXY_CRT_FILE);
       String nonProxyHosts =
           (String) connectionPropertiesMap.get(SFSessionProperty.NON_PROXY_HOSTS);
       String proxyProtocol = (String) connectionPropertiesMap.get(SFSessionProperty.PROXY_PROTOCOL);
@@ -546,6 +547,7 @@ public HttpClientSettingsKey getHttpClientKey() throws SnowflakeSQLException {
               proxyPassword,
               proxyProtocol,
               userAgentSuffix,
+              proxyCrtFile,
               gzipDisabled);
 
       logHttpClientInitInfo(ocspAndProxyAndGzipKey);
@@ -563,6 +565,7 @@ public HttpClientSettingsKey getHttpClientKey() throws SnowflakeSQLException {
       String httpsProxyHost = systemGetProperty("https.proxyHost");
       String httpsProxyPort = systemGetProperty("https.proxyPort");
       String httpsProxyUser = systemGetProperty("https.proxyUser");
+      String proxyCrtFile = systemGetProperty("https.proxyCrtFile");
       String httpsProxyPassword = systemGetProperty("https.proxyPassword");
       String httpProxyProtocol = systemGetProperty("http.proxyProtocol");
       String noProxy = systemGetEnv("NO_PROXY");
@@ -627,6 +630,7 @@ public HttpClientSettingsKey getHttpClientKey() throws SnowflakeSQLException {
                   httpsProxyPassword,
                   "https",
                   userAgentSuffix,
+                  proxyCrtFile,
                   gzipDisabled);
           logHttpClientInitInfo(ocspAndProxyAndGzipKey);
         } else if (proxyProtocol.equals("http")
@@ -650,6 +654,7 @@ public HttpClientSettingsKey getHttpClientKey() throws SnowflakeSQLException {
                   httpProxyPassword,
                   "http",
                   userAgentSuffix,
+                  proxyCrtFile,
                   gzipDisabled);
           logHttpClientInitInfo(ocspAndProxyAndGzipKey);
         } else {

src/main/java/net/snowflake/client/core/SFSession.java

@@ -654,6 +654,8 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
         .setToken((String) connectionPropertiesMap.get(SFSessionProperty.TOKEN))
         .setPasscodeInPassword(passcodeInPassword)
         .setPasscode((String) connectionPropertiesMap.get(SFSessionProperty.PASSCODE))
+            //or from sessionParametersMap
+        .setAdditionalHttpHeadersForSnowsight(getHttpHeaders((String)connectionPropertiesMap.get(SFSessionProperty.AdditionalHttpHeaders)))
         .setConnectionTimeout(
             connectionPropertiesMap.get(SFSessionProperty.HTTP_CLIENT_CONNECTION_TIMEOUT) != null
                 ? Duration.ofMillis(
@@ -791,6 +793,20 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
     logger.info("Session {} opened in {} ms.", getSessionId(), stopwatch.elapsedMillis());
   }
 
+  private Map<String, String> getHttpHeaders(String headers) {
+    if(headers!=null && !headers.isEmpty()) {
+      Map<String, String> headersMap = new HashMap<>();
+      for (String headerKeyPair : headers.split(";")) {
+        String[] split = headerKeyPair.split(":");
+        if (split.length >= 2) {
+          headersMap.put(split[0], split[1]);
+        }
+      }
+      return headersMap;
+    }
+    return null;
+  }
+
   /**
    * If authenticator is null and private key is specified, jdbc will assume key pair authentication
    *

src/main/java/net/snowflake/client/core/SFSessionProperty.java

@@ -48,6 +48,7 @@ public enum SFSessionProperty {
   PROXY_HOST("proxyHost", false, String.class),
   PROXY_PORT("proxyPort", false, String.class),
   PROXY_USER("proxyUser", false, String.class),
+  PROXY_CRT_FILE("proxyCrtFile", false, String.class),
   PROXY_PASSWORD("proxyPassword", false, String.class),
   NON_PROXY_HOSTS("nonProxyHosts", false, String.class),
   PROXY_PROTOCOL("proxyProtocol", false, String.class),
@@ -69,6 +70,8 @@ public enum SFSessionProperty {
   // Adds a suffix to the user agent header in the http requests made by the jdbc driver
   USER_AGENT_SUFFIX("user_agent_suffix", false, String.class),
 
+  AdditionalHttpHeaders("additional_http_headers", false, String.class),
+
   CLIENT_OUT_OF_BAND_TELEMETRY_ENABLED(
       "CLIENT_OUT_OF_BAND_TELEMETRY_ENABLED", false, Boolean.class),
   GZIP_DISABLED("gzipDisabled", false, Boolean.class),

src/main/java/net/snowflake/client/core/SFTrustManager.java

@@ -15,10 +15,8 @@
 import com.fasterxml.jackson.databind.node.JsonNodeType;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.google.common.base.Strings;
-import java.io.ByteArrayOutputStream;
-import java.io.File;
-import java.io.IOException;
-import java.io.OutputStream;
+
+import java.io.*;
 import java.math.BigInteger;
 import java.net.URI;
 import java.net.URISyntaxException;
@@ -47,17 +45,15 @@
 import java.util.TimeZone;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
-import javax.net.ssl.KeyManagerFactory;
-import javax.net.ssl.SSLEngine;
-import javax.net.ssl.TrustManager;
-import javax.net.ssl.TrustManagerFactory;
-import javax.net.ssl.X509ExtendedTrustManager;
-import javax.net.ssl.X509TrustManager;
+import javax.net.ssl.*;
+
 import net.snowflake.client.jdbc.OCSPErrorCode;
 import net.snowflake.client.log.SFLogger;
 import net.snowflake.client.log.SFLoggerFactory;
 import net.snowflake.client.util.DecorrelatedJitterBackoff;
 import net.snowflake.client.util.SFPair;
+import nl.altindag.ssl.SSLFactory;
+import nl.altindag.ssl.pem.util.PemUtils;
 import org.apache.commons.codec.binary.Base64;
 import org.apache.commons.io.IOUtils;
 import org.apache.http.HttpHost;
@@ -694,6 +690,7 @@ private void checkNewOCSPEndpointAvailability() {
    */
   private X509TrustManager getTrustManager(String algorithm) {
     try {
+      //todo need make sure I load key to this instance
       TrustManagerFactory factory = TrustManagerFactory.getInstance(algorithm);
       factory.init((KeyStore) null);
       X509TrustManager ret = null;

src/main/java/net/snowflake/client/jdbc/DefaultSFConnectionHandler.java

@@ -335,6 +335,7 @@ private void initSessionProperties(SnowflakeConnectString conStr, String appID,
         boolean booleanV = v0 instanceof Boolean ? (Boolean) v0 : Boolean.parseBoolean((String) v0);
         sfSession.setSfSQLMode(booleanV);
       }
+      //here we add properties
       sfSession.addSFSessionProperty(property.getKey(), property.getValue());
     }
 

src/main/java/net/snowflake/client/jdbc/SnowflakeUtil.java

@@ -756,6 +756,7 @@ public static HttpClientSettingsKey convertProxyPropertiesToHttpClientKey(
         String proxyUser = info.getProperty(SFSessionProperty.PROXY_USER.getPropertyKey());
         String proxyPassword = info.getProperty(SFSessionProperty.PROXY_PASSWORD.getPropertyKey());
         String nonProxyHosts = info.getProperty(SFSessionProperty.NON_PROXY_HOSTS.getPropertyKey());
+        String proxyCrtFile = info.getProperty(SFSessionProperty.PROXY_CRT_FILE.getPropertyKey());
         String proxyProtocol = info.getProperty(SFSessionProperty.PROXY_PROTOCOL.getPropertyKey());
         String userAgentSuffix =
             info.getProperty(SFSessionProperty.USER_AGENT_SUFFIX.getPropertyKey());
@@ -776,6 +777,7 @@ public static HttpClientSettingsKey convertProxyPropertiesToHttpClientKey(
             proxyPassword,
             proxyProtocol,
             userAgentSuffix,
+            proxyCrtFile,
             gzipDisabled);
       }
     }