/*
 * 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.pulsar.client.impl;

import io.netty.channel.Channel;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelInitializer;
import io.netty.channel.socket.SocketChannel;
import io.netty.handler.flush.FlushConsolidationHandler;
import io.netty.handler.proxy.Socks5ProxyHandler;
import io.netty.handler.ssl.SslContext;
import io.netty.handler.ssl.SslHandler;
import io.netty.handler.ssl.SslProvider;
import java.net.InetSocketAddress;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.pulsar.client.api.AuthenticationDataProvider;
import org.apache.pulsar.client.api.PulsarClientException;
import org.apache.pulsar.client.impl.conf.ClientConfigurationData;
import org.apache.pulsar.client.util.ObjectCache;
import org.apache.pulsar.common.protocol.ByteBufPair;
import org.apache.pulsar.common.protocol.Commands;
import org.apache.pulsar.common.protocol.FrameDecoderUtil;
import org.apache.pulsar.common.util.SecurityUtility;
import org.apache.pulsar.common.util.keystoretls.NettySSLContextAutoRefreshBuilder;
import org.apache.pulsar.common.util.netty.NettyFutureUtil;

@Slf4j
public class PulsarChannelInitializer extends ChannelInitializer<SocketChannel> {

    public static final String TLS_HANDLER = "tls";

    private final Supplier<ClientCnx> clientCnxSupplier;
    @Getter
    private final boolean tlsEnabled;
    private final boolean tlsHostnameVerificationEnabled;
    private final boolean tlsEnabledWithKeyStore;
    private final InetSocketAddress socks5ProxyAddress;
    private final String socks5ProxyUsername;
    private final String socks5ProxyPassword;
    private final ClientConfigurationData conf;
    private Map<String, Supplier<SslContext>> sslContextSupplierMap;
    private Map<String, NettySSLContextAutoRefreshBuilder> nettySSLContextAutoRefreshBuilderMap;

    private static final long TLS_CERTIFICATE_CACHE_MILLIS = TimeUnit.MINUTES.toMillis(1);

    public PulsarChannelInitializer(ClientConfigurationData conf, Supplier<ClientCnx> clientCnxSupplier)
            throws Exception {
        super();
        this.clientCnxSupplier = clientCnxSupplier;
        this.tlsEnabled = conf.isUseTls();
        this.tlsHostnameVerificationEnabled = conf.isTlsHostnameVerificationEnable();
        this.socks5ProxyAddress = conf.getSocks5ProxyAddress();
        this.socks5ProxyUsername = conf.getSocks5ProxyUsername();
        this.socks5ProxyPassword = conf.getSocks5ProxyPassword();

        this.tlsEnabledWithKeyStore = conf.isUseKeyStoreTls();
        this.conf = conf.clone();
        this.sslContextSupplierMap = new ConcurrentHashMap<>();
        this.nettySSLContextAutoRefreshBuilderMap = new ConcurrentHashMap<>();
    }

    @Override
    public void initChannel(SocketChannel ch) throws Exception {
        ch.pipeline().addLast("consolidation", new FlushConsolidationHandler(1024, true));

        // Setup channel except for the SsHandler for TLS enabled connections
        ch.pipeline().addLast("ByteBufPairEncoder", ByteBufPair.getEncoder(tlsEnabled));
        FrameDecoderUtil.addFrameDecoder(ch.pipeline(), Commands.DEFAULT_MAX_MESSAGE_SIZE);
        ChannelHandler clientCnx = clientCnxSupplier.get();
        ch.pipeline().addLast("handler", clientCnx);
    }

    private NettySSLContextAutoRefreshBuilder getNettySSLContextAutoRefreshBuilder(String host)
            throws PulsarClientException {
        if (tlsEnabledWithKeyStore) {
            AuthenticationDataProvider authData1 = conf.getAuthentication().getAuthData(host);
            if (StringUtils.isBlank(conf.getTlsTrustStorePath())) {
                throw new PulsarClientException("Failed to create TLS context, the tlsTrustStorePath"
                        + " need to be configured if useKeyStoreTls enabled");
            }
            return nettySSLContextAutoRefreshBuilderMap.computeIfAbsent(host,
                    key -> new NettySSLContextAutoRefreshBuilder(
                            conf.getSslProvider(),
                            conf.isTlsAllowInsecureConnection(),
                            conf.getTlsTrustStoreType(),
                            conf.getTlsTrustStorePath(),
                            conf.getTlsTrustStorePassword(),
                            conf.getTlsKeyStoreType(),
                            conf.getTlsKeyStorePath(),
                            conf.getTlsKeyStorePassword(),
                            conf.getTlsCiphers(),
                            conf.getTlsProtocols(),
                            TLS_CERTIFICATE_CACHE_MILLIS,
                            authData1));
        }
        throw new PulsarClientException(
                "Failed to create TLS context, the tlsEnabledWithKeyStore need to be true");
    }

    private Supplier<SslContext> getSslContextSupplier(String host) {
        return sslContextSupplierMap.computeIfAbsent(host, key -> new ObjectCache<>(() -> {
            try {
                SslProvider sslProvider = null;
                if (conf.getSslProvider() != null) {
                    sslProvider = SslProvider.valueOf(conf.getSslProvider());
                }

                // Set client certificate if available
                AuthenticationDataProvider authData = conf.getAuthentication().getAuthData(host);
                if (authData.hasDataForTls()) {
                    return authData.getTlsTrustStoreStream() == null
                            ? SecurityUtility.createNettySslContextForClient(
                            sslProvider,
                            conf.isTlsAllowInsecureConnection(),
                            conf.getTlsTrustCertsFilePath(),
                            authData.getTlsCertificates(),
                            authData.getTlsPrivateKey(),
                            conf.getTlsCiphers(),
                            conf.getTlsProtocols())
                            : SecurityUtility.createNettySslContextForClient(sslProvider,
                            conf.isTlsAllowInsecureConnection(),
                            authData.getTlsTrustStoreStream(),
                            authData.getTlsCertificates(), authData.getTlsPrivateKey(),
                            conf.getTlsCiphers(),
                            conf.getTlsProtocols());
                } else {
                    return SecurityUtility.createNettySslContextForClient(
                            sslProvider,
                            conf.isTlsAllowInsecureConnection(),
                            conf.getTlsTrustCertsFilePath(),
                            conf.getTlsCertificateFilePath(),
                            conf.getTlsKeyFilePath(),
                            conf.getTlsCiphers(),
                            conf.getTlsProtocols());
                }
            } catch (Exception e) {
                throw new RuntimeException("Failed to create TLS context", e);
            }
        }, TLS_CERTIFICATE_CACHE_MILLIS, TimeUnit.MILLISECONDS));
    }

    /**
     * Initialize TLS for a channel. Should be invoked before the channel is connected to the remote address.
     *
     * @param ch      the channel
     * @param sniHost the value of this argument will be passed as peer host and port when creating the SSLEngine (which
     *                in turn will use these values to set SNI header when doing the TLS handshake). Cannot be
     *                <code>null</code>.
     * @return a {@link CompletableFuture} that completes when the TLS is set up.
     */
    CompletableFuture<Channel> initTls(Channel ch, InetSocketAddress sniHost) {
        Objects.requireNonNull(ch, "A channel is required");
        Objects.requireNonNull(sniHost, "A sniHost is required");
        if (!tlsEnabled) {
            throw new IllegalStateException("TLS is not enabled in client configuration");
        }
        CompletableFuture<Channel> initTlsFuture = new CompletableFuture<>();
        ch.eventLoop().execute(() -> {
            try {
                SslHandler handler = tlsEnabledWithKeyStore
                        ? new SslHandler(getNettySSLContextAutoRefreshBuilder(sniHost.getHostName()).get()
                        .createSSLEngine(sniHost.getHostString(), sniHost.getPort()))
                        : getSslContextSupplier(sniHost.getHostName()).get()
                        .newHandler(ch.alloc(), sniHost.getHostString(), sniHost.getPort());

                if (tlsHostnameVerificationEnabled) {
                    SecurityUtility.configureSSLHandler(handler);
                }

                ch.pipeline().addFirst(TLS_HANDLER, handler);
                initTlsFuture.complete(ch);
            } catch (Throwable t) {
                initTlsFuture.completeExceptionally(t);
            }
        });

        return initTlsFuture;
    }

    CompletableFuture<Channel> initSocks5IfConfig(Channel ch) {
        CompletableFuture<Channel> initSocks5Future = new CompletableFuture<>();
        if (socks5ProxyAddress != null) {
            ch.eventLoop().execute(() -> {
                try {
                    Socks5ProxyHandler socks5ProxyHandler =
                            new Socks5ProxyHandler(socks5ProxyAddress, socks5ProxyUsername, socks5ProxyPassword);
                    ch.pipeline().addFirst(socks5ProxyHandler.protocol(), socks5ProxyHandler);
                    initSocks5Future.complete(ch);
                } catch (Throwable t) {
                    initSocks5Future.completeExceptionally(t);
                }
            });
        } else {
            initSocks5Future.complete(ch);
        }

        return initSocks5Future;
    }

    CompletableFuture<Channel> initializeClientCnx(Channel ch,
                                                   InetSocketAddress logicalAddress,
                                                   InetSocketAddress unresolvedPhysicalAddress) {
        return NettyFutureUtil.toCompletableFuture(ch.eventLoop().submit(() -> {
            final ClientCnx cnx = (ClientCnx) ch.pipeline().get("handler");

            if (cnx == null) {
                throw new IllegalStateException("Missing ClientCnx. This should not happen.");
            }

            if (!logicalAddress.equals(unresolvedPhysicalAddress)) {
                // We are connecting through a proxy. We need to set the target broker in the ClientCnx object so that
                // it can be specified when sending the CommandConnect.
                cnx.setTargetBroker(logicalAddress);
            }

            cnx.setRemoteHostName(unresolvedPhysicalAddress.getHostString());

            return ch;
        }));
    }
}

