feat: enhance Kernel to support memory limits

Author: BSchneppe

core/src/main/java/org/extism/sdk/chicory/Kernel.java

@@ -1,12 +1,14 @@
 package org.extism.sdk.chicory;
 
+import com.dylibso.chicory.runtime.ByteArrayMemory;
 import com.dylibso.chicory.runtime.ExportFunction;
 import com.dylibso.chicory.runtime.HostFunction;
 import com.dylibso.chicory.runtime.Instance;
 import com.dylibso.chicory.runtime.Machine;
 import com.dylibso.chicory.wasm.Parser;
 import com.dylibso.chicory.wasm.WasmModule;
 import com.dylibso.chicory.wasm.types.FunctionType;
+import com.dylibso.chicory.wasm.types.MemoryLimits;
 import com.dylibso.chicory.wasm.types.ValType;
 
 import java.util.List;
@@ -38,11 +40,11 @@ public class Kernel {
     final ExportFunction memoryBytes;
 
     public Kernel() {
-        this(null);
+        this(null, null);
     }
 
-    Kernel(Function<Instance, Machine> machineFactory) {
-        Instance kernel = instance(machineFactory);
+    Kernel(Function<Instance, Machine> machineFactory, MemoryLimits memoryLimits) {
+        Instance kernel = instance(machineFactory, memoryLimits);
         instanceMemory = kernel.memory();
         alloc = kernel.export("alloc");
         free = kernel.export("free");
@@ -66,13 +68,17 @@ public Kernel() {
         memoryBytes = kernel.export("memory_bytes");
     }
 
-    private static Instance instance(Function<Instance, Machine> machineFactory) {
+    private static Instance instance(Function<Instance, Machine> machineFactory, MemoryLimits memoryLimits) {
         var kernelStream = Kernel.class.getClassLoader().getResourceAsStream("extism-runtime.wasm");
         WasmModule module = Parser.parse(kernelStream);
         if (machineFactory != null && machineFactory instanceof CachedAotMachineFactory) {
             ((CachedAotMachineFactory) machineFactory).compile(module);
         }
-        return Instance.builder(module).withMachineFactory(machineFactory).build();
+        var builder = Instance.builder(module).withMachineFactory(machineFactory);
+        if (memoryLimits != null) {
+            builder.withMemoryFactory(ByteArrayMemory::new).withMemoryLimits(memoryLimits);
+        }
+        return builder.build();
     }
 
     public void setInput(byte[] input) {

core/src/main/java/org/extism/sdk/chicory/Linker.java

@@ -57,7 +57,7 @@ Plugin link() {
         httpConfig = options.httpConfig;
 
         // Register the HostEnv exports.
-        var hostEnv = new HostEnv(new Kernel(machineFactory), config, allowedHosts, enableHttpResponseHeaders, httpConfig, logger);
+        var hostEnv = new HostEnv(new Kernel(machineFactory, options.memoryLimits), config, allowedHosts, enableHttpResponseHeaders, httpConfig, logger);
         dg.registerFunctions(hostEnv.toHostFunctions());
 
         // Register the WASI host functions.

core/src/test/java/org/extism/sdk/chicory/e2e/PluginTest.java

@@ -1,9 +1,13 @@
 package org.extism.sdk.chicory.e2e;
 
+import com.dylibso.chicory.runtime.ByteArrayMemory;
+import com.dylibso.chicory.runtime.Instance;
 import com.dylibso.chicory.wasi.WasiOptions;
 import com.dylibso.chicory.wasm.UninstantiableException;
+import com.dylibso.chicory.wasm.types.MemoryLimits;
 import com.google.common.jimfs.Configuration;
 import com.google.common.jimfs.Jimfs;
+import java.lang.reflect.Field;
 import junit.framework.TestCase;
 import org.extism.sdk.chicory.ExtismFunction;
 import org.extism.sdk.chicory.ExtismFunctionException;
@@ -21,6 +25,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import org.extism.sdk.chicory.http.HttpConfig;
 
 public class PluginTest extends TestCase {
 
@@ -72,6 +77,29 @@ public void testGreetFailingMemory() {
         }
     }
 
+    public void testFailingMemoryHog()
+        throws IOException, NoSuchFieldException, IllegalAccessException {
+        byte[] bytes =
+            PluginTest.class.getClassLoader().getResourceAsStream("memory-leaker/memory-leaker.wasm")
+                .readAllBytes();
+        var wasm = ManifestWasm.fromBytes(bytes).build();
+        int memorySize = 1 << 8;
+        var manifest = Manifest.ofWasms(wasm).withOptions(
+            new Manifest.Options().withAoT().withWasi(WasiOptions.builder().build()).withAoT(true)
+                .withHttpConfig(
+                    HttpConfig.builder().withClientAdapter(() -> null).withJsonCodec(() -> null)
+                        .build()).withMemoryLimits(memorySize, memorySize)).build();
+        Plugin plugin = Plugin.ofManifest(manifest).build();
+        try {
+            plugin.call("leak", new byte[0]);
+        } catch (ExtismFunctionException ex) {
+            assertTrue(ex.getMessage().contains("out of memory"));
+            assertEquals(memorySize,plugin.memory().memory().initialPages());
+            assertEquals(memorySize,plugin.memory().memory().maximumPages());
+        }
+
+    }
+
     public void testCountVowels() {
         var url = "https://github.com/extism/plugins/releases/download/v1.1.1/count_vowels.wasm";
         var wasm = ManifestWasm.fromUrl(url).build();