Prevent packets from having dangerously sized collections

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

@@ -11,6 +11,7 @@
 import net.minecraft.client.multiplayer.ServerData;
 import net.minecraft.network.protocol.game.ClientboundLoginPacket;
 import net.minecraft.network.protocol.game.ClientboundRespawnPacket;
+import net.minecraft.world.level.ChunkPos;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.jetbrains.annotations.NotNull;
@@ -225,23 +226,21 @@ public void handleRegionTimestamps(ClientboundRegionTimestampsPacket packet, Syn
 		if (!dimension.dimension.location().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) {
@@ -262,6 +261,7 @@ public void handleCatchupData(ClientboundChunkTimestampsResponsePacket packet) {
 		dimensionState.addCatchupChunks(packet.chunks);
 	}
 
+	private record Pos2D(int x, int z) {}
 	public void requestCatchupData(List<CatchupChunk> chunks) {
 		if (chunks == null || chunks.isEmpty()) {
 			debugLog("not requesting more catchup: null/empty");
@@ -275,8 +275,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.get(0).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())
+				));
+			}
 		}
 	}
 

mod/common/src/main/java/gjum/minecraft/mapsync/common/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()
+      )
+    );
   }
 }

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

@@ -5,8 +5,6 @@
 import io.netty.buffer.ByteBuf;
 import org.jetbrains.annotations.NotNull;
 
-import java.util.List;
-
 /**
  * You send this in response to a {@link ClientboundRegionTimestampsPacket},
  * listing all the regions you'd like the server to elaborate on. You should
@@ -16,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());
   }
 }

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,24 @@ 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 << 4,
+        maxChunkX = minChunkX + 16;
+    const minChunkZ = regionZ << 4,
+        maxChunkZ = minChunkZ + 16;
     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)
+        .orderBy("ts", "desc")
         .execute();
 }
 

server/src/main.ts

@@ -6,6 +6,7 @@ import { CatchupRequestPacket } from "./protocol/CatchupRequestPacket";
 import { ChunkTilePacket } from "./protocol/ChunkTilePacket";
 import { TcpClient, TcpServer } from "./server";
 import { RegionCatchupPacket } from "./protocol/RegionCatchupPacket";
+import { CatchupChunk, RegionChunkTimestamps } from "./model";
 
 let config: metadata.Config = null!;
 Promise.resolve().then(async () => {
@@ -48,11 +49,15 @@ export class Main {
         // TODO check version, mc server, user access
 
         const timestamps = await database.getRegionTimestamps(client.world!);
-        client.send({
-            type: "RegionTimestamps",
-            world: client.world!,
-            regions: timestamps,
-        });
+        await Promise.allSettled(
+            timestamps.map((chunk) =>
+                client.send({
+                    type: "RegionTimestamps",
+                    world: client.world!,
+                    region: chunk,
+                }),
+            ),
+        );
     }
 
     handleClientDisconnected(client: ProtocolClient) {}
@@ -151,7 +156,8 @@ export class Main {
 
         const chunks = await database.getChunkTimestamps(
             pkt.world,
-            pkt.regions,
+            pkt.regionX,
+            pkt.regionZ
         );
         if (chunks.length)
             client.send({ type: "Catchup", world: pkt.world, chunks });

server/src/model.ts

@@ -14,3 +14,9 @@ export interface Pos2D {
     readonly x: number;
     readonly z: number;
 }
+
+export interface RegionChunkTimestamps {
+    readonly regionX: number;
+    readonly regionZ: number;
+    readonly chunks: Array<CatchupChunk>;
+}

server/src/protocol/RegionCatchupPacket.ts

@@ -1,23 +1,19 @@
 import { BufReader } from "./BufReader";
-import { type Pos2D } from "../model";
 
 export interface RegionCatchupPacket {
     type: "RegionCatchup";
     world: string;
-    regions: Pos2D[];
+    regionX: number;
+    regionZ: number;
 }
 
 export namespace RegionCatchupPacket {
     export function decode(reader: BufReader): RegionCatchupPacket {
-        let world = reader.readString();
-        const len = reader.readInt16();
-        const regions: Pos2D[] = [];
-        for (let i = 0; i < len; i++) {
-            regions.push({
-                x: reader.readInt16(),
-                z: reader.readInt16(),
-            });
-        }
-        return { type: "RegionCatchup", world, regions };
+        return {
+            type: "RegionCatchup",
+            world: reader.readString(),
+            regionX: reader.readInt16(),
+            regionZ: reader.readInt16(),
+        };
     }
 }

server/src/protocol/RegionTimestampsPacket.ts

@@ -4,19 +4,15 @@ import { CatchupRegion } from "../model";
 export interface RegionTimestampsPacket {
     type: "RegionTimestamps";
     world: string;
-    regions: Array<CatchupRegion>;
+    region: CatchupRegion;
 }
 
 export namespace RegionTimestampsPacket {
-    export function encode(pkt: RegionTimestampsPacket, writer: BufWriter) {
-        writer.writeString(pkt.world);
-        writer.writeInt16(pkt.regions.length);
-        console.log("Sending regions " + JSON.stringify(pkt.regions));
-        for (let i = 0; i < pkt.regions.length; i++) {
-            let region = pkt.regions[i];
-            writer.writeInt16(region.regionX);
-            writer.writeInt16(region.regionZ);
-            writer.writeInt64(region.timestamp);
-        }
+    export function encode(packet: RegionTimestampsPacket, writer: BufWriter) {
+        writer.writeString(packet.world);
+        console.log(`Sending region for [${packet.world}]`, packet);
+        writer.writeInt16(packet.region.regionX);
+        writer.writeInt16(packet.region.regionZ);
+        writer.writeInt64(packet.region.timestamp);
     }
 }

server/src/server.ts

@@ -108,12 +108,10 @@ export class TcpClient {
 
                     // prevent Out of Memory
                     if (frameSize > this.maxFrameSize) {
-                        return this.kick(
-                            "Frame too large: " +
-                                frameSize +
-                                " have " +
-                                accBuf.length,
+                        this.kick(
+                            `Frame's length [${frameSize}] is too large, max is [${this.maxFrameSize}]!`,
                         );
+                        return;
                     }
 
                     if (accBuf.length < 4 + frameSize) return; // wait for more data