/*
 * 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.ratis.netty;

import org.apache.ratis.security.SecurityTestUtils;
import org.apache.ratis.security.TlsConf;
import org.apache.ratis.thirdparty.io.netty.bootstrap.Bootstrap;
import org.apache.ratis.thirdparty.io.netty.bootstrap.ServerBootstrap;
import org.apache.ratis.thirdparty.io.netty.buffer.ByteBuf;
import org.apache.ratis.thirdparty.io.netty.buffer.Unpooled;
import org.apache.ratis.thirdparty.io.netty.channel.ChannelFuture;
import org.apache.ratis.thirdparty.io.netty.channel.ChannelHandlerContext;
import org.apache.ratis.thirdparty.io.netty.channel.ChannelInboundHandler;
import org.apache.ratis.thirdparty.io.netty.channel.ChannelInboundHandlerAdapter;
import org.apache.ratis.thirdparty.io.netty.channel.ChannelInitializer;
import org.apache.ratis.thirdparty.io.netty.channel.ChannelOption;
import org.apache.ratis.thirdparty.io.netty.channel.ChannelPipeline;
import org.apache.ratis.thirdparty.io.netty.channel.EventLoopGroup;
import org.apache.ratis.thirdparty.io.netty.channel.nio.NioEventLoopGroup;
import org.apache.ratis.thirdparty.io.netty.channel.socket.SocketChannel;
import org.apache.ratis.thirdparty.io.netty.channel.socket.nio.NioSocketChannel;
import org.apache.ratis.thirdparty.io.netty.handler.logging.LogLevel;
import org.apache.ratis.thirdparty.io.netty.handler.logging.LoggingHandler;
import org.apache.ratis.thirdparty.io.netty.handler.ssl.IdentityCipherSuiteFilter;
import org.apache.ratis.thirdparty.io.netty.handler.ssl.SslContext;
import org.apache.ratis.thirdparty.io.netty.handler.ssl.SslProvider;
import org.apache.ratis.util.JavaUtils;
import org.apache.ratis.util.NettyUtils;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Closeable;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;
import java.util.Queue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;


/**
 * Testing {@link TlsConf} and the security related utility methods in {@link NettyUtils}.
 */
public class TestTlsConfWithNetty {
  private static final Logger LOG = LoggerFactory.getLogger(TestTlsConfWithNetty.class);

  static String buffer2String(ByteBuf buf){
    try {
      return buf.toString(StandardCharsets.UTF_8);
    } finally {
      buf.release();
    }
  }

  static ByteBuf unpooledBuffer(String s) {
    final ByteBuf buf = Unpooled.buffer();
    buf.writeBytes(s.getBytes(StandardCharsets.UTF_8));
    return buf;
  }

  static int randomPort() {
    final int port = 50000 + ThreadLocalRandom.current().nextInt(10000);
    LOG.info("randomPort: {}", port);
    return port;
  }

  @Test
  public void testNoSsl() throws Exception {
    runTest(randomPort(), null, null);
  }

  @Test
  public void testSsl() throws Exception {
    final TlsConf serverTlsConfig = SecurityTestUtils.newServerTlsConfig(true);
    final TlsConf clientTlsConfig = SecurityTestUtils.newClientTlsConfig(true);
    runTest(randomPort(), serverTlsConfig, clientTlsConfig);
  }

  /** RSA-certificate cipher, broadly supported on the JDK provider for TLSv1.2. */
  private static final String VALID_CIPHER = "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256";
  private static final String UNSUPPORTED_CIPHER = "TLS_UNSUPPORTED_BOGUS_CIPHER";
  private static final List<String> TLS_1_2 = Collections.singletonList("TLSv1.2");

  static TlsConf newServerTlsConfig(SslProvider provider, List<String> protocols, List<String> cipherSuites) {
    return TlsConf.newBuilder()
        .setName("server")
        .setPrivateKey(new TlsConf.PrivateKeyConf(SecurityTestUtils.getResource("ssl/server.pem")))
        .setKeyCertificates(new TlsConf.CertificatesConf(SecurityTestUtils.getResource("ssl/server.crt")))
        .setSslProvider(provider)
        .setProtocols(protocols)
        .setCipherSuites(cipherSuites)
        .build();
  }

  static TlsConf newClientTlsConfig(SslProvider provider, List<String> protocols, List<String> cipherSuites) {
    return TlsConf.newBuilder()
        .setName("client")
        .setTrustCertificates(new TlsConf.CertificatesConf(SecurityTestUtils.getResource("ssl/ca.crt")))
        .setSslProvider(provider)
        .setProtocols(protocols)
        .setCipherSuites(cipherSuites)
        .build();
  }

  /**
   * The configured cipher suites must actually be applied to the {@link SslContext}; previously
   * they were silently dropped on the Netty path. See RATIS-2598.
   */
  @Test
  public void testConfiguredCipherSuitesAreApplied() {
    final SslContext ctx = NettyUtils.buildSslContextForServer(
        newServerTlsConfig(SslProvider.JDK, TLS_1_2, Collections.singletonList(VALID_CIPHER)));
    Assertions.assertEquals(Collections.singletonList(VALID_CIPHER), ctx.cipherSuites());
  }

  /**
   * An unsupported cipher suite must be filtered out rather than passed through verbatim; otherwise
   * building the {@link SslContext} fails and crashes the server. See RATIS-2598.
   */
  @Test
  public void testUnsupportedCipherSuiteIsFilteredOut() {
    final SslContext ctx = NettyUtils.buildSslContextForServer(
        newServerTlsConfig(SslProvider.JDK, TLS_1_2, Arrays.asList(UNSUPPORTED_CIPHER, VALID_CIPHER)));
    Assertions.assertEquals(Collections.singletonList(VALID_CIPHER), ctx.cipherSuites());
  }

  /** The configured cipher suite filter must be passed to the Netty {@link SslContext}. See RATIS-2602. */
  @Test
  public void testIdentityCipherSuiteFilterIsApplied() {
    final List<String> ciphers = Arrays.asList(UNSUPPORTED_CIPHER, VALID_CIPHER);
    final TlsConf conf = TlsConf.newBuilder()
        .setName("server")
        .setPrivateKey(new TlsConf.PrivateKeyConf(SecurityTestUtils.getResource("ssl/server.pem")))
        .setKeyCertificates(new TlsConf.CertificatesConf(SecurityTestUtils.getResource("ssl/server.crt")))
        .setSslProvider(SslProvider.JDK)
        .setProtocols(TLS_1_2)
        .setCipherSuites(ciphers)
        .setCipherSuiteFilter(IdentityCipherSuiteFilter.INSTANCE)
        .build();

    final SslContext ctx = NettyUtils.buildSslContextForServer(conf);
    Assertions.assertEquals(ciphers, ctx.cipherSuites());
  }

  /**
   * End-to-end handshake with explicit protocol and cipher configuration, including an unsupported
   * cipher that must be filtered out without breaking the handshake.
   */
  @Test
  public void testSslWithProtocolAndCipher() throws Exception {
    final List<String> ciphers = Arrays.asList(UNSUPPORTED_CIPHER, VALID_CIPHER);
    final TlsConf serverTlsConfig = newServerTlsConfig(SslProvider.JDK, TLS_1_2, ciphers);
    final TlsConf clientTlsConfig = newClientTlsConfig(SslProvider.JDK, TLS_1_2, ciphers);
    runTest(randomPort(), serverTlsConfig, clientTlsConfig);
  }

  /** An unset JSSE provider name resolves to null (no provider override). See RATIS-2598. */
  @Test
  public void testNoJsseProviderReturnsNull() {
    Assertions.assertNull(NettyUtils.getJsseProvider(TlsConf.newBuilder().build()));
  }

  /** An unknown JSSE provider name must fail fast rather than be silently ignored. */
  @Test
  public void testUnknownJsseProviderThrows() {
    final TlsConf conf = TlsConf.newBuilder().setJsseProviderName("NoSuchJsseProvider").build();
    Assertions.assertThrows(IllegalArgumentException.class, () -> NettyUtils.getJsseProvider(conf));
  }

  static void runTest(int port, TlsConf serverSslConf, TlsConf clientSslConf) throws Exception {
    final SslContext serverSslContext = serverSslConf == null? null
        : NettyUtils.buildSslContextForServer(serverSslConf);
    final SslContext clientSslContext = clientSslConf == null? null
        : NettyUtils.buildSslContextForClient(clientSslConf);

    final String message = "Hey, how are you?";
    final String[] words = message.split(" ");
    try (NettyTestServer server = new NettyTestServer(port, serverSslContext);
         NettyTestClient client = new NettyTestClient("localhost", port, clientSslContext)) {
      final List<CompletableFuture<String>> replyFutures = new ArrayList<>();
      for(String word : words) {
        final ByteBuf buf = unpooledBuffer(word + " ");
        final CompletableFuture<String> f = client.writeAndFlush(buf);
        replyFutures.add(f);
      }
      for(int i = 0; i < replyFutures.size(); i++) {
        final CompletableFuture<String> future = replyFutures.get(i);
        final String reply = future.get(3, TimeUnit.SECONDS);
        LOG.info(reply);
        Assertions.assertEquals(NettyTestServer.toReply(words[i]), reply);
      }
    }
  }

  static class NettyTestServer implements Closeable {
    private static final Logger LOG = LoggerFactory.getLogger(NettyTestServer.class);

    static final String CLASS_NAME = JavaUtils.getClassSimpleName(NettyTestServer.class);

    private final EventLoopGroup bossGroup = NettyUtils.newEventLoopGroup(
        CLASS_NAME + "-bossGroup", 3, true);
    private final EventLoopGroup workerGroup = NettyUtils.newEventLoopGroup(
        CLASS_NAME + "-workerGroup", 3, true);
    private final ChannelFuture channelFuture;

    public NettyTestServer(int port, SslContext sslContext) {
      this.channelFuture = new ServerBootstrap()
          .group(bossGroup, workerGroup)
          .channel(NettyUtils.getServerChannelClass(bossGroup))
          .handler(new LoggingHandler(getClass(), LogLevel.INFO))
          .childHandler(newChannelInitializer(sslContext))
          .bind(port)
          .syncUninterruptibly();
    }

    private ChannelInitializer<SocketChannel> newChannelInitializer(SslContext sslContext){
      return new ChannelInitializer<SocketChannel>(){
        @Override
        public void initChannel(SocketChannel ch) {
          final ChannelPipeline p = ch.pipeline();
          if (sslContext != null) {
            p.addLast("ssl", sslContext.newHandler(ch.alloc()));
          }
          p.addLast(newServerHandler());
        }
      };
    }

    private ChannelInboundHandler newServerHandler(){
      return new ChannelInboundHandlerAdapter() {
        @Override
        public void channelRead(ChannelHandlerContext ctx, Object obj) {
          if (obj instanceof ByteBuf) {
            final String s = buffer2String((ByteBuf) obj);
            LOG.info("channelRead: " + s);
            for(String word : s.split(" ")) {
              ctx.writeAndFlush(unpooledBuffer(toReply(word) + " "));
            }
          }
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable throwable) {
          LOG.error(NettyTestServer.this.getClass().getSimpleName() + ": exceptionCaught", throwable);
          ctx.close();
        }
      };
    }

    static String toReply(String request) {
      return "[" + request + "]";
    }


    @Override
    public void close() {
      channelFuture.channel().close();
      bossGroup.shutdownGracefully();
      workerGroup.shutdownGracefully();
    }
  }

  static class NettyTestClient implements Closeable {
    private static final Logger LOG = LoggerFactory.getLogger(NettyTestClient.class);

    private final EventLoopGroup workerGroup = new NioEventLoopGroup(3);
    private final ChannelFuture channelFuture;


    private final Queue<CompletableFuture<String>> queue = new LinkedList<>();

    public NettyTestClient(String host, int port, SslContext sslContext) {
      this.channelFuture = new Bootstrap()
          .group(workerGroup)
          .channel(NioSocketChannel.class)
          .handler(new LoggingHandler(getClass(), LogLevel.INFO))
          .handler(newChannelInitializer(sslContext, host, port))
          .option(ChannelOption.SO_KEEPALIVE, true)
          .option(ChannelOption.TCP_NODELAY, true)
          .connect(host, port)
          .syncUninterruptibly();
    }

    public CompletableFuture<String> writeAndFlush(ByteBuf buf) {
      final CompletableFuture<String> reply = new CompletableFuture<>();
      queue.offer(reply);
      channelFuture.channel().writeAndFlush(buf);
      return reply;
    }

    private ChannelInitializer<SocketChannel> newChannelInitializer(SslContext sslContext, String host, int port) {
      return new ChannelInitializer<SocketChannel>(){
        @Override
        public void initChannel(SocketChannel ch) {
          ChannelPipeline p = ch.pipeline();
          if (sslContext != null) {
            p.addLast("ssl", sslContext.newHandler(ch.alloc(), host, port));
          }
          p.addLast(getClientHandler());
        }
      };
    }

    private ChannelInboundHandler getClientHandler(){
      return new ChannelInboundHandlerAdapter(){
        @Override
        public void channelRead(ChannelHandlerContext ctx, Object obj) {
          final String s = buffer2String((ByteBuf) obj);
          LOG.info("received: " + s);
          for(String word : s.split(" ")) {
            queue.remove().complete(word);
          }
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
          LOG.error(NettyTestClient.this.getClass().getSimpleName() + ": exceptionCaught", cause);
          ctx.close();
        }
      };
    }

    @Override
    public void close() {
      channelFuture.channel().close();
      workerGroup.shutdownGracefully();
    }
  }
}
