cc.agentx.server.net.nio.XConnectHandler.java Source code

Java tutorial

Introduction

Here is the source code for cc.agentx.server.net.nio.XConnectHandler.java

Source

/*
 * Copyright 2017 ZhangJiupeng
 *
 * Licensed 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 cc.agentx.server.net.nio;

import cc.agentx.protocol.request.XRequest;
import cc.agentx.protocol.request.XRequestWrapper;
import cc.agentx.server.Configuration;
import cc.agentx.server.cache.DnsCache;
import cc.agentx.wrapper.Wrapper;
import io.netty.bootstrap.Bootstrap;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.*;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.util.ReferenceCountUtil;
import io.netty.util.concurrent.Future;
import io.netty.util.concurrent.FutureListener;
import io.netty.util.concurrent.Promise;
import io.netty.util.internal.logging.InternalLogger;
import io.netty.util.internal.logging.InternalLoggerFactory;

import java.io.ByteArrayOutputStream;
import java.net.UnknownHostException;
import java.util.Arrays;

@ChannelHandler.Sharable
public final class XConnectHandler extends ChannelInboundHandlerAdapter {
    private static final InternalLogger log = InternalLoggerFactory.getInstance(XConnectHandler.class);

    private final Bootstrap bootstrap = new Bootstrap();
    private final ByteArrayOutputStream tailDataBuffer;
    private final XRequestWrapper requestWrapper;
    private final boolean exposedRequest;
    private final Wrapper wrapper;

    private boolean requestParsed;

    public XConnectHandler() {
        this.tailDataBuffer = new ByteArrayOutputStream();
        Configuration config = Configuration.INSTANCE;
        this.requestWrapper = config.getXRequestWrapper();
        this.exposedRequest = requestWrapper.exposeRequest();
        this.wrapper = config.getWrapper();
    }

    @Override
    public void channelRead(ChannelHandlerContext ctx, Object msg) {
        try {
            ByteBuf byteBuf = (ByteBuf) msg;
            if (!byteBuf.hasArray()) {
                byte[] bytes = new byte[byteBuf.readableBytes()];
                byteBuf.getBytes(0, bytes);

                if (!requestParsed) {
                    if (!exposedRequest) {
                        bytes = wrapper.unwrap(bytes);
                        if (bytes == null) {
                            log.info("\tClient -> Proxy           \tHalf Request");
                            return;
                        }
                    }
                    XRequest xRequest = requestWrapper.parse(bytes);
                    String host = xRequest.getHost();
                    int port = xRequest.getPort();
                    int dataLength = xRequest.getSubsequentDataLength();
                    if (dataLength > 0) {
                        byte[] tailData = Arrays.copyOfRange(bytes, bytes.length - dataLength, bytes.length);
                        if (exposedRequest) {
                            tailData = wrapper.unwrap(tailData);
                            if (tailData != null) {
                                tailDataBuffer.write(tailData, 0, tailData.length);
                            }
                        } else {
                            tailDataBuffer.write(tailData, 0, tailData.length);
                        }
                    }
                    log.info("\tClient -> Proxy           \tTarget {}:{}{}", host, port,
                            DnsCache.isCached(host) ? " [Cached]" : "");
                    if (xRequest.getAtyp() == XRequest.Type.DOMAIN) {
                        try {
                            host = DnsCache.get(host);
                            if (host == null) {
                                host = xRequest.getHost();
                            }
                        } catch (UnknownHostException e) {
                            log.warn("\tClient <- Proxy           \tBad DNS! ({})", e.getMessage());
                            ctx.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE);
                            return;
                        }
                    }

                    Promise<Channel> promise = ctx.executor().newPromise();
                    promise.addListener(new FutureListener<Channel>() {
                        @Override
                        public void operationComplete(final Future<Channel> future) throws Exception {
                            final Channel outboundChannel = future.getNow();
                            if (future.isSuccess()) {
                                // handle tail
                                byte[] tailData = tailDataBuffer.toByteArray();
                                tailDataBuffer.close();
                                if (tailData.length > 0) {
                                    log.info("\tClient ==========> Target \tSend Tail [{} bytes]", tailData.length);
                                }
                                outboundChannel
                                        .writeAndFlush((tailData.length > 0) ? Unpooled.wrappedBuffer(tailData)
                                                : Unpooled.EMPTY_BUFFER)
                                        .addListener(channelFuture -> {
                                            // task handover
                                            outboundChannel.pipeline()
                                                    .addLast(new XRelayHandler(ctx.channel(), wrapper, false));
                                            ctx.pipeline()
                                                    .addLast(new XRelayHandler(outboundChannel, wrapper, true));
                                            ctx.pipeline().remove(XConnectHandler.this);
                                        });

                            } else {
                                if (ctx.channel().isActive()) {
                                    ctx.writeAndFlush(Unpooled.EMPTY_BUFFER)
                                            .addListener(ChannelFutureListener.CLOSE);
                                }
                            }
                        }
                    });

                    final String finalHost = host;
                    bootstrap.group(ctx.channel().eventLoop()).channel(NioSocketChannel.class)
                            .option(ChannelOption.CONNECT_TIMEOUT_MILLIS, 10000)
                            .option(ChannelOption.SO_KEEPALIVE, true)
                            .handler(new XPingHandler(promise, System.currentTimeMillis())).connect(host, port)
                            .addListener(new ChannelFutureListener() {
                                @Override
                                public void operationComplete(ChannelFuture future) throws Exception {
                                    if (!future.isSuccess()) {
                                        if (ctx.channel().isActive()) {
                                            log.warn("\tClient <- Proxy           \tBad Ping! ({}:{})", finalHost,
                                                    port);
                                            ctx.writeAndFlush(Unpooled.EMPTY_BUFFER)
                                                    .addListener(ChannelFutureListener.CLOSE);
                                        }
                                    }
                                }
                            });

                    requestParsed = true;
                } else {
                    bytes = wrapper.unwrap(bytes);
                    if (bytes != null)
                        tailDataBuffer.write(bytes, 0, bytes.length);
                }
            }
        } finally {
            ReferenceCountUtil.release(msg);
        }
    }

    @Override
    public void channelActive(ChannelHandlerContext ctx) {
        ctx.writeAndFlush(Unpooled.EMPTY_BUFFER);
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
        if (ctx.channel().isActive()) {
            ctx.writeAndFlush(Unpooled.EMPTY_BUFFER).addListener(ChannelFutureListener.CLOSE);
        }
        cause.printStackTrace();
        log.warn("\tBad Connection! ({})", cause.getMessage());
    }
}