Stop reload task when there are no external references to the managers

The reload task will shut down the executor service when the GC has
determined that there are no other references to the key/trust manager.

2 files changed, 126 insertions, 18 deletions

diff --git a/security-utils/src/main/java/com/yahoo/security/tls/ReloadingTlsContext.java b/security-utils/src/main/java/com/yahoo/security/tls/ReloadingTlsContext.java
index debf14a27f8..16f66f91da6 100644
--- a/security-utils/src/main/java/com/yahoo/security/tls/ReloadingTlsContext.java
+++ b/security-utils/src/main/java/com/yahoo/security/tls/ReloadingTlsContext.java
@@ -15,6 +15,7 @@ import javax.net.ssl.SSLParameters;
 import javax.net.ssl.X509ExtendedTrustManager;
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.lang.ref.WeakReference;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.security.KeyStore;
@@ -22,6 +23,7 @@ import java.time.Duration;
 import java.util.Set;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -37,37 +39,31 @@ public class ReloadingTlsContext implements TlsContext {
 
     private static final Logger log = Logger.getLogger(ReloadingTlsContext.class.getName());
 
-    private final Path tlsOptionsConfigFile;
     private final TlsContext tlsContext;
-    private final MutableX509TrustManager trustManager = new MutableX509TrustManager();
-    private final MutableX509KeyManager keyManager = new MutableX509KeyManager();
-    private final ScheduledExecutorService scheduler =
-            Executors.newSingleThreadScheduledExecutor(runnable -> {
-                Thread thread = new Thread(runnable, "tls-context-reloader");
-                thread.setDaemon(true);
-                return thread;
-            });
+    private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(new ReloaderThreadFactory());
 
     public ReloadingTlsContext(Path tlsOptionsConfigFile, AuthorizationMode mode) {
-        this.tlsOptionsConfigFile = tlsOptionsConfigFile;
         TransportSecurityOptions options = TransportSecurityOptions.fromJsonFile(tlsOptionsConfigFile);
-        reloadCryptoMaterial(options, trustManager, keyManager);
+        MutableX509TrustManager trustManager = new MutableX509TrustManager();
+        MutableX509KeyManager keyManager = new MutableX509KeyManager();
+        reloadTrustManager(options, trustManager);
+        reloadKeyManager(options, keyManager);
         this.tlsContext = createDefaultTlsContext(options, mode, trustManager, keyManager);
-        this.scheduler.scheduleAtFixedRate(new CryptoMaterialReloader(),
+        this.scheduler.scheduleAtFixedRate(new CryptoMaterialReloader(tlsOptionsConfigFile, scheduler, trustManager, keyManager),
                                            UPDATE_PERIOD.getSeconds()/*initial delay*/,
                                            UPDATE_PERIOD.getSeconds(),
                                            TimeUnit.SECONDS);
     }
 
-    private static void reloadCryptoMaterial(TransportSecurityOptions options,
-                                             MutableX509TrustManager trustManager,
-                                             MutableX509KeyManager keyManager) {
+    private static void reloadTrustManager(TransportSecurityOptions options, MutableX509TrustManager trustManager) {
         if (options.getCaCertificatesFile().isPresent()) {
             trustManager.updateTruststore(loadTruststore(options.getCaCertificatesFile().get()));
         } else {
             trustManager.useDefaultTruststore();
         }
+    }
 
+    private static void reloadKeyManager(TransportSecurityOptions options, MutableX509KeyManager keyManager) {
         if (options.getPrivateKeyFile().isPresent() && options.getCertificatesFile().isPresent()) {
             keyManager.updateKeystore(loadKeystore(options.getPrivateKeyFile().get(), options.getCertificatesFile().get()), new char[0]);
         } else {
@@ -122,21 +118,63 @@ public class ReloadingTlsContext implements TlsContext {
     public void close() {
         try {
             scheduler.shutdownNow();
-            scheduler.awaitTermination(5, TimeUnit.SECONDS);
+            if (!scheduler.awaitTermination(10, TimeUnit.SECONDS)) {
+                throw new RuntimeException("Unable to shutdown executor before timeout");
+            }
         } catch (InterruptedException e) {
             throw new RuntimeException(e);
         }
     }
 
-    private class CryptoMaterialReloader implements Runnable {
+    // Note: no reference to outer class (directly or indirectly) to ensure trust/key managers are eventually GCed once
+    // there are no more use of the outer class and the underlying SSLContext
+    private static class CryptoMaterialReloader implements Runnable {
+
+        final Path tlsOptionsConfigFile;
+        final ScheduledExecutorService scheduler;
+        final WeakReference<MutableX509TrustManager> trustManager;
+        final WeakReference<MutableX509KeyManager> keyManager;
+
+        CryptoMaterialReloader(Path tlsOptionsConfigFile,
+                               ScheduledExecutorService scheduler,
+                               MutableX509TrustManager trustManager,
+                               MutableX509KeyManager keyManager) {
+            this.tlsOptionsConfigFile = tlsOptionsConfigFile;
+            this.scheduler = scheduler;
+            this.trustManager = new WeakReference<>(trustManager);
+            this.keyManager = new WeakReference<>(keyManager);
+        }
+
         @Override
         public void run() {
             try {
-                reloadCryptoMaterial(TransportSecurityOptions.fromJsonFile(tlsOptionsConfigFile), trustManager, keyManager);
+                MutableX509TrustManager trustManager = this.trustManager.get();
+                MutableX509KeyManager keyManager = this.keyManager.get();
+                if (trustManager == null && keyManager == null) {
+                    scheduler.shutdown();
+                    return;
+                }
+                TransportSecurityOptions options = TransportSecurityOptions.fromJsonFile(tlsOptionsConfigFile);
+                if (trustManager != null) {
+                    reloadTrustManager(options, trustManager);
+                }
+                if (keyManager != null) {
+                    reloadKeyManager(options, keyManager);
+                }
             } catch (Throwable t) {
                 log.log(Level.SEVERE, String.format("Failed to reload crypto material (path='%s'): %s", tlsOptionsConfigFile, t.getMessage()), t);
             }
         }
     }
 
+    // Static class to ensure no reference to outer class is contained
+    private static class ReloaderThreadFactory implements ThreadFactory {
+        @Override
+        public Thread newThread(Runnable r) {
+            Thread thread = new Thread(r, "tls-context-reloader");
+            thread.setDaemon(true);
+            return thread;
+        }
+    }
+
 }
diff --git a/security-utils/src/test/java/com/yahoo/security/tls/ReloadingTlsContextTest.java b/security-utils/src/test/java/com/yahoo/security/tls/ReloadingTlsContextTest.java
new file mode 100644
index 00000000000..f991f86fdce
--- /dev/null
+++ b/security-utils/src/test/java/com/yahoo/security/tls/ReloadingTlsContextTest.java
@@ -0,0 +1,70 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.security.tls;
+
+import com.yahoo.security.KeyUtils;
+import com.yahoo.security.X509CertificateBuilder;
+import com.yahoo.security.X509CertificateUtils;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import javax.net.ssl.SSLEngine;
+import javax.security.auth.x500.X500Principal;
+import java.io.IOException;
+import java.math.BigInteger;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.security.KeyPair;
+import java.security.cert.X509Certificate;
+
+import static com.yahoo.security.KeyAlgorithm.EC;
+import static com.yahoo.security.SignatureAlgorithm.SHA256_WITH_ECDSA;
+import static java.time.Instant.EPOCH;
+import static java.time.temporal.ChronoUnit.DAYS;
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * @author bjorncs
+ */
+public class ReloadingTlsContextTest {
+
+    @Rule
+    public TemporaryFolder tempDirectory = new TemporaryFolder();
+
+    @Test
+    public void can_create_sslcontext_from_credentials() throws IOException, InterruptedException {
+        KeyPair keyPair = KeyUtils.generateKeypair(EC);
+        Path privateKeyFile = tempDirectory.newFile().toPath();
+        Files.writeString(privateKeyFile, KeyUtils.toPem(keyPair.getPrivate()));
+
+        X509Certificate certificate = X509CertificateBuilder
+                .fromKeypair(keyPair, new X500Principal("CN=dummy"), EPOCH, EPOCH.plus(1, DAYS), SHA256_WITH_ECDSA, BigInteger.ONE)
+                .build();
+        Path certificateChainFile = tempDirectory.newFile().toPath();
+        String certificatePem = X509CertificateUtils.toPem(certificate);
+        Files.writeString(certificateChainFile, certificatePem);
+
+        Path caCertificatesFile = tempDirectory.newFile().toPath();
+        Files.writeString(caCertificatesFile, certificatePem);
+
+        TransportSecurityOptions options = new TransportSecurityOptions.Builder()
+                .withCertificates(certificateChainFile, privateKeyFile)
+                .withCaCertificates(caCertificatesFile)
+                .build();
+
+        Path optionsFile = tempDirectory.newFile().toPath();
+        options.toJsonFile(optionsFile);
+
+        try (TlsContext tlsContext = new ReloadingTlsContext(optionsFile, AuthorizationMode.ENFORCE)) {
+            SSLEngine sslEngine = tlsContext.createSslEngine();
+            assertThat(sslEngine).isNotNull();
+            String[] enabledCiphers = sslEngine.getEnabledCipherSuites();
+            assertThat(enabledCiphers).isNotEmpty();
+            assertThat(enabledCiphers).isSubsetOf(DefaultTlsContext.ALLOWED_CIPHER_SUITES.toArray(new String[0]));
+
+            String[] enabledProtocols = sslEngine.getEnabledProtocols();
+            assertThat(enabledProtocols).contains("TLSv1.2");
+        }
+    }
+
+}
\ No newline at end of file