content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.bookkeeper.tls;
  
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertFalse;
  import static org.junit.Assert.assertNotNull;
  import static org.junit.Assert.assertNull;
  import static org.junit.Assert.assertTrue;
  import static org.junit.Assert.fail;
  import static org.junit.Assume.assumeTrue;
  
  import java.io.File;
  import java.io.IOException;
  import java.nio.file.Paths;
  import java.security.cert.Certificate;
  import java.security.cert.X509Certificate;
  import java.util.Arrays;
  import java.util.Collection;
  import java.util.Enumeration;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Set;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicInteger;
  import org.apache.bookkeeper.auth.AuthCallbacks;
  import org.apache.bookkeeper.auth.AuthToken;
  import org.apache.bookkeeper.auth.BookieAuthProvider;
  import org.apache.bookkeeper.auth.ClientAuthProvider;
  import org.apache.bookkeeper.client.BKException;
  import org.apache.bookkeeper.client.BookKeeper;
  import org.apache.bookkeeper.client.BookKeeper.DigestType;
  import org.apache.bookkeeper.client.BookKeeperAdmin;
  import org.apache.bookkeeper.client.BookKeeperClientStats;
  import org.apache.bookkeeper.client.BookKeeperTestClient;
  import org.apache.bookkeeper.client.LedgerEntry;
  import org.apache.bookkeeper.client.LedgerHandle;
  import org.apache.bookkeeper.client.api.LedgerMetadata;
  import org.apache.bookkeeper.conf.ClientConfiguration;
  import org.apache.bookkeeper.conf.ServerConfiguration;
  import org.apache.bookkeeper.net.BookieId;
  import org.apache.bookkeeper.proto.BookieConnectionPeer;
  import org.apache.bookkeeper.proto.BookieServer;
  import org.apache.bookkeeper.proto.ClientConnectionPeer;
  import org.apache.bookkeeper.proto.TestPerChannelBookieClient;
  import org.apache.bookkeeper.test.BookKeeperClusterTestCase;
  import org.apache.bookkeeper.test.TestStatsProvider;
  import org.apache.bookkeeper.tls.TLSContextFactory.KeyStoreType;
  import org.apache.bookkeeper.util.IOUtils;
  import org.apache.bookkeeper.util.TestUtils;
  import org.apache.commons.io.FileUtils;
  import org.junit.After;
  import org.junit.Assert;
  import org.junit.Before;
  import org.junit.Test;
  import org.junit.runner.RunWith;
  import org.junit.runners.Parameterized;
  import org.junit.runners.Parameterized.Parameters;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * Tests with TLS enabled.
   */
  @RunWith(Parameterized.class)
  public class TestTLS extends BookKeeperClusterTestCase {
  
      private static final Logger LOG = LoggerFactory.getLogger(TestPerChannelBookieClient.class);
  
      private static boolean secureClientSideChannel = false;
      private static Collection<Object> secureClientSideChannelPrincipals = null;
  
      private static boolean secureBookieSideChannel = false;
      private static Collection<Object> secureBookieSideChannelPrincipals = null;
  
      private KeyStoreType clientKeyStoreFormat;
      private KeyStoreType clientTrustStoreFormat;
      private KeyStoreType serverKeyStoreFormat;
      private KeyStoreType serverTrustStoreFormat;
      private final boolean useV2Protocol;
  
      @Parameters
      public static Collection<Object[]> data() {
          return Arrays.asList(new Object[][] {
                   { "JKS", "JKS", false },
                   { "PEM", "PEM", false },
                   { "PEM", "PEM", true },
                   { "PKCS12", "PKCS12", false },
                   { "JKS", "PEM", false },
                   { "PEM", "PKCS12", false },
                   { "PKCS12", "JKS", false }
              });
      }
      public TestTLS(String keyStoreFormat,
                     String trustStoreFormat,
                     boolean useV2Protocol) {
          super(3);
          this.clientKeyStoreFormat = KeyStoreType.valueOf(keyStoreFormat);
          this.clientTrustStoreFormat = KeyStoreType.valueOf(trustStoreFormat);
          this.serverKeyStoreFormat = KeyStoreType.valueOf(keyStoreFormat);
          this.serverTrustStoreFormat = KeyStoreType.valueOf(trustStoreFormat);
          this.useV2Protocol = useV2Protocol;
      }
  
      private String getResourcePath(String resource) throws Exception {
          return Paths.get(this.getClass().getClassLoader().getResource(resource).toURI()).toString();
      }
  
      @Before
      @Override
      public void setUp() throws Exception {
          /* client configuration */
          baseClientConf.setTLSProviderFactoryClass(TLSContextFactory.class.getName());
          baseClientConf.setTLSClientAuthentication(true);
          baseClientConf.setUseV2WireProtocol(useV2Protocol);
          baseClientConf.setLimitStatsLogging(false);
  
          switch (clientKeyStoreFormat) {
          case PEM:
              baseClientConf.setTLSKeyStoreType("PEM");
              baseClientConf.setTLSKeyStore(getResourcePath("client-key.pem"));
              baseClientConf.setTLSCertificatePath(getResourcePath("client-cert.pem"));
  
              break;
          case JKS:
              baseClientConf.setTLSKeyStoreType("JKS");
              baseClientConf.setTLSKeyStore(getResourcePath("client-key.jks"));
              baseClientConf.setTLSKeyStorePasswordPath(getResourcePath("keyStoreClientPassword.txt"));
  
              break;
          case PKCS12:
              baseClientConf.setTLSKeyStoreType("PKCS12");
              baseClientConf.setTLSKeyStore(getResourcePath("client-key.p12"));
              baseClientConf.setTLSKeyStorePasswordPath(getResourcePath("keyStoreClientPassword.txt"));
  
              break;
          default:
              throw new Exception("Invalid client keystore format" + clientKeyStoreFormat);
          }
  
          switch (clientTrustStoreFormat) {
          case PEM:
              baseClientConf.setTLSTrustStoreType("PEM");
              baseClientConf.setTLSTrustStore(getResourcePath("server-cert.pem"));
  
              break;
          case JKS:
              baseClientConf.setTLSTrustStoreType("JKS");
              baseClientConf.setTLSTrustStore(getResourcePath("server-key.jks"));
              baseClientConf.setTLSTrustStorePasswordPath(getResourcePath("keyStoreServerPassword.txt"));
  
              break;
          case PKCS12:
              baseClientConf.setTLSTrustStoreType("PKCS12");
              baseClientConf.setTLSTrustStore(getResourcePath("server-key.p12"));
              baseClientConf.setTLSTrustStorePasswordPath(getResourcePath("keyStoreServerPassword.txt"));
  
              break;
          default:
              throw new Exception("Invalid client keystore format" + clientTrustStoreFormat);
          }
  
          /* server configuration */
          baseConf.setTLSProviderFactoryClass(TLSContextFactory.class.getName());
          baseConf.setTLSClientAuthentication(true);
  
          switch (serverKeyStoreFormat) {
          case PEM:
              baseConf.setTLSKeyStoreType("PEM");
              baseConf.setTLSKeyStore(getResourcePath("server-key.pem"));
              baseConf.setTLSCertificatePath(getResourcePath("server-cert.pem"));
  
              break;
          case JKS:
              baseConf.setTLSKeyStoreType("JKS");
              baseConf.setTLSKeyStore(getResourcePath("server-key.jks"));
              baseConf.setTLSKeyStorePasswordPath(getResourcePath("keyStoreServerPassword.txt"));
  
              break;
          case PKCS12:
              baseConf.setTLSKeyStoreType("PKCS12");
              baseConf.setTLSKeyStore(getResourcePath("server-key.p12"));
              baseConf.setTLSKeyStorePasswordPath(getResourcePath("keyStoreServerPassword.txt"));
  
              break;
          default:
              throw new Exception("Invalid server keystore format" + serverKeyStoreFormat);
          }
  
          switch (serverTrustStoreFormat) {
          case PEM:
              baseConf.setTLSTrustStoreType("PEM");
              baseConf.setTLSTrustStore(getResourcePath("client-cert.pem"));
  
              break;
          case JKS:
              baseConf.setTLSTrustStoreType("JKS");
              baseConf.setTLSTrustStore(getResourcePath("client-key.jks"));
              baseConf.setTLSTrustStorePasswordPath(getResourcePath("keyStoreClientPassword.txt"));
  
              break;
  
          case PKCS12:
              baseConf.setTLSTrustStoreType("PKCS12");
              baseConf.setTLSTrustStore(getResourcePath("client-key.p12"));
              baseConf.setTLSTrustStorePasswordPath(getResourcePath("keyStoreClientPassword.txt"));
  
              break;
          default:
              throw new Exception("Invalid server keystore format" + serverTrustStoreFormat);
          }
  
          super.setUp();
      }
  
      @After
      @Override
      public void tearDown() throws Exception {
          super.tearDown();
      }
  
      /**
       * Verify the BouncyCastleProvider Name is expected.
       */
      @Test
      public void testGetBouncyCastleProviderName() throws Exception {
          String bcName = TLSContextFactory.getProvider().getName();
          Assert.assertEquals(bcName, TLSContextFactory.BC_FIPS);
      }
  
      /**
       * Verify that a server will not start if tls is enabled but no cert is specified.
       */
      @Test
      public void testStartTLSServerNoKeyStore() throws Exception {
          ServerConfiguration bookieConf = newServerConfiguration().setTLSKeyStore(null);
  
          try {
              startAndAddBookie(bookieConf);
              fail("Shouldn't have been able to start");
          } catch (SecurityException se) {
              assertTrue(true);
          }
      }
  
      /**
       * Verify handshake failure with a bad cert.
       */
      @Test
      public void testKeyMismatchFailure() throws Exception {
          // Valid test case only for PEM format keys
          assumeTrue(serverKeyStoreFormat == KeyStoreType.PEM);
  
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf);
  
          // restart a bookie with bad cert
          int restartBookieIdx = 0;
          ServerConfiguration bookieConf = confByIndex(restartBookieIdx)
                  .setTLSCertificatePath(getResourcePath("client-cert.pem"));
          killBookie(restartBookieIdx);
          LOG.info("Sleeping for 1s before restarting bookie with bad cert");
          Thread.sleep(1000);
          startAndAddBookie(bookieConf);
  
          // Create ledger and write entries
          BookKeeper client = new BookKeeper(clientConf);
          byte[] passwd = "testPassword".getBytes();
          int numEntries = 2;
          byte[] testEntry = "testEntry".getBytes();
  
          try (LedgerHandle lh = client.createLedger(numBookies, numBookies, DigestType.CRC32, passwd)) {
              for (int i = 0; i <= numEntries; i++) {
                  lh.addEntry(testEntry);
              }
              fail("Should have failed with not enough bookies to write");
          } catch (BKException.BKNotEnoughBookiesException bke) {
              // expected
          }
      }
  
      /**
       * Verify that a server will not start if ssl is enabled but the cert password is incorrect.
       */
      @Test
      public void testStartTLSServerBadPassword() throws Exception {
          ServerConfiguration bookieConf = newServerConfiguration().setTLSKeyStorePasswordPath("badpassword");
          try {
              startAndAddBookie(bookieConf);
              fail("Shouldn't have been able to start");
          } catch (SecurityException se) {
              assertTrue(true);
          }
      }
  
      private LedgerMetadata testClient(BookKeeper client, int clusterSize) throws Exception {
          byte[] passwd = "testPassword".getBytes();
          int numEntries = 100;
          long lid;
          byte[] testEntry = "testEntry".getBytes();
          try (LedgerHandle lh = client.createLedger(clusterSize, clusterSize, DigestType.CRC32, passwd)) {
              for (int i = 0; i <= numEntries; i++) {
                  lh.addEntry(testEntry);
              }
              lid = lh.getId();
          }
          try (LedgerHandle lh = client.openLedger(lid, DigestType.CRC32, passwd)) {
              Enumeration<LedgerEntry> entries = lh.readEntries(0, numEntries);
              while (entries.hasMoreElements()) {
                  LedgerEntry e = entries.nextElement();
                  assertTrue("Entry contents incorrect", Arrays.equals(e.getEntry(), testEntry));
              }
              BookKeeperAdmin admin = new BookKeeperAdmin(client, baseClientConf);
              return admin.getLedgerMetadata(lh);
          }
      }
  
      private LedgerMetadata testClient(ClientConfiguration conf, int clusterSize) throws Exception {
          try (BookKeeper client = new BookKeeper(conf)) {
              return testClient(client, clusterSize);
          }
      }
  
      /**
       * Verify the basic use of TLS. TLS client, TLS servers.
       */
      @Test
      public void testConnectToTLSClusterTLSClient() throws Exception {
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf);
          testClient(clientConf, numBookies);
      }
  
      /**
       * Verify the basic use of TLS. TLS client, TLS servers with LocalTransport.
       */
      @Test
      public void testConnectToLocalTLSClusterTLSClient() throws Exception {
          restartBookies(c -> {
                  c.setDisableServerSocketBind(true);
                  c.setEnableLocalTransport(true);
                  return c;
              });
  
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf);
          testClient(clientConf, numBookies);
      }
  
      /**
       * Verify Bookie refreshes certs at configured duration.
       */
      @Test
      public void testRefreshDurationForBookieCerts() throws Exception {
          assumeTrue(serverKeyStoreFormat == KeyStoreType.PEM);
          String originalTlsKeyFilePath = confByIndex(0).getTLSKeyStore();
          String invalidServerKey = getResourcePath("client-key.pem");
          File originalTlsCertFile = new File(originalTlsKeyFilePath);
          File newTlsKeyFile = IOUtils.createTempFileAndDeleteOnExit(originalTlsKeyFilePath, "refresh");
          // clean up temp file even if test fails
          newTlsKeyFile.deleteOnExit();
          File invalidServerKeyFile = new File(invalidServerKey);
          // copy invalid cert to new temp file
          FileUtils.copyFile(invalidServerKeyFile, newTlsKeyFile);
          long refreshDurationInSec = 1;
          restartBookies(c -> {
                  c.setTLSCertFilesRefreshDurationSeconds(1);
                  c.setTLSKeyStore(newTlsKeyFile.getAbsolutePath());
                  return c;
              });
  
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf);
          try {
              testClient(clientConf, numBookies);
              Assert.fail("Should have fail due to invalid cert");
          } catch (Exception e) {
              // Ok.
          }
  
          // Sleep so, cert file can be refreshed
          Thread.sleep(refreshDurationInSec * 1000 + 1000);
  
          // copy valid key-file at given new location
          FileUtils.copyFile(originalTlsCertFile, newTlsKeyFile);
          newTlsKeyFile.setLastModified(System.currentTimeMillis() + 1000);
          // client should be successfully able to add entries over tls
          testClient(clientConf, numBookies);
          newTlsKeyFile.delete();
      }
  
      /**
       * Verify Bookkeeper-client refreshes certs at configured duration.
       */
      @Test
      public void testRefreshDurationForBookkeeperClientCerts() throws Exception {
          assumeTrue(serverKeyStoreFormat == KeyStoreType.PEM);
  
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf);
          String originalTlsCertFilePath = baseClientConf.getTLSCertificatePath();
          String invalidClientCert = getResourcePath("server-cert.pem");
          File originalTlsCertFile = new File(originalTlsCertFilePath);
          File newTlsCertFile = IOUtils.createTempFileAndDeleteOnExit(originalTlsCertFilePath, "refresh");
          // clean up temp file even if test fails
          newTlsCertFile.deleteOnExit();
          File invalidClientCertFile = new File(invalidClientCert);
          // copy invalid cert to new temp file
          FileUtils.copyFile(invalidClientCertFile, newTlsCertFile);
          long refreshDurationInSec = 2;
          clientConf.setTLSCertFilesRefreshDurationSeconds(1);
          clientConf.setTLSCertificatePath(newTlsCertFile.getAbsolutePath());
  
          // create a bookkeeper-client
          try (BookKeeper client = new BookKeeper(clientConf)) {
              byte[] testEntry = "testEntry".getBytes();
              byte[] passwd = "testPassword".getBytes();
              int totalAddEntries = 1;
              CountDownLatch latch = new CountDownLatch(totalAddEntries);
              AtomicInteger result = new AtomicInteger(-1);
              LedgerHandle lh = client.createLedger(1, 1, DigestType.CRC32, passwd);
  
              for (int i = 0; i <= totalAddEntries; i++) {
                  lh.asyncAddEntry(testEntry, (rc, lgh, entryId, ctx) -> {
                      result.set(rc);
                      latch.countDown();
                  }, null);
              }
              latch.await(1, TimeUnit.SECONDS);
              Assert.assertNotEquals(result.get(), BKException.Code.OK);
  
              // Sleep so, cert file can be refreshed
              Thread.sleep(refreshDurationInSec * 1000 + 1000);
  
              // copy valid key-file at given new location
              FileUtils.copyFile(originalTlsCertFile, newTlsCertFile);
              newTlsCertFile.setLastModified(System.currentTimeMillis() + 1000);
              // client should be successfully able to add entries over tls
              CountDownLatch latchWithValidCert = new CountDownLatch(totalAddEntries);
              AtomicInteger validCertResult = new AtomicInteger(-1);
              lh = client.createLedger(1, 1, DigestType.CRC32, passwd);
              for (int i = 0; i <= totalAddEntries; i++) {
                  lh.asyncAddEntry(testEntry, (rc, lgh, entryId, ctx) -> {
                      validCertResult.set(rc);
                      latchWithValidCert.countDown();
                  }, null);
              }
              latchWithValidCert.await(1, TimeUnit.SECONDS);
              Assert.assertEquals(validCertResult.get(), BKException.Code.OK);
              newTlsCertFile.delete();
          }
      }
  
      /**
       * Multiple clients, some with TLS, and some without TLS.
       */
      @Test
      public void testConnectToTLSClusterMixedClient() throws Exception {
          ClientConfiguration confWithTLS = new ClientConfiguration(baseClientConf);
          testClient(confWithTLS, numBookies);
  
          ClientConfiguration confNoTLS = new ClientConfiguration(baseClientConf);
          confNoTLS.setTLSProviderFactoryClass(null);
          testClient(confNoTLS, numBookies);
      }
  
      /**
       * Verify the basic use of TLS. TLS client, TLS servers. No Mutual Authentication.
       */
      @Test
      public void testConnectToTLSClusterTLSClientWithTLSNoAuthentication() throws Exception {
          restartBookies(c -> {
                  c.setTLSClientAuthentication(false);
                  return c;
              });
  
          ClientConfiguration conf = new ClientConfiguration(baseClientConf);
          testClient(conf, numBookies);
      }
  
      /**
       * Verify the basic use of TLS. TLS client, TLS servers with mutual Auth.
       */
      @Test
      public void testConnectToTLSClusterTLSClientWithAuthentication() throws Exception {
          ClientConfiguration conf = new ClientConfiguration(baseClientConf);
          try {
              testClient(conf, numBookies);
          } catch (BKException.BKNotEnoughBookiesException nnbe) {
              fail("Client should be able to connect to bookie");
          }
      }
  
      /**
       * Verify that a client without tls enabled can connect to a cluster with TLS.
       */
      @Test
      public void testConnectToTLSClusterNonTLSClient() throws Exception {
          ClientConfiguration conf = new ClientConfiguration(baseClientConf);
          conf.setTLSProviderFactoryClass(null);
          try {
              testClient(conf, numBookies);
          } catch (BKException.BKNotEnoughBookiesException nnbe) {
              fail("non tls client should be able to connect to tls enabled bookies");
          }
      }
  
      /**
       * Verify that a client will fail to connect to a server if it has asked for TLS, but it is not available.
       */
      @Test
      public void testClientWantsTLSNoServersHaveIt() throws Exception {
          restartBookies(c -> {
                  c.setTLSProviderFactoryClass(null);
                  return c;
              });
  
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf);
          try {
              testClient(clientConf, numBookies);
              fail("Shouldn't be able to connect");
          } catch (BKException.BKNotEnoughBookiesException nnbe) {
              // correct response
          }
      }
  
      /**
       * Verify that a client will be able to connect to a bookie cluster if it has asked for TLS, and there are enough
       * bookies with TLS enabled in the cluster, although few bookies do not have TLS enabled.
       */
      @Test
      public void testTLSClientButOnlyFewTLSServers() throws Exception {
          // disable TLS on initial set of bookies
          restartBookies(c -> {
                  c.setTLSProviderFactoryClass(null);
                  return c;
              });
  
          // add two bookies which support TLS
          baseConf.setTLSProviderFactoryClass(TLSContextFactory.class.getName());
  
          Set<Integer> tlsBookiePorts = new HashSet<>();
          tlsBookiePorts.add(startNewBookie());
          tlsBookiePorts.add(startNewBookie());
  
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf);
          LedgerMetadata metadata = testClient(clientConf, 2);
          assertFalse(metadata.getAllEnsembles().isEmpty());
          Collection<? extends List<BookieId>> ensembles = metadata.getAllEnsembles().values();
          try (BookKeeper client = new BookKeeper(clientConf)) {
              for (List<BookieId> bookies : ensembles) {
                  for (BookieId bookieAddress : bookies) {
                      int port = client.getBookieAddressResolver().resolve(bookieAddress).getPort();
                      assertTrue(tlsBookiePorts.contains(port));
                  }
              }
          }
      }
  
      /**
       * Verify that a client-side Auth plugin can access server certificates.
       */
      @Test
      public void testClientAuthPlugin() throws Exception {
          secureClientSideChannel = false;
          secureClientSideChannelPrincipals = null;
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf);
  
          clientConf.setClientAuthProviderFactoryClass(AllowOnlyBookiesWithX509Certificates.class.getName());
  
          testClient(clientConf, numBookies);
          assertTrue(secureClientSideChannel);
          assertNotNull(secureClientSideChannelPrincipals);
          assertTrue(!secureClientSideChannelPrincipals.isEmpty());
          assertTrue(secureClientSideChannelPrincipals.iterator().next() instanceof Certificate);
          Certificate cert = (Certificate) secureClientSideChannelPrincipals.iterator().next();
          assertTrue(cert instanceof X509Certificate);
      }
  
      /**
       * Verify that a bookie-side Auth plugin can access server certificates.
       */
      @Test
      public void testBookieAuthPluginRequireClientTLSAuthentication() throws Exception {
          restartBookies(c -> {
                  c.setBookieAuthProviderFactoryClass(
                          AllowOnlyClientsWithX509Certificates.class.getName());
                  return c;
              });
  
          secureBookieSideChannel = false;
          secureBookieSideChannelPrincipals = null;
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf);
  
          testClient(clientConf, numBookies);
          assertTrue(secureBookieSideChannel);
          assertNotNull(secureBookieSideChannelPrincipals);
          assertTrue(!secureBookieSideChannelPrincipals.isEmpty());
          assertTrue(secureBookieSideChannelPrincipals.iterator().next() instanceof Certificate);
          Certificate cert = (Certificate) secureBookieSideChannelPrincipals.iterator().next();
          assertTrue(cert instanceof X509Certificate);
      }
  
      /**
       * Verify that a bookie-side Auth plugin can access server certificates over LocalTransport.
       */
      @Test
      public void testBookieAuthPluginRequireClientTLSAuthenticationLocal() throws Exception {
          restartBookies(c -> {
                  c.setBookieAuthProviderFactoryClass(
                          AllowOnlyClientsWithX509Certificates.class.getName());
                  c.setDisableServerSocketBind(true);
                  c.setEnableLocalTransport(true);
                  return c;
              });
  
          secureBookieSideChannel = false;
          secureBookieSideChannelPrincipals = null;
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf);
  
          testClient(clientConf, numBookies);
          assertTrue(secureBookieSideChannel);
          assertNotNull(secureBookieSideChannelPrincipals);
          assertTrue(!secureBookieSideChannelPrincipals.isEmpty());
          assertTrue(secureBookieSideChannelPrincipals.iterator().next() instanceof Certificate);
          Certificate cert = (Certificate) secureBookieSideChannelPrincipals.iterator().next();
          assertTrue(cert instanceof X509Certificate);
      }
  
      /**
       * Verify that given role in client certificate is checked when BookieAuthZFactory is set.
       * Positive test case where all given roles are present.
       * If authorization fails unexpectedly, we catch the UnauthorizedAccessException and fail.
       * Otherwise we exit the test and mark it as success
       */
      @Test
      public void testRoleBasedAuthZInCertificate() throws Exception {
          restartBookies(serverConf -> {
              serverConf.setBookieAuthProviderFactoryClass(BookieAuthZFactory.class.getCanonicalName());
              serverConf.setAuthorizedRoles("testRole,testRole1");
              return serverConf;
          });
  
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf);
  
          try {
              testClient(clientConf, numBookies);
          } catch (BKException.BKUnauthorizedAccessException bke) {
              fail("Could not verify given role.");
          }
      }
  
      /**
       * Verify that a bookie-side Auth plugin can access server certificates.
       */
      @Test
      public void testBookieAuthPluginDenyAccesstoClientWithoutTLSAuthentication() throws Exception {
          restartBookies(c -> {
                  c.setTLSClientAuthentication(false);
                  c.setBookieAuthProviderFactoryClass(
                          AllowOnlyClientsWithX509Certificates.class.getName());
                  return c;
              });
  
          secureBookieSideChannel = false;
          secureBookieSideChannelPrincipals = null;
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf);
          clientConf.setTLSClientAuthentication(false);
  
          try {
              testClient(clientConf, numBookies);
              fail("Shouldn't be able to connect");
          } catch (BKException.BKUnauthorizedAccessException authFailed) {
          } catch (BKException.BKNotEnoughBookiesException notEnoughBookiesException) {
              if (!useV2Protocol) {
                  fail("Unexpected exception occurred.");
              }
          }
  
          assertTrue(secureBookieSideChannel);
          assertNotNull(secureBookieSideChannelPrincipals);
          assertTrue(secureBookieSideChannelPrincipals.isEmpty());
      }
  
      /**
       * Verify that a bookie-side Auth plugin can access server certificates over LocalTransport.
       */
      @Test
      public void testBookieAuthPluginDenyAccessToClientWithoutTLSAuthenticationLocal() throws Exception {
          restartBookies(c -> {
                  c.setTLSClientAuthentication(false);
                  c.setBookieAuthProviderFactoryClass(
                          AllowOnlyClientsWithX509Certificates.class.getName());
                  c.setDisableServerSocketBind(true);
                  c.setEnableLocalTransport(true);
                  return c;
              });
  
          secureBookieSideChannel = false;
          secureBookieSideChannelPrincipals = null;
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf);
          clientConf.setTLSClientAuthentication(false);
  
          try {
              testClient(clientConf, numBookies);
              fail("Shouldn't be able to connect");
          } catch (BKException.BKUnauthorizedAccessException authFailed) {
          } catch (BKException.BKNotEnoughBookiesException notEnoughBookiesException) {
              if (!useV2Protocol) {
                  fail("Unexpected exception occurred.");
              }
          }
  
          assertTrue(secureBookieSideChannel);
          assertNotNull(secureBookieSideChannelPrincipals);
          assertTrue(secureBookieSideChannelPrincipals.isEmpty());
      }
  
      /**
       * Verify that a bookie-side Auth plugin can access server certificates.
       */
      @Test
      public void testBookieAuthPluginDenyAccessToClientWithoutTLS() throws Exception {
          restartBookies(c -> {
                  c.setBookieAuthProviderFactoryClass(
                          AllowOnlyClientsWithX509Certificates.class.getName());
                  return c;
              });
  
          secureBookieSideChannel = false;
          secureBookieSideChannelPrincipals = null;
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf);
          clientConf.setTLSProviderFactoryClass(null);
  
          try {
              testClient(clientConf, numBookies);
              fail("Shouldn't be able to connect");
          } catch (BKException.BKUnauthorizedAccessException authFailed) {
          } catch (BKException.BKNotEnoughBookiesException notEnoughBookiesException) {
              if (!useV2Protocol) {
                  fail("Unexpected exception occurred.");
              }
          }
  
          assertFalse(secureBookieSideChannel);
          assertNull(secureBookieSideChannelPrincipals);
      }
  
      private static class AllowOnlyBookiesWithX509Certificates implements ClientAuthProvider.Factory {
  
          @Override
          public String getPluginName() {
              return "tls";
          }
  
          @Override
          public void init(ClientConfiguration conf) {
          }
  
          @Override
          public ClientAuthProvider newProvider(ClientConnectionPeer addr,
                  final AuthCallbacks.GenericCallback<Void> completeCb) {
              return new ClientAuthProvider() {
  
                  AuthCallbacks.GenericCallback<AuthToken> completeCallback;
  
                  @Override
                  public void init(AuthCallbacks.GenericCallback<AuthToken> cb) {
                      this.completeCallback = cb;
                  }
  
                  @Override
                  public void onProtocolUpgrade() {
                      secureClientSideChannel = addr.isSecure();
                      secureClientSideChannelPrincipals = addr.getProtocolPrincipals();
                      Collection<Object> certificates = addr.getProtocolPrincipals();
                      if (addr.isSecure() && !certificates.isEmpty()) {
                          assertTrue(certificates.iterator().next() instanceof X509Certificate);
                          completeCallback.operationComplete(BKException.Code.OK, AuthToken.NULL);
                      } else {
                          completeCallback.operationComplete(BKException.Code.UnauthorizedAccessException,
                                  AuthToken.NULL);
                      }
                  }
  
                  @Override
                  public void process(AuthToken m, AuthCallbacks.GenericCallback<AuthToken> cb) {
                  }
              };
          }
      }
  
      private static class AllowOnlyClientsWithX509Certificates implements BookieAuthProvider.Factory {
  
          @Override
          public String getPluginName() {
              return "tls";
          }
  
          @Override
          public void init(ServerConfiguration conf) throws IOException {
          }
  
          @Override
          public BookieAuthProvider newProvider(BookieConnectionPeer addr,
                  final AuthCallbacks.GenericCallback<Void> completeCb) {
              return new BookieAuthProvider() {
  
                  AuthCallbacks.GenericCallback<Void> completeCallback = completeCb;
  
                  @Override
                  public void onProtocolUpgrade() {
                      secureBookieSideChannel = addr.isSecure();
                      secureBookieSideChannelPrincipals = addr.getProtocolPrincipals();
                      Collection<Object> certificates = addr.getProtocolPrincipals();
                      if (addr.isSecure() && !certificates.isEmpty()) {
                          assertTrue(certificates.iterator().next() instanceof X509Certificate);
                          completeCallback.operationComplete(BKException.Code.OK, null);
                      } else {
                          completeCallback.operationComplete(BKException.Code.UnauthorizedAccessException, null);
                      }
                  }
  
                  @Override
                  public void process(AuthToken m, AuthCallbacks.GenericCallback<AuthToken> cb) {
                  }
              };
          }
      }
  
      /**
       * Verify that a client will fail to connect to a server if it has asked for TLS, but it is not available. Verify
       * that if there are enough TLS servers to fill the ensemble, it will eventually use those rather than the non-TLS
       */
      @Test
      public void testMixedCluster() throws Exception {
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf);
          int origNumBookies = numBookies;
  
          ServerConfiguration bookieConf = newServerConfiguration();
          bookieConf.setTLSProviderFactoryClass(TLSContextFactory.class.getName());
          startAndAddBookie(bookieConf);
          testClient(clientConf, origNumBookies + 1);
      }
  
      /**
       * Verify that if the server hangs while an TLS client is trying to connect, the client can continue.
       */
      @Test
      public void testHungServer() throws Exception {
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf);
          CountDownLatch latch = new CountDownLatch(1);
          sleepBookie(getBookie(0), latch);
          try {
              testClient(clientConf, numBookies);
              fail("Shouldn't be able to connect");
          } catch (BKException.BKNotEnoughBookiesException nnbe) {
              // correct response
          }
          LOG.info("latch countdown");
          latch.countDown();
      }
  
      /**
       * Verify TLS and non-TLS channel counters.
       */
      @Test
      public void testTLSChannelCounters() throws Exception {
          ClientConfiguration tlsClientConf = new ClientConfiguration(baseClientConf)
                  .setNumChannelsPerBookie(1);
          ClientConfiguration nonTlsClientConf = new ClientConfiguration(baseClientConf)
                  .setNumChannelsPerBookie(1)
                  .setTLSProviderFactoryClass(null);
  
          TestStatsProvider tlsStatsProvider = new TestStatsProvider();
          TestStatsProvider nonTlsStatsProvider = new TestStatsProvider();
          BookKeeperTestClient tlsClient = new BookKeeperTestClient(tlsClientConf, tlsStatsProvider);
          BookKeeperTestClient nonTlsClient = new BookKeeperTestClient(nonTlsClientConf, nonTlsStatsProvider);
  
          // IO load from clients
          testClient(tlsClient, numBookies);
          testClient(nonTlsClient, numBookies);
  
          // verify stats
          for (int i = 0; i < numBookies; i++) {
              BookieServer bookie = serverByIndex(i);
              StringBuilder nameBuilder = new StringBuilder(BookKeeperClientStats.CHANNEL_SCOPE)
                      .append(".")
                      .append("bookie_")
                      .append(TestUtils.buildStatsCounterPathFromBookieID(bookie.getBookieId()))
                      .append(".");
              // check stats on TLS enabled client
              assertEquals("Mismatch TLS channel count", 1,
                      tlsClient.getTestStatsProvider().getCounter(nameBuilder
                      + BookKeeperClientStats.ACTIVE_TLS_CHANNEL_COUNTER).get().longValue());
              assertEquals("TLS handshake failure unexpected", 0,
                      tlsClient.getTestStatsProvider().getCounter(nameBuilder
                      + BookKeeperClientStats.FAILED_TLS_HANDSHAKE_COUNTER).get().longValue());
              assertEquals("Mismatch non-TLS channel count", 0,
                      tlsClient.getTestStatsProvider().getCounter(nameBuilder
                      + BookKeeperClientStats.ACTIVE_NON_TLS_CHANNEL_COUNTER).get().longValue());
              assertEquals("Connection failures unexpected", 0,
                      tlsClient.getTestStatsProvider().getCounter(nameBuilder
                      + BookKeeperClientStats.FAILED_CONNECTION_COUNTER).get().longValue());
  
              // check stats on non-TLS enabled client
              assertEquals("Mismatch TLS channel count", 0,
                      nonTlsClient.getTestStatsProvider().getCounter(nameBuilder
                      + BookKeeperClientStats.ACTIVE_TLS_CHANNEL_COUNTER).get().longValue());
              assertEquals("TLS handshake failure unexpected", 0,
                      nonTlsClient.getTestStatsProvider().getCounter(nameBuilder
                      + BookKeeperClientStats.FAILED_TLS_HANDSHAKE_COUNTER).get().longValue());
              assertEquals("Mismatch non-TLS channel count", 1,
                      nonTlsClient.getTestStatsProvider().getCounter(nameBuilder
                      + BookKeeperClientStats.ACTIVE_NON_TLS_CHANNEL_COUNTER).get().longValue());
              assertEquals("Connection failures unexpected", 0,
                      nonTlsClient.getTestStatsProvider().getCounter(nameBuilder
                      + BookKeeperClientStats.FAILED_CONNECTION_COUNTER).get().longValue());
  
              bookie.shutdown();
              assertEquals("Mismatch TLS channel count", 0,
                      tlsClient.getTestStatsProvider().getCounter(nameBuilder
                      + BookKeeperClientStats.ACTIVE_TLS_CHANNEL_COUNTER).get().longValue());
              assertEquals("Mismatch non-TLS channel count", 0,
                      nonTlsClient.getTestStatsProvider().getCounter(nameBuilder
                      + BookKeeperClientStats.ACTIVE_NON_TLS_CHANNEL_COUNTER).get().longValue());
  
          }
      }
  
      /**
       * Verify handshake failure due to missing entry in trust store.
       */
      @Test
      public void testHandshakeFailure() throws Exception {
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf)
                  .setNumChannelsPerBookie(1);
  
          // restart a bookie with wrong trust store
          int restartBookieIdx = 0;
          ServerConfiguration badBookieConf = confByIndex(restartBookieIdx);
  
          switch (serverTrustStoreFormat) {
              case PEM:
                  badBookieConf.setTLSTrustStore(getResourcePath("server-cert.pem"));
                  break;
              case JKS:
                  badBookieConf.setTLSTrustStore(getResourcePath("server-key.jks"))
                          .setTLSTrustStorePasswordPath(getResourcePath("keyStoreServerPassword.txt"));
                  break;
              case PKCS12:
                  badBookieConf.setTLSTrustStore(getResourcePath("server-key.p12"))
                          .setTLSTrustStorePasswordPath(getResourcePath("keyStoreServerPassword.txt"));
                  break;
              default:
                  throw new Exception("Unrecognized trust store format: " + serverTrustStoreFormat);
          }
  
  
          killBookie(restartBookieIdx);
          LOG.info("Sleeping for 1s before restarting bookie with bad cert");
          Thread.sleep(1000);
          BookieServer bookie = startAndAddBookie(badBookieConf).getServer();
          // Create ledger and write entries
          TestStatsProvider testStatsProvider = new TestStatsProvider();
          BookKeeperTestClient client = new BookKeeperTestClient(clientConf, testStatsProvider);
          byte[] passwd = "testPassword".getBytes();
          int numEntries = 2;
          byte[] testEntry = "testEntry".getBytes();
  
          // should fail to write entries whey WQ == AQ == 3
          try (LedgerHandle lh = client.createLedger(numBookies, numBookies, numBookies, DigestType.CRC32, passwd)) {
              for (int i = 0; i <= numEntries; i++) {
                  lh.addEntry(testEntry);
              }
              fail("Should have failed with not enough bookies to write");
          } catch (BKException.BKNotEnoughBookiesException bke) {
              // expected
          }
  
          // check failed handshake counter
          StringBuilder nameBuilder = new StringBuilder(BookKeeperClientStats.CHANNEL_SCOPE)
                  .append(".")
                  .append("bookie_")
                  .append(TestUtils.buildStatsCounterPathFromBookieID(bookie.getBookieId()))
                  .append(".");
          assertEquals("TLS handshake failure expected", 1,
                  client.getTestStatsProvider().getCounter(nameBuilder
                  + BookKeeperClientStats.FAILED_TLS_HANDSHAKE_COUNTER).get().longValue());
      }
  
      /**
       * Verify that a client fails to connect to bookie if hostname verification
       * fails.
       */
      @Test
      public void testClientAuthPluginWithHostnameVerificationEnabled() throws Exception {
          secureClientSideChannel = false;
          secureClientSideChannelPrincipals = null;
          ClientConfiguration clientConf = new ClientConfiguration(baseClientConf);
  
          clientConf.setClientAuthProviderFactoryClass(AllowOnlyBookiesWithX509Certificates.class.getName());
          clientConf.setHostnameVerificationEnabled(true);
          try {
              testClient(clientConf, numBookies);
              fail("should have failed with unauthorized exception");
          } catch (BKException.BKNotEnoughBookiesException e) {
              // Ok.
          }
      }
  }
  