Fix & improve PacketContainer serialization & cloning (#1794)

src/main/java/com/comphenix/protocol/events/PacketContainer.java

@@ -17,37 +17,49 @@
 
 package com.comphenix.protocol.events;
 
-import javax.annotation.Nullable;
 import java.io.IOException;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
 import java.io.Serializable;
-import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.function.Function;
 
 import com.comphenix.protocol.PacketType;
 import com.comphenix.protocol.injector.StructureCache;
 import com.comphenix.protocol.reflect.FuzzyReflection;
 import com.comphenix.protocol.reflect.ObjectWriter;
 import com.comphenix.protocol.reflect.StructureModifier;
-import com.comphenix.protocol.reflect.cloning.*;
+import com.comphenix.protocol.reflect.accessors.Accessors;
+import com.comphenix.protocol.reflect.accessors.ConstructorAccessor;
+import com.comphenix.protocol.reflect.accessors.MethodAccessor;
+import com.comphenix.protocol.reflect.cloning.AggregateCloner;
 import com.comphenix.protocol.reflect.cloning.AggregateCloner.BuilderParameters;
+import com.comphenix.protocol.reflect.cloning.BukkitCloner;
+import com.comphenix.protocol.reflect.cloning.Cloner;
+import com.comphenix.protocol.reflect.cloning.CollectionCloner;
+import com.comphenix.protocol.reflect.cloning.FieldCloner;
+import com.comphenix.protocol.reflect.cloning.GuavaOptionalCloner;
+import com.comphenix.protocol.reflect.cloning.ImmutableDetector;
+import com.comphenix.protocol.reflect.cloning.JavaOptionalCloner;
 import com.comphenix.protocol.reflect.fuzzy.FuzzyMethodContract;
+import com.comphenix.protocol.reflect.instances.DefaultInstances;
 import com.comphenix.protocol.reflect.instances.MinecraftGenerator;
 import com.comphenix.protocol.utility.MinecraftMethods;
 import com.comphenix.protocol.utility.MinecraftReflection;
 import com.comphenix.protocol.utility.MinecraftVersion;
 import com.comphenix.protocol.wrappers.Converters;
-import com.google.common.base.Function;
 import com.google.common.collect.Sets;
-
 import io.netty.buffer.ByteBuf;
-import io.netty.buffer.UnpooledByteBufAllocator;
+import io.netty.util.ReferenceCountUtil;
+import javax.annotation.Nullable;
 
 /**
  * Represents a Minecraft packet indirectly.
@@ -61,8 +73,7 @@ public class PacketContainer extends AbstractStructure implements Serializable {
 	private PacketType type;
 
 	// Support for serialization
-	private static ConcurrentMap<Class<?>, Method> writeMethods = new ConcurrentHashMap<>();
+	private static final Map<PacketType, Function<Object, Object>> PACKET_DESERIALIZER_METHODS = new ConcurrentHashMap<>();
-	private static ConcurrentMap<Class<?>, Method> readMethods = new ConcurrentHashMap<>();
 
 	// Used to clone packets
 	private static final AggregateCloner DEEP_CLONER = AggregateCloner
@@ -217,28 +228,36 @@ public class PacketContainer extends AbstractStructure implements Serializable {
 	 * @return A deep copy of the current packet.
 	 */
 	public PacketContainer deepClone() {
-		Object clonedPacket = null;
+		Object handle = this.getHandle();
+		PacketType packetType = this.getType();
+		if (handle == null || packetType == null) {
+			// nothing to clone, just carry on (this should normally not happen)
+			return this;
+		}
 
-		if (!FAST_CLONE_UNSUPPORTED.contains(type)) {
+		// try fast cloning first
+		if (!FAST_CLONE_UNSUPPORTED.contains(packetType)) {
 			try {
-				clonedPacket = DEEP_CLONER.clone(getHandle());
+				Object cloned = DEEP_CLONER.clone(handle);
+				return new PacketContainer(packetType, cloned);
 			} catch (Exception ex) {
-				FAST_CLONE_UNSUPPORTED.add(type);
+				FAST_CLONE_UNSUPPORTED.add(packetType);
 			}
 		}
 
 		// Fall back on the slower alternative method of reading and writing back the packet
-		if (clonedPacket == null) {
+		Object serialized = this.serializeToBuffer();
-			clonedPacket = SerializableCloner.clone(this).getHandle();
+		Object deserialized = deserializeFromBuffer(packetType, serialized);
-		}
 
-		return new PacketContainer(getType(), clonedPacket);
+		// ensure that we don't leak memory
+		ReferenceCountUtil.safeRelease(serialized);
+		return new PacketContainer(packetType, deserialized);
 	}
 
 	// To save space, we'll skip copying the inflated buffers in packet 51 and 56
-	private static Function<BuilderParameters, Cloner> getSpecializedDeepClonerFactory() {
+	private static com.google.common.base.Function<BuilderParameters, Cloner> getSpecializedDeepClonerFactory() {
 		// Look at what you've made me do Java, look at it!!
-		return new Function<BuilderParameters, Cloner>() {
+		return new com.google.common.base.Function<BuilderParameters, Cloner>() {
 			@Override
 			public Cloner apply(@Nullable BuilderParameters param) {
 				return new FieldCloner(param.getAggregateCloner(), param.getInstanceProvider()) {{
@@ -262,17 +281,17 @@ public class PacketContainer extends AbstractStructure implements Serializable {
 		// Default serialization
 		output.defaultWriteObject();
 
-		// We'll take care of NULL packets as well
+		// serialize the packet
-		output.writeBoolean(handle != null);
+		ByteBuf buffer = (ByteBuf) this.serializeToBuffer();
-
+		if (buffer != null) {
-		try {
+			output.writeBoolean(true);
-			ByteBuf buffer = createPacketBuffer();
-			MinecraftMethods.getPacketWriteByteBufMethod().invoke(handle, buffer);
-
 			output.writeInt(buffer.readableBytes());
 			buffer.readBytes(output, buffer.readableBytes());
-		} catch (IllegalArgumentException e) {
+
-			throw new IOException("Minecraft packet doesn't support DataOutputStream", e);
+			// ensure that we don't leak memory
+			ReferenceCountUtil.safeRelease(buffer);
+		} else {
+			output.writeBoolean(false);
 		}
 	}
 
@@ -280,68 +299,101 @@ public class PacketContainer extends AbstractStructure implements Serializable {
 		// Default deserialization
 		input.defaultReadObject();
 
-		// Get structure modifier
+		// Deserialize the packet from the stream (if present)
-		structureModifier = StructureCache.getStructure(type);
+		this.structureModifier = StructureCache.getStructure(this.type);
-
-	    // Don't read NULL packets
 		if (input.readBoolean()) {
-			ByteBuf buffer = createPacketBuffer();
+			int dataLength = input.readInt();
-			buffer.writeBytes(input, input.readInt());
 
-	    	// Create a default instance of the packet
+			ByteBuf byteBuf = (ByteBuf) MinecraftReflection.createPacketDataSerializer(dataLength);
-			if (MinecraftVersion.CAVES_CLIFFS_1.atOrAbove()) {
+			while (true) {
-				Object serializer = MinecraftReflection.getPacketDataSerializer(buffer);
+				// ObjectInputStream only reads a specific amount of bytes before moving the cursor forwards and
+				// allows reading the next byte chunk. So we need to read until the data is gone from the stream and
+				// fully transferred into the buffer.
+				int transferredBytes = byteBuf.writeBytes(input, dataLength);
 
-				try {
+				// check if we reached the end of the stream, or if the stream has no more data available
-					handle = type.getPacketClass()
+				dataLength -= transferredBytes;
-							.getConstructor(MinecraftReflection.getPacketDataSerializerClass())
+				if (dataLength <= 0 || transferredBytes <= 0) {
-							.newInstance(serializer);
+					break;
-				} catch (NoSuchMethodException | IllegalAccessException | InstantiationException ex) {
-					// they might have a static method to create them instead
-					Method method = FuzzyReflection.fromClass(type.getPacketClass(), true)
-							.getMethod(FuzzyMethodContract
-									.newBuilder()
-									.requireModifier(Modifier.STATIC)
-									.returnTypeExact(type.getPacketClass())
-									.parameterExactArray(MinecraftReflection.getPacketDataSerializerClass())
-									.build());
-					try {
-						handle = method.invoke(null, serializer);
-					} catch (ReflectiveOperationException exception) {
-						throw new RuntimeException("Failed to construct packet for " + type, exception);
-					}
-				} catch (InvocationTargetException ex) {
-					throw new RuntimeException("Unable to clone packet " + type + " using constructor", ex.getCause());
-				}
-			} else {
-				handle = StructureCache.newPacket(type);
-
-				// Call the read method
-				try {
-					MinecraftMethods.getPacketReadByteBufMethod().invoke(handle, buffer);
-				} catch (IllegalArgumentException e) {
-					throw new IOException("Minecraft packet doesn't support DataInputStream", e);
 				}
 			}
 
-			// And we're done
+			// deserialize & ensure that we don't leak memory
-			structureModifier = structureModifier.withTarget(handle);
+			Object packet = deserializeFromBuffer(this.type, byteBuf);
+			ReferenceCountUtil.safeRelease(byteBuf);
+
+			this.handle = packet;
+			this.structureModifier = this.structureModifier.withTarget(packet);
 		}
 	}
 
 	/**
 	 * Construct a new packet data serializer.
 	 * @return The packet data serializer.
+	 * @deprecated use {@link MinecraftReflection#createPacketDataSerializer(int)} instead
 	 */
+	@Deprecated
 	public static ByteBuf createPacketBuffer() {
-		ByteBuf buffer = UnpooledByteBufAllocator.DEFAULT.buffer();
+		return (ByteBuf) MinecraftReflection.createPacketDataSerializer(0);
-		Class<?> packetSerializer = MinecraftReflection.getPacketDataSerializerClass();
-
-		try {
-			return (ByteBuf) packetSerializer.getConstructor(ByteBuf.class).newInstance(buffer);
-		} catch (Exception e) {
-			throw new RuntimeException("Cannot construct packet serializer.", e);
 	}
+
+	// ---- Cloning
+
+	public static Object deserializeFromBuffer(PacketType packetType, Object buffer) {
+		if (buffer == null) {
+			return null;
+		}
+
+		Function<Object, Object> deserializer = PACKET_DESERIALIZER_METHODS.computeIfAbsent(packetType, type -> {
+			if (MinecraftVersion.CAVES_CLIFFS_1.atOrAbove()) {
+				// best guess - a constructor which takes a buffer as the only argument
+				ConstructorAccessor bufferConstructor = Accessors.getConstructorAccessorOrNull(
+						type.getPacketClass(),
+						MinecraftReflection.getPacketDataSerializerClass());
+				if (bufferConstructor != null) {
+					return bufferConstructor::invoke;
+				}
+
+				// they might have a static method to create them instead
+				List<Method> methods = FuzzyReflection.fromClass(type.getPacketClass(), true)
+						.getMethodList(FuzzyMethodContract.newBuilder()
+								.requireModifier(Modifier.STATIC)
+								.returnTypeExact(type.getPacketClass())
+								.parameterExactArray(MinecraftReflection.getPacketDataSerializerClass())
+								.build());
+				if (!methods.isEmpty()) {
+					MethodAccessor accessor = Accessors.getMethodAccessor(methods.get(0));
+					return buf -> accessor.invoke(null, buf);
+				}
+			}
+
+			// try to construct a packet instance using a no-args constructor and invoke the read method
+			MethodAccessor readMethod = MinecraftMethods.getPacketReadByteBufMethod();
+			Objects.requireNonNull(readMethod,
+					"Unable to find the Packet#read(ByteBuf) method, cannot deserialize " + type);
+
+			Object checkInstance = DefaultInstances.DEFAULT.create(type.getPacketClass());
+			Objects.requireNonNull(checkInstance, "Unable to construct empty packet, cannot deserialize " + type);
+
+			// okay, Packet#read exists
+			return buf -> {
+				Object packet = DefaultInstances.DEFAULT.create(type.getPacketClass());
+				readMethod.invoke(packet, buf);
+				return packet;
+			};
+		});
+		return deserializer.apply(buffer);
+	}
+
+	public Object serializeToBuffer() {
+		Object handle = this.getHandle();
+		if (handle == null) {
+			return null;
+		}
+
+		Object targetBuffer = MinecraftReflection.createPacketDataSerializer(0);
+		MinecraftMethods.getPacketWriteByteBufMethod().invoke(handle, targetBuffer);
+		return targetBuffer;
 	}
 
 	// ---- Metadata

src/main/java/com/comphenix/protocol/utility/MinecraftReflection.java

@@ -35,9 +35,13 @@ import com.comphenix.protocol.reflect.FuzzyReflection;
 import com.comphenix.protocol.reflect.accessors.Accessors;
 import com.comphenix.protocol.reflect.accessors.FieldAccessor;
 import com.comphenix.protocol.reflect.accessors.MethodAccessor;
-import com.comphenix.protocol.reflect.fuzzy.*;
+import com.comphenix.protocol.reflect.fuzzy.AbstractFuzzyMatcher;
+import com.comphenix.protocol.reflect.fuzzy.FuzzyClassContract;
+import com.comphenix.protocol.reflect.fuzzy.FuzzyFieldContract;
+import com.comphenix.protocol.reflect.fuzzy.FuzzyMatchers;
+import com.comphenix.protocol.reflect.fuzzy.FuzzyMethodContract;
 import com.comphenix.protocol.wrappers.EnumWrappers;
-
+import io.netty.buffer.Unpooled;
 import org.bukkit.Bukkit;
 import org.bukkit.Material;
 import org.bukkit.Server;
@@ -1434,6 +1438,7 @@ public final class MinecraftReflection {
 	 */
 	public static Object getPacketDataSerializer(Object buffer) {
 		try {
+			// TODO: move this to MinecraftMethods, or at least, cache the constructor accessor
 			Class<?> packetSerializer = getPacketDataSerializerClass();
 			return packetSerializer.getConstructor(getByteBufClass()).newInstance(buffer);
 		} catch (Exception e) {
@@ -1441,6 +1446,16 @@ public final class MinecraftReflection {
 		}
 	}
 
+	public static Object createPacketDataSerializer(int initialSize) {
+		// validate the initial size
+		if (initialSize <= 0) {
+			initialSize = 256;
+		}
+
+		Object buffer = Unpooled.buffer(initialSize);
+		return getPacketDataSerializer(buffer);
+	}
+
 	public static Class<?> getNbtTagTypes() {
 		return getMinecraftClass("nbt.NBTTagTypes", "NBTTagTypes");
 	}

src/test/java/com/comphenix/protocol/events/PacketContainerTest.java

@@ -15,18 +15,16 @@
  */
 package com.comphenix.protocol.events;
 
-import static com.comphenix.protocol.utility.TestUtils.assertItemCollectionsEqual;
+import java.lang.reflect.Array;
-import static com.comphenix.protocol.utility.TestUtils.assertItemsEqual;
+import java.lang.reflect.Field;
-import static com.comphenix.protocol.utility.TestUtils.equivalentItem;
+import java.lang.reflect.Modifier;
-import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import java.util.ArrayList;
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import java.util.List;
-import static org.junit.jupiter.api.Assertions.assertFalse;
+import java.util.Objects;
-import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import java.util.Optional;
-import static org.junit.jupiter.api.Assertions.assertNotSame;
+import java.util.Set;
-import static org.junit.jupiter.api.Assertions.assertNull;
+import java.util.UUID;
-import static org.junit.jupiter.api.Assertions.assertSame;
+import java.util.concurrent.ThreadLocalRandom;
-import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import com.comphenix.protocol.BukkitInitialization;
 import com.comphenix.protocol.PacketType;
@@ -36,6 +34,7 @@ import com.comphenix.protocol.reflect.FuzzyReflection;
 import com.comphenix.protocol.reflect.StructureModifier;
 import com.comphenix.protocol.reflect.accessors.Accessors;
 import com.comphenix.protocol.reflect.accessors.FieldAccessor;
+import com.comphenix.protocol.reflect.cloning.SerializableCloner;
 import com.comphenix.protocol.utility.MinecraftReflection;
 import com.comphenix.protocol.wrappers.BlockPosition;
 import com.comphenix.protocol.wrappers.BukkitConverters;
@@ -55,22 +54,13 @@ import com.comphenix.protocol.wrappers.WrappedDataWatcher.Registry;
 import com.comphenix.protocol.wrappers.WrappedDataWatcher.WrappedDataWatcherObject;
 import com.comphenix.protocol.wrappers.WrappedEnumEntityUseAction;
 import com.comphenix.protocol.wrappers.WrappedGameProfile;
-import com.comphenix.protocol.wrappers.WrappedSaltedSignature;
 import com.comphenix.protocol.wrappers.WrappedRegistry;
+import com.comphenix.protocol.wrappers.WrappedSaltedSignature;
 import com.comphenix.protocol.wrappers.WrappedWatchableObject;
 import com.comphenix.protocol.wrappers.nbt.NbtCompound;
 import com.comphenix.protocol.wrappers.nbt.NbtFactory;
 import com.google.common.collect.Lists;
-import java.lang.reflect.Array;
+import io.netty.buffer.ByteBuf;
-import java.lang.reflect.Field;
-import java.lang.reflect.Modifier;
-import java.time.Instant;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Objects;
-import java.util.Optional;
-import java.util.Set;
-import java.util.UUID;
 import net.md_5.bungee.api.chat.BaseComponent;
 import net.md_5.bungee.api.chat.ClickEvent;
 import net.md_5.bungee.api.chat.ComponentBuilder;
@@ -101,6 +91,19 @@ import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 
+import static com.comphenix.protocol.utility.TestUtils.assertItemCollectionsEqual;
+import static com.comphenix.protocol.utility.TestUtils.assertItemsEqual;
+import static com.comphenix.protocol.utility.TestUtils.equivalentItem;
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertNotSame;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
 public class PacketContainerTest {
 
 	private static BaseComponent[] TEST_COMPONENT;
@@ -404,6 +407,32 @@ public class PacketContainerTest {
 		assertFalse(pos.isInsideBlock());
 	}
 
+	@Test
+	public void testBigPacketSerialization() {
+		PacketContainer payload = new PacketContainer(PacketType.Play.Server.CUSTOM_PAYLOAD);
+		payload.getMinecraftKeys().write(0, new com.comphenix.protocol.wrappers.MinecraftKey("test"));
+
+		byte[] randomData = new byte[8192];
+		ThreadLocalRandom.current().nextBytes(randomData);
+
+		ByteBuf serializer = (ByteBuf) MinecraftReflection.createPacketDataSerializer(randomData.length);
+		serializer.writeBytes(randomData);
+
+		payload.getModifier().withType(MinecraftReflection.getPacketDataSerializerClass()).write(0, serializer);
+
+		PacketContainer cloned = SerializableCloner.clone(payload);
+		com.comphenix.protocol.wrappers.MinecraftKey clonedKey = cloned.getMinecraftKeys().read(0);
+
+		byte[] clonedData = new byte[randomData.length];
+		ByteBuf clonedBuffer = (ByteBuf) cloned.getModifier()
+				.withType(MinecraftReflection.getPacketDataSerializerClass())
+				.read(0);
+		clonedBuffer.readBytes(clonedData);
+
+		assertEquals("minecraft:test", clonedKey.getFullKey());
+		assertArrayEquals(randomData, clonedData);
+	}
+
 	@Test
 	public void testIntList() {
 		PacketContainer destroy = new PacketContainer(PacketType.Play.Server.ENTITY_DESTROY);