com.github.zk1931.jzab.NettyTransport.java Source code

Java tutorial

Introduction

Here is the source code for com.github.zk1931.jzab.NettyTransport.java

Source

/**
 * Licensed to the zk1931 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 com.github.zk1931.jzab;

import com.google.protobuf.InvalidProtocolBufferException;
import io.netty.bootstrap.Bootstrap;
import io.netty.bootstrap.ServerBootstrap;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelFutureListener;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
import io.netty.channel.ChannelInitializer;
import io.netty.channel.ChannelOption;
import io.netty.channel.DefaultFileRegion;
import io.netty.channel.EventLoopGroup;
import io.netty.channel.FileRegion;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.codec.ByteToMessageDecoder;
import io.netty.handler.codec.LengthFieldPrepender;
import io.netty.handler.ssl.SslHandler;
import io.netty.handler.stream.ChunkedFile;
import io.netty.handler.stream.ChunkedWriteHandler;
import io.netty.util.AttributeKey;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.RandomAccessFile;
import java.nio.ByteBuffer;
import java.security.GeneralSecurityException;
import java.security.KeyStore;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.concurrent.BlockingDeque;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.TimeUnit;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLEngine;
import javax.net.ssl.TrustManagerFactory;
import com.github.zk1931.jzab.proto.ZabMessage.Message;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import static com.github.zk1931.jzab.proto.ZabMessage.Message.MessageType;

/**
 * Netty-based transport.
 */
class NettyTransport extends Transport {
    private static final Logger LOG = LoggerFactory.getLogger(NettyTransport.class);
    static final AttributeKey<String> SENDER = AttributeKey.valueOf("remote");

    private final String hostPort;
    private final EventLoopGroup bossGroup = new NioEventLoopGroup();
    private final EventLoopGroup workerGroup = new NioEventLoopGroup();
    Channel channel;
    private final File keyStore;
    private final char[] keyStorePassword;
    private final File trustStore;
    private final char[] trustStorePassword;
    private SSLContext clientContext;
    private SSLContext serverContext;
    private final File dir;

    // remote id => sender map.
    ConcurrentMap<String, Sender> senders = new ConcurrentHashMap<>();

    // remote id => incoming channel context
    ConcurrentMap<String, ChannelHandlerContext> receivers = new ConcurrentHashMap<>();

    public NettyTransport(String hostPort, final Receiver receiver, final File dir)
            throws InterruptedException, GeneralSecurityException, IOException {
        this(hostPort, receiver, new ZabConfig.SslParameters(), dir);
    }

    /**
     * Constructs a NettyTransport object.
     *
     * @param hostPort "hostname:port" string. The netty transport binds to the
     *                 port specified in the string.
     * @param receiver receiver callback.
     * @param sslParam Ssl parameters.
     * @param dir the directory used to store the received file.
     */
    public NettyTransport(String hostPort, final Receiver receiver, ZabConfig.SslParameters sslParam,
            final File dir) throws InterruptedException, GeneralSecurityException, IOException {
        super(receiver);
        this.keyStore = sslParam.getKeyStore();
        this.trustStore = sslParam.getTrustStore();
        this.keyStorePassword = sslParam.getKeyStorePassword() != null
                ? sslParam.getKeyStorePassword().toCharArray()
                : null;
        this.trustStorePassword = sslParam.getTrustStorePassword() != null
                ? sslParam.getTrustStorePassword().toCharArray()
                : null;
        this.dir = dir;
        if (isSslEnabled()) {
            initSsl();
        }

        this.hostPort = hostPort;
        String[] address = hostPort.split(":", 2);
        int port = Integer.parseInt(address[1]);
        ServerBootstrap b = new ServerBootstrap();
        b.group(bossGroup, workerGroup).channel(NioServerSocketChannel.class).option(ChannelOption.SO_BACKLOG, 128)
                .option(ChannelOption.SO_REUSEADDR, true).childOption(ChannelOption.SO_KEEPALIVE, true)
                .childOption(ChannelOption.TCP_NODELAY, true).childHandler(new ChannelInitializer<SocketChannel>() {
                    @Override
                    public void initChannel(SocketChannel ch) throws Exception {
                        if (isSslEnabled()) {
                            SSLEngine engine = serverContext.createSSLEngine();
                            engine.setUseClientMode(false);
                            engine.setNeedClientAuth(true);
                            ch.pipeline().addLast(new SslHandler(engine));
                        }
                        // Incoming handlers
                        ch.pipeline().addLast(new MainHandler());
                        ch.pipeline().addLast(new ServerHandshakeHandler());
                        ch.pipeline().addLast(new NotifyHandler());
                        ch.pipeline().addLast(new ServerErrorHandler());
                    }
                });

        // Travis build fails once in a while because it fails to bind to a port.
        // This is most likely a transient failure. Retry binding for 5 times with
        // 1 second sleep in between before giving up.
        int bindRetryCount = 5;
        for (int i = 0;; i++) {
            try {
                channel = b.bind(port).sync().channel();
                LOG.info("Server started: {}", hostPort);
                return;
            } catch (Exception ex) {
                if (i >= bindRetryCount) {
                    throw ex;
                }
                LOG.debug("Failed to bind to {}. Retrying after 1 second.", hostPort);
                Thread.sleep(1000);
            }
        }
    }

    private boolean isSslEnabled() {
        return keyStore != null && trustStore != null;
    }

    private void initSsl() throws IOException, GeneralSecurityException {
        String kmAlgorithm = KeyManagerFactory.getDefaultAlgorithm();
        String tmAlgorithm = TrustManagerFactory.getDefaultAlgorithm();
        // TODO make the protocol and keystore type configurable.
        String protocol = "TLS";
        KeyStore ks = KeyStore.getInstance("JKS");
        KeyStore ts = KeyStore.getInstance("JKS");
        try (FileInputStream keyStoreStream = new FileInputStream(keyStore);
                FileInputStream trustStoreStream = new FileInputStream(trustStore)) {
            ks.load(keyStoreStream, keyStorePassword);
            ts.load(trustStoreStream, trustStorePassword);
        }
        KeyManagerFactory kmf = KeyManagerFactory.getInstance(kmAlgorithm);
        TrustManagerFactory tmf = TrustManagerFactory.getInstance(tmAlgorithm);
        kmf.init(ks, keyStorePassword);
        tmf.init(ts);
        serverContext = SSLContext.getInstance(protocol);
        clientContext = SSLContext.getInstance(protocol);
        serverContext.init(kmf.getKeyManagers(), tmf.getTrustManagers(), null);
        clientContext.init(kmf.getKeyManagers(), tmf.getTrustManagers(), null);
    }

    /**
     * Destroys the transport.
     */
    @Override
    public void shutdown() throws InterruptedException {
        try {
            channel.close();
            for (Map.Entry<String, Sender> entry : senders.entrySet()) {
                entry.getValue().shutdown();
            }
            senders.clear();

            for (Map.Entry<String, ChannelHandlerContext> entry : receivers.entrySet()) {
                entry.getValue().close();
            }
            receivers.clear();

        } finally {
            try {
                long quietPeriodSec = 0;
                long timeoutSec = 10;
                io.netty.util.concurrent.Future wf = workerGroup.shutdownGracefully(quietPeriodSec, timeoutSec,
                        TimeUnit.SECONDS);
                io.netty.util.concurrent.Future bf = bossGroup.shutdownGracefully(quietPeriodSec, timeoutSec,
                        TimeUnit.SECONDS);
                wf.await();
                bf.await();
                LOG.debug("Shutdown complete: {}", hostPort);
            } catch (InterruptedException ex) {
                LOG.debug("Interrupted while shutting down NioEventLoopGroup", ex);
            }
        }
    }

    /**
     * Handles server-side handshake.
     */
    private class ServerHandshakeHandler extends SimpleChannelInboundHandler<Message> {
        @Override
        public void channelRead0(ChannelHandlerContext ctx, Message msg) throws Exception {
            // Make sure it's a handshake message.
            if (msg.getType() != MessageType.HANDSHAKE) {
                LOG.debug("The first message from {} was not a handshake", ctx.channel().remoteAddress());
                ctx.close();
                return;
            }
            String remoteId = msg.getHandshake().getNodeId();

            ChannelHandlerContext current = receivers.putIfAbsent(remoteId, ctx);
            if (current != null) {
                LOG.warn("Rejecting a duplicate connection from {}", remoteId);
                ctx.close();
                return;
            }
            LOG.debug("{} accepted a connection from {}", hostPort, remoteId);

            // Attach the Sender to this channel. Subsequent handlers use this
            // information to determine origins of messages.
            ctx.channel().attr(SENDER).set(remoteId);
            ctx.pipeline().remove(this);
        }
    }

    private class MainHandler extends ByteToMessageDecoder {
        private FileReceiver fileReceiver = null;

        private Message decodeToMessage(ByteBuf in) {
            if (in.readableBytes() < 4) {
                return null;
            }
            in.markReaderIndex();
            int messageLength = in.readInt();
            if (in.readableBytes() < messageLength) {
                in.resetReaderIndex();
                return null;
            }
            byte[] buffer = new byte[messageLength];
            in.readBytes(buffer);
            try {
                Message msg = Message.parseFrom(buffer);
                return msg;
            } catch (InvalidProtocolBufferException e) {
                LOG.error("Exception when parse protocol buffer.", e);
                Message msg = MessageBuilder.buildInvalidMessage(buffer);
                return msg;
            }
        }

        @Override
        protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) throws Exception {
            if (fileReceiver == null) {
                Message msg = decodeToMessage(in);
                if (msg == null) {
                    return;
                } else if (msg.getType() == MessageType.FILE_HEADER) {
                    LOG.debug("Got FILE_HEADER.");
                    fileReceiver = new FileReceiver(msg.getFileHeader().getLength());
                } else {
                    out.add(msg);
                }
            } else {
                fileReceiver.process(in);
                if (fileReceiver.isDone()) {
                    String filePath = fileReceiver.file.getPath();
                    Message msg = MessageBuilder.buildFileReceived(filePath);
                    out.add(msg);
                    // Resets it to null to switch back to normal decode mode.
                    fileReceiver = null;
                }
            }
        }

        class FileReceiver {
            final long fileLength;
            long receivedLength = 0;
            final File file;
            final FileOutputStream fout;

            public FileReceiver(long length) throws IOException {
                this.file = File.createTempFile("transport", "", dir);
                this.fileLength = length;
                this.fout = new FileOutputStream(this.file);
            }

            public void process(ByteBuf in) throws IOException {
                long readableBytes = in.readableBytes();
                long remainingBytes = fileLength - receivedLength;
                long bytesToRead = (remainingBytes < readableBytes) ? remainingBytes : readableBytes;
                byte[] buffer = new byte[(int) bytesToRead];
                in.readBytes(buffer);
                fout.write(buffer);
                receivedLength += bytesToRead;
                if (receivedLength == fileLength) {
                    fout.getChannel().force(false);
                    fout.close();
                }
            }

            boolean isDone() {
                return receivedLength == fileLength;
            }
        }
    }

    private class NotifyHandler extends SimpleChannelInboundHandler<Message> {
        @Override
        public void channelRead0(ChannelHandlerContext ctx, Message msg) {
            String remoteId = ctx.channel().attr(NettyTransport.SENDER).get();
            receiver.onReceived(remoteId, msg);
        }
    }

    /**
     * Handles server-side errors.
     */
    private class ServerErrorHandler extends ChannelInboundHandlerAdapter {
        @Override
        public void channelInactive(ChannelHandlerContext ctx) throws Exception {
            String remote = ctx.channel().attr(NettyTransport.SENDER).get();
            if (remote != null) {
                // remote might be null, for example if client connection got closed
                // because SSL handshake failed.
                LOG.debug("Disconnected from client {}", remote);
            }
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
            // Don't handle errors here. Call ctx.close() and let channelInactive()
            // handle all the errrors.
            LOG.debug("Caught an exception", cause);
            ctx.close();
        }
    }

    /**
     * Handles client-side errors.
     */
    private class ClientErrorHandler extends ChannelInboundHandlerAdapter {
        @Override
        public void channelInactive(ChannelHandlerContext ctx) throws Exception {
            String remote = ctx.channel().attr(NettyTransport.SENDER).get();
            LOG.debug("Disconnected from server {}", remote);

            // close incoming connection synchronously to make sure we don't receive
            // any more messages from this channel.
            ChannelHandlerContext remoteChannel = receivers.get(remote);
            if (remoteChannel != null) {
                remoteChannel.close().awaitUninterruptibly();
            }

            receiver.onDisconnected(remote);
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
            // Don't handle errors here. Call ctx.close() and let channelInactive()
            // handle all the errrors.
            LOG.debug("Caught an exception", cause);
            ctx.close();
        }
    }

    @Override
    public void send(final String destination, Message message) {
        if (destination.equals(hostPort)) {
            // The message is being sent to itself. Don't bother going over TCP.
            // Directly call onReceived.
            receiver.onReceived(destination, message);
            return;
        }
        ByteBuffer bytes = ByteBuffer.wrap(message.toByteArray());
        Sender currentSender = senders.get(destination);
        if (currentSender != null) {
            currentSender.requests.add(bytes);
        } else {
            // no connection exists.
            LOG.debug("No connection from {} to {}. Creating a new one", hostPort, destination);
            Sender newSender = new Sender(hostPort, destination);
            currentSender = senders.putIfAbsent(destination, newSender);
            if (currentSender == null) {
                newSender.requests.add(bytes);
                newSender.startHandshake();
            } else {
                currentSender.requests.add(bytes);
            }
        }
    }

    @Override
    public void send(final String destination, File file) {
        if (destination.equals(hostPort)) {
            LOG.error("Can't send file to itself.");
            throw new RuntimeException("Can't send file to itself.");
        }
        Sender currentSender = senders.get(destination);
        if (currentSender != null) {
            currentSender.requests.add(file);
        } else {
            // no connection exists.
            LOG.debug("No connection from {} to {}. Creating a new one", hostPort, destination);
            Sender newSender = new Sender(hostPort, destination);
            currentSender = senders.putIfAbsent(destination, newSender);
            if (currentSender == null) {
                newSender.requests.add(file);
                newSender.startHandshake();
            } else {
                currentSender.requests.add(file);
            }
        }
    }

    @Override
    public void clear(String destination) {
        LOG.debug("Closing the connection to {}", destination);
        Sender sender = senders.remove(destination);

        if (sender != null && sender.channel != null) {
            try {
                sender.channel.pipeline().remove("clientError");
            } catch (NoSuchElementException ex) {
                LOG.debug("Channel pipeline has already been cleared");
            }
            sender.shutdown();
        }

        // close incoming connection synchronously to make sure we don't receive
        // any more messages from this channel.
        ChannelHandlerContext receiver = receivers.remove(destination);
        if (receiver != null) {
            receiver.close().awaitUninterruptibly();
        }
    }

    /**
     * sender thread.
     */
    private class Sender implements Callable<Void> {
        final String destination;
        private Bootstrap bootstrap = null;
        private Channel channel;
        Future<Void> future;
        BlockingDeque<Object> requests = new LinkedBlockingDeque<>();

        public Sender(final String source, final String destination) {
            this.destination = destination;
            bootstrap = new Bootstrap();
            bootstrap.group(workerGroup);
            bootstrap.channel(NioSocketChannel.class);
            bootstrap.option(ChannelOption.CONNECT_TIMEOUT_MILLIS, 1000);
            bootstrap.option(ChannelOption.SO_KEEPALIVE, true);
            bootstrap.option(ChannelOption.TCP_NODELAY, true);
            bootstrap.handler(new ChannelInitializer<SocketChannel>() {
                @Override
                public void initChannel(SocketChannel ch) throws Exception {
                    if (isSslEnabled()) {
                        SSLEngine engine = serverContext.createSSLEngine();
                        engine.setUseClientMode(true);
                        ch.pipeline().addLast(new SslHandler(engine));
                    }
                    // Inbound handlers.
                    ch.pipeline().addLast("clientError", new ClientErrorHandler());
                    // Outbound handlers.
                    ch.pipeline().addLast("frameEncoder", new LengthFieldPrepender(4));
                }
            });
        }

        public void startHandshake() {
            String[] address = destination.split(":", 2);
            String host = address[0];
            int port = Integer.parseInt(address[1]);
            LOG.debug("host: {}, port: {}", host, port);
            bootstrap.connect(host, port).addListener(new ChannelFutureListener() {
                @Override
                public void operationComplete(ChannelFuture cfuture) {
                    if (cfuture.isSuccess()) {
                        LOG.debug("{} connected to {}. Sending a connection header", hostPort, destination);
                        channel = cfuture.channel();
                        // Associates the Sender with channel once the connection
                        // has been established.
                        channel.attr(SENDER).set(destination);
                        Message msg = MessageBuilder.buildHandshake(hostPort);
                        ByteBuffer bb = ByteBuffer.wrap(msg.toByteArray());
                        channel.writeAndFlush(Unpooled.wrappedBuffer(bb));
                        start();
                    } else {
                        LOG.debug("Failed to connect to {}: {}", destination, cfuture.cause().getMessage());
                        shutdown();
                        receiver.onDisconnected(destination);
                    }
                }
            });
        }

        void sendFile(File file) throws Exception {
            long length = file.length();
            LOG.debug("Got request of sending file {} of length {}.", file, length);
            Message handshake = MessageBuilder.buildFileHeader(length);
            byte[] bytes = handshake.toByteArray();
            // Sends HANDSHAKE first before transferring actual file data, the
            // HANDSHAKE will tell the peer's channel to prepare for the file
            // transferring.
            channel.writeAndFlush(Unpooled.wrappedBuffer(bytes)).sync();
            ChannelHandler prepender = channel.pipeline().get("frameEncoder");
            // Removes length prepender, we don't need this handler for file
            // transferring.
            channel.pipeline().remove(prepender);
            // Adds ChunkedWriteHandler for file transferring.
            ChannelHandler cwh = new ChunkedWriteHandler();
            channel.pipeline().addLast(cwh);
            // Begins file transferring.
            RandomAccessFile raf = new RandomAccessFile(file, "r");
            if (channel.pipeline().get(SslHandler.class) != null) {
                // Zero-Copy file transferring is not supported for ssl.
                channel.writeAndFlush(new ChunkedFile(raf, 0, length, 8912));
            } else {
                // Use Zero-Copy file transferring in non-ssl mode.
                FileRegion region = new DefaultFileRegion(raf.getChannel(), 0, length);
                channel.writeAndFlush(region);
            }
            // Restores pipeline to original state.
            channel.pipeline().remove(cwh);
            channel.pipeline().addLast("frameEncoder", prepender);
        }

        @Override
        public Void call() throws Exception {
            LOG.debug("Started the sender: {} => {}", hostPort, destination);
            try {
                while (true) {
                    Object req = requests.take();
                    if (req instanceof ByteBuffer) {
                        ByteBuffer buf = (ByteBuffer) req;
                        channel.writeAndFlush(Unpooled.wrappedBuffer(buf));
                    } else if (req instanceof File) {
                        File file = (File) req;
                        sendFile(file);
                    } else if (req instanceof Shutdown) {
                        LOG.debug("Got shutdown request.");
                        break;
                    }
                }
            } catch (InterruptedException ex) {
                LOG.debug("Sender to {} got interrupted", destination);
                return null;
            } catch (Exception ex) {
                LOG.warn("Sender failed with an exception", ex);
                throw ex;
            } finally {
                channel.close();
            }
            return null;
        }

        public synchronized void start() {
            ExecutorService es = Executors.newSingleThreadExecutor();
            future = es.submit(this);
            es.shutdown();
        }

        public void shutdown() {
            LOG.debug("Shutting down the sender: {} => {}", hostPort, destination);
            try {
                if (future != null) {
                    try {
                        this.requests.add(new Shutdown());
                        future.get();
                    } catch (InterruptedException | ExecutionException ex) {
                        LOG.debug("Ignore the exception", ex);
                    }
                }
                if (channel != null) {
                    channel.close().syncUninterruptibly();
                }
            } catch (RejectedExecutionException ex) {
                LOG.debug("Ignoring rejected execution exception", ex);
            }
        }

        class Shutdown {
            // We use it to shutdown the sender thread.
        }
    }
}