// Copyright 2017 The Nomulus Authors. All Rights Reserved.
//
// Licensed 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 google.registry.networking.handler;

import static com.google.common.truth.Truth.assertThat;
import static google.registry.networking.handler.SslInitializerTestUtils.getKeyPair;
import static google.registry.networking.handler.SslInitializerTestUtils.setUpSslChannel;
import static google.registry.networking.handler.SslInitializerTestUtils.signKeyPair;
import static google.registry.networking.handler.SslInitializerTestUtils.verifySslException;
import static google.registry.networking.handler.SslServerInitializer.CLIENT_CERTIFICATE_PROMISE_KEY;
import static google.registry.util.DateTimeUtils.minusDays;
import static google.registry.util.DateTimeUtils.plusDays;

import com.google.common.base.Suppliers;
import com.google.common.collect.ImmutableList;
import google.registry.testing.FakeClock;
import google.registry.util.SelfSignedCaCertificate;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelInitializer;
import io.netty.channel.ChannelPipeline;
import io.netty.channel.embedded.EmbeddedChannel;
import io.netty.channel.local.LocalAddress;
import io.netty.channel.local.LocalChannel;
import io.netty.handler.ssl.OpenSsl;
import io.netty.handler.ssl.SslContextBuilder;
import io.netty.handler.ssl.SslHandler;
import io.netty.handler.ssl.SslProvider;
import java.nio.channels.ClosedChannelException;
import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.cert.CertificateException;
import java.security.cert.CertificateExpiredException;
import java.security.cert.CertificateNotYetValidException;
import java.security.cert.X509Certificate;
import java.time.Instant;
import java.util.Collections;
import java.util.List;
import java.util.stream.Stream;
import javax.net.ssl.SSLEngine;
import javax.net.ssl.SSLException;
import javax.net.ssl.SSLHandshakeException;
import javax.net.ssl.SSLParameters;
import javax.net.ssl.SSLSession;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;

/**
 * Unit tests for {@link SslServerInitializer}.
 *
 * <p>To validate that the handler accepts & rejects connections as expected, a test server and a
 * test client are spun up, and both connect to the {@link LocalAddress} within the JVM. This avoids
 * the overhead of routing traffic through the network layer, even if it were to go through
 * loopback. It also alleviates the need to pick a free port to use.
 *
 * <p>The local addresses used in each test method must be different, otherwise tests run in
 * parallel may interfere with each other.
 */
class SslServerInitializerTest {

  /** Fake host to test if the SSL engine gets the correct peer host. */
  private static final String SSL_HOST = "www.example.tld";

  /** Fake port to test if the SSL engine gets the correct peer port. */
  private static final int SSL_PORT = 12345;

  @RegisterExtension NettyExtension nettyExtension = new NettyExtension();

  @SuppressWarnings("unused")
  static Stream<Arguments> provideTestCombinations() {
    Stream.Builder<Arguments> args = Stream.builder();
    // We do our best effort to test all available SSL providers.
    args.add(Arguments.of(SslProvider.JDK));
    if (OpenSsl.isAvailable()) {
      args.add(Arguments.of(SslProvider.OPENSSL));
    }
    return args.build();
  }

  private ChannelHandler getServerHandler(
      boolean requireClientCert,
      boolean validateClientCert,
      SslProvider sslProvider,
      PrivateKey privateKey,
      X509Certificate... certificates) {
    return new SslServerInitializer<LocalChannel>(
        requireClientCert,
        validateClientCert,
        sslProvider,
        Suppliers.ofInstance(privateKey),
        Suppliers.ofInstance(ImmutableList.copyOf(certificates)));
  }

  private ChannelHandler getClientHandler(
      SslProvider sslProvider,
      X509Certificate trustedCertificate,
      PrivateKey privateKey,
      X509Certificate certificate,
      String protocol,
      List<String> cipher) {
    return new ChannelInitializer<LocalChannel>() {
      @Override
      protected void initChannel(LocalChannel ch) throws Exception {
        SslContextBuilder sslContextBuilder =
            SslContextBuilder.forClient()
                .trustManager(trustedCertificate)
                .sslProvider(sslProvider)
                .ciphers(cipher);
        if (protocol != null) {
          sslContextBuilder.protocols(protocol);
        }
        if (privateKey != null && certificate != null) {
          sslContextBuilder.keyManager(privateKey, certificate);
        }
        SslHandler sslHandler =
            sslContextBuilder.build().newHandler(ch.alloc(), SSL_HOST, SSL_PORT);

        // Enable hostname verification.
        SSLEngine sslEngine = sslHandler.engine();
        SSLParameters sslParameters = sslEngine.getSSLParameters();
        sslParameters.setEndpointIdentificationAlgorithm("HTTPS");
        sslEngine.setSSLParameters(sslParameters);

        ch.pipeline().addLast(sslHandler);
      }
    };
  }

  private ChannelHandler getClientHandler(
      SslProvider sslProvider,
      X509Certificate trustedCertificate,
      PrivateKey privateKey,
      X509Certificate certificate) {
    return getClientHandler(sslProvider, trustedCertificate, privateKey, certificate, null, null);
  }

  @ParameterizedTest
  @MethodSource("provideTestCombinations")
  void testSuccess_swappedInitializerWithSslHandler(SslProvider sslProvider) throws Exception {
    SelfSignedCaCertificate ssc = SelfSignedCaCertificate.create(SSL_HOST, new FakeClock());
    SslServerInitializer<EmbeddedChannel> sslServerInitializer =
        new SslServerInitializer<>(
            true,
            false,
            sslProvider,
            Suppliers.ofInstance(ssc.key()),
            Suppliers.ofInstance(ImmutableList.of(ssc.cert())));
    EmbeddedChannel channel = new EmbeddedChannel();
    ChannelPipeline pipeline = channel.pipeline();
    pipeline.addLast(sslServerInitializer);
    ChannelHandler firstHandler = pipeline.first();
    assertThat(firstHandler.getClass()).isEqualTo(SslHandler.class);
    SslHandler sslHandler = (SslHandler) firstHandler;
    assertThat(sslHandler.engine().getNeedClientAuth()).isTrue();
    assertThat(channel.isActive()).isTrue();
  }

  @ParameterizedTest
  @MethodSource("provideTestCombinations")
  void testSuccess_trustAnyClientCert(SslProvider sslProvider) throws Exception {
    SelfSignedCaCertificate serverSsc = SelfSignedCaCertificate.create(SSL_HOST, new FakeClock());
    LocalAddress localAddress = new LocalAddress("TRUST_ANY_CLIENT_CERT_" + sslProvider);

    nettyExtension.setUpServer(
        localAddress,
        getServerHandler(true, false, sslProvider, serverSsc.key(), serverSsc.cert()));
    SelfSignedCaCertificate clientSsc = SelfSignedCaCertificate.create(new FakeClock());
    nettyExtension.setUpClient(
        localAddress,
        getClientHandler(sslProvider, serverSsc.cert(), clientSsc.key(), clientSsc.cert()));

    SSLSession sslSession = setUpSslChannel(nettyExtension.getClientChannel(), serverSsc.cert());
    nettyExtension.assertThatMessagesWork();

    // Verify that the SSL session gets the client cert. Note that this SslSession is for the client
    // channel, therefore its local certificates are the remote certificates of the SslSession for
    // the server channel, and vice versa.
    assertThat(sslSession.getLocalCertificates()).asList().containsExactly(clientSsc.cert());
    assertThat(sslSession.getPeerCertificates()).asList().containsExactly(serverSsc.cert());
  }

  @ParameterizedTest
  @MethodSource("provideTestCombinations")
  void testFailure_cipherNotAccepted(SslProvider sslProvider) throws Exception {
    SelfSignedCaCertificate serverSsc = SelfSignedCaCertificate.create(SSL_HOST, new FakeClock());
    LocalAddress localAddress = new LocalAddress("CIPHER_NOT_ACCEPTED_" + sslProvider);

    nettyExtension.setUpServer(
        localAddress, getServerHandler(true, true, sslProvider, serverSsc.key(), serverSsc.cert()));
    SelfSignedCaCertificate clientSsc =
        SelfSignedCaCertificate.create(
            "CLIENT", minusDays(Instant.now(), 2), plusDays(Instant.now(), 1));
    nettyExtension.setUpClient(
        localAddress,
        getClientHandler(
            sslProvider,
            serverSsc.cert(),
            clientSsc.key(),
            clientSsc.cert(),
            "TLSv1.2",
            Collections.singletonList("TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA")));

    verifySslException(
        nettyExtension.getServerChannel(),
        channel -> channel.attr(CLIENT_CERTIFICATE_PROMISE_KEY).get().get(),
        SSLHandshakeException.class);
  }

  @ParameterizedTest
  @MethodSource("provideTestCombinations")
  void testSuccess_someCiphersNotAccepted(SslProvider sslProvider) throws Exception {
    SelfSignedCaCertificate serverSsc = SelfSignedCaCertificate.create(SSL_HOST, new FakeClock());
    LocalAddress localAddress = new LocalAddress("SOME_CIPHERS_NOT_ACCEPTED_" + sslProvider);

    nettyExtension.setUpServer(
        localAddress,
        new SslServerInitializer<LocalChannel>(
            true,
            true,
            sslProvider,
            Suppliers.ofInstance(serverSsc.key()),
            Suppliers.ofInstance(ImmutableList.of(serverSsc.cert()))));
    SelfSignedCaCertificate clientSsc =
        SelfSignedCaCertificate.create(
            "CLIENT", minusDays(Instant.now(), 2), plusDays(Instant.now(), 1));
    nettyExtension.setUpClient(
        localAddress,
        getClientHandler(
            sslProvider,
            serverSsc.cert(),
            clientSsc.key(),
            clientSsc.cert(),
            "TLSv1.2",
            ImmutableList.of(
                "TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA",
                "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256", // Only accepted cipher
                "TLS_RSA_WITH_AES_256_CBC_SHA")));

    SSLSession sslSession = setUpSslChannel(nettyExtension.getClientChannel(), serverSsc.cert());
    nettyExtension.assertThatMessagesWork();

    assertThat(sslSession.getLocalCertificates()).asList().containsExactly(clientSsc.cert());
    assertThat(sslSession.getPeerCertificates()).asList().containsExactly(serverSsc.cert());
    assertThat(sslSession.getCipherSuite()).isEqualTo("TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256");
  }

  @ParameterizedTest
  @MethodSource("provideTestCombinations")
  void testFailure_protocolNotAccepted(SslProvider sslProvider) throws Exception {
    SelfSignedCaCertificate serverSsc = SelfSignedCaCertificate.create(SSL_HOST, new FakeClock());
    LocalAddress localAddress = new LocalAddress("PROTOCOL_NOT_ACCEPTED_" + sslProvider);

    nettyExtension.setUpServer(
        localAddress, getServerHandler(true, true, sslProvider, serverSsc.key(), serverSsc.cert()));
    SelfSignedCaCertificate clientSsc =
        SelfSignedCaCertificate.create(
            "CLIENT", minusDays(Instant.now(), 2), plusDays(Instant.now(), 1));
    nettyExtension.setUpClient(
        localAddress,
        getClientHandler(
            sslProvider, serverSsc.cert(), clientSsc.key(), clientSsc.cert(), "TLSv1.1", null));

    // In JDK v11.0.11 and above, TLS 1.1 is not supported anymore, in which case attempting to
    // connect with TLS 1.1 results in a ClosedChannelException instead of a SSLHandShakeException.
    // See https://www.oracle.com/java/technologies/javase/11-0-11-relnotes.html#JDK-8202343
    Class<? extends Exception> rootCause =
        sslProvider == SslProvider.JDK
            ? ClosedChannelException.class
            : SSLHandshakeException.class;

    verifySslException(
        nettyExtension.getServerChannel(),
        channel -> channel.attr(CLIENT_CERTIFICATE_PROMISE_KEY).get().get(),
        rootCause);
  }

  @ParameterizedTest
  @MethodSource("provideTestCombinations")
  void testFailure_clientCertExpired(SslProvider sslProvider) throws Exception {
    SelfSignedCaCertificate serverSsc = SelfSignedCaCertificate.create(SSL_HOST, new FakeClock());
    LocalAddress localAddress = new LocalAddress("CLIENT_CERT_EXPIRED_" + sslProvider);

    nettyExtension.setUpServer(
        localAddress, getServerHandler(true, true, sslProvider, serverSsc.key(), serverSsc.cert()));
    SelfSignedCaCertificate clientSsc =
        SelfSignedCaCertificate.create(
            "CLIENT", minusDays(Instant.now(), 2), minusDays(Instant.now(), 1));
    nettyExtension.setUpClient(
        localAddress,
        getClientHandler(sslProvider, serverSsc.cert(), clientSsc.key(), clientSsc.cert()));

    verifySslException(
        nettyExtension.getServerChannel(),
        channel -> channel.attr(CLIENT_CERTIFICATE_PROMISE_KEY).get().get(),
        CertificateExpiredException.class);
  }

  @ParameterizedTest
  @MethodSource("provideTestCombinations")
  void testFailure_clientCertNotYetValid(SslProvider sslProvider) throws Exception {
    SelfSignedCaCertificate serverSsc = SelfSignedCaCertificate.create(SSL_HOST, new FakeClock());
    LocalAddress localAddress = new LocalAddress("CLIENT_CERT_EXPIRED_" + sslProvider);

    nettyExtension.setUpServer(
        localAddress, getServerHandler(true, true, sslProvider, serverSsc.key(), serverSsc.cert()));
    SelfSignedCaCertificate clientSsc =
        SelfSignedCaCertificate.create(
            "CLIENT", plusDays(Instant.now(), 1), plusDays(Instant.now(), 2));
    nettyExtension.setUpClient(
        localAddress,
        getClientHandler(sslProvider, serverSsc.cert(), clientSsc.key(), clientSsc.cert()));

    verifySslException(
        nettyExtension.getServerChannel(),
        channel -> channel.attr(CLIENT_CERTIFICATE_PROMISE_KEY).get().get(),
        CertificateNotYetValidException.class);
  }

  @ParameterizedTest
  @MethodSource("provideTestCombinations")
  void testSuccess_doesNotRequireClientCert(SslProvider sslProvider) throws Exception {
    SelfSignedCaCertificate serverSsc = SelfSignedCaCertificate.create(SSL_HOST, new FakeClock());
    LocalAddress localAddress = new LocalAddress("DOES_NOT_REQUIRE_CLIENT_CERT_" + sslProvider);

    nettyExtension.setUpServer(
        localAddress,
        getServerHandler(false, false, sslProvider, serverSsc.key(), serverSsc.cert()));
    nettyExtension.setUpClient(
        localAddress, getClientHandler(sslProvider, serverSsc.cert(), null, null));

    SSLSession sslSession = setUpSslChannel(nettyExtension.getClientChannel(), serverSsc.cert());
    nettyExtension.assertThatMessagesWork();

    // Verify that the SSL session does not contain any client cert. Note that this SslSession is
    // for the client channel, therefore its local certificates are the remote certificates of the
    // SslSession for the server channel, and vice versa.
    assertThat(sslSession.getLocalCertificates()).isNull();
    assertThat(sslSession.getPeerCertificates()).asList().containsExactly(serverSsc.cert());
  }

  @ParameterizedTest
  @MethodSource("provideTestCombinations")
  void testSuccess_CertSignedByOtherCa(SslProvider sslProvider) throws Exception {
    // The self-signed cert of the CA.
    SelfSignedCaCertificate caSsc = SelfSignedCaCertificate.create(new FakeClock());
    KeyPair keyPair = getKeyPair();
    X509Certificate serverCert = signKeyPair(caSsc, keyPair, SSL_HOST);
    LocalAddress localAddress = new LocalAddress("CERT_SIGNED_BY_OTHER_CA_" + sslProvider);

    nettyExtension.setUpServer(
        localAddress,
        getServerHandler(
            true,
            false,
            sslProvider,
            keyPair.getPrivate(),
            // Serving both the server cert, and the CA cert
            serverCert,
            caSsc.cert()));
    SelfSignedCaCertificate clientSsc = SelfSignedCaCertificate.create(new FakeClock());
    nettyExtension.setUpClient(
        localAddress,
        getClientHandler(
            sslProvider,
            // Client trusts the CA cert
            caSsc.cert(),
            clientSsc.key(),
            clientSsc.cert()));

    SSLSession sslSession =
        setUpSslChannel(nettyExtension.getClientChannel(), serverCert, caSsc.cert());
    nettyExtension.assertThatMessagesWork();

    assertThat(sslSession.getLocalCertificates()).asList().containsExactly(clientSsc.cert());
    assertThat(sslSession.getPeerCertificates())
        .asList()
        .containsExactly(serverCert, caSsc.cert())
        .inOrder();
  }

  @ParameterizedTest
  @MethodSource("provideTestCombinations")
  void testFailure_requireClientCertificate(SslProvider sslProvider) throws Exception {
    SelfSignedCaCertificate serverSsc = SelfSignedCaCertificate.create(SSL_HOST, new FakeClock());
    LocalAddress localAddress = new LocalAddress("REQUIRE_CLIENT_CERT_" + sslProvider);

    nettyExtension.setUpServer(
        localAddress,
        getServerHandler(true, false, sslProvider, serverSsc.key(), serverSsc.cert()));
    nettyExtension.setUpClient(
        localAddress,
        getClientHandler(
            sslProvider,
            serverSsc.cert(),
            // No client cert/private key used.
            null,
            null));

    // When the server rejects the client during handshake due to lack of client certificate, both
    // should throw exceptions.
    nettyExtension.assertThatServerRootCause().isInstanceOf(SSLHandshakeException.class);
    nettyExtension.assertThatClientRootCause().isInstanceOf(SSLException.class);
    assertThat(nettyExtension.getClientChannel().isActive()).isFalse();
  }

  @ParameterizedTest
  @MethodSource("provideTestCombinations")
  void testFailure_wrongHostnameInCertificate(SslProvider sslProvider) throws Exception {
    SelfSignedCaCertificate serverSsc =
        SelfSignedCaCertificate.create("wrong.com", new FakeClock());
    LocalAddress localAddress = new LocalAddress("WRONG_HOSTNAME_" + sslProvider);

    nettyExtension.setUpServer(
        localAddress,
        getServerHandler(true, false, sslProvider, serverSsc.key(), serverSsc.cert()));
    SelfSignedCaCertificate clientSsc = SelfSignedCaCertificate.create(new FakeClock());
    nettyExtension.setUpClient(
        localAddress,
        getClientHandler(sslProvider, serverSsc.cert(), clientSsc.key(), clientSsc.cert()));

    // When the client rejects the server cert due to the wrong hostname, both the server and the
    // client throw exceptions.
    nettyExtension.assertThatClientRootCause().isInstanceOf(CertificateException.class);
    nettyExtension.assertThatClientRootCause().hasMessageThat().contains(SSL_HOST);
    nettyExtension.assertThatServerRootCause().isInstanceOf(SSLException.class);
    assertThat(nettyExtension.getClientChannel().isActive()).isFalse();
  }

  /** Returns true if v1 is larger or equals to v2. */
  private static boolean compareSemanticVersion(
      ImmutableList<Integer> v1, ImmutableList<Integer> v2) {
    for (int i : ImmutableList.of(0, 1, 2)) {
      if (v1.get(i) > v2.get(i)) {
        return true;
      }
      if (v1.get(i) < v2.get(i)) {
        return false;
      }
    }
    return true;
  }
}
