From e8857192d4b15b01de62e6a7591410797e6c09f9 Mon Sep 17 00:00:00 2001
From: Minecraftschurli <minecraftschurli@gmail.com>
Date: Sun, 31 Dec 2023 12:16:23 +0100
Subject: [PATCH] Sanitize registry snapshot

---
 .../payload/FrozenRegistryPayload.java        |  4 +-
 .../neoforge/registries/RegistrySnapshot.java | 88 ++++++++++---------
 2 files changed, 50 insertions(+), 42 deletions(-)

diff --git a/src/main/java/net/neoforged/neoforge/network/payload/FrozenRegistryPayload.java b/src/main/java/net/neoforged/neoforge/network/payload/FrozenRegistryPayload.java
index f91ccc455b..7515e60b29 100644
--- a/src/main/java/net/neoforged/neoforge/network/payload/FrozenRegistryPayload.java
+++ b/src/main/java/net/neoforged/neoforge/network/payload/FrozenRegistryPayload.java
@@ -24,13 +24,13 @@ public record FrozenRegistryPayload(ResourceLocation registryName, RegistrySnaps
     public static final ResourceLocation ID = new ResourceLocation(NeoForgeVersion.MOD_ID, "frozen_registry");
 
     public FrozenRegistryPayload(FriendlyByteBuf buf) {
-        this(buf.readResourceLocation(), RegistrySnapshot.read(buf));
+        this(buf.readResourceLocation(), new RegistrySnapshot(buf));
     }
 
     @Override
     public void write(FriendlyByteBuf buf) {
         buf.writeResourceLocation(registryName());
-        buf.writeBytes(snapshot().getPacketData());
+        snapshot().write(buf);
     }
 
     @Override
diff --git a/src/main/java/net/neoforged/neoforge/registries/RegistrySnapshot.java b/src/main/java/net/neoforged/neoforge/registries/RegistrySnapshot.java
index 06d6d5512d..ce46db76b3 100644
--- a/src/main/java/net/neoforged/neoforge/registries/RegistrySnapshot.java
+++ b/src/main/java/net/neoforged/neoforge/registries/RegistrySnapshot.java
@@ -31,7 +31,7 @@ public class RegistrySnapshot {
     @Nullable
     private final Registry<?> fullBackup;
     @Nullable
-    private FriendlyByteBuf binary = null;
+    private byte[] binary = null;
 
     /**
      * Creates a blank snapshot to populate.
@@ -43,7 +43,7 @@ private RegistrySnapshot() {
     /**
      * Creates a registry snapshot based on the given registry.
      *
-     * @param registry the registry to snapshot
+     * @param registry the registry to snapshot.
      * @param full     if {@code true}, all entries will be stored in this snapshot.
      *                 These entries are never saved to disk nor sent to the client.
      * @param <T>      the registry type
@@ -66,56 +66,64 @@ public <T> RegistrySnapshot(Registry<T> registry, boolean full) {
         }
     }
 
-    public Int2ObjectSortedMap<ResourceLocation> getIds() {
-        return this.idsView;
-    }
-
-    public Map<ResourceLocation, ResourceLocation> getAliases() {
-        return this.aliasesView;
-    }
+    /**
+     * Creates a registry snapshot from the received buffer.
+     * 
+     * @param buf the buffer containing the data of the received snapshot.
+     */
+    public RegistrySnapshot(FriendlyByteBuf buf) {
+        this();
+        int len = buf.readVarInt();
+        for (int x = 0; x < len; x++)
+            this.ids.put(buf.readVarInt(), buf.readResourceLocation());
 
-    @SuppressWarnings("unchecked")
-    @Nullable
-    public <T> Registry<T> getFullBackup() {
-        return (Registry<T>) this.fullBackup;
+        len = buf.readVarInt();
+        for (int x = 0; x < len; x++)
+            this.aliases.put(buf.readResourceLocation(), buf.readResourceLocation());
     }
 
-    public synchronized FriendlyByteBuf getPacketData() {
+    /**
+     * Write the registry snapshot to the given buffer and cache the binary data.
+     * 
+     * @param buf the buffer to write to.
+     */
+    public synchronized void write(FriendlyByteBuf buf) {
         if (this.binary == null) {
             FriendlyByteBuf pkt = new FriendlyByteBuf(Unpooled.buffer());
 
-            pkt.writeVarInt(this.ids.size());
-            this.ids.forEach((k, v) -> {
-                pkt.writeVarInt(k);
-                pkt.writeResourceLocation(v);
-            });
-
-            pkt.writeVarInt(this.aliases.size());
-            this.aliases.forEach((k, v) -> {
-                pkt.writeResourceLocation(k);
-                pkt.writeResourceLocation(v);
-            });
+            try {
+                pkt.writeVarInt(this.ids.size());
+                this.ids.forEach((k, v) -> {
+                    pkt.writeVarInt(k);
+                    pkt.writeResourceLocation(v);
+                });
+
+                pkt.writeVarInt(this.aliases.size());
+                this.aliases.forEach((k, v) -> {
+                    pkt.writeResourceLocation(k);
+                    pkt.writeResourceLocation(v);
+                });
+            } finally {
+                pkt.release();
+            }
 
-            this.binary = pkt;
+            this.binary = pkt.array();
         }
 
-        return new FriendlyByteBuf(this.binary.slice());
+        buf.writeBytes(this.binary);
     }
 
-    public static RegistrySnapshot read(@Nullable FriendlyByteBuf buf) {
-        if (buf == null)
-            return new RegistrySnapshot();
-
-        RegistrySnapshot ret = new RegistrySnapshot();
-
-        int len = buf.readVarInt();
-        for (int x = 0; x < len; x++)
-            ret.ids.put(buf.readVarInt(), buf.readResourceLocation());
+    public Int2ObjectSortedMap<ResourceLocation> getIds() {
+        return this.idsView;
+    }
 
-        len = buf.readVarInt();
-        for (int x = 0; x < len; x++)
-            ret.aliases.put(buf.readResourceLocation(), buf.readResourceLocation());
+    public Map<ResourceLocation, ResourceLocation> getAliases() {
+        return this.aliasesView;
+    }
 
-        return ret;
+    @SuppressWarnings("unchecked")
+    @Nullable
+    public <T> Registry<T> getFullBackup() {
+        return (Registry<T>) this.fullBackup;
     }
 }