org.rx.socks.tcp.TcpServer Maven / Gradle / Ivy

package org.rx.socks.tcp;

import io.netty.bootstrap.ServerBootstrap;
import io.netty.channel.*;
import io.netty.handler.codec.compression.ZlibCodecFactory;
import io.netty.handler.codec.compression.ZlibWrapper;
import io.netty.handler.codec.serialization.ClassResolvers;
import io.netty.handler.codec.serialization.ObjectDecoder;
import io.netty.handler.codec.serialization.ObjectEncoder;
import io.netty.handler.ssl.SslContext;
import io.netty.handler.ssl.SslContextBuilder;
import io.netty.handler.ssl.util.SelfSignedCertificate;
import lombok.*;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.rx.core.*;
import org.rx.socks.Sockets;
import org.rx.socks.tcp.packet.ErrorPacket;
import org.rx.socks.tcp.packet.HandshakePacket;

import java.io.Serializable;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.BiConsumer;

import static org.rx.core.Contract.*;

@Slf4j
@RequiredArgsConstructor
public class TcpServer extends Disposable implements EventTarget> {
    //not shareable
    private class PacketServerHandler extends ChannelInboundHandlerAdapter {
        private SessionClient client;

        @Override
        public void channelActive(ChannelHandlerContext ctx) {
//        super.channelActive(ctx);
            log.debug("serverActive {}", ctx.channel().remoteAddress());
            if (getClientSize() > getCapacity()) {
                log.warn("Not enough capacity");
                Sockets.closeOnFlushed(ctx.channel());
                return;
            }
            client = new SessionClient<>(ctx, getStateType() == null ? null : Reflects.newInstance(getStateType()));
            PackEventArgs args = new PackEventArgs<>(client, null);
            raiseEvent(onConnected, args);
            if (args.isCancel()) {
                log.warn("Close client");
                Sockets.closeOnFlushed(ctx.channel());
            }
        }

        @Override
        public void channelRead(ChannelHandlerContext ctx, Object msg) {
            log.debug("serverRead {} {}", ctx.channel().remoteAddress(), msg.getClass());

            Serializable pack;
            if ((pack = as(msg, Serializable.class)) == null) {
                ctx.writeAndFlush(new ErrorPacket("Packet discard"));
                Sockets.closeOnFlushed(ctx.channel());
                return;
            }
            if (tryAs(pack, HandshakePacket.class, p -> {
                client.setGroupId(p.getGroupId());
                addClient(client);
            })) {
                return;
            }

            if (client.getGroupId() == null) {
                log.warn("ServerHandshake fail");
                return;
            }
            //异步避免阻塞
            Tasks.run(() -> raiseEvent(onReceive, new PackEventArgs<>(client, pack)));
        }

        @Override
        public void channelInactive(ChannelHandlerContext ctx) {
            log.debug("serverInactive {}", ctx.channel().remoteAddress());
            removeClient(client);
            raiseEvent(onDisconnected, new PackEventArgs<>(client, null));
        }

        @Override
        public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
            log.error("serverCaught {}", ctx.channel().remoteAddress(), cause);
            if (!ctx.channel().isActive()) {
                return;
            }

            ErrorEventArgs args = new ErrorEventArgs<>(client, cause);
            try {
                raiseEvent(onError, args);
            } catch (Exception e) {
                log.error("serverCaught", e);
            }
            if (args.isCancel()) {
                return;
            }
            Sockets.closeOnFlushed(ctx.channel());
        }
    }

    public volatile BiConsumer, PackEventArgs> onConnected, onDisconnected, onSend, onReceive;
    public volatile BiConsumer, ErrorEventArgs> onError;
    public volatile BiConsumer, EventArgs> onClosed;
    @Getter
    private final TcpConfig config;
    @Getter
    private final Class stateType;
    private ServerBootstrap bootstrap;
    private SslContext sslCtx;
    private final Map>> clients = new ConcurrentHashMap<>();
    @Getter
    private volatile boolean isStarted;
    @Getter
    @Setter
    private volatile int capacity = 1000000;

    public int getClientSize() {
        return clients.size();
    }

    @Override
    protected void freeObjects() {
        Sockets.closeBootstrap(bootstrap);
        isStarted = false;
        raiseEvent(onClosed, EventArgs.Empty);
    }

    public void start() {
        start(false);
    }

    @SneakyThrows
    public void start(boolean waitClose) {
        if (isStarted) {
            throw new InvalidOperationException("Server has started");
        }

        if (config.isEnableSsl()) {
            SelfSignedCertificate ssc = new SelfSignedCertificate();
            sslCtx = SslContextBuilder.forServer(ssc.certificate(), ssc.privateKey()).build();
        }
        bootstrap = Sockets.serverBootstrap(1, config.getWorkThread(), config.getMemoryMode(), channel -> {
            ChannelPipeline pipeline = channel.pipeline();
            if (sslCtx != null) {
                pipeline.addLast(sslCtx.newHandler(channel.alloc()));
            }
            if (config.isEnableCompress()) {
                pipeline.addLast(ZlibCodecFactory.newZlibEncoder(ZlibWrapper.GZIP));
                pipeline.addLast(ZlibCodecFactory.newZlibDecoder(ZlibWrapper.GZIP));
            }
            pipeline.addLast(new ObjectEncoder(),
                    new ObjectDecoder(ClassResolvers.weakCachingConcurrentResolver(TcpConfig.class.getClassLoader())),
                    new PacketServerHandler());
        }).option(ChannelOption.SO_REUSEADDR, true)
                .option(ChannelOption.CONNECT_TIMEOUT_MILLIS, config.getConnectTimeout());
        ChannelFuture future = bootstrap.bind(config.getEndpoint()).addListeners(Sockets.FireExceptionThenCloseOnFailure, f -> {
            if (!f.isSuccess()) {
                return;
            }
            isStarted = true;
            log.debug("Listened on port {}..", config.getEndpoint());
        });
        if (waitClose) {
            future.channel().closeFuture().sync();
        }
    }

    protected void addClient(SessionClient client) {
        require(client.getGroupId());

        clients.computeIfAbsent(client.getGroupId(), k -> Collections.synchronizedSet(new HashSet<>())).add(client);
    }

    protected void removeClient(SessionClient client) {
        Set> set = getClients(client.getGroupId());
        if (set.removeIf(p -> p == client) && set.isEmpty()) {
            clients.remove(client.getGroupId());
        }
    }

    public Set> getClients(String groupId) {
        require(groupId);

        Set> set = clients.get(groupId);
        if (CollectionUtils.isEmpty(set)) {
            throw new InvalidOperationException(String.format("GroupId %s not found", groupId));
        }
        return set;
    }

    public SessionClient getClient(String groupId, ChannelId channelId) {
        SessionClient client = NQuery.of(getClients(groupId)).where(p -> p.getId().equals(channelId)).singleOrDefault();
        if (client == null) {
            throw new InvalidOperationException(String.format("GroupId %s with ClientId %s not found", groupId, channelId));
        }
        return client;
    }

    public void send(SessionClient client, Serializable pack) {
        checkNotClosed();
        require(client, pack);

        PackEventArgs args = new PackEventArgs<>(client, pack);
        raiseEvent(onSend, args);
        if (args.isCancel()) {
            return;
        }
        client.ctx.writeAndFlush(pack).addListener(ChannelFutureListener.FIRE_EXCEPTION_ON_FAILURE);
    }
}