A fix for some TLS issues in the MongoDB IO

mransley · Nov 28, 2019 · e097e89 · e097e89
1 parent 438055c
commit e097e89
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 12 deletions.
diff --git a/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java b/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/MongoDbIO.java
@@ -39,6 +39,7 @@
 import java.util.List;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
+import javax.net.ssl.SSLContext;
 import org.apache.beam.sdk.annotations.Experimental;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.SerializableCoder;
@@ -347,14 +348,19 @@ public void populateDisplayData(DisplayData.Builder builder) {
  }
 
  private static MongoClientOptions.Builder getOptions(
- int maxConnectionIdleTime, boolean sslEnabled, boolean sslInvalidHostNameAllowed) {
+ int maxConnectionIdleTime,
+ boolean sslEnabled,
+ boolean sslInvalidHostNameAllowed,
+ boolean ignoreSSLCertificate) {
  MongoClientOptions.Builder optionsBuilder = new MongoClientOptions.Builder();
  optionsBuilder.maxConnectionIdleTime(maxConnectionIdleTime);
  if (sslEnabled) {
- optionsBuilder
- .sslEnabled(sslEnabled)
- .sslInvalidHostNameAllowed(sslInvalidHostNameAllowed)
- .sslContext(SSLUtils.ignoreSSLCertificate());
+ optionsBuilder.sslEnabled(sslEnabled).sslInvalidHostNameAllowed(sslInvalidHostNameAllowed);
+ if (ignoreSSLCertificate) {
+ SSLContext sslContext = SSLUtils.ignoreSSLCertificate();
+ optionsBuilder.sslContext(sslContext);
+ optionsBuilder.socketFactory(sslContext.getSocketFactory());
+ }
  }
  return optionsBuilder;
  }
@@ -396,7 +402,8 @@ long getDocumentCount() {
  getOptions(
  spec.maxConnectionIdleTime(),
  spec.sslEnabled(),
- spec.sslInvalidHostNameAllowed())))) {
+ spec.sslInvalidHostNameAllowed(),
+ spec.ignoreSSLCertificate())))) {
  return getDocumentCount(mongoClient, spec.database(), spec.collection());
  } catch (Exception e) {
  return -1;
@@ -424,7 +431,8 @@ public long getEstimatedSizeBytes(PipelineOptions pipelineOptions) {
  getOptions(
  spec.maxConnectionIdleTime(),
  spec.sslEnabled(),
- spec.sslInvalidHostNameAllowed())))) {
+ spec.sslInvalidHostNameAllowed(),
+ spec.ignoreSSLCertificate())))) {
  return getEstimatedSizeBytes(mongoClient, spec.database(), spec.collection());
  }
  }
@@ -452,7 +460,8 @@ public List<BoundedSource<Document>> split(
  getOptions(
  spec.maxConnectionIdleTime(),
  spec.sslEnabled(),
- spec.sslInvalidHostNameAllowed())))) {
+ spec.sslInvalidHostNameAllowed(),
+ spec.ignoreSSLCertificate())))) {
  MongoDatabase mongoDatabase = mongoClient.getDatabase(spec.database());
 
  List<Document> splitKeys;
@@ -743,7 +752,8 @@ private MongoClient createClient(Read spec) {
  getOptions(
  spec.maxConnectionIdleTime(),
  spec.sslEnabled(),
- spec.sslInvalidHostNameAllowed())));
+ spec.sslInvalidHostNameAllowed(),
+ spec.ignoreSSLCertificate())));
  }
  }
 
@@ -925,7 +935,8 @@ public void createMongoClient() {
  getOptions(
  spec.maxConnectionIdleTime(),
  spec.sslEnabled(),
- spec.sslInvalidHostNameAllowed())));
+ spec.sslInvalidHostNameAllowed(),
+ spec.ignoreSSLCertificate())));
  }
 
  @StartBundle

diff --git a/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/SSLUtils.java b/sdks/java/io/mongodb/src/main/java/org/apache/beam/sdk/io/mongodb/SSLUtils.java
@@ -19,7 +19,6 @@
 
 import java.security.KeyStore;
 import java.security.cert.X509Certificate;
-import javax.net.ssl.HttpsURLConnection;
 import javax.net.ssl.KeyManagerFactory;
 import javax.net.ssl.SSLContext;
 import javax.net.ssl.TrustManager;
@@ -55,7 +54,6 @@ static SSLContext ignoreSSLCertificate() {
  // Install the all-trusting trust manager
  SSLContext sc = SSLContext.getInstance("TLS");
  sc.init(null, trustAllCerts, new java.security.SecureRandom());
- HttpsURLConnection.setDefaultSSLSocketFactory(sc.getSocketFactory());
 
  KeyStore ks = KeyStore.getInstance("JKS");
  ks.load(