me.xingrz.prox.tcp.tunnel.Tunnel.java Source code

Introduction

Here is the source code for me.xingrz.prox.tcp.tunnel.Tunnel.java

Source

/*
 * Copyright (C) 2015 XiNGRZ <chenxingyu92@gmail.com>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, write to the Free Software Foundation, Inc.,
 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */

package me.xingrz.prox.tcp.tunnel;

import org.apache.commons.io.IOUtils;

import java.io.Closeable;
import java.io.IOException;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;

import me.xingrz.prox.logging.FormattingLogger;
import me.xingrz.prox.selectable.Readable;
import me.xingrz.prox.selectable.Writable;

/**
 * Tunnel 
 *  Tunnel ? {@link java.nio.channels.SocketChannel}
 *  Tunnel ? {@link #setBrother(Tunnel)} 
 */
public abstract class Tunnel implements Closeable, Readable, Writable {

    protected final FormattingLogger logger;

    protected final Selector selector;
    protected final SocketChannel channel;

    protected Tunnel brother;

    private ByteBuffer receiving = ByteBuffer.allocate(0xFFFF);
    private ByteBuffer remaining = ByteBuffer.allocate(0xFFFF);

    private volatile boolean closed;

    public Tunnel(Selector selector, SocketChannel channel, String sessionKey) {
        this.selector = selector;
        this.channel = channel;
        this.logger = getLogger(sessionKey);
    }

    protected abstract FormattingLogger getLogger(String sessionKey);

    public Socket socket() {
        return channel.socket();
    }

    /**
     * @param brother ? Tunnel
     */
    public final void setBrother(Tunnel brother) {
        this.brother = brother;
    }

    /**
     * ?
     */
    public final void beginReceiving() {
        try {
            if (channel.isBlocking()) {
                channel.configureBlocking(false);
            }

            channel.register(selector, SelectionKey.OP_READ, this);
        } catch (IOException e) {
            logger.w(e, "Failed to begin receiving, close");
            IOUtils.closeQuietly(this);
        }
    }

    /**
     *  Selector  OP_READ 
     *
     * @param key Selection Key
     */
    @Override
    public final void onReadable(SelectionKey key) {
        if (closed || receiving == null) {
            return;
        }

        receiving.clear();

        int read;

        try {
            read = channel.read(receiving);
        } catch (IOException e) {
            logger.w(e, "Failed to read from channel, terminated");
            IOUtils.closeQuietly(this);
            return;
        }

        if (read == -1) {
            IOUtils.closeQuietly(this);
            return;
        }

        receiving.flip();

        // 
        if (!receiving.hasRemaining()) {
            return;
        }

        // ???
        if (afterReceived(receiving)) {
            return;
        }

        // ??
        if (!receiving.hasRemaining()) {
            return;
        }

        brother.beforeSending(receiving);

        if (!brother.write(receiving)) {
            logger.v("Brother not ready for receiving, canceled");
            key.cancel();
        }
    }

    /**
     * Tunnel  {@link #channel} ?
     *
     * @param buffer ?
     * @return {@code true} ?{@code false} ?? {@link #brother}
     */
    protected abstract boolean afterReceived(ByteBuffer buffer);

    /**
     * ?? Tunnel ?????
     *
     * @param buffer ????
     */
    protected abstract void beforeSending(ByteBuffer buffer);

    /**
     * ? Tunnel  {@link #channel} ???
     *
     * @param buffer ????
     * @return ????
     */
    public boolean write(ByteBuffer buffer) {
        return writeInternal(buffer);
    }

    protected final boolean writeInternal(ByteBuffer buffer) {
        if (closed || remaining == null) {
            return false;
        }

        try {
            while (buffer.hasRemaining()) {
                if (channel.write(buffer) == 0) {
                    break;
                }
            }
        } catch (IOException e) {
            logger.w(e, "Failed writing to " + channelToString());
            return false;
        }

        if (buffer.hasRemaining()) {
            keepRemaining(buffer);
            return false;
        } else {
            return true;
        }
    }

    protected void keepRemaining(ByteBuffer buffer) {
        remaining.clear();
        remaining.put(buffer);
        remaining.flip();

        // ??????
        if (!channel.isConnected()) {
            return;
        }

        try {
            channel.register(selector, SelectionKey.OP_WRITE, this);
        } catch (ClosedChannelException e) {
            logger.w(e, "Failed to register OP_WRITE since channel is closed");
            IOUtils.closeQuietly(this);
        }
    }

    /**
     *  Selector  OP_WRITE 
     *
     * @param key Selection Key
     */
    @Override
    public final void onWritable(SelectionKey key) {
        if (closed) {
            return;
        }

        beforeSending(remaining);
        writeRemaining();
        key.cancel();
    }

    protected void writeRemaining() {
        try {
            while (remaining.hasRemaining()) {
                if (channel.write(remaining) == 0) {
                    break;
                }
            }
        } catch (IOException e) {
            logger.w(e, "Failed writing remaining data, closed");
            IOUtils.closeQuietly(this);
            return;
        }

        brother.beginReceiving();
    }

    /**
     *  Tunnel ?
     */
    @Override
    public void close() {
        closeInternal(true);
    }

    private void closeInternal(boolean withBrother) {
        if (!closed) {
            closed = true;

            IOUtils.closeQuietly(channel);

            if (withBrother && brother != null) {
                brother.closeInternal(false);
            }

            remaining = null;
            receiving = null;

            brother = null;

            onClose();
        }
    }

    /**
     * Tunnel 
     */
    protected abstract void onClose();

    protected final String channelToString() {
        return String.format("Channel[%s:%d -> %s:%d]", channel.socket().getLocalAddress().getHostAddress(),
                channel.socket().getLocalPort(), channel.socket().getInetAddress().getHostAddress(),
                channel.socket().getPort());
    }

}