Prevent packets from having dangerously sized collections (Electric Boogaloo)

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

@@ -20,6 +20,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 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 +32,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 +274,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) {
@@ -322,8 +322,27 @@ public void requestCatchupData(List<CatchupChunk> chunks) {
 			list.add(chunk);
 		}
 		for (List<CatchupChunk> chunksForServer : byServer.values()) {
-			SyncClient client = chunksForServer.get(0).syncClient;
-			client.send(new ServerboundCatchupRequestPacket(chunksForServer));
+			final SyncClient client = chunksForServer.getFirst().syncClient;
+
+			final Map<RegionPos, Map<ChunkPos, CatchupChunk>> regionChunkRequests = new HashMap<>();
+			for (final CatchupChunk chunk : chunksForServer) {
+				final ChunkPos chunkPos = chunk.chunkPos();
+                final Map<ChunkPos, CatchupChunk> regionChunks = regionChunkRequests.computeIfAbsent(
+					RegionPos.forChunkPos(chunkPos),
+					(regionPos) -> new HashMap<>()
+				);
+				final CatchupChunk existingCatchup = regionChunks.get(chunkPos);
+				if (existingCatchup != null && existingCatchup.timestamp() > chunk.timestamp()) {
+					continue;
+				}
+				regionChunks.put(chunkPos, chunk);
+			}
+
+			for (final Map<ChunkPos, CatchupChunk> catchupChunks : regionChunkRequests.values()) {
+				client.send(new ServerboundCatchupRequestPacket(
+					new ArrayList<>(catchupChunks.values())
+				));
+			}
 		}
 	}
 

mapsync-mod/src/main/java/gjum/minecraft/mapsync/mod/net/packet/ClientboundRegionTimestampsPacket.java

@@ -14,35 +14,29 @@ public class ClientboundRegionTimestampsPacket implements Packet {
 
 	private final String dimension;
 
-	private final RegionTimestamp[] timestamps;
+	private final RegionTimestamp timestamp;
 
-	public ClientboundRegionTimestampsPacket(String dimension, RegionTimestamp[] timestamps) {
+	public ClientboundRegionTimestampsPacket(String dimension, RegionTimestamp timestamp) {
 		this.dimension = dimension;
-		this.timestamps = timestamps;
+		this.timestamp = timestamp;
 	}
 
 	public String getDimension() {
 		return dimension;
 	}
 
-	public RegionTimestamp[] getTimestamps() {
-		return timestamps;
+	public RegionTimestamp getTimestamp() {
+		return timestamp;
 	}
 
 	public static Packet read(ByteBuf buf) {
-		String dimension = Packet.readUtf8String(buf);
-
-		short totalRegions = buf.readShort();
-		RegionTimestamp[] timestamps = new RegionTimestamp[totalRegions];
-		// row = x
-		for (short i = 0; i < totalRegions; i++) {
-			short regionX = buf.readShort();
-			short regionZ = buf.readShort();
-
-			long timestamp = buf.readLong();
-			timestamps[i] = new RegionTimestamp(regionX, regionZ, timestamp);
-		}
-
-		return new ClientboundRegionTimestampsPacket(dimension, timestamps);
+		return new ClientboundRegionTimestampsPacket(
+			Packet.readUtf8String(buf),
+			new RegionTimestamp(
+				buf.readShort(),
+				buf.readShort(),
+				buf.readLong()
+			)
+		);
 	}
 }

mapsync-mod/src/main/java/gjum/minecraft/mapsync/mod/net/packet/ServerboundChunkTimestampsRequestPacket.java

@@ -3,7 +3,6 @@
 import gjum.minecraft.mapsync.mod.data.RegionPos;
 import gjum.minecraft.mapsync.mod.net.Packet;
 import io.netty.buffer.ByteBuf;
-import java.util.List;
 import org.jetbrains.annotations.NotNull;
 
 /**
@@ -15,20 +14,17 @@ public class ServerboundChunkTimestampsRequestPacket implements Packet {
 	public static final int PACKET_ID = 8;
 
 	private final String dimension;
-	private final List<RegionPos> regions;
+	private final RegionPos region;
 
-	public ServerboundChunkTimestampsRequestPacket(String dimension, List<RegionPos> regions) {
+	public ServerboundChunkTimestampsRequestPacket(String dimension, RegionPos region) {
 		this.dimension = dimension;
-		this.regions = regions;
+		this.region = region;
 	}
 
 	@Override
 	public void write(@NotNull ByteBuf buf) {
 		Packet.writeUtf8String(buf, dimension);
-		buf.writeShort(regions.size());
-		for (var region : regions) {
-			buf.writeShort(region.x());
-			buf.writeShort(region.z());
-		}
+		buf.writeShort(region.x());
+		buf.writeShort(region.z());
 	}
 }

mapsync-mod/src/main/java/gjum/minecraft/mapsync/mod/net/packet/ServerboundHandshakePacket.java

@@ -14,21 +14,21 @@ public class ServerboundHandshakePacket implements Packet {
 	public final @NotNull String modVersion;
 	public final @NotNull String username;
 	public final @NotNull String gameAddress;
-	public final @NotNull String world;
+	public final @NotNull String dimension;
 
-	public ServerboundHandshakePacket(@NotNull String modVersion, @NotNull String username, @NotNull String gameAddress, @NotNull String world) {
+	public ServerboundHandshakePacket(@NotNull String modVersion, @NotNull String username, @NotNull String gameAddress, @NotNull String dimension) {
 		this.modVersion = modVersion;
 		this.username = username;
 		this.gameAddress = gameAddress;
-		this.world = world;
+		this.dimension = dimension;
 	}
 
 	@Override
 	public void write(@NotNull ByteBuf out) {
 		Packet.writeUtf8String(out, modVersion);
 		Packet.writeUtf8String(out, username);
 		Packet.writeUtf8String(out, gameAddress);
-		Packet.writeUtf8String(out, world);
+		Packet.writeUtf8String(out, dimension);
 	}
 
 	@Override

mapsync-server/src/cli.ts

@@ -5,6 +5,7 @@ import * as metadata from "./metadata";
 import { TcpServer } from "./server";
 
 import * as database from "./database";
+import { ClientboundRegionTimestampsPacket } from "./protocol";
 
 //idk where these come from lol
 interface TerminalExtras {
@@ -143,18 +144,25 @@ async function handle_input(input: string): Promise<void> {
             return;
         }
 
-        const world = client.world;
-        if (!world) {
+        const dimension = client.dimension;
+        if (!dimension) {
             console.log("Client has no world yet");
             return;
         }
 
-        const regions = await database.getRegionTimestamps(world);
-        await client.send({
-            type: "RegionTimestamps",
-            world,
-            regions,
-        });
+        const regions = await database.getRegionTimestamps(client.dimension!);
+        await Promise.allSettled(
+            regions.map((region) =>
+                client.send(
+                    new ClientboundRegionTimestampsPacket(
+                        client.dimension!,
+                        region.regionX,
+                        region.regionZ,
+                        region.timestamp,
+                    ),
+                ),
+            ),
+        );
     } else if (command === "kick") {
         const target = extras.trim(); // IGN or UUID
 

mapsync-server/src/database.ts

@@ -1,7 +1,6 @@
 import * as kysely from "kysely";
 import sqlite from "better-sqlite3";
 import { DATA_FOLDER } from "./metadata";
-import { type Pos2D } from "./model";
 
 let database: kysely.Kysely<Database> | null = null;
 
@@ -98,33 +97,29 @@ export function getRegionTimestamps(dimension: string) {
 /**
  * Converts an array of region coords into an array of timestamped chunk coords.
  */
-export async function getChunkTimestamps(dimension: string, regions: Pos2D[]) {
+export async function getChunkTimestamps(
+    dimension: string,
+    regionX: number,
+    regionZ: number,
+) {
+    const minChunkX = regionX << 5,
+        maxChunkX = minChunkX + 32;
+    const minChunkZ = regionZ << 5,
+        maxChunkZ = minChunkZ + 32;
     return get()
-        .with("regions", (db) =>
-            db
-                .selectFrom("player_chunk")
-                .select([
-                    (eb) =>
-                        kysely.sql<string>`(cast(floor(${eb.ref(
-                            "chunk_x",
-                        )} / 32.0) as int) || '_' || cast(floor(${eb.ref(
-                            "chunk_z",
-                        )} / 32.0) as int))`.as("region"),
-                    "chunk_x as x",
-                    "chunk_z as z",
-                    (eb) => eb.fn.max("ts").as("timestamp"),
-                ])
-                .where("world", "=", dimension)
-                .groupBy(["x", "z"]),
-        )
-        .selectFrom("regions")
-        .select(["x as chunkX", "z as chunkZ", "timestamp"])
-        .where(
-            "region",
-            "in",
-            regions.map((region) => region.x + "_" + region.z),
-        )
-        .orderBy("timestamp", "desc")
+        .selectFrom("player_chunk")
+        .select([
+            "chunk_x as chunkX",
+            "chunk_z as chunkZ",
+            (eb) => eb.fn.max("ts").as("timestamp"),
+        ])
+        .where("world", "=", dimension)
+        .where("chunk_x", ">=", minChunkX)
+        .where("chunk_x", "<", maxChunkX)
+        .where("chunk_z", ">=", minChunkZ)
+        .where("chunk_z", "<", maxChunkZ)
+        .groupBy(["chunk_x", "chunk_z"])
+        .orderBy("ts", "desc")
         .execute();
 }