Merge pull request #66 from ruivieira-forking/release-0.12.0-rc0-tls

[cherry-pick] feat: Add SSL context loading
opendatahub-io · Aug 1, 2024 · 245615a · 245615a
2 parents 325ce12 + a645c37
commit 245615a
Showing 1 changed file with 35 additions and 6 deletions.
diff --git a/src/main/java/com/ibm/watson/modelmesh/ModelMesh.java b/src/main/java/com/ibm/watson/modelmesh/ModelMesh.java
@@ -101,17 +101,16 @@
 
 import javax.annotation.concurrent.GuardedBy;
 import javax.net.ssl.SSLContext;
-import java.io.File;
-import java.io.IOException;
-import java.io.InterruptedIOException;
-import java.io.UncheckedIOException;
+import javax.net.ssl.TrustManagerFactory;
+import java.io.*;
 import java.lang.management.ManagementFactory;
 import java.lang.management.MemoryMXBean;
 import java.lang.management.MemoryUsage;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.net.URI;
 import java.nio.channels.ClosedByInterruptException;
+import java.security.KeyStore;
 import java.security.NoSuchAlgorithmException;
 import java.text.ParseException;
 import java.text.SimpleDateFormat;
@@ -432,6 +431,34 @@ public abstract class ModelMesh extends ThriftService
         }
     }
 
+    private static final String SSL_TRUSTSTORE_PATH_PROPERTY = "watson.ssl.truststore.path";
+    private static final String SSL_TRUSTSTORE_PASSWORD_PROPERTY = "watson.ssl.truststore.password";
+
+    private static SSLContext sslContext = null;
+
+    private static SSLContext loadSSLContext() throws Exception {
+        if (sslContext == null) {
+            final String trustStorePath = System.getProperty(SSL_TRUSTSTORE_PATH_PROPERTY);
+            final String trustStorePassword = System.getProperty(SSL_TRUSTSTORE_PASSWORD_PROPERTY);
+
+            if (trustStorePath == null || trustStorePassword == null) {
+                throw new IllegalArgumentException("Truststore settings not found in system properties");
+            }
+
+            final KeyStore trustStore = KeyStore.getInstance("JKS");
+            try (FileInputStream trustStoreStream = new FileInputStream(trustStorePath)) {
+                trustStore.load(trustStoreStream, trustStorePassword.toCharArray());
+            }
+
+            final TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
+            trustManagerFactory.init(trustStore);
+
+            sslContext = SSLContext.getInstance("TLS");
+            sslContext.init(null, trustManagerFactory.getTrustManagers(), null);
+        }
+        return sslContext;
+    }
+
     private PayloadProcessor initPayloadProcessor() {
         String payloadProcessorsDefinitions = getStringParameter(MM_PAYLOAD_PROCESSORS, null);
         logger.info("Parsing PayloadProcessor definition '{}'", payloadProcessorsDefinitions);
@@ -445,14 +472,16 @@ private PayloadProcessor initPayloadProcessor() {
                     String modelId = uri.getQuery();
                     String method = uri.getFragment();
                     if ("http".equals(processorName)) {
+                        logger.info("Initializing HTTP payload processor");
                         processor = new RemotePayloadProcessor(uri);
                     } else if ("https".equals(processorName)) {
                         SSLContext sslContext;
                         try {
-                            sslContext = SSLContext.getDefault();
-                        } catch (NoSuchAlgorithmException missingAlgorithmException) {
+                            sslContext = loadSSLContext();
+                        } catch (Exception missingAlgorithmException) {
                             throw new UncheckedIOException(new IOException(missingAlgorithmException));
                         }
+                        logger.info("Initializing HTTPS payload processor");
                         processor = new RemotePayloadProcessor(uri, sslContext, sslContext.getDefaultSSLParameters());
                     } else if ("logger".equals(processorName)) {
                         processor = new LoggingPayloadProcessor();