BukkitEncodeHandler.java

package us.myles.ViaVersion.bukkit.handlers;

import io.netty.buffer.ByteBuf;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.MessageToByteEncoder;
import us.myles.ViaVersion.api.PacketWrapper;
import us.myles.ViaVersion.api.data.UserConnection;
import us.myles.ViaVersion.api.type.Type;
import us.myles.ViaVersion.bukkit.util.NMSUtil;
import us.myles.ViaVersion.exception.CancelException;
import us.myles.ViaVersion.handlers.ChannelHandlerContextWrapper;
import us.myles.ViaVersion.handlers.ViaHandler;
import us.myles.ViaVersion.packets.Direction;
import us.myles.ViaVersion.protocols.base.ProtocolInfo;
import us.myles.ViaVersion.util.PipelineUtil;

import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;

public class BukkitEncodeHandler extends MessageToByteEncoder implements ViaHandler {
    private static Field versionField = null;

    static {
        try {
            versionField = NMSUtil.nms("PacketEncoder").getDeclaredField("version");
            versionField.setAccessible(true);
        } catch (Exception e) {
            // Not compat version
        }
    }

    private final UserConnection info;
    private final MessageToByteEncoder minecraftEncoder;

    public BukkitEncodeHandler(UserConnection info, MessageToByteEncoder minecraftEncoder) {
        this.info = info;
        this.minecraftEncoder = minecraftEncoder;
    }

    @Override
    public void handlerAdded(ChannelHandlerContext ctx) throws Exception {
        if (versionField != null) {
            versionField.set(minecraftEncoder, versionField.get(this));
        }
    }

    @Override
    protected void encode(ChannelHandlerContext ctx, Object o, ByteBuf buf) throws Exception {
        info.incrementSent();

        if (isTransformationRequired() && info.isActive()) {
            this.doTransform(ctx, o, buf);
        } else {
            // Directly encode into the original buffer
            this.encodeMinecraftBuffer(ctx, o, buf);
        }
    }

    private boolean isTransformationRequired() {
        ProtocolInfo protocolInfo = info.get(ProtocolInfo.class);
        int serverProtocol = protocolInfo.getServerProtocolVersion();
        int clientProtocol = protocolInfo.getProtocolVersion();
        return clientProtocol != serverProtocol;
    }

    private void doTransform(ChannelHandlerContext ctx, Object o, ByteBuf buf) throws Exception {
        ByteBuf source;
        if (!(o instanceof ByteBuf)) {
            source = encodeMinecraftBuffer(ctx, o);
        } else {
            source = ((ByteBuf) o).retain();
        }

        try {
            transform(source, buf);
        } finally {
            source.release();
        }
    }

    private ByteBuf encodeMinecraftBuffer(ChannelHandlerContext ctx, Object o) throws Exception {
        ByteBuf dest = ctx.alloc().ioBuffer();
        try {
            this.encodeMinecraftBuffer(ctx, o, dest);
        } catch (Exception e) {
            dest.release();
            throw e;
        }
        return dest;
    }

    private void encodeMinecraftBuffer(ChannelHandlerContext ctx, Object o, ByteBuf buf) throws Exception {
        // call minecraft encoder
        try {
            PipelineUtil.callEncode(this.minecraftEncoder, new ChannelHandlerContextWrapper(ctx, this), o, buf);
        } catch (InvocationTargetException e) {
            if (e.getCause() instanceof Exception) {
                throw (Exception) e.getCause();
            }
        }
    }

    @Override
    public void transform(ByteBuf source) throws Exception {
        // Provided because other things rely on this API - the internal translation done by this class does NOT rely
        // on making a copy of the source packet.

        if (source.readableBytes() == 0) {
            return; // Someone Already Decoded It!
        }

        ByteBuf oldPacket = source.copy();
        source.clear();
        try {
            this.transform(oldPacket, source);
        } finally {
            oldPacket.release();
        }
    }

    private void transform(ByteBuf source, ByteBuf destination) throws Exception {
        // Handle ID
        int id = Type.VAR_INT.read(source);

        // Transform
        PacketWrapper wrapper = new PacketWrapper(id, source, info);
        ProtocolInfo protInfo = info.get(ProtocolInfo.class);
        protInfo.getPipeline().transform(Direction.OUTGOING, protInfo.getState(), wrapper);
        wrapper.writeToBuffer(destination);
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        if (PipelineUtil.containsCause(cause, CancelException.class)) return;
        super.exceptionCaught(ctx, cause);
    }
}