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.celeborn.common.network.ssl;
  
  import static org.apache.celeborn.common.network.ssl.SslSampleConfigs.*;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertFalse;
  import static org.junit.Assert.assertThrows;
  
  import java.io.File;
  import java.io.FileOutputStream;
  import java.io.IOException;
  import java.io.OutputStream;
  import java.nio.file.Files;
  import java.security.KeyPair;
  import java.security.KeyStore;
  import java.security.cert.X509Certificate;
  import java.util.HashMap;
  import java.util.Map;
  
  import org.junit.Test;
  
  public class ReloadingX509TrustManagerSuiteJ {
  
    /**
     * Waits until reload count hits the requested value, sleeping 100ms at a time. If the maximum
     * number of attempts is hit, throws a RuntimeException
     *
     * @param tm the trust manager to wait for
     * @param count The count to wait for
     * @param attempts The number of attempts to wait for
     */
    private void waitForReloadCount(ReloadingX509TrustManager tm, int count, int attempts)
        throws InterruptedException {
      if (tm.reloadCount > count) {
        throw new IllegalStateException(
            "Passed invalid count "
                + count
                + " to waitForReloadCount, already have "
                + tm.reloadCount);
      }
      for (int i = 0; i < attempts; i++) {
        if (tm.reloadCount >= count) {
          return;
        }
        // Adapted from SystemClock.waitTillTime
        long startTime = System.currentTimeMillis();
        long targetTime = startTime + 100;
        long currentTime = startTime;
        while (currentTime < targetTime) {
          long sleepTime = Math.min(10, targetTime - currentTime);
          Thread.sleep(sleepTime);
          currentTime = System.currentTimeMillis();
        }
      }
      throw new IllegalStateException("Trust store not reloaded after " + attempts + " attempts!");
    }
  
    /**
     * Waits until we make some number of attempts to reload, and verifies that the actual reload
     * count did not change
     *
     * @param tm the trust manager to wait for
     * @param attempts The number of attempts to wait for
     */
    private void waitForNoReload(ReloadingX509TrustManager tm, int attempts)
        throws InterruptedException {
      int oldReloadCount = tm.reloadCount;
      int checkCount = tm.needsReloadCheckCounts;
      int target = checkCount + attempts;
      while (checkCount < target) {
        Thread.sleep(100);
        checkCount = tm.needsReloadCheckCounts;
      }
      assertEquals(oldReloadCount, tm.reloadCount);
    }
  
    /**
     * Tests to ensure that loading a missing trust-store fails
     *
     * @throws Exception
     */
    @Test
    public void testLoadMissingTrustStore() throws Exception {
      File trustStore = new File("testmissing.jks");
      assertFalse(trustStore.exists());
  
      assertThrows(
          IOException.class,
          () -> new ReloadingX509TrustManager(KeyStore.getDefaultType(), trustStore, "password", 10));
    }
  
    /**
     * Tests to ensure that loading a corrupt trust-store fails
     *
     * @throws Exception
     */
    @Test
    public void testLoadCorruptTrustStore() throws Exception {
      File corruptStore = File.createTempFile("truststore-corrupt", "jks");
      corruptStore.deleteOnExit();
      OutputStream os = new FileOutputStream(corruptStore);
      os.write(1);
      os.close();
  
      assertThrows(
          IOException.class,
          () ->
              new ReloadingX509TrustManager(KeyStore.getDefaultType(), corruptStore, "password", 10));
    }
  
    /**
     * Tests that we successfully reload when a file is updated
     *
     * @throws Exception
     */
    @Test
    public void testReload() throws Exception {
      KeyPair kp = generateKeyPair("RSA");
      X509Certificate cert1 = generateCertificate("CN=Cert1", kp, 30, "SHA1withRSA");
      X509Certificate cert2 = generateCertificate("CN=Cert2", kp, 30, "SHA1withRSA");
      File trustStore = File.createTempFile("testreload", "jks");
      trustStore.deleteOnExit();
      createTrustStore(trustStore, "password", "cert1", cert1);
  
      ReloadingX509TrustManager tm = new ReloadingX509TrustManager("jks", trustStore, "password", 1);
      assertEquals(1, tm.getReloadInterval());
      assertEquals(0, tm.reloadCount);
      try {
        tm.init();
        assertEquals(1, tm.getAcceptedIssuers().length);
        // At this point we haven't reloaded, just the initial load
        assertEquals(0, tm.reloadCount);
  
        // Wait so that the file modification time is different
        Thread.sleep((tm.getReloadInterval() + 1000));
  
        // Add another cert
        Map<String, X509Certificate> certs = new HashMap<String, X509Certificate>();
        certs.put("cert1", cert1);
        certs.put("cert2", cert2);
        createTrustStore(trustStore, "password", certs);
  
        // Wait up to 10s until we reload
        waitForReloadCount(tm, 1, 100);
  
        assertEquals(2, tm.getAcceptedIssuers().length);
      } finally {
        tm.destroy();
        trustStore.delete();
      }
    }
  
    /**
     * Tests that we keep old certs if the trust store goes missing
     *
     * @throws Exception
     */
    @Test
    public void testReloadMissingTrustStore() throws Exception {
      KeyPair kp = generateKeyPair("RSA");
      X509Certificate cert1 = generateCertificate("CN=Cert1", kp, 30, "SHA1withRSA");
      File trustStore = new File("testmissing.jks");
      trustStore.deleteOnExit();
      assertFalse(trustStore.exists());
      createTrustStore(trustStore, "password", "cert1", cert1);
  
      ReloadingX509TrustManager tm = new ReloadingX509TrustManager("jks", trustStore, "password", 1);
      assertEquals(0, tm.reloadCount);
      try {
        tm.init();
        assertEquals(1, tm.getAcceptedIssuers().length);
        X509Certificate cert = tm.getAcceptedIssuers()[0];
        trustStore.delete();
  
        // Wait for up to 5s - we should *not* reload
        waitForNoReload(tm, 50);
  
        assertEquals(1, tm.getAcceptedIssuers().length);
        assertEquals(cert, tm.getAcceptedIssuers()[0]);
      } finally {
        tm.destroy();
      }
    }
  
    /**
     * Tests that we keep old certs if the new truststore is corrupt
     *
     * @throws Exception
     */
    @Test
    public void testReloadCorruptTrustStore() throws Exception {
      KeyPair kp = generateKeyPair("RSA");
      X509Certificate cert1 = generateCertificate("CN=Cert1", kp, 30, "SHA1withRSA");
      File corruptStore = File.createTempFile("truststore-corrupt", "jks");
      corruptStore.deleteOnExit();
      createTrustStore(corruptStore, "password", "cert1", cert1);
  
      ReloadingX509TrustManager tm =
          new ReloadingX509TrustManager("jks", corruptStore, "password", 1);
      assertEquals(0, tm.reloadCount);
      try {
        tm.init();
        assertEquals(1, tm.getAcceptedIssuers().length);
        X509Certificate cert = tm.getAcceptedIssuers()[0];
  
        OutputStream os = new FileOutputStream(corruptStore);
        os.write(1);
        os.close();
        corruptStore.setLastModified(System.currentTimeMillis() - 1000);
  
        // Wait for up to 5s - we should *not* reload
        waitForNoReload(tm, 50);
  
        assertEquals(1, tm.getAcceptedIssuers().length);
        assertEquals(cert, tm.getAcceptedIssuers()[0]);
      } finally {
        tm.destroy();
        corruptStore.delete();
      }
    }
  
    /**
     * Tests that we successfully reload when the trust store is a symlink and we update the contents
     * of the pointed-to file or we update the file it points to.
     *
     * @throws Exception
     */
    @Test
    public void testReloadSymlink() throws Exception {
      KeyPair kp = generateKeyPair("RSA");
      X509Certificate cert1 = generateCertificate("CN=Cert1", kp, 30, "SHA1withRSA");
      X509Certificate cert2 = generateCertificate("CN=Cert2", kp, 30, "SHA1withRSA");
      X509Certificate cert3 = generateCertificate("CN=Cert3", kp, 30, "SHA1withRSA");
  
      File trustStore1 = File.createTempFile("testreload", "jks");
      trustStore1.deleteOnExit();
      createTrustStore(trustStore1, "password", "cert1", cert1);
  
      File trustStore2 = File.createTempFile("testreload", "jks");
      Map<String, X509Certificate> certs = new HashMap<String, X509Certificate>();
      certs.put("cert1", cert1);
      certs.put("cert2", cert2);
      createTrustStore(trustStore2, "password", certs);
  
      File trustStoreSymlink = File.createTempFile("testreloadsymlink", "jks");
      trustStoreSymlink.delete();
      Files.createSymbolicLink(trustStoreSymlink.toPath(), trustStore1.toPath());
  
      ReloadingX509TrustManager tm =
          new ReloadingX509TrustManager("jks", trustStoreSymlink, "password", 1);
      assertEquals(1, tm.getReloadInterval());
      assertEquals(0, tm.reloadCount);
      try {
        tm.init();
        assertEquals(1, tm.getAcceptedIssuers().length);
        // At this point we haven't reloaded, just the initial load
        assertEquals(0, tm.reloadCount);
  
        // Repoint to trustStore2, which has another cert
        trustStoreSymlink.delete();
        Files.createSymbolicLink(trustStoreSymlink.toPath(), trustStore2.toPath());
  
        // Wait up to 10s until we reload
        waitForReloadCount(tm, 1, 100);
  
        assertEquals(2, tm.getAcceptedIssuers().length);
  
        // Add another cert
        certs.put("cert3", cert3);
        createTrustStore(trustStore2, "password", certs);
  
        // Wait up to 10s until we reload
        waitForReloadCount(tm, 2, 100);
  
        assertEquals(3, tm.getAcceptedIssuers().length);
      } finally {
        tm.destroy();
        trustStore1.delete();
        trustStore2.delete();
        trustStoreSymlink.delete();
      }
    }
  }
  