Add mutable x509 key manager

Add a x509 key manager where certificates can be updated while the
manager is in use.

2 files changed, 171 insertions, 0 deletions

diff --git a/security-utils/src/main/java/com/yahoo/security/tls/MutableX509KeyManager.java b/security-utils/src/main/java/com/yahoo/security/tls/MutableX509KeyManager.java
new file mode 100644
index 00000000000..936dbe61ad6
--- /dev/null
+++ b/security-utils/src/main/java/com/yahoo/security/tls/MutableX509KeyManager.java
@@ -0,0 +1,106 @@
+// 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 javax.net.ssl.SSLEngine;
+import javax.net.ssl.X509ExtendedKeyManager;
+import java.net.Socket;
+import java.security.KeyStore;
+import java.security.Principal;
+import java.security.PrivateKey;
+import java.security.cert.X509Certificate;
+import java.util.WeakHashMap;
+
+/**
+ * A {@link X509ExtendedKeyManager} which can be updated with new certificate chain and private key while in use.
+ *
+ * The implementations assumes that aliases are retrieve from the same thread as the certificate chain and private key.
+ * This is case for OpenJDK 11.
+ *
+ * @author bjorncs
+ */
+public class MutableX509KeyManager extends X509ExtendedKeyManager {
+
+    // Not using ThreadLocal as we want the x509 key manager instances to be collected
+    // when either the thread dies or the MutableX509KeyManager instance is collected (latter not the case for ThreadLocal).
+    private final WeakHashMap<Thread, X509ExtendedKeyManager> threadLocalManager = new WeakHashMap<>();
+    private volatile X509ExtendedKeyManager currentManager;
+
+    public MutableX509KeyManager(KeyStore keystore, char[] password) {
+        this.currentManager = KeyManagerUtils.createDefaultX509KeyManager(keystore, password);
+    }
+
+    public MutableX509KeyManager() {
+        this.currentManager = KeyManagerUtils.createDefaultX509KeyManager();
+    }
+
+    public void updateKeystore(KeyStore keystore, char[] password) {
+        this.currentManager = KeyManagerUtils.createDefaultX509KeyManager(keystore, password);
+    }
+
+    public void useDefaultKeystore() {
+        this.currentManager = KeyManagerUtils.createDefaultX509KeyManager();
+    }
+
+    @Override
+    public String[] getServerAliases(String keyType, Principal[] issuers) {
+        return updateAndGetThreadLocalManager()
+                .getServerAliases(keyType, issuers);
+    }
+
+    @Override
+    public String[] getClientAliases(String keyType, Principal[] issuers) {
+        return updateAndGetThreadLocalManager()
+                .getClientAliases(keyType, issuers);
+    }
+
+    @Override
+    public String chooseServerAlias(String keyType, Principal[] issuers, Socket socket) {
+        return updateAndGetThreadLocalManager()
+                .chooseServerAlias(keyType, issuers, socket);
+    }
+
+    @Override
+    public String chooseClientAlias(String[] keyType, Principal[] issuers, Socket socket) {
+        return updateAndGetThreadLocalManager()
+                .chooseClientAlias(keyType, issuers, socket);
+    }
+
+    @Override
+    public String chooseEngineServerAlias(String keyType, Principal[] issuers, SSLEngine engine) {
+        return updateAndGetThreadLocalManager()
+                .chooseEngineServerAlias(keyType, issuers, engine);
+    }
+
+    @Override
+    public String chooseEngineClientAlias(String[] keyType, Principal[] issuers, SSLEngine engine) {
+        return updateAndGetThreadLocalManager()
+                .chooseEngineClientAlias(keyType, issuers, engine);
+    }
+
+    private X509ExtendedKeyManager updateAndGetThreadLocalManager() {
+        X509ExtendedKeyManager currentManager = this.currentManager;
+        threadLocalManager.put(Thread.currentThread(), currentManager);
+        return currentManager;
+    }
+
+    @Override
+    public X509Certificate[] getCertificateChain(String alias) {
+        return getThreadLocalManager()
+                .getCertificateChain(alias);
+    }
+
+    @Override
+    public PrivateKey getPrivateKey(String alias) {
+        return getThreadLocalManager()
+                .getPrivateKey(alias);
+    }
+
+    private X509ExtendedKeyManager getThreadLocalManager() {
+        X509ExtendedKeyManager manager = threadLocalManager.get(Thread.currentThread());
+        if (manager == null) {
+            throw new IllegalStateException("Methods to retrieve valid aliases has not been called previously from this thread");
+        }
+        return manager;
+    }
+
+}
diff --git a/security-utils/src/test/java/com/yahoo/security/tls/MutableX509KeyManagerTest.java b/security-utils/src/test/java/com/yahoo/security/tls/MutableX509KeyManagerTest.java
new file mode 100644
index 00000000000..30e54d3c09d
--- /dev/null
+++ b/security-utils/src/test/java/com/yahoo/security/tls/MutableX509KeyManagerTest.java
@@ -0,0 +1,65 @@
+// 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.KeyStoreBuilder;
+import com.yahoo.security.KeyStoreType;
+import com.yahoo.security.KeyUtils;
+import com.yahoo.security.SignatureAlgorithm;
+import com.yahoo.security.X509CertificateBuilder;
+import org.junit.Test;
+
+import javax.security.auth.x500.X500Principal;
+import java.math.BigInteger;
+import java.security.KeyPair;
+import java.security.KeyStore;
+import java.security.Principal;
+import java.security.cert.X509Certificate;
+import java.time.Instant;
+
+import static java.time.temporal.ChronoUnit.DAYS;
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * @author bjorncs
+ */
+public class MutableX509KeyManagerTest {
+
+    private static final X500Principal SUBJECT = new X500Principal("CN=dummy");
+
+    @Test
+    public void key_manager_can_be_updated_with_new_certificate() {
+        KeyPair keyPair = KeyUtils.generateKeypair(KeyAlgorithm.EC);
+
+        BigInteger serialNumberInitialCertificate = BigInteger.ONE;
+        KeyStore initialKeystore = generateKeystore(keyPair, serialNumberInitialCertificate);
+
+        MutableX509KeyManager keyManager = new MutableX509KeyManager(initialKeystore, new char[0]);
+
+        String[] initialAliases = keyManager.getClientAliases(keyPair.getPublic().getAlgorithm(), new Principal[]{SUBJECT});
+        assertThat(initialAliases).hasSize(1);
+        X509Certificate[] certChain = keyManager.getCertificateChain(initialAliases[0]);
+        assertThat(certChain).hasSize(1);
+        assertThat(certChain[0].getSerialNumber()).isEqualTo(serialNumberInitialCertificate);
+
+        BigInteger serialNumberUpdatedCertificate = BigInteger.TWO;
+        KeyStore updatedKeystore = generateKeystore(keyPair, serialNumberUpdatedCertificate);
+        keyManager.updateKeystore(updatedKeystore, new char[0]);
+
+        String[] updatedAliases = keyManager.getClientAliases(keyPair.getPublic().getAlgorithm(), new Principal[]{SUBJECT});
+        assertThat(updatedAliases).hasSize(1);
+        X509Certificate[] updatedCertChain = keyManager.getCertificateChain(updatedAliases[0]);
+        assertThat(updatedCertChain).hasSize(1);
+        assertThat(updatedCertChain[0].getSerialNumber()).isEqualTo(serialNumberUpdatedCertificate);
+    }
+
+    private static KeyStore generateKeystore(KeyPair keyPair, BigInteger serialNumber) {
+        X509Certificate certificate = X509CertificateBuilder.fromKeypair(
+                keyPair, SUBJECT, Instant.EPOCH, Instant.EPOCH.plus(1, DAYS), SignatureAlgorithm.SHA256_WITH_ECDSA, serialNumber)
+                .build();
+        return KeyStoreBuilder.withType(KeyStoreType.PKCS12)
+                .withKeyEntry("default", keyPair.getPrivate(), certificate)
+                .build();
+    }
+
+}
\ No newline at end of file