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

import static org.apache.bookkeeper.auth.AuthProviderFactoryFactory.AUTHENTICATION_DISABLED_PLUGIN_NAME;

import com.google.protobuf.ByteString;
import io.netty.buffer.ByteBuf;
import io.netty.channel.Channel;
import io.netty.channel.ChannelDuplexHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
import io.netty.channel.ChannelPromise;
import java.io.IOException;
import java.net.SocketAddress;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicLong;
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.proto.BookkeeperProtocol.AuthMessage;
import org.apache.bookkeeper.util.ByteBufList;
import org.apache.bookkeeper.util.NettyChannelUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

class AuthHandler {
    static final Logger LOG = LoggerFactory.getLogger(AuthHandler.class);

    static class ServerSideHandler extends ChannelInboundHandlerAdapter {
        volatile boolean authenticated = false;
        final BookieAuthProvider.Factory authProviderFactory;
        final BookieConnectionPeer connectionPeer;
        BookieAuthProvider authProvider;

        ServerSideHandler(BookieConnectionPeer connectionPeer, BookieAuthProvider.Factory authProviderFactory) {
            this.authProviderFactory = authProviderFactory;
            this.connectionPeer = connectionPeer;
            authProvider = null;
        }

        @Override
        public void channelActive(ChannelHandlerContext ctx) throws Exception {
            authProvider = authProviderFactory.newProvider(connectionPeer, new AuthHandshakeCompleteCallback());
            super.channelActive(ctx);
        }

        @Override
        public void channelInactive(ChannelHandlerContext ctx) throws Exception {
            if (authProvider != null) {
                authProvider.close();
            }
            super.channelInactive(ctx);
        }

        @Override
        public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
            if (authProvider == null) {
                // close the channel, authProvider should only be
                // null if the other end of line is an InetSocketAddress
                // anything else is strange, and we don't want to deal
                // with it
                ctx.channel().close();
                return;
            }

            if (authenticated) {
                super.channelRead(ctx, msg);
            } else if (msg instanceof BookieProtocol.AuthRequest) { // pre-PB-client
                BookieProtocol.AuthRequest req = (BookieProtocol.AuthRequest) msg;
                assert (req.getOpCode() == BookieProtocol.AUTH);
                if (checkAuthPlugin(req.getAuthMessage(), ctx.channel())) {
                    byte[] payload = req
                        .getAuthMessage()
                        .getPayload()
                        .toByteArray();
                    authProvider.process(AuthToken.wrap(payload),
                                new AuthResponseCallbackLegacy(req, ctx.channel()));
                } else {
                    ctx.channel().close();
                }
            } else if (msg instanceof BookieProtocol.Request) {
                BookieProtocol.Request req = (BookieProtocol.Request) msg;
                if (req.getOpCode() == BookieProtocol.ADDENTRY) {
                    final BookieProtocol.AddResponse response = BookieProtocol.AddResponse.create(
                            req.getProtocolVersion(), BookieProtocol.EUA,
                            req.getLedgerId(), req.getEntryId());
                    NettyChannelUtil.writeAndFlushWithVoidPromise(ctx.channel(), response);
                } else if (req.getOpCode() == BookieProtocol.READENTRY) {
                    final BookieProtocol.ReadResponse response = new BookieProtocol.ReadResponse(
                            req.getProtocolVersion(), BookieProtocol.EUA,
                            req.getLedgerId(), req.getEntryId());
                    NettyChannelUtil.writeAndFlushWithVoidPromise(ctx.channel(), response);
                } else {
                    ctx.channel().close();
                }
            } else if (msg instanceof BookkeeperProtocol.Request) { // post-PB-client
                BookkeeperProtocol.Request req = (BookkeeperProtocol.Request) msg;
                if (req.getHeader().getOperation() == BookkeeperProtocol.OperationType.AUTH
                        && req.hasAuthRequest()
                        && checkAuthPlugin(req.getAuthRequest(), ctx.channel())) {
                    byte[] payload = req
                        .getAuthRequest()
                        .getPayload()
                        .toByteArray();
                    authProvider.process(AuthToken.wrap(payload),
                            new AuthResponseCallback(req, ctx.channel(), authProviderFactory.getPluginName()));
                } else if (req.getHeader().getOperation() == BookkeeperProtocol.OperationType.START_TLS
                        && req.hasStartTLSRequest()) {
                    super.channelRead(ctx, msg);
                } else {
                    BookkeeperProtocol.Response.Builder builder = BookkeeperProtocol.Response.newBuilder()
                        .setHeader(req.getHeader())
                        .setStatus(BookkeeperProtocol.StatusCode.EUA);

                    NettyChannelUtil.writeAndFlushWithVoidPromise(ctx.channel(), builder.build());
                }
            } else {
                // close the channel, junk coming over it
                ctx.channel().close();
            }
        }

        private boolean checkAuthPlugin(AuthMessage am, final Channel src) {
            if (!am.hasAuthPluginName() || !am.getAuthPluginName().equals(authProviderFactory.getPluginName())) {
                LOG.error("Received message from incompatible auth plugin. Local = {}, Remote = {}, Channel = {}",
                        authProviderFactory.getPluginName(), am.getAuthPluginName(), src);
                return false;
            }
            return true;
        }

        public boolean isAuthenticated() {
            return authenticated;
        }

        static class AuthResponseCallbackLegacy implements AuthCallbacks.GenericCallback<AuthToken> {
            final BookieProtocol.AuthRequest req;
            final Channel channel;

            AuthResponseCallbackLegacy(BookieProtocol.AuthRequest req, Channel channel) {
                this.req = req;
                this.channel = channel;
            }

            @Override
            public void operationComplete(int rc, AuthToken newam) {
                if (rc != BKException.Code.OK) {
                    LOG.error("Error processing auth message, closing connection");
                    channel.close();
                    return;
                }
                AuthMessage message = AuthMessage.newBuilder().setAuthPluginName(req.authMessage.getAuthPluginName())
                        .setPayload(ByteString.copyFrom(newam.getData())).build();
                final BookieProtocol.AuthResponse response =
                        new BookieProtocol.AuthResponse(req.getProtocolVersion(), message);
                NettyChannelUtil.writeAndFlushWithVoidPromise(channel, response);
            }
        }

        static class AuthResponseCallback implements AuthCallbacks.GenericCallback<AuthToken> {
            final BookkeeperProtocol.Request req;
            final Channel channel;
            final String pluginName;

            AuthResponseCallback(BookkeeperProtocol.Request req, Channel channel, String pluginName) {
                this.req = req;
                this.channel = channel;
                this.pluginName = pluginName;
            }

            @Override
            public void operationComplete(int rc, AuthToken newam) {
                BookkeeperProtocol.Response.Builder builder = BookkeeperProtocol.Response.newBuilder()
                        .setHeader(req.getHeader());

                if (rc != BKException.Code.OK) {
                    LOG.error("Error processing auth message, closing connection");

                    builder.setStatus(BookkeeperProtocol.StatusCode.EUA);
                    NettyChannelUtil.writeAndFlushWithClosePromise(
                            channel, builder.build()
                    );
                    return;
                } else {
                    AuthMessage message = AuthMessage.newBuilder().setAuthPluginName(pluginName)
                            .setPayload(ByteString.copyFrom(newam.getData())).build();
                    builder.setStatus(BookkeeperProtocol.StatusCode.EOK).setAuthResponse(message);
                    NettyChannelUtil.writeAndFlushWithVoidPromise(
                            channel, builder.build()
                    );
                }
            }
        }

        class AuthHandshakeCompleteCallback implements AuthCallbacks.GenericCallback<Void> {
            @Override
            public void operationComplete(int rc, Void v) {
                if (rc == BKException.Code.OK) {
                    authenticated = true;
                    LOG.info("Authentication success on server side");
                } else {
                    if (LOG.isDebugEnabled()) {
                        LOG.debug("Authentication failed on server side");
                    }
                }
            }
        }
    }

    static class ClientSideHandler extends ChannelDuplexHandler {
        volatile boolean authenticated = false;
        final ClientAuthProvider.Factory authProviderFactory;
        ClientAuthProvider authProvider;
        final AtomicLong transactionIdGenerator;
        final Queue<Object> waitingForAuth = new ConcurrentLinkedQueue<>();
        final ClientConnectionPeer connectionPeer;

        private final boolean isUsingV2Protocol;

        public ClientAuthProvider getAuthProvider() {
            return authProvider;
        }

        ClientSideHandler(ClientAuthProvider.Factory authProviderFactory, AtomicLong transactionIdGenerator,
                ClientConnectionPeer connectionPeer, boolean isUsingV2Protocol) {
            this.authProviderFactory = authProviderFactory;
            this.transactionIdGenerator = transactionIdGenerator;
            this.connectionPeer = connectionPeer;
            authProvider = null;
            this.isUsingV2Protocol = isUsingV2Protocol;
        }

        @Override
        public void channelActive(ChannelHandlerContext ctx) throws Exception {
            authProvider = authProviderFactory.newProvider(connectionPeer, new AuthHandshakeCompleteCallback(ctx));
            authProvider.init(new AuthRequestCallback(ctx, authProviderFactory.getPluginName()));

            super.channelActive(ctx);
        }

        @Override
        public void channelInactive(ChannelHandlerContext ctx) throws Exception {
            if (authProvider != null) {
                authProvider.close();
            }
            super.channelInactive(ctx);
        }

        @Override
        public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
            assert (authProvider != null);

            if (authenticated) {
                super.channelRead(ctx, msg);
            } else if (msg instanceof BookkeeperProtocol.Response) {
                BookkeeperProtocol.Response resp = (BookkeeperProtocol.Response) msg;
                if (null == resp.getHeader().getOperation()) {
                    LOG.info("dropping received malformed message {} from bookie {}", msg, ctx.channel());
                    // drop the message without header
                } else {
                    switch (resp.getHeader().getOperation()) {
                    case START_TLS:
                        super.channelRead(ctx, msg);
                        break;
                    case AUTH:
                        if (resp.getStatus() != BookkeeperProtocol.StatusCode.EOK) {
                            authenticationError(ctx, resp.getStatus().getNumber());
                        } else {
                            assert (resp.hasAuthResponse());
                            BookkeeperProtocol.AuthMessage am = resp.getAuthResponse();
                            if (AUTHENTICATION_DISABLED_PLUGIN_NAME.equals(am.getAuthPluginName())){
                                SocketAddress remote = ctx.channel().remoteAddress();
                                LOG.info("Authentication is not enabled."
                                    + "Considering this client {} authenticated", remote);
                                AuthHandshakeCompleteCallback cb = new AuthHandshakeCompleteCallback(ctx);
                                cb.operationComplete(BKException.Code.OK, null);
                                return;
                            }
                            byte[] payload = am.getPayload().toByteArray();
                            authProvider.process(AuthToken.wrap(payload), new AuthRequestCallback(ctx,
                                authProviderFactory.getPluginName()));
                        }
                        break;
                    default:
                        LOG.warn("dropping received message {} from bookie {}", msg, ctx.channel());
                        // else just drop the message,
                        // we're not authenticated so nothing should be coming through
                        break;
                    }
                }
            } else if (msg instanceof BookieProtocol.Response) {
                BookieProtocol.Response resp = (BookieProtocol.Response) msg;
                switch (resp.opCode) {
                case BookieProtocol.AUTH:
                    if (resp.errorCode != BookieProtocol.EOK) {
                        authenticationError(ctx, resp.errorCode);
                    } else {
                        BookkeeperProtocol.AuthMessage am = ((BookieProtocol.AuthResponse) resp).authMessage;
                        if (AUTHENTICATION_DISABLED_PLUGIN_NAME.equals(am.getAuthPluginName())) {
                            SocketAddress remote = ctx.channel().remoteAddress();
                            LOG.info("Authentication is not enabled."
                                    + "Considering this client {} authenticated", remote);
                            AuthHandshakeCompleteCallback cb = new AuthHandshakeCompleteCallback(ctx);
                            cb.operationComplete(BKException.Code.OK, null);
                            return;
                        }
                        byte[] payload = am.getPayload().toByteArray();
                        authProvider.process(AuthToken.wrap(payload), new AuthRequestCallback(ctx,
                                authProviderFactory.getPluginName()));
                    }
                    break;
                default:
                    LOG.warn("dropping received message {} from bookie {}", msg, ctx.channel());
                    // else just drop the message, we're not authenticated so nothing should be coming
                    // through
                    break;
                }
            }
        }

        @Override
        public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {
            synchronized (this) {
                if (authenticated) {
                    super.write(ctx, msg, promise);
                    super.flush(ctx);
                } else if (msg instanceof BookkeeperProtocol.Request) {
                    // let auth messages through, queue the rest
                    BookkeeperProtocol.Request req = (BookkeeperProtocol.Request) msg;
                    if (req.getHeader().getOperation()
                            == BookkeeperProtocol.OperationType.AUTH
                        || req.getHeader().getOperation() == BookkeeperProtocol.OperationType.START_TLS) {
                        super.write(ctx, msg, promise);
                        super.flush(ctx);
                    } else {
                        addMsgAndPromiseToQueue(msg, promise);
                    }
                } else if (msg instanceof BookieProtocol.Request) {
                    // let auth messages through, queue the rest
                    BookieProtocol.Request req = (BookieProtocol.Request) msg;
                    if (BookieProtocol.AUTH == req.getOpCode()) {
                        super.write(ctx, msg, promise);
                        super.flush(ctx);
                    } else {
                        addMsgAndPromiseToQueue(msg, promise);
                    }
                } else if (msg instanceof ByteBuf || msg instanceof ByteBufList) {
                    addMsgAndPromiseToQueue(msg, promise);
                } else {
                    LOG.info("[{}] dropping write of message {}", ctx.channel(), msg);
                }
            }
        }

        // Add the message and the associated promise to the queue.
        // The promise is added to the same queue as the message without an additional wrapper object so
        // that object allocations can be avoided. A similar solution is used in Netty codebase.
        private void addMsgAndPromiseToQueue(Object msg, ChannelPromise promise) {
            waitingForAuth.add(msg);
            if (promise != null && !promise.isVoid()) {
                waitingForAuth.add(promise);
            }
        }

        long newTxnId() {
            return transactionIdGenerator.incrementAndGet();
        }

        void authenticationError(ChannelHandlerContext ctx, int errorCode) {
            LOG.error("Error processing auth message, erroring connection {}", errorCode);
            ctx.fireExceptionCaught(new AuthenticationException("Auth failed with error " + errorCode));
        }

        class AuthRequestCallback implements AuthCallbacks.GenericCallback<AuthToken> {
            Channel channel;
            ChannelHandlerContext ctx;
            String pluginName;

            AuthRequestCallback(ChannelHandlerContext ctx, String pluginName) {
                this.channel = ctx.channel();
                this.ctx = ctx;
                this.pluginName = pluginName;
            }

            @Override
            public void operationComplete(int rc, AuthToken newam) {
                if (rc != BKException.Code.OK) {
                    authenticationError(ctx, rc);
                    return;
                }

                AuthMessage message = AuthMessage.newBuilder().setAuthPluginName(pluginName)
                        .setPayload(ByteString.copyFrom(newam.getData())).build();

                if (isUsingV2Protocol) {
                    final BookieProtocol.AuthRequest msg =
                            new BookieProtocol.AuthRequest(BookieProtocol.CURRENT_PROTOCOL_VERSION, message);
                    NettyChannelUtil.writeAndFlushWithVoidPromise(channel, msg);
                } else {
                    // V3 protocol
                    BookkeeperProtocol.BKPacketHeader header = BookkeeperProtocol.BKPacketHeader.newBuilder()
                            .setVersion(BookkeeperProtocol.ProtocolVersion.VERSION_THREE)
                            .setOperation(BookkeeperProtocol.OperationType.AUTH).setTxnId(newTxnId()).build();
                    BookkeeperProtocol.Request.Builder builder = BookkeeperProtocol.Request.newBuilder()
                            .setHeader(header)
                            .setAuthRequest(message);
                    NettyChannelUtil.writeAndFlushWithVoidPromise(channel, builder.build());
                }
            }
        }

        class AuthHandshakeCompleteCallback implements AuthCallbacks.GenericCallback<Void> {
            ChannelHandlerContext ctx;

            AuthHandshakeCompleteCallback(ChannelHandlerContext ctx) {
                this.ctx = ctx;
            }

            @Override
            public void operationComplete(int rc, Void v) {
                if (rc == BKException.Code.OK) {
                    synchronized (this) {
                        authenticated = true;
                        while (true) {
                            Object msg = waitingForAuth.poll();
                            if (msg == null) {
                                break;
                            }
                            ChannelPromise promise;
                            // check if the message has an associated promise as the next element in the queue
                            if (waitingForAuth.peek() instanceof ChannelPromise) {
                                promise = (ChannelPromise) waitingForAuth.poll();
                            } else {
                                promise = ctx.voidPromise();
                            }
                            ctx.writeAndFlush(msg, promise);
                        }
                    }
                } else {
                    LOG.warn("Client authentication failed");
                    authenticationError(ctx, rc);
                }
            }
        }
    }

    @SuppressWarnings("serial")
    static class AuthenticationException extends IOException {
        AuthenticationException(String reason) {
            super(reason);
        }
    }
}
