Remove size limit for WASI random_get

dylibso · Sep 6, 2024 · 5eb6670 · 5eb6670
1 parent c05fd30
commit 5eb6670
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 7 deletions.
diff --git a/wasi/README.md b/wasi/README.md
@@ -73,7 +73,7 @@ We also have a table:
 | poll_oneoff             | ❌         |                                                                                                                                                                                        |
 | proc_exit               | ✅         |                                                                                                                                                                                        |
 | proc_raise              | 💀        |                                                                                                                                                                                        |
-| random_get              | ✅         | Max size here is `100_000` bytes                                                                                                                                                       |
+| random_get              | ✅         |                                                                                                                                                                                        |
 | reset_adapter_state     | ❌         |                                                                                                                                                                                        |
 | sched_yield             | ✅         |                                                                                                                                                                                        |
 | set_allocation_state    | ❌         |                                                                                                                                                                                        |

diff --git a/wasi/src/main/java/com/dylibso/chicory/wasi/WasiPreview1.java b/wasi/src/main/java/com/dylibso/chicory/wasi/WasiPreview1.java
@@ -23,6 +23,7 @@
 import com.dylibso.chicory.wasi.Descriptors.InStream;
 import com.dylibso.chicory.wasi.Descriptors.OutStream;
 import com.dylibso.chicory.wasi.Descriptors.PreopenedDirectory;
+import com.dylibso.chicory.wasm.exceptions.ChicoryException;
 import java.io.Closeable;
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -1079,13 +1080,21 @@ public int randomGet(Memory memory, int buf, int bufLen) {
         if (bufLen < 0) {
             return wasiResult(WasiErrno.EINVAL);
         }
-        if (bufLen >= 100_000) {
-            throw new WASMRuntimeException("random_get: bufLen must be < 100_000");
-        }
 
-        byte[] data = new byte[bufLen];
-        random.nextBytes(data);
-        memory.write(buf, data);
+        byte[] data = new byte[min(bufLen, 4096)];
+        int written = 0;
+        while (written < bufLen) {
+            if (Thread.currentThread().isInterrupted()) {
+                throw new ChicoryException("Thread interrupted");
+            }
+            int size = min(data.length, bufLen - written);
+            if (size < data.length) {
+                data = new byte[size];
+            }
+            random.nextBytes(data);
+            memory.write(buf + written, data, 0, size);
+            written += size;
+        }
         return wasiResult(WasiErrno.ESUCCESS);
     }
 

diff --git a/wasi/src/test/java/com/dylibso/chicory/wasi/WasiPreview1Test.java b/wasi/src/test/java/com/dylibso/chicory/wasi/WasiPreview1Test.java
@@ -1,18 +1,22 @@
 package com.dylibso.chicory.wasi;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import com.dylibso.chicory.log.Logger;
 import com.dylibso.chicory.log.SystemLogger;
 import com.dylibso.chicory.runtime.HostImports;
 import com.dylibso.chicory.runtime.Instance;
+import com.dylibso.chicory.runtime.Memory;
 import com.dylibso.chicory.wasm.Module;
 import com.dylibso.chicory.wasm.Parser;
+import com.dylibso.chicory.wasm.types.MemoryLimits;
 import com.dylibso.chicory.wasm.types.Value;
 import java.io.ByteArrayInputStream;
 import java.util.List;
+import java.util.Random;
 import org.junit.jupiter.api.Test;
 
 public class WasiPreview1Test {
@@ -123,4 +127,25 @@ public void shouldRunWasiDemoDotnetModule() throws Exception {
 
         assertEquals("Hello, Wasi Console!\n", fakeStdout.output());
     }
+
+    @Test
+    public void wasiRandom() {
+        var seed = 0x12345678;
+        var wasi =
+                new WasiPreview1(
+                        this.logger, WasiOptions.builder().withRandom(new Random(seed)).build());
+
+        var memory = new Memory(new MemoryLimits(8, 8));
+        assertEquals(0, wasi.randomGet(memory, 0, 123_456));
+        assertEquals(0, wasi.randomGet(memory, 222_222, 87_654));
+
+        var random = new Random(seed);
+        byte[] first = new byte[123_456];
+        random.nextBytes(first);
+        byte[] second = new byte[87_654];
+        random.nextBytes(second);
+
+        assertArrayEquals(first, memory.readBytes(0, 123_456));
+        assertArrayEquals(second, memory.readBytes(222_222, 87_654));
+    }
 }