Write more TLS config and a key store to be used by ZooKeeper server

2 files changed, 189 insertions, 59 deletions

diff --git a/zookeeper-server/zookeeper-server-3.5/src/main/java/com/yahoo/vespa/zookeeper/VespaZooKeeperServerImpl.java b/zookeeper-server/zookeeper-server-3.5/src/main/java/com/yahoo/vespa/zookeeper/VespaZooKeeperServerImpl.java
index c4e1f8130a1..ed8be3ad7f4 100644
--- a/zookeeper-server/zookeeper-server-3.5/src/main/java/com/yahoo/vespa/zookeeper/VespaZooKeeperServerImpl.java
+++ b/zookeeper-server/zookeeper-server-3.5/src/main/java/com/yahoo/vespa/zookeeper/VespaZooKeeperServerImpl.java
@@ -5,13 +5,26 @@ import com.google.inject.Inject;
 import com.yahoo.cloud.config.ZookeeperServerConfig;
 import com.yahoo.component.AbstractComponent;
 import com.yahoo.log.LogLevel;
+import com.yahoo.security.KeyStoreBuilder;
+import com.yahoo.security.KeyStoreType;
+import com.yahoo.security.KeyStoreUtils;
+import com.yahoo.security.KeyUtils;
+import com.yahoo.security.X509CertificateUtils;
 import com.yahoo.security.tls.TlsContext;
+import com.yahoo.security.tls.TransportSecurityOptions;
+import com.yahoo.security.tls.TransportSecurityUtils;
+import com.yahoo.text.Utf8;
 
 import static com.yahoo.vespa.defaults.Defaults.getDefaults;
 
 import java.io.File;
 import java.io.FileWriter;
 import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.security.PrivateKey;
+import java.security.cert.X509Certificate;
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
@@ -31,14 +44,19 @@ public class VespaZooKeeperServerImpl extends AbstractComponent implements Runna
     static final String ZOOKEEPER_JUTE_MAX_BUFFER = "jute.maxbuffer";
     private final Thread zkServerThread;
     private final ZookeeperServerConfig zookeeperServerConfig;
+    private final String configFilePath;
+    private final String jksKeyStoreFilePath;
 
-    VespaZooKeeperServerImpl(ZookeeperServerConfig zookeeperServerConfig, boolean startServer) {
+    VespaZooKeeperServerImpl(ZookeeperServerConfig zookeeperServerConfig, boolean startServer,
+                            Optional<TransportSecurityOptions> transportSecurityOptions) {
         this.zookeeperServerConfig = zookeeperServerConfig;
         System.setProperty("zookeeper.jmx.log4j.disable", "true");
         System.setProperty("zookeeper.snapshot.trust.empty", Boolean.valueOf(zookeeperServerConfig.trustEmptySnapshot()).toString());
         System.setProperty(ZOOKEEPER_JUTE_MAX_BUFFER, Integer.valueOf(zookeeperServerConfig.juteMaxBuffer()).toString());
 
-        writeConfigToDisk(zookeeperServerConfig);
+        configFilePath = getDefaults().underVespaHome(zookeeperServerConfig.zooKeeperConfigFile());
+        jksKeyStoreFilePath = getDefaults().underVespaHome(zookeeperServerConfig.jksKeyStoreFile());
+        writeConfigToDisk(zookeeperServerConfig, transportSecurityOptions);
         zkServerThread = new Thread(this, "zookeeper server");
         if (startServer) {
             zkServerThread.start();
@@ -47,21 +65,30 @@ public class VespaZooKeeperServerImpl extends AbstractComponent implements Runna
 
     @Inject
     public VespaZooKeeperServerImpl(ZookeeperServerConfig zookeeperServerConfig) {
-        this(zookeeperServerConfig, true);
+        this(zookeeperServerConfig, true, TransportSecurityUtils.getOptions());
     }
 
-    private void writeConfigToDisk(ZookeeperServerConfig config) {
-       String configFilePath = getDefaults().underVespaHome(config.zooKeeperConfigFile());
-       new File(configFilePath).getParentFile().mkdirs();
-       try (FileWriter writer = new FileWriter(configFilePath)) {
-           writer.write(transformConfigToString(config));
-           writeMyIdFile(config);
-       } catch (IOException e) {
-           throw new RuntimeException("Error writing zookeeper config", e);
-       }
+    private void writeConfigToDisk(ZookeeperServerConfig config, Optional<TransportSecurityOptions> transportSecurityOptions) {
+        new File(configFilePath).getParentFile().mkdirs();
+
+        try {
+            writeZooKeeperConfigFile(zookeeperServerConfig, transportSecurityOptions);
+            writeMyIdFile(config);
+            transportSecurityOptions.ifPresent(this::writeJksKeystore);
+        } catch (IOException e) {
+            throw new RuntimeException("Error writing zookeeper config", e);
+        }
    }
 
-    private String transformConfigToString(ZookeeperServerConfig config) {
+    private void writeZooKeeperConfigFile(ZookeeperServerConfig config,
+                                          Optional<TransportSecurityOptions> transportSecurityOptions) throws IOException {
+        try (FileWriter writer = new FileWriter(configFilePath)) {
+            writer.write(transformConfigToString(config, transportSecurityOptions));
+        }
+    }
+
+    private String transformConfigToString(ZookeeperServerConfig config,
+                                           Optional<TransportSecurityOptions> transportSecurityOptions) {
         StringBuilder sb = new StringBuilder();
         sb.append("tickTime=").append(config.tickTime()).append("\n");
         sb.append("initLimit=").append(config.initLimit()).append("\n");
@@ -80,12 +107,16 @@ public class VespaZooKeeperServerImpl extends AbstractComponent implements Runna
         sb.append("serverCnxnFactory=org.apache.zookeeper.server.NettyServerCnxnFactory").append("\n");
         ensureThisServerIsRepresented(config.myid(), config.server());
         config.server().forEach(server -> addServerToCfg(sb, server));
-        sb.append(createTlsQuorumConfig(getEnvironmentVariable("VESPA_TLS_FOR_ZOOKEEPER_QUORUM_COMMUNICATION")
-                                                .orElse(config.tlsForQuorumCommunication().name())));
+        sb.append(createTlsQuorumConfig(config, transportSecurityOptions));
         return sb.toString();
     }
 
-    private String createTlsQuorumConfig(String tlsSetting) {
+    private String createTlsQuorumConfig(ZookeeperServerConfig config, Optional<TransportSecurityOptions> transportSecurityOptions) {
+        String tlsSetting = getEnvironmentVariable("VESPA_TLS_FOR_ZOOKEEPER_QUORUM_COMMUNICATION")
+                .orElse(config.tlsForQuorumCommunication().name());
+        if (transportSecurityOptions.isEmpty() && !tlsSetting.equals("OFF"))
+            throw new RuntimeException("Could not retrieve transport security options");
+
         StringBuilder sb = new StringBuilder();
 
         // Common config
@@ -119,6 +150,15 @@ public class VespaZooKeeperServerImpl extends AbstractComponent implements Runna
         sb.append("sslQuorum=").append(sslQuorum).append("\n");
         sb.append("portUnification=").append(portUnification).append("\n");
 
+        transportSecurityOptions.ifPresent(options -> {
+            sb.append("ssl.quorum.keyStore.location=").append(jksKeyStoreFilePath).append("\n");
+            sb.append("ssl.quorum.keyStore.type=JKS\n");
+
+            Path caCertificatesFile = options.getCaCertificatesFile().orElseThrow(() -> new RuntimeException("Could not find ca certificates file"));
+            sb.append("ssl.quorum.trustStore.location=").append(caCertificatesFile).append("\n");
+            sb.append("ssl.quorum.trustStore.type=PEM\n");
+        });
+
         return sb.toString();
     }
 
@@ -130,6 +170,25 @@ public class VespaZooKeeperServerImpl extends AbstractComponent implements Runna
         }
     }
 
+    private void writeJksKeystore(TransportSecurityOptions options) {
+        Path privateKeyFile = options.getPrivateKeyFile().orElseThrow(() -> new RuntimeException("Could not find private key file"));
+        Path certificatesFile = options.getCertificatesFile().orElseThrow(() -> new RuntimeException("Could not find certificates file"));
+
+        PrivateKey privateKey;
+        List<X509Certificate> certificates;
+        try {
+            privateKey = KeyUtils.fromPemEncodedPrivateKey(Utf8.toString(Files.readAllBytes(privateKeyFile)));
+            certificates = X509CertificateUtils.certificateListFromPem(Utf8.toString(Files.readAllBytes(certificatesFile)));
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+        KeyStoreBuilder keyStoreBuilder = KeyStoreBuilder
+                .withType(KeyStoreType.JKS)
+                .withKeyEntry("foo", privateKey, certificates);
+
+        KeyStoreUtils.writeKeyStoreToFile(keyStoreBuilder.build(), Paths.get(jksKeyStoreFilePath));
+    }
+
     private void ensureThisServerIsRepresented(int myid, List<ZookeeperServerConfig.Server> servers) {
         boolean found = false;
         for (ZookeeperServerConfig.Server server : servers) {
diff --git a/zookeeper-server/zookeeper-server-3.5/src/test/java/com/yahoo/vespa/zookeeper/VespaZooKeeperServerImplTest.java b/zookeeper-server/zookeeper-server-3.5/src/test/java/com/yahoo/vespa/zookeeper/VespaZooKeeperServerImplTest.java
index fcca602a2ed..953a85aeaaf 100644
--- a/zookeeper-server/zookeeper-server-3.5/src/test/java/com/yahoo/vespa/zookeeper/VespaZooKeeperServerImplTest.java
+++ b/zookeeper-server/zookeeper-server-3.5/src/test/java/com/yahoo/vespa/zookeeper/VespaZooKeeperServerImplTest.java
@@ -3,30 +3,59 @@ package com.yahoo.vespa.zookeeper;
 
 import com.yahoo.cloud.config.ZookeeperServerConfig;
 import com.yahoo.io.IOUtils;
+import com.yahoo.security.KeyUtils;
+import com.yahoo.security.X509CertificateBuilder;
+import com.yahoo.security.X509CertificateUtils;
+import com.yahoo.security.tls.TransportSecurityOptions;
+import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
 
+import javax.security.auth.x500.X500Principal;
 import java.io.File;
 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 java.util.Optional;
 
+import static com.yahoo.cloud.config.ZookeeperServerConfig.TlsForQuorumCommunication.Enum.PORT_UNIFICATION;
+import static com.yahoo.cloud.config.ZookeeperServerConfig.TlsForQuorumCommunication.Enum.TLS_ONLY;
+import static com.yahoo.cloud.config.ZookeeperServerConfig.TlsForQuorumCommunication.Enum.TLS_WITH_PORT_UNIFICATION;
+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.hamcrest.CoreMatchers.is;
 import static org.junit.Assert.assertThat;
 import static com.yahoo.vespa.defaults.Defaults.getDefaults;
+import static org.junit.Assert.assertTrue;
 
 /**
  * Tests the zookeeper server.
  */
 public class VespaZooKeeperServerImplTest {
 
+    private File cfgFile;
+    private File idFile;
+    private File jksKeyStoreFile;
+
     @Rule
     public TemporaryFolder folder = new TemporaryFolder();
 
+    @Before
+    public void setup() throws IOException {
+        cfgFile = folder.newFile();
+        idFile = folder.newFile();
+        jksKeyStoreFile = folder.newFile();
+    }
+
     @Test
     public void config_is_written_correctly_when_one_server() throws IOException {
-        File cfgFile = folder.newFile();
-        File idFile = folder.newFile();
-        ZookeeperServerConfig.Builder builder = createConfigBuilderForSingleHost(cfgFile, idFile);
+        ZookeeperServerConfig.Builder builder = createConfigBuilderForSingleHost(cfgFile, idFile, jksKeyStoreFile);
         createServer(builder);
         validateConfigFileSingleHost(cfgFile);
         validateIdFile(idFile, "");
@@ -34,8 +63,6 @@ public class VespaZooKeeperServerImplTest {
 
     @Test
     public void config_is_written_correctly_when_multiple_servers() throws IOException {
-        File cfgFile = folder.newFile();
-        File idFile = folder.newFile();
         ZookeeperServerConfig.Builder builder = new ZookeeperServerConfig.Builder();
         builder.zooKeeperConfigFile(cfgFile.getAbsolutePath());
         builder.server(newServer(0, "foo", 123, 321));
@@ -50,45 +77,50 @@ public class VespaZooKeeperServerImplTest {
 
     @Test
     public void config_is_written_correctly_with_tls_for_quorum_communication_port_unification() throws IOException {
-        File cfgFile = folder.newFile();
-        File idFile = folder.newFile();
-        ZookeeperServerConfig.Builder builder = createConfigBuilderForSingleHost(cfgFile, idFile);
-        builder.tlsForQuorumCommunication(ZookeeperServerConfig.TlsForQuorumCommunication.Enum.PORT_UNIFICATION);
-        createServer(builder);
-        validateConfigFilePortUnification(cfgFile);
+        ZookeeperServerConfig.Builder builder = createConfigBuilderForSingleHost(cfgFile, idFile, jksKeyStoreFile);
+        builder.tlsForQuorumCommunication(PORT_UNIFICATION);
+        Optional<TransportSecurityOptions> transportSecurityOptions = createTransportSecurityOptions();
+        createServer(builder, transportSecurityOptions);
+        validateConfigFilePortUnification(cfgFile, jksKeyStoreFile, transportSecurityOptions.get().getCaCertificatesFile().get().toFile());
+        validateThatJksKeyStoreFileExists(jksKeyStoreFile);
     }
 
     @Test
     public void config_is_written_correctly_with_tls_for_quorum_communication_tls_with_port_unification() throws IOException {
-        File cfgFile = folder.newFile();
-        File idFile = folder.newFile();
-        ZookeeperServerConfig.Builder builder = createConfigBuilderForSingleHost(cfgFile, idFile);
-        builder.tlsForQuorumCommunication(ZookeeperServerConfig.TlsForQuorumCommunication.Enum.TLS_WITH_PORT_UNIFICATION);
-        createServer(builder);
-        validateConfigFileTlsWithPortUnification(cfgFile);
+        ZookeeperServerConfig.Builder builder = createConfigBuilderForSingleHost(cfgFile, idFile, jksKeyStoreFile);
+        builder.tlsForQuorumCommunication(TLS_WITH_PORT_UNIFICATION);
+        Optional<TransportSecurityOptions> transportSecurityOptions = createTransportSecurityOptions();
+        createServer(builder, transportSecurityOptions);
+        validateConfigFileTlsWithPortUnification(cfgFile, jksKeyStoreFile, transportSecurityOptions.get().getCaCertificatesFile().get().toFile());
+        validateThatJksKeyStoreFileExists(jksKeyStoreFile);
     }
 
     @Test
     public void config_is_written_correctly_with_tls_for_quorum_communication_tls_only() throws IOException {
-        File cfgFile = folder.newFile();
-        File idFile = folder.newFile();
-        ZookeeperServerConfig.Builder builder = createConfigBuilderForSingleHost(cfgFile, idFile);
-        builder.tlsForQuorumCommunication(ZookeeperServerConfig.TlsForQuorumCommunication.Enum.TLS_ONLY);
-        createServer(builder);
-        validateConfigFileTlsOnly(cfgFile);
+        ZookeeperServerConfig.Builder builder = createConfigBuilderForSingleHost(cfgFile, idFile, jksKeyStoreFile);
+        builder.tlsForQuorumCommunication(TLS_ONLY);
+        Optional<TransportSecurityOptions> transportSecurityOptions = createTransportSecurityOptions();
+        createServer(builder, transportSecurityOptions);
+        validateConfigFileTlsOnly(cfgFile, jksKeyStoreFile, transportSecurityOptions.get().getCaCertificatesFile().get().toFile());
+        validateThatJksKeyStoreFileExists(jksKeyStoreFile);
     }
 
-    private ZookeeperServerConfig.Builder createConfigBuilderForSingleHost(File cfgFile, File idFile) {
+    private ZookeeperServerConfig.Builder createConfigBuilderForSingleHost(File cfgFile, File idFile, File jksKeyStoreFile) {
         ZookeeperServerConfig.Builder builder = new ZookeeperServerConfig.Builder();
         builder.zooKeeperConfigFile(cfgFile.getAbsolutePath());
         builder.myidFile(idFile.getAbsolutePath());
         builder.server(newServer(0, "foo", 123, 321));
         builder.myid(0);
+        builder.jksKeyStoreFile(jksKeyStoreFile.getAbsolutePath());
         return builder;
     }
 
     private void createServer(ZookeeperServerConfig.Builder builder) {
-        new VespaZooKeeperServerImpl(new ZookeeperServerConfig(builder), false);
+        createServer(builder, Optional.empty());
+    }
+
+    private void createServer(ZookeeperServerConfig.Builder builder, Optional<TransportSecurityOptions> options) {
+        new VespaZooKeeperServerImpl(new ZookeeperServerConfig(builder), false, options);
     }
 
     @Test(expected = RuntimeException.class)
@@ -135,18 +167,28 @@ public class VespaZooKeeperServerImplTest {
     }
 
     private String commonConfig() {
-       return "tickTime=2000\n" +
-              "initLimit=20\n" +
-              "syncLimit=15\n" +
-              "maxClientCnxns=0\n" +
-              "snapCount=50000\n" +
-              "dataDir=" + getDefaults().underVespaHome("var/zookeeper") + "\n" +
-              "clientPort=2181\n" +
-              "autopurge.purgeInterval=1\n" +
-              "autopurge.snapRetainCount=15\n" +
-              "4lw.commands.whitelist=conf,cons,crst,dirs,dump,envi,mntr,ruok,srst,srvr,stat,wchs\n" +
-              "admin.enableServer=false\n" +
-              "serverCnxnFactory=org.apache.zookeeper.server.NettyServerCnxnFactory\n";
+        return "tickTime=2000\n" +
+               "initLimit=20\n" +
+               "syncLimit=15\n" +
+               "maxClientCnxns=0\n" +
+               "snapCount=50000\n" +
+               "dataDir=" + getDefaults().underVespaHome("var/zookeeper") + "\n" +
+               "clientPort=2181\n" +
+               "autopurge.purgeInterval=1\n" +
+               "autopurge.snapRetainCount=15\n" +
+               "4lw.commands.whitelist=conf,cons,crst,dirs,dump,envi,mntr,ruok,srst,srvr,stat,wchs\n" +
+               "admin.enableServer=false\n" +
+               "serverCnxnFactory=org.apache.zookeeper.server.NettyServerCnxnFactory\n";
+    }
+
+    private String keyStoreAndTrustStoreConfig(File jksKeyStoreFilePath, File caCertificatesFilePath) {
+        StringBuilder sb = new StringBuilder();
+
+        sb.append("ssl.quorum.keyStore.location=").append(jksKeyStoreFilePath.getAbsolutePath()).append("\n");
+        sb.append("ssl.quorum.keyStore.type=JKS\n");
+        sb.append("ssl.quorum.trustStore.location=").append(caCertificatesFilePath.getAbsolutePath()).append("\n");
+        sb.append("ssl.quorum.trustStore.type=PEM\n");
+        return sb.toString();
     }
 
     private void validateConfigFileSingleHost(File cfgFile) throws IOException {
@@ -179,33 +221,36 @@ public class VespaZooKeeperServerImplTest {
         validateConfigFile(cfgFile, expected);
     }
 
-    private void validateConfigFilePortUnification(File cfgFile) throws IOException {
+    private void validateConfigFilePortUnification(File cfgFile, File jksKeyStoreFile, File caCertificatesFile) throws IOException {
         String expected =
                 commonConfig() +
                 "server.0=foo:321:123\n" +
                 commonTlsConfig() +
                 "sslQuorum=false\n" +
-                "portUnification=true\n";
+                "portUnification=true\n" +
+                keyStoreAndTrustStoreConfig(jksKeyStoreFile, caCertificatesFile);
         validateConfigFile(cfgFile, expected);
     }
 
-    private void validateConfigFileTlsWithPortUnification(File cfgFile) throws IOException {
+    private void validateConfigFileTlsWithPortUnification(File cfgFile, File jksKeyStoreFile, File caCertificatesFile) throws IOException {
         String expected =
                 commonConfig() +
                 "server.0=foo:321:123\n" +
                 commonTlsConfig() +
                 "sslQuorum=true\n" +
-                "portUnification=true\n";
+                "portUnification=true\n" +
+                keyStoreAndTrustStoreConfig(jksKeyStoreFile, caCertificatesFile);
         validateConfigFile(cfgFile, expected);
     }
 
-    private void validateConfigFileTlsOnly(File cfgFile) throws IOException {
+    private void validateConfigFileTlsOnly(File cfgFile, File jksKeyStoreFile, File caCertificatesFile) throws IOException {
         String expected =
                 commonConfig() +
                 "server.0=foo:321:123\n" +
                 commonTlsConfig() +
                 "sslQuorum=true\n" +
-                "portUnification=false\n";
+                "portUnification=false\n" +
+                keyStoreAndTrustStoreConfig(jksKeyStoreFile, caCertificatesFile);
         validateConfigFile(cfgFile, expected);
     }
 
@@ -213,4 +258,30 @@ public class VespaZooKeeperServerImplTest {
         String actual = IOUtils.readFile(cfgFile);
         assertThat(actual, is(expected));
     }
+
+    private void validateThatJksKeyStoreFileExists(File cfgFile) {
+        assertTrue(cfgFile.exists() && cfgFile.canRead());
+    }
+
+    private Optional<TransportSecurityOptions> createTransportSecurityOptions() throws IOException {
+        KeyPair keyPair = KeyUtils.generateKeypair(EC);
+        Path privateKeyFile = folder.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 = folder.newFile().toPath();
+        String certificatePem = X509CertificateUtils.toPem(certificate);
+        Files.writeString(certificateChainFile, certificatePem);
+
+        Path caCertificatesFile = folder.newFile().toPath();
+        Files.writeString(caCertificatesFile, certificatePem);
+
+        return Optional.of(new TransportSecurityOptions.Builder()
+                                   .withCertificates(certificateChainFile, privateKeyFile)
+                                   .withCaCertificates(caCertificatesFile)
+                                   .build());
+    }
+
 }