Add x509 key manager that regularly updates cert chain from PEM files

3 files changed, 239 insertions, 0 deletions

diff --git a/security-utils/pom.xml b/security-utils/pom.xml
index 0a26c73cf70..10dec598915 100644
--- a/security-utils/pom.xml
+++ b/security-utils/pom.xml
@@ -54,6 +54,11 @@
       <artifactId>assertj-core</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-core</artifactId>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
   <build>
     <plugins>
diff --git a/security-utils/src/main/java/com/yahoo/security/tls/AutoReloadingX509KeyManager.java b/security-utils/src/main/java/com/yahoo/security/tls/AutoReloadingX509KeyManager.java
new file mode 100644
index 00000000000..0dae185995c
--- /dev/null
+++ b/security-utils/src/main/java/com/yahoo/security/tls/AutoReloadingX509KeyManager.java
@@ -0,0 +1,150 @@
+// 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.KeyStoreBuilder;
+import com.yahoo.security.KeyStoreType;
+import com.yahoo.security.KeyUtils;
+import com.yahoo.security.X509CertificateUtils;
+
+import javax.net.ssl.SSLEngine;
+import javax.net.ssl.X509ExtendedKeyManager;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.net.Socket;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.security.KeyStore;
+import java.security.Principal;
+import java.security.PrivateKey;
+import java.security.cert.X509Certificate;
+import java.time.Duration;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+/**
+ * A {@link X509ExtendedKeyManager} that reloads the certificate and private key from file regularly.
+ *
+ * @author bjorncs
+ */
+public class AutoReloadingX509KeyManager extends X509ExtendedKeyManager implements AutoCloseable {
+
+    private static final Duration UPDATE_PERIOD = Duration.ofHours(1);
+
+    private static final Logger log = Logger.getLogger(AutoReloadingX509KeyManager.class.getName());
+
+    private final MutableX509KeyManager mutableX509KeyManager;
+    private final ScheduledExecutorService scheduler;
+    private final Path privateKeyFile;
+    private final Path certificatesFile;
+
+    private AutoReloadingX509KeyManager(Path privateKeyFile, Path certificatesFile) {
+        this(privateKeyFile, certificatesFile, createDefaultScheduler());
+    }
+
+    AutoReloadingX509KeyManager(Path privateKeyFile, Path certificatesFile, ScheduledExecutorService scheduler) {
+            this.privateKeyFile = privateKeyFile;
+            this.certificatesFile = certificatesFile;
+            this.scheduler = scheduler;
+            this.mutableX509KeyManager = new MutableX509KeyManager(createKeystore(privateKeyFile, certificatesFile), new char[0]);
+            scheduler.scheduleAtFixedRate(
+                    new KeyManagerReloader(), UPDATE_PERIOD.getSeconds()/*initial delay*/, UPDATE_PERIOD.getSeconds(), TimeUnit.SECONDS);
+        }
+
+    public static AutoReloadingX509KeyManager fromPemFiles(Path privateKeyFile, Path certificatesFile) {
+        return new AutoReloadingX509KeyManager(privateKeyFile, certificatesFile);
+    }
+
+    private static KeyStore createKeystore(Path privateKey, Path certificateChain) {
+        try {
+            return KeyStoreBuilder.withType(KeyStoreType.PKCS12)
+                    .withKeyEntry(
+                            "default",
+                            KeyUtils.fromPemEncodedPrivateKey(Files.readString(privateKey)),
+                            X509CertificateUtils.certificateListFromPem(Files.readString(certificateChain)))
+                    .build();
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
+    private static ScheduledExecutorService createDefaultScheduler() {
+        return Executors.newSingleThreadScheduledExecutor(runnable -> {
+            Thread thread = new Thread(runnable, "auto-reloading-x509-key-manager");
+            thread.setDaemon(true);
+            return thread;
+        });
+    }
+
+    private class KeyManagerReloader implements Runnable {
+        @Override
+        public void run() {
+            try {
+                log.log(Level.FINE, () -> String.format("Reloading key and certificate chain (private-key='%s', certificates='%s')", privateKeyFile, certificatesFile));
+                mutableX509KeyManager.updateKeystore(createKeystore(privateKeyFile, certificatesFile), new char[0]);
+            } catch (Throwable t) {
+                log.log(Level.SEVERE,
+                        String.format("Failed to load X509 key manager (private-key='%s', certificates='%s'): %s",
+                                      privateKeyFile, certificatesFile, t.getMessage()),
+                        t);
+            }
+        }
+    }
+
+    @Override
+    public void close() {
+        try {
+            scheduler.shutdownNow();
+            scheduler.awaitTermination(5, TimeUnit.SECONDS);
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    //
+    // Methods from X509ExtendedKeyManager
+    //
+
+    @Override
+    public String[] getServerAliases(String keyType, Principal[] issuers) {
+        return mutableX509KeyManager.getServerAliases(keyType, issuers);
+    }
+
+    @Override
+    public String[] getClientAliases(String keyType, Principal[] issuers) {
+        return mutableX509KeyManager.getClientAliases(keyType, issuers);
+    }
+
+    @Override
+    public String chooseServerAlias(String keyType, Principal[] issuers, Socket socket) {
+        return mutableX509KeyManager.chooseServerAlias(keyType, issuers, socket);
+    }
+
+    @Override
+    public String chooseClientAlias(String[] keyType, Principal[] issuers, Socket socket) {
+        return mutableX509KeyManager.chooseClientAlias(keyType, issuers, socket);
+    }
+
+    @Override
+    public String chooseEngineServerAlias(String keyType, Principal[] issuers, SSLEngine engine) {
+        return mutableX509KeyManager.chooseEngineServerAlias(keyType, issuers, engine);
+    }
+
+    @Override
+    public String chooseEngineClientAlias(String[] keyType, Principal[] issuers, SSLEngine engine) {
+        return mutableX509KeyManager.chooseEngineClientAlias(keyType, issuers, engine);
+    }
+
+    @Override
+    public X509Certificate[] getCertificateChain(String alias) {
+        return mutableX509KeyManager.getCertificateChain(alias);
+    }
+
+    @Override
+    public PrivateKey getPrivateKey(String alias) {
+        return mutableX509KeyManager.getPrivateKey(alias);
+    }
+
+}
diff --git a/security-utils/src/test/java/com/yahoo/security/tls/AutoReloadingX509KeyManagerTest.java b/security-utils/src/test/java/com/yahoo/security/tls/AutoReloadingX509KeyManagerTest.java
new file mode 100644
index 00000000000..139d5313074
--- /dev/null
+++ b/security-utils/src/test/java/com/yahoo/security/tls/AutoReloadingX509KeyManagerTest.java
@@ -0,0 +1,84 @@
+// 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.KeyAlgorithm;
+import com.yahoo.security.KeyUtils;
+import com.yahoo.security.SignatureAlgorithm;
+import com.yahoo.security.X509CertificateBuilder;
+import com.yahoo.security.X509CertificateUtils;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
+
+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.Principal;
+import java.security.cert.X509Certificate;
+import java.time.Instant;
+import java.util.concurrent.ScheduledExecutorService;
+
+import static java.time.temporal.ChronoUnit.DAYS;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.anyLong;
+import static org.mockito.Mockito.verify;
+
+/**
+ * @author bjorncs
+ */
+public class AutoReloadingX509KeyManagerTest {
+    private static final X500Principal SUBJECT = new X500Principal("CN=dummy");
+
+    @Rule
+    public TemporaryFolder tempDirectory = new TemporaryFolder();
+
+    @Test
+    public void crypto_material_is_reloaded_when_scheduler_task_is_executed() throws IOException {
+        KeyPair keyPair = KeyUtils.generateKeypair(KeyAlgorithm.EC);
+        Path privateKeyFile = tempDirectory.newFile().toPath();
+        Files.writeString(privateKeyFile, KeyUtils.toPem(keyPair.getPrivate()));
+
+        Path certificateFile = tempDirectory.newFile().toPath();
+        BigInteger serialNumberInitialCertificate = BigInteger.ONE;
+        X509Certificate initialCertificate = generateCertificate(keyPair, serialNumberInitialCertificate);
+        Files.writeString(certificateFile, X509CertificateUtils.toPem(initialCertificate));
+
+        ScheduledExecutorService scheduler = Mockito.mock(ScheduledExecutorService.class);
+        ArgumentCaptor<Runnable> updaterTaskCaptor = ArgumentCaptor.forClass(Runnable.class);
+
+        AutoReloadingX509KeyManager keyManager = new AutoReloadingX509KeyManager(privateKeyFile, certificateFile, scheduler);
+        verify(scheduler).scheduleAtFixedRate(updaterTaskCaptor.capture(), anyLong(), anyLong(), any());
+
+        String[] initialAliases = keyManager.getClientAliases(keyPair.getPublic().getAlgorithm(), new Principal[]{SUBJECT});
+        X509Certificate[] certChain = keyManager.getCertificateChain(initialAliases[0]);
+        assertThat(certChain).hasSize(1);
+        assertThat(certChain[0].getSerialNumber()).isEqualTo(serialNumberInitialCertificate);
+
+        BigInteger serialNumberUpdatedCertificate = BigInteger.TWO;
+        X509Certificate updatedCertificate = generateCertificate(keyPair, serialNumberUpdatedCertificate);
+        Files.writeString(certificateFile, X509CertificateUtils.toPem(updatedCertificate));
+
+        updaterTaskCaptor.getValue().run(); // run update task in ReloadingX509KeyManager
+
+        String[] updatedAliases = keyManager.getClientAliases(keyPair.getPublic().getAlgorithm(), new Principal[]{SUBJECT});
+        X509Certificate[] updatedCertChain = keyManager.getCertificateChain(updatedAliases[0]);
+        assertThat(updatedCertChain).hasSize(1);
+        assertThat(updatedCertChain[0].getSerialNumber()).isEqualTo(serialNumberUpdatedCertificate);
+    }
+
+    private static X509Certificate generateCertificate(KeyPair keyPair, BigInteger serialNumber) {
+        return X509CertificateBuilder.fromKeypair(keyPair,
+                                                  SUBJECT,
+                                                  Instant.EPOCH,
+                                                  Instant.EPOCH.plus(1, DAYS),
+                                                  SignatureAlgorithm.SHA256_WITH_ECDSA,
+                                                  serialNumber)
+                .build();
+    }
+}
\ No newline at end of file