Prevent packets from having dangerously sized collections (Electric Boogaloo)

mapsync-mod/build.gradle.kts

@@ -35,7 +35,7 @@ dependencies {
 
 	libs.voxelmap.also {
 		modCompileOnly(it)
-		//modLocalDep(it) // Uncomment to test VoxelMap
+		modLocalDep(it) // Uncomment to test VoxelMap
 	}
 	libs.journeymap.also {
 		modCompileOnly(it)

mapsync-mod/src/main/java/gjum/minecraft/mapsync/mod/Cartography.java

@@ -3,6 +3,7 @@
 import gjum.minecraft.mapsync.mod.data.BlockColumn;
 import gjum.minecraft.mapsync.mod.data.BlockInfo;
 import gjum.minecraft.mapsync.mod.data.ChunkTile;
+import gjum.minecraft.mapsync.mod.net.buffers.BufferWriter;
 import gjum.minecraft.mapsync.mod.utils.Shortcuts;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
@@ -14,6 +15,7 @@
 import net.minecraft.world.level.LightLayer;
 import net.minecraft.world.level.chunk.LevelChunk;
 import net.minecraft.world.level.levelgen.Heightmap;
+import org.apache.commons.lang3.function.Failable;
 
 public class Cartography {
 	public static ChunkTile chunkTileFromLevel(Level level, int cx, int cz) {
@@ -35,7 +37,7 @@ public static ChunkTile chunkTileFromLevel(Level level, int cx, int cz) {
 		// TODO speedup: don't serialize twice (once here, once later when writing to network)
 		final byte[] dataHash; {
 			final ByteBuf columnsBuf = Unpooled.buffer();
-			ChunkTile.writeColumns(columns, columnsBuf);
+			Failable.run(() -> ChunkTile.writeColumns(columns, new BufferWriter(columnsBuf)));
 			final MessageDigest md = Shortcuts.shaHash();
 			md.update(columnsBuf.nioBuffer());
 			dataHash = md.digest();

mapsync-mod/src/main/java/gjum/minecraft/mapsync/mod/CatchupLogic.java

@@ -71,7 +71,7 @@ synchronized void maybeRequestMoreCatchup() {
 			// if none get received within a second (all outdated etc.) then request more anyway
 			tsRequestMore = now + 1000;
 			var chunksToRequest = pollCatchupChunks(WATERMARK_REQUEST_MORE);
-			getMod().requestCatchupData(chunksToRequest);
+			getMod().requestCatchupData(dimensionState, chunksToRequest);
 		}
 	}
 

mapsync-mod/src/main/java/gjum/minecraft/mapsync/mod/MapSyncMod.java

@@ -13,13 +13,17 @@
 import gjum.minecraft.mapsync.mod.net.packet.ClientboundRegionTimestampsPacket;
 import gjum.minecraft.mapsync.mod.net.packet.ServerboundCatchupRequestPacket;
 import gjum.minecraft.mapsync.mod.net.packet.ServerboundChunkTimestampsRequestPacket;
+import it.unimi.dsi.fastutil.objects.Object2LongArrayMap;
+import it.unimi.dsi.fastutil.objects.Object2LongMap;
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.IdentityHashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.stream.Collectors;
 import net.fabricmc.api.ClientModInitializer;
 import net.fabricmc.fabric.api.client.event.lifecycle.v1.ClientTickEvents;
@@ -31,6 +35,7 @@
 import net.minecraft.network.protocol.game.ClientboundLoginPacket;
 import net.minecraft.network.protocol.game.ClientboundRespawnPacket;
 import net.minecraft.resources.Identifier;
+import net.minecraft.world.level.ChunkPos;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.jetbrains.annotations.NotNull;
@@ -272,23 +277,21 @@ public void handleRegionTimestamps(ClientboundRegionTimestampsPacket packet, Syn
 		if (!dimension.dimension.identifier().toString().equals(packet.getDimension())) {
 			return;
 		}
-		var outdatedRegions = new ArrayList<RegionPos>();
-		for (var regionTs : packet.getTimestamps()) {
-			var regionPos = new RegionPos(regionTs.x(), regionTs.z());
-			long oldestChunkTs = dimension.getOldestChunkTsInRegion(regionPos);
-			boolean requiresUpdate = regionTs.timestamp() > oldestChunkTs;
-
-			debugLog("region " + regionPos
-					+ (requiresUpdate ? " requires update." : " is up to date.")
-					+ " oldest client chunk ts: " + oldestChunkTs
-					+ ", newest server chunk ts: " + regionTs.timestamp());
-
-			if (requiresUpdate) {
-				outdatedRegions.add(regionPos);
-			}
-		}
 
-		client.send(new ServerboundChunkTimestampsRequestPacket(packet.getDimension(), outdatedRegions));
+		var regionTs = packet.getTimestamp();
+
+		var regionPos = new RegionPos(regionTs.x(), regionTs.z());
+		long oldestChunkTs = dimension.getOldestChunkTsInRegion(regionPos);
+		boolean requiresUpdate = regionTs.timestamp() > oldestChunkTs;
+
+		debugLog("region " + regionPos
+				+ (requiresUpdate ? " requires update." : " is up to date.")
+				+ " oldest client chunk ts: " + oldestChunkTs
+				+ ", newest server chunk ts: " + regionTs.timestamp());
+
+		if (requiresUpdate) {
+			client.send(new ServerboundChunkTimestampsRequestPacket(packet.getDimension(), regionPos));
+		}
 	}
 
 	public void handleSharedChunk(ChunkTile chunkTile) {
@@ -305,25 +308,44 @@ public void handleSharedChunk(ChunkTile chunkTile) {
 	public void handleCatchupData(ClientboundChunkTimestampsResponsePacket packet) {
 		var dimensionState = getDimensionState();
 		if (dimensionState == null) return;
-		debugLog("received catchup: " + packet.chunks.size() + " " + packet.chunks.get(0).syncClient.address);
-		dimensionState.addCatchupChunks(packet.chunks);
+		debugLog("received catchup: " + packet.chunks().size() + " " + packet.chunks().get(0).syncClient.address);
+		dimensionState.addCatchupChunks(packet.chunks());
 	}
 
-	public void requestCatchupData(List<CatchupChunk> chunks) {
+	public void requestCatchupData(
+		final @NotNull DimensionState dimensionState,
+		final List<@NotNull CatchupChunk> chunks
+	) {
 		if (chunks == null || chunks.isEmpty()) {
 			debugLog("not requesting more catchup: null/empty");
 			return;
 		}
-
-		debugLog("requesting more catchup: " + chunks.size());
-		var byServer = new HashMap<String, List<CatchupChunk>>();
-		for (CatchupChunk chunk : chunks) {
-			var list = byServer.computeIfAbsent(chunk.syncClient.address, (a) -> new ArrayList<>());
-			list.add(chunk);
+		debugLog("requesting %d more catchup chunks".formatted(
+			chunks.size()
+		));
+		final var catchupChunksBySyncServer = new IdentityHashMap<SyncClient, List<CatchupChunk>>();
+		for (final CatchupChunk chunk : chunks) {
+			catchupChunksBySyncServer
+				.computeIfAbsent(chunk.syncClient, (key) -> new ArrayList<>())
+				.add(chunk);
 		}
-		for (List<CatchupChunk> chunksForServer : byServer.values()) {
-			SyncClient client = chunksForServer.get(0).syncClient;
-			client.send(new ServerboundCatchupRequestPacket(chunksForServer));
+		for (final var byServerEntry : catchupChunksBySyncServer.entrySet()) {
+			final SyncClient syncConnection = byServerEntry.getKey();
+			final Map<RegionPos, Object2LongMap<ChunkPos>> regionChunkRequests = new HashMap<>();
+			for (final CatchupChunk catchupChunk : byServerEntry.getValue()) {
+				regionChunkRequests
+					.computeIfAbsent(RegionPos.forChunkPos(catchupChunk.chunkPos()), (regionPos) -> new Object2LongArrayMap<>())
+					.mergeLong(catchupChunk.chunkPos(), catchupChunk.timestamp(), Math::max);
+			}
+			for (final var byRegionEntry : regionChunkRequests.entrySet()) {
+				final RegionPos regionPos = byRegionEntry.getKey();
+				syncConnection.send(new ServerboundCatchupRequestPacket(
+					dimensionState.dimension.identifier(),
+					(short) regionPos.x(),
+					(short) regionPos.z(),
+					byRegionEntry.getValue()
+				));
+			}
 		}
 	}
 

mapsync-mod/src/main/java/gjum/minecraft/mapsync/mod/data/BlockColumn.java

@@ -2,7 +2,8 @@
 
 import static gjum.minecraft.mapsync.mod.Utils.getBiomeRegistry;
 
-import io.netty.buffer.ByteBuf;
+import gjum.minecraft.mapsync.mod.net.buffers.BufferReader;
+import gjum.minecraft.mapsync.mod.net.buffers.BufferWriter;
 import java.util.ArrayList;
 import java.util.List;
 import net.minecraft.world.level.biome.Biome;
@@ -12,27 +13,27 @@ public record BlockColumn(
 		int light,
 		List<BlockInfo> layers
 ) {
-	public void write(ByteBuf buf) {
-		buf.writeShort(getBiomeRegistry().getId(biome));
-		buf.writeByte(light);
+	public void write(BufferWriter writer) throws Exception {
+		writer.writeUnt16(getBiomeRegistry().getId(biome));
+		writer.writeUnt8(light);
 		// write at most 127 layers, and always include the bottom layer
-		buf.writeByte(Math.min(127, layers.size()));
+		writer.writeUnt8(Math.min(127, layers.size()));
 		int i = 0;
 		for (BlockInfo layer : layers) {
 			if (++i == 127) break;
-			layer.write(buf);
+			layer.write(writer);
 		}
-		if (i == 127) layers.get(layers.size() - 1).write(buf);
+		if (i == 127) layers.getLast().write(writer);
 	}
 
-	public static BlockColumn fromBuf(ByteBuf buf) {
-		int biomeId = buf.readUnsignedShort();
+	public static BlockColumn read(BufferReader reader) throws Exception {
+		int biomeId = reader.readUnt16();
 		Biome biome = getBiomeRegistry().byId(biomeId);
-		int light = buf.readUnsignedByte();
-		int numLayers = buf.readUnsignedByte();
+		int light = reader.readUnt8();
+		int numLayers = reader.readUnt8();
 		var layers = new ArrayList<BlockInfo>(numLayers);
 		for (int i = 0; i < numLayers; i++) {
-			layers.add(BlockInfo.fromBuf(buf));
+			layers.add(BlockInfo.read(reader));
 		}
 		return new BlockColumn(biome, light, layers);
 	}

mapsync-mod/src/main/java/gjum/minecraft/mapsync/mod/data/BlockInfo.java

@@ -1,18 +1,19 @@
 package gjum.minecraft.mapsync.mod.data;
 
-import io.netty.buffer.ByteBuf;
+import gjum.minecraft.mapsync.mod.net.buffers.BufferReader;
+import gjum.minecraft.mapsync.mod.net.buffers.BufferWriter;
 import net.minecraft.world.level.block.Block;
 import net.minecraft.world.level.block.state.BlockState;
 
 public record BlockInfo(int y, BlockState state) {
-	public void write(ByteBuf buf) {
-		buf.writeShort(y);
-		buf.writeShort(Block.BLOCK_STATE_REGISTRY.getId(state)); // we can assume this never becomes large enough to overflow
+	public void write(BufferWriter writer) throws Exception {
+		writer.writeInt16((short) y);
+		writer.writeUnt16(Block.BLOCK_STATE_REGISTRY.getId(state)); // we can assume this never becomes large enough to overflow
 	}
 
-	public static BlockInfo fromBuf(ByteBuf in) {
-		int y = in.readShort();
-		int stateId = in.readUnsignedShort();
+	public static BlockInfo read(BufferReader reader) throws Exception {
+		int y = reader.readInt16();
+		int stateId = reader.readUnt16();
 		return new BlockInfo(y, Block.BLOCK_STATE_REGISTRY.byId(stateId));
 	}
 }

mapsync-mod/src/main/java/gjum/minecraft/mapsync/mod/data/ChunkTile.java

@@ -1,9 +1,9 @@
 package gjum.minecraft.mapsync.mod.data;
 
-import gjum.minecraft.mapsync.mod.net.Packet;
+import gjum.minecraft.mapsync.mod.net.buffers.BufferReader;
+import gjum.minecraft.mapsync.mod.net.buffers.BufferWriter;
 import gjum.minecraft.mapsync.mod.utils.Assertions;
 import gjum.minecraft.mapsync.mod.utils.MagicValues;
-import io.netty.buffer.ByteBuf;
 import net.minecraft.core.registries.Registries;
 import net.minecraft.resources.ResourceKey;
 import net.minecraft.world.level.ChunkPos;
@@ -18,48 +18,48 @@ public record ChunkTile(
 		BlockColumn[] columns
 ) {
 	public ChunkTile {
-		Assertions.assertNotNull("dataHash", dataHash);
-		Assertions.assertLength("dataHash", dataHash.length, MagicValues.SHA1_HASH_LENGTH);
+		Assertions.assertNotNull(dataHash);
+		Assertions.assertLength(dataHash.length, MagicValues.SHA1_HASH_LENGTH);
 	}
 
 	public ChunkPos chunkPos() {
 		return new ChunkPos(x, z);
 	}
 
-	public void write(ByteBuf buf) {
-		writeMetadata(buf);
-		writeColumns(columns, buf);
+	public void write(BufferWriter writer) throws Exception {
+		writeMetadata(writer);
+		writeColumns(columns, writer);
 	}
 
 	/**
 	 * without columns
 	 */
-	public void writeMetadata(ByteBuf buf) {
-		Packet.writeResourceKey(buf, dimension);
-		buf.writeInt(x);
-		buf.writeInt(z);
-		buf.writeLong(timestamp);
-		buf.writeShort(dataVersion);
-		buf.writeBytes(dataHash);
+	public void writeMetadata(BufferWriter writer) throws Exception {
+		writer.writeString(dimension.identifier().toString());
+		writer.writeInt32(x);
+		writer.writeInt32(z);
+		writer.writeInt64(timestamp);
+		writer.writeUnt16(dataVersion);
+		writer.writeBytes(dataHash);
 	}
 
-	public static void writeColumns(BlockColumn[] columns, ByteBuf buf) {
+	public static void writeColumns(BlockColumn[] columns, BufferWriter writer) throws Exception {
 		// TODO compress
 		for (BlockColumn column : columns) {
-			column.write(buf);
+			column.write(writer);
 		}
 	}
 
-	public static ChunkTile fromBuf(ByteBuf buf) {
-		final ResourceKey<Level> dimension = Packet.readResourceKey(buf, Registries.DIMENSION);
-		int x = buf.readInt();
-		int z = buf.readInt();
-		long timestamp = buf.readLong();
-		int dataVersion = buf.readUnsignedShort();
-		byte[] hash = Packet.readByteArrayOfSize(buf, MagicValues.SHA1_HASH_LENGTH);
+	public static ChunkTile read(BufferReader reader) throws Exception {
+		final ResourceKey<Level> dimension = reader.readResourceKey(Registries.DIMENSION);
+		int x = reader.readInt32();
+		int z = reader.readInt32();
+		long timestamp = reader.readInt64();
+		int dataVersion = reader.readUnt16();
+		byte[] hash = reader.readBytesOfLength(MagicValues.SHA1_HASH_LENGTH);
 		var columns = new BlockColumn[256];
 		for (int i = 0; i < 256; i++) {
-			columns[i] = BlockColumn.fromBuf(buf);
+			columns[i] = BlockColumn.read(reader);
 		}
 		return new ChunkTile(dimension, x, z, timestamp, dataVersion, hash, columns);
 	}

mapsync-mod/src/main/java/gjum/minecraft/mapsync/mod/net/ClientHandler.java

@@ -34,7 +34,7 @@ public void channelRead(ChannelHandlerContext ctx, Object packet) {
 			} else if (packet instanceof ClientboundRegionTimestampsPacket pktRegionTimestamps) {
 				getMod().handleRegionTimestamps(pktRegionTimestamps, client);
 			} else if (packet instanceof ClientboundChunkTimestampsResponsePacket pktCatchup) {
-				for (CatchupChunk chunk : pktCatchup.chunks) {
+				for (CatchupChunk chunk : pktCatchup.chunks()) {
 					chunk.syncClient = this.client;
 				}
 				getMod().handleCatchupData((ClientboundChunkTimestampsResponsePacket) packet);

mapsync-mod/src/main/java/gjum/minecraft/mapsync/mod/net/ClientboundPacketDecoder.java

@@ -1,5 +1,6 @@
 package gjum.minecraft.mapsync.mod.net;
 
+import gjum.minecraft.mapsync.mod.net.buffers.BufferReader;
 import gjum.minecraft.mapsync.mod.net.packet.ChunkTilePacket;
 import gjum.minecraft.mapsync.mod.net.packet.ClientboundChunkTimestampsResponsePacket;
 import gjum.minecraft.mapsync.mod.net.packet.ClientboundEncryptionRequestPacket;
@@ -11,19 +12,19 @@
 import org.jetbrains.annotations.Nullable;
 
 public class ClientboundPacketDecoder extends ReplayingDecoder<Void> {
-	public static @Nullable Packet constructServerPacket(int id, ByteBuf buf) {
-		if (id == ChunkTilePacket.PACKET_ID) return ChunkTilePacket.read(buf);
-		if (id == ClientboundEncryptionRequestPacket.PACKET_ID) return ClientboundEncryptionRequestPacket.read(buf);
-		if (id == ClientboundChunkTimestampsResponsePacket.PACKET_ID) return ClientboundChunkTimestampsResponsePacket.read(buf);
-		if (id == ClientboundRegionTimestampsPacket.PACKET_ID) return ClientboundRegionTimestampsPacket.read(buf);
+	public static @Nullable Packet constructServerPacket(int id, BufferReader reader) throws Exception {
+		if (id == ChunkTilePacket.PACKET_ID) return ChunkTilePacket.read(reader);
+		if (id == ClientboundEncryptionRequestPacket.PACKET_ID) return ClientboundEncryptionRequestPacket.read(reader);
+		if (id == ClientboundChunkTimestampsResponsePacket.PACKET_ID) return ClientboundChunkTimestampsResponsePacket.read(reader);
+		if (id == ClientboundRegionTimestampsPacket.PACKET_ID) return ClientboundRegionTimestampsPacket.read(reader);
 		return null;
 	}
 
 	@Override
 	protected void decode(ChannelHandlerContext ctx, ByteBuf buf, List<Object> out) {
 		try {
 			byte id = buf.readByte();
-			final Packet packet = constructServerPacket(id, buf);
+			final Packet packet = constructServerPacket(id, new BufferReader(buf));
 			if (packet == null) {
 				SyncClient.logger.error("[ServerPacketDecoder] " +
 						"Unknown server packet id " + id + " 0x" + Integer.toHexString(id));