From d8d568daca1308236503deffb5ed3eaf18d18ddd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Florian=20N=C3=BCcke?= <florian@nuecke.de>
Date: Mon, 5 Oct 2020 11:26:27 +0200
Subject: [PATCH] Added utility class for limiting memory allocated for VMs.

---
 src/main/java/li/cil/oc2/Config.java          |  5 +
 .../li/cil/oc2/common/sandbox/Allocator.java  | 99 +++++++++++++++++++
 .../cil/oc2/common/sandbox/package-info.java  |  7 ++
 3 files changed, 111 insertions(+)
 create mode 100644 src/main/java/li/cil/oc2/Config.java
 create mode 100644 src/main/java/li/cil/oc2/common/sandbox/Allocator.java
 create mode 100644 src/main/java/li/cil/oc2/common/sandbox/package-info.java

diff --git a/src/main/java/li/cil/oc2/Config.java b/src/main/java/li/cil/oc2/Config.java
new file mode 100644
index 00000000..bec96830
--- /dev/null
+++ b/src/main/java/li/cil/oc2/Config.java
@@ -0,0 +1,5 @@
+package li.cil.oc2;
+
+public final class Config {
+    public static final long maxAllocatedData = 512 * 1024 * 1024;
+}
diff --git a/src/main/java/li/cil/oc2/common/sandbox/Allocator.java b/src/main/java/li/cil/oc2/common/sandbox/Allocator.java
new file mode 100644
index 00000000..28f9c0fe
--- /dev/null
+++ b/src/main/java/li/cil/oc2/common/sandbox/Allocator.java
@@ -0,0 +1,99 @@
+package li.cil.oc2.common.sandbox;
+
+import li.cil.oc2.Config;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import java.util.HashMap;
+import java.util.UUID;
+
+/**
+ * Cooperative memory allocation limit enforcement.
+ * <p>
+ * Call sites must be cooperative and only free claimed memory when actually being sure the
+ * allocated memory associated with the claim will be garbage collected.
+ */
+public final class Allocator {
+    private static final Logger LOGGER = LogManager.getLogger();
+
+    private static final HashMap<UUID, Allocation> ALLOCATIONS = new HashMap<>();
+    private static long allocated;
+
+    /**
+     * Creates a new handle that can be used to claim memory.
+     *
+     * @return a new handle.
+     */
+    public static UUID createHandle() {
+        return UUID.randomUUID();
+    }
+
+    /**
+     * Tries to claim the specified amount of memory using the specified {@code handle}.
+     * <p>
+     * Claimed memory <em>must</em> be returned using {@link #freeMemory(UUID)} to prevent leaks.
+     *
+     * @param handle the handle to use for claiming memory.
+     * @param size   the amount of memory to claim.
+     * @return {@code true} if the memory was successfully claimed; {@code false} otherwise.
+     */
+    public static boolean claimMemory(final UUID handle, final int size) {
+        if (!checkArgs(handle, size)) {
+            return false;
+        }
+        if (size != 0) {
+            ALLOCATIONS.put(handle, new Allocation(size));
+        }
+        return true;
+    }
+
+    /**
+     * Frees memory that was claimed using the specified handle.
+     * <p>
+     * Using this if there was no memory claimed with ths handle or if the handle has already been
+     * freed does nothing.
+     *
+     * @param handle the handle to release the claimed memory for.
+     */
+    public static void freeMemory(final UUID handle) {
+        final Allocation allocation = ALLOCATIONS.remove(handle);
+        if (allocation != null) {
+            allocated -= allocation.size;
+        }
+    }
+
+    /**
+     * Clears all remaining allocations and logs their stack traces.
+     */
+    public static void resetAndCheckLeaks() {
+        if (allocated > 0) {
+            LOGGER.error("Not all memory was released; leaked allocation stack traces follow.");
+            for (final Allocation allocation : ALLOCATIONS.values()) {
+                LOGGER.error(allocation.stacktrace);
+            }
+        }
+
+        ALLOCATIONS.clear();
+        allocated = 0;
+    }
+
+    private static boolean checkArgs(final UUID handle, final int size) {
+        if (ALLOCATIONS.containsKey(handle)) {
+            throw new IllegalStateException("Handle is already in use. It must be freed before it can be reused.");
+        }
+        if (size < 0) {
+            throw new IllegalArgumentException();
+        }
+        return Config.maxAllocatedData - size >= allocated;
+    }
+
+    private static final class Allocation {
+        public final int size;
+        private final StackTraceElement[] stacktrace;
+
+        private Allocation(final int size) {
+            this.size = size;
+            this.stacktrace = new Throwable().getStackTrace();
+        }
+    }
+}
diff --git a/src/main/java/li/cil/oc2/common/sandbox/package-info.java b/src/main/java/li/cil/oc2/common/sandbox/package-info.java
new file mode 100644
index 00000000..293198af
--- /dev/null
+++ b/src/main/java/li/cil/oc2/common/sandbox/package-info.java
@@ -0,0 +1,7 @@
+@ParametersAreNonnullByDefault
+@MethodsReturnNonnullByDefault
+package li.cil.oc2.common.sandbox;
+
+import mcp.MethodsReturnNonnullByDefault;
+
+import javax.annotation.ParametersAreNonnullByDefault;
\ No newline at end of file