Use AutoReloadingX509KeyManager in SiaIdentityProvider

The SSLContext instance provided by SiaIdentityProvider
can now be reused as the underlying key manager will automatically
update the key material every hour.

2 files changed, 16 insertions, 57 deletions

diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaIdentityProvider.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaIdentityProvider.java
index 2b0e50ed982..cab28e55b21 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaIdentityProvider.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaIdentityProvider.java
@@ -3,23 +3,17 @@ package com.yahoo.vespa.athenz.identity;
 
 import com.google.inject.Inject;
 import com.yahoo.component.AbstractComponent;
-import com.yahoo.log.LogLevel;
-import com.yahoo.vespa.athenz.api.AthenzIdentity;
-import com.yahoo.vespa.athenz.api.AthenzService;
 import com.yahoo.security.KeyStoreType;
 import com.yahoo.security.SslContextBuilder;
+import com.yahoo.security.tls.AutoReloadingX509KeyManager;
+import com.yahoo.vespa.athenz.api.AthenzIdentity;
+import com.yahoo.vespa.athenz.api.AthenzService;
 import com.yahoo.vespa.athenz.utils.SiaUtils;
 
 import javax.net.ssl.SSLContext;
 import java.io.File;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.time.Duration;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.ScheduledThreadPoolExecutor;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicReference;
-import java.util.logging.Logger;
 
 /**
  * A {@link ServiceIdentityProvider} that provides the credentials stored on file system.
@@ -29,24 +23,19 @@ import java.util.logging.Logger;
  */
 public class SiaIdentityProvider extends AbstractComponent implements ServiceIdentityProvider {
 
-    private static final Logger log = Logger.getLogger(SiaIdentityProvider.class.getName());
-
-    private static final Duration REFRESH_INTERVAL = Duration.ofHours(1);
-
-    private final AtomicReference<SSLContext> sslContext = new AtomicReference<>();
+    private final AutoReloadingX509KeyManager keyManager;
+    private final SSLContext sslContext;
     private final AthenzIdentity service;
     private final File privateKeyFile;
     private final File certificateFile;
     private final File trustStoreFile;
-    private final ScheduledExecutorService scheduler;
 
     @Inject
     public SiaIdentityProvider(SiaProviderConfig config) {
         this(new AthenzService(config.athenzDomain(), config.athenzService()),
              SiaUtils.getPrivateKeyFile(Paths.get(config.keyPathPrefix()), new AthenzService(config.athenzDomain(), config.athenzService())).toFile(),
              SiaUtils.getCertificateFile(Paths.get(config.keyPathPrefix()), new AthenzService(config.athenzDomain(), config.athenzService())).toFile(),
-             new File(config.trustStorePath()),
-             createScheduler());
+             new File(config.trustStorePath()));
     }
 
     public SiaIdentityProvider(AthenzIdentity service,
@@ -55,30 +44,19 @@ public class SiaIdentityProvider extends AbstractComponent implements ServiceIde
         this(service,
              SiaUtils.getPrivateKeyFile(siaPath, service).toFile(),
              SiaUtils.getCertificateFile(siaPath, service).toFile(),
-             trustStoreFile,
-             createScheduler());
+             trustStoreFile);
     }
 
     public SiaIdentityProvider(AthenzIdentity service,
                                File privateKeyFile,
                                File certificateFile,
-                               File trustStoreFile,
-                               ScheduledExecutorService scheduler) {
+                               File trustStoreFile) {
         this.service = service;
         this.privateKeyFile = privateKeyFile;
         this.certificateFile = certificateFile;
         this.trustStoreFile = trustStoreFile;
-        this.scheduler = scheduler;
-        this.sslContext.set(createIdentitySslContext());
-        scheduler.scheduleAtFixedRate(this::reloadSslContext, REFRESH_INTERVAL.toMinutes(), REFRESH_INTERVAL.toMinutes(), TimeUnit.MINUTES);
-    }
-
-    private static ScheduledThreadPoolExecutor createScheduler() {
-        return new ScheduledThreadPoolExecutor(1, runnable -> {
-            Thread thread = new Thread(runnable);
-            thread.setName("sia-identity-provider-sslcontext-updater");
-            return thread;
-        });
+        this.keyManager = AutoReloadingX509KeyManager.fromPemFiles(privateKeyFile.toPath(), certificateFile.toPath());
+        this.sslContext = createIdentitySslContext(keyManager, trustStoreFile.toPath());
     }
 
     @Override
@@ -88,34 +66,18 @@ public class SiaIdentityProvider extends AbstractComponent implements ServiceIde
 
     @Override
     public SSLContext getIdentitySslContext() {
-        return sslContext.get();
+        return sslContext;
     }
 
-    private SSLContext createIdentitySslContext() {
+    private static SSLContext createIdentitySslContext(AutoReloadingX509KeyManager keyManager, Path trustStoreFile) {
         return new SslContextBuilder()
-                .withTrustStore(trustStoreFile.toPath(), KeyStoreType.JKS)
-                .withKeyStore(privateKeyFile.toPath(), certificateFile.toPath())
+                .withTrustStore(trustStoreFile, KeyStoreType.JKS)
+                .withKeyManager(keyManager)
                 .build();
     }
 
-    private void reloadSslContext() {
-        log.log(LogLevel.DEBUG, "Updating SSLContext for identity " + service.getFullName());
-        try {
-            SSLContext sslContext = createIdentitySslContext();
-            this.sslContext.set(sslContext);
-        } catch (Exception e) {
-            log.log(LogLevel.SEVERE, "Failed to update SSLContext: " + e.getMessage(), e);
-        }
-    }
-
-
     @Override
     public void deconstruct() {
-        try {
-            scheduler.shutdownNow();
-            scheduler.awaitTermination(90, TimeUnit.SECONDS);
-        } catch (InterruptedException e) {
-            throw new RuntimeException(e);
-        }
+        keyManager.close();
     }
 }
diff --git a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/identity/SiaIdentityProviderTest.java b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/identity/SiaIdentityProviderTest.java
index 0195d6000e1..31152a4602f 100644
--- a/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/identity/SiaIdentityProviderTest.java
+++ b/vespa-athenz/src/test/java/com/yahoo/vespa/athenz/identity/SiaIdentityProviderTest.java
@@ -24,10 +24,8 @@ import java.security.KeyStore;
 import java.security.cert.X509Certificate;
 import java.time.Duration;
 import java.time.Instant;
-import java.util.concurrent.ScheduledExecutorService;
 
 import static org.junit.Assert.assertNotNull;
-import static org.mockito.Mockito.mock;
 
 /**
  * @author bjorncs
@@ -55,8 +53,7 @@ public class SiaIdentityProviderTest {
                         new AthenzService("domain", "service-name"),
                         keyFile,
                         certificateFile,
-                        trustStoreFile,
-                        mock(ScheduledExecutorService.class));
+                        trustStoreFile);
 
         assertNotNull(provider.getIdentitySslContext());
     }