From e36ba049216bff011ad17508552113a9f782ceb9 Mon Sep 17 00:00:00 2001
From: Moulberry <james.jenour@protonmail.com>
Date: Sat, 28 Oct 2023 09:44:12 -0400
Subject: [PATCH] Avoid mods causing double free when BufferBuilder leak fix is
 enabled

---
 .../BufferBuilderMixin.java                   |  5 +-
 .../modernfix/render/UnsafeBufferHelper.java  | 47 +++++++++++++++++++
 2 files changed, 49 insertions(+), 3 deletions(-)
 create mode 100644 common/src/main/java/org/embeddedt/modernfix/render/UnsafeBufferHelper.java

diff --git a/common/src/main/java/org/embeddedt/modernfix/common/mixin/bugfix/buffer_builder_leak/BufferBuilderMixin.java b/common/src/main/java/org/embeddedt/modernfix/common/mixin/bugfix/buffer_builder_leak/BufferBuilderMixin.java
index 3a6e190b..7249216d 100644
--- a/common/src/main/java/org/embeddedt/modernfix/common/mixin/bugfix/buffer_builder_leak/BufferBuilderMixin.java
+++ b/common/src/main/java/org/embeddedt/modernfix/common/mixin/bugfix/buffer_builder_leak/BufferBuilderMixin.java
@@ -2,7 +2,7 @@ package org.embeddedt.modernfix.common.mixin.bugfix.buffer_builder_leak;
 
 import com.mojang.blaze3d.vertex.BufferBuilder;
 import org.embeddedt.modernfix.ModernFix;
-import org.lwjgl.system.MemoryUtil;
+import org.embeddedt.modernfix.render.UnsafeBufferHelper;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
 
@@ -12,7 +12,6 @@ import java.nio.ByteBuffer;
 public class BufferBuilderMixin {
     @Shadow private ByteBuffer buffer;
 
-    private static final MemoryUtil.MemoryAllocator ALLOCATOR = MemoryUtil.getAllocator(false);
     private static boolean leakReported = false;
 
     @Override
@@ -25,7 +24,7 @@ public class BufferBuilderMixin {
                     leakReported = true;
                     ModernFix.LOGGER.warn("One or more BufferBuilders have been leaked, ModernFix will attempt to correct this.");
                 }
-                ALLOCATOR.free(MemoryUtil.memAddress0(buf));
+                UnsafeBufferHelper.free(buf);
                 buffer = null;
             }
         } finally {
diff --git a/common/src/main/java/org/embeddedt/modernfix/render/UnsafeBufferHelper.java b/common/src/main/java/org/embeddedt/modernfix/render/UnsafeBufferHelper.java
new file mode 100644
index 00000000..acfeee83
--- /dev/null
+++ b/common/src/main/java/org/embeddedt/modernfix/render/UnsafeBufferHelper.java
@@ -0,0 +1,47 @@
+package org.embeddedt.modernfix.render;
+
+import org.embeddedt.modernfix.ModernFix;
+import org.lwjgl.system.MemoryUtil;
+import sun.misc.Unsafe;
+
+import java.lang.reflect.Field;
+import java.nio.ByteBuffer;
+
+/**
+ * Helper that frees ByteBuffers allocated by BufferBuilders, and nulls out the address pointer
+ * to prevent double frees.
+ *
+ * @author Moulberry
+ */
+public class UnsafeBufferHelper {
+    private static final MemoryUtil.MemoryAllocator ALLOCATOR = MemoryUtil.getAllocator(false);
+
+    private static sun.misc.Unsafe UNSAFE = null;
+    private static long ADDRESS = -1;
+
+    static {
+        try {
+            final Field theUnsafe = Unsafe.class.getDeclaredField("theUnsafe");
+            theUnsafe.setAccessible(true);
+            UNSAFE = (Unsafe)theUnsafe.get(null);
+
+            final Field addressField = MemoryUtil.class.getDeclaredField("ADDRESS");
+            addressField.setAccessible(true);
+            ADDRESS = addressField.getLong(null);
+        } catch(Throwable t) {
+            ModernFix.LOGGER.error("Could load unsafe/buffer address", t);
+        }
+    }
+
+    public static void free(ByteBuffer buf) {
+        if(UNSAFE != null && ADDRESS >= 0) {
+            // set the address to 0 to prevent double free
+            long address = UNSAFE.getAndSetLong(buf, ADDRESS, 0);
+            if(address != 0) {
+                ALLOCATOR.free(address);
+            }
+        } else {
+            ALLOCATOR.free(MemoryUtil.memAddress0(buf));
+        }
+    }
+}