From 372f6fe26eb2387bebb3f1d1e8277e2a29ce467b Mon Sep 17 00:00:00 2001
From: Marcin Kolny <mkolny@amazon.com>
Date: Mon, 9 Jan 2023 11:07:52 +0000
Subject: [PATCH] Update simple test

 * don't assume the return value of the `__wasi_thread_spawn` API should be zero; assert on thread ID
 * implement wasi_thread_start in ASM to avoid potential stack corruption
---
 test/build.sh                        |  2 +-
 test/testsuite/thread_spawn-simple.c | 41 +++++++++++++++++-----------
 test/testsuite/wasi_thread_spawn.S   | 28 +++++++++++++++++++
 3 files changed, 54 insertions(+), 17 deletions(-)
 create mode 100644 test/testsuite/wasi_thread_spawn.S

diff --git a/test/build.sh b/test/build.sh
index 01188f5..960174e 100755
--- a/test/build.sh
+++ b/test/build.sh
@@ -7,6 +7,6 @@ for input in testsuite/*.c; do
 
   if [ "$input" -nt "$output" ]; then
     echo "Compiling $input"
-    $CC "$input" -o "$output"
+    $CC "$input" testsuite/wasi_thread_spawn.S -o "$output"
   fi
 done
diff --git a/test/testsuite/thread_spawn-simple.c b/test/testsuite/thread_spawn-simple.c
index 671f95b..2df1013 100644
--- a/test/testsuite/thread_spawn-simple.c
+++ b/test/testsuite/thread_spawn-simple.c
@@ -1,9 +1,12 @@
 #include <assert.h>
+#include <stdlib.h>
 #include <wasi/api.h>
 
 static const int64_t SECOND = 1000 * 1000 * 1000;
 
-typedef struct {
+typedef struct
+{
+    char *stack;
     int th_ready;
     int th_continue;
     int th_done;
@@ -12,8 +15,7 @@ typedef struct {
     int value;
 } shared_t;
 
-__attribute__((export_name("wasi_thread_start"))) void
-wasi_thread_start(int thread_id, int *start_arg)
+void __wasi_thread_start_C(int thread_id, int *start_arg)
 {
     shared_t *data = (shared_t *)start_arg;
 
@@ -21,7 +23,8 @@ wasi_thread_start(int thread_id, int *start_arg)
     __builtin_wasm_memory_atomic_notify(&data->th_ready, 1);
 
     // so we can have all the threads alive at the same time
-    if (__builtin_wasm_memory_atomic_wait32(&data->th_continue, 0, SECOND) == 2) {
+    if (__builtin_wasm_memory_atomic_wait32(&data->th_continue, 0, SECOND) == 2)
+    {
         data->failed = 1;
         return;
     }
@@ -35,36 +38,42 @@ wasi_thread_start(int thread_id, int *start_arg)
     __builtin_wasm_memory_atomic_notify(&data->th_done, 1);
 }
 
-int
-main(int argc, char **argv)
+int main(int argc, char **argv)
 {
-    shared_t data[3] = { 0 };
+    shared_t data[3] = {0};
+    int tid[3];
     int data_count = sizeof(data) / sizeof(data[0]);
     int i, j;
 
-    for (i = 0; i < data_count; i++) {
+    for (i = 0; i < data_count; i++)
+    {
+        data[i].stack = malloc(128);
         data[i].value = 52;
-        assert(__wasi_thread_spawn(&data[i]) == 0);
+        tid[i] = __wasi_thread_spawn(&data[i]);
+        assert(tid[i] > 0);
         assert(__builtin_wasm_memory_atomic_wait32(&data[i].th_ready, 0,
-                                                   SECOND)
-               != 2); // not a timeout
+                                                   SECOND) != 2); // not a timeout
     }
 
-    for (i = 0; i < data_count; i++) {
+    for (i = 0; i < data_count; i++)
+    {
         __builtin_wasm_memory_atomic_notify(&data[i].th_continue, 1);
     }
 
-    for (i = 0; i < data_count; i++) {
+    for (i = 0; i < data_count; i++)
+    {
         assert(__builtin_wasm_memory_atomic_wait32(&data[i].th_done, 0,
-                                                   SECOND)
-               != 2); // not a timeout
+                                                   SECOND) != 2); // not a timeout
+        assert(data[i].tid == tid[i]);
         assert(data[i].value == 60);
 
-        for (j = i + 1; j < data_count; j++) {
+        for (j = i + 1; j < data_count; j++)
+        {
             assert(data[i].tid != data[j].tid);
         }
 
         assert(data[i].failed == 0);
+        free(data[i].stack);
     }
 
     return 0;
diff --git a/test/testsuite/wasi_thread_spawn.S b/test/testsuite/wasi_thread_spawn.S
new file mode 100644
index 0000000..f9dd2cc
--- /dev/null
+++ b/test/testsuite/wasi_thread_spawn.S
@@ -0,0 +1,28 @@
+# A copy of the wasi-libc implementation:
+# https://github.com/WebAssembly/wasi-libc/blob/main/libc-top-half/musl/src/thread/wasm32/wasi_thread_start.s
+	.text
+
+	.export_name	wasi_thread_start, wasi_thread_start
+
+	.globaltype	__stack_pointer, i32
+	.functype	__wasi_thread_start_C (i32, i32) -> ()
+
+	.hidden	wasi_thread_start
+	.globl	wasi_thread_start
+	.type	wasi_thread_start,@function
+
+wasi_thread_start:
+	.functype	wasi_thread_start (i32, i32) -> ()
+
+	# Set up the minimum C environment.
+	# Note: offsetof(start_arg, stack) == 0
+	local.get   1  # start_arg
+	i32.load    0  # stack
+	global.set  __stack_pointer
+
+	# Make the C function do the rest of work.
+	local.get   0  # tid
+	local.get   1  # start_arg
+	call __wasi_thread_start_C
+
+	end_function
\ No newline at end of file