fix: stack-depth guard and shadow-stack restore for deep recursion (follow-up to #179)

templates/Makefile.template

@@ -78,7 +78,7 @@ ifdef EMSCRIPTEN
 		-I$(LIBPG_QUERY_DIR)/vendor \
 		-L$(LIBPG_QUERY_DIR) \
 		-sEXPORTED_FUNCTIONS="{{EXPORTED_FUNCTIONS}}" \
-		-sEXPORTED_RUNTIME_METHODS="['lengthBytesUTF8','stringToUTF8','getValue','UTF8ToString','HEAPU8','HEAPU32']" \
+		-sEXPORTED_RUNTIME_METHODS="['lengthBytesUTF8','stringToUTF8','getValue','UTF8ToString','HEAPU8','HEAPU32','stackSave','stackRestore']" \
 		-sEXPORT_NAME="$(WASM_MODULE_NAME)" \
 		-sENVIRONMENT="web,node,worker" \
         -sASSERTIONS=0 \

templates/full/index.ts

@@ -177,9 +177,14 @@ instantiate();
 function callWasm(fn: () => number): number {
   const proc = (globalThis as any).process;
   const exitCode = proc?.exitCode;
+  const module = wasmModule;
+  const stackPointer = module.stackSave();
   try {
     return fn();
   } catch (error) {
+    // The throw unwound the JS frames but not the C shadow stack; give it
+    // back so later sync calls on this instance do not start half-way down it.
+    module.stackRestore(stackPointer);
     failures++;
     stale = true;
     if (++staleFailures >= MAX_STALE_FAILURES) wasmModule = undefined;

templates/full/wasm_wrapper.c

@@ -1,10 +1,21 @@
 #include "pg_query.h"
 #include "protobuf/pg_query.pb-c.h"
 #include <emscripten.h>
+#include <emscripten/stack.h>
 #include <stdlib.h>
 #include <string.h>
 #include <stdio.h>
 
+/* libpg_query's check_stack_depth() needs a reference point and a limit per
+ * entry: the shadow stack left after the previous call, minus a 512 KiB
+ * margin for the frames between here and the recursion. */
+extern void pg_query_set_stack_limit(size_t bytes) __attribute__((weak));
+static void set_stack_guard(void) {
+    if (!pg_query_set_stack_limit) return; /* libpg_query without the guard (upstream pins) */
+    size_t free_bytes = emscripten_stack_get_free();
+    pg_query_set_stack_limit(free_bytes > 524288 ? free_bytes - 524288 : 0);
+}
+
 static int validate_input(const char* input) {
     return input != NULL && strlen(input) > 0;
 }
@@ -28,6 +39,7 @@ static void* safe_malloc(size_t size) {
 
 EMSCRIPTEN_KEEPALIVE
 char* wasm_parse_query(const char* input) {
+    set_stack_guard();
     if (!validate_input(input)) {
         return safe_strdup("Invalid input: query cannot be null or empty");
     }
@@ -47,6 +59,7 @@ char* wasm_parse_query(const char* input) {
 
 EMSCRIPTEN_KEEPALIVE
 PgQueryParseResult* wasm_parse_query_raw(const char* input) {
+    set_stack_guard();
     PgQueryParseResult* result = (PgQueryParseResult*)malloc(sizeof(PgQueryParseResult));
     if (!result) {
         return NULL;
@@ -66,6 +79,7 @@ void wasm_free_parse_result(PgQueryParseResult* result) {
 
 EMSCRIPTEN_KEEPALIVE
 char* wasm_parse_plpgsql(const char* input) {
+    set_stack_guard();
     if (!validate_input(input)) {
         return safe_strdup("Invalid input: query cannot be null or empty");
     }
@@ -106,6 +120,7 @@ char* wasm_parse_plpgsql(const char* input) {
 
 EMSCRIPTEN_KEEPALIVE
 char* wasm_fingerprint(const char* input) {
+    set_stack_guard();
     if (!validate_input(input)) {
         return safe_strdup("Invalid input: query cannot be null or empty");
     }
@@ -125,6 +140,7 @@ char* wasm_fingerprint(const char* input) {
 
 EMSCRIPTEN_KEEPALIVE
 char* wasm_normalize_query(const char* input) {
+    set_stack_guard();
     if (!validate_input(input)) {
         return safe_strdup("Invalid input: query cannot be null or empty");
     }
@@ -165,6 +181,7 @@ typedef struct {
 
 EMSCRIPTEN_KEEPALIVE
 WasmDetailedResult* wasm_parse_query_detailed(const char* input) {
+    set_stack_guard();
     WasmDetailedResult* result = safe_malloc(sizeof(WasmDetailedResult));
     if (!result) {
         return NULL;
@@ -365,6 +382,7 @@ static char* build_scan_json(PgQuery__ScanResult *scan_result, const char* origi
 
 EMSCRIPTEN_KEEPALIVE
 char* wasm_scan(const char* input) {
+    set_stack_guard();
     if (!validate_input(input)) {
         return safe_strdup("Invalid input: query cannot be null or empty");
     }

templates/index.ts

@@ -165,9 +165,14 @@ instantiate();
 function callWasm(fn: () => number): number {
   const proc = (globalThis as any).process;
   const exitCode = proc?.exitCode;
+  const module = wasmModule;
+  const stackPointer = module.stackSave();
   try {
     return fn();
   } catch (error) {
+    // The throw unwound the JS frames but not the C shadow stack; give it
+    // back so later sync calls on this instance do not start half-way down it.
+    module.stackRestore(stackPointer);
     failures++;
     stale = true;
     if (++staleFailures >= MAX_STALE_FAILURES) wasmModule = undefined;

templates/wasm_wrapper.c

@@ -1,8 +1,19 @@
 #include <emscripten.h>
+#include <emscripten/stack.h>
 #include <pg_query.h>
 #include <stdlib.h>
 #include <string.h>
 
+/* libpg_query's check_stack_depth() needs a reference point and a limit per
+ * entry: the shadow stack left after the previous call, minus a 512 KiB
+ * margin for the frames between here and the recursion. */
+extern void pg_query_set_stack_limit(size_t bytes) __attribute__((weak));
+static void set_stack_guard(void) {
+    if (!pg_query_set_stack_limit) return; /* libpg_query without the guard (upstream pins) */
+    size_t free_bytes = emscripten_stack_get_free();
+    pg_query_set_stack_limit(free_bytes > 524288 ? free_bytes - 524288 : 0);
+}
+
 static int validate_input(const char* input) {
     return input != NULL && strlen(input) > 0;
 }
@@ -18,6 +29,7 @@ static void* safe_malloc(size_t size) {
 // Raw struct access functions for parse
 EMSCRIPTEN_KEEPALIVE
 PgQueryParseResult* wasm_parse_query_raw(const char* input) {
+    set_stack_guard();
     if (!validate_input(input)) {
         return NULL;
     }

versions/13/Makefile

@@ -82,7 +82,7 @@ ifdef EMSCRIPTEN
 		-I$(LIBPG_QUERY_DIR)/vendor \
 		-L$(LIBPG_QUERY_DIR) \
 		-sEXPORTED_FUNCTIONS="['_malloc','_free','_wasm_parse_query_raw','_wasm_free_parse_result']" \
-		-sEXPORTED_RUNTIME_METHODS="['lengthBytesUTF8','stringToUTF8','getValue','UTF8ToString','HEAPU8','HEAPU32']" \
+		-sEXPORTED_RUNTIME_METHODS="['lengthBytesUTF8','stringToUTF8','getValue','UTF8ToString','HEAPU8','HEAPU32','stackSave','stackRestore']" \
 		-sEXPORT_NAME="$(WASM_MODULE_NAME)" \
 		-sENVIRONMENT="web,node,worker" \
         -sASSERTIONS=0 \

versions/13/src/index.ts

@@ -172,9 +172,14 @@ instantiate();
 function callWasm(fn: () => number): number {
   const proc = (globalThis as any).process;
   const exitCode = proc?.exitCode;
+  const module = wasmModule;
+  const stackPointer = module.stackSave();
   try {
     return fn();
   } catch (error) {
+    // The throw unwound the JS frames but not the C shadow stack; give it
+    // back so later sync calls on this instance do not start half-way down it.
+    module.stackRestore(stackPointer);
     failures++;
     stale = true;
     if (++staleFailures >= MAX_STALE_FAILURES) wasmModule = undefined;

versions/13/src/wasm_wrapper.c

@@ -6,10 +6,21 @@
  */
 
 #include <emscripten.h>
+#include <emscripten/stack.h>
 #include <pg_query.h>
 #include <stdlib.h>
 #include <string.h>
 
+/* libpg_query's check_stack_depth() needs a reference point and a limit per
+ * entry: the shadow stack left after the previous call, minus a 512 KiB
+ * margin for the frames between here and the recursion. */
+extern void pg_query_set_stack_limit(size_t bytes) __attribute__((weak));
+static void set_stack_guard(void) {
+    if (!pg_query_set_stack_limit) return; /* libpg_query without the guard (upstream pins) */
+    size_t free_bytes = emscripten_stack_get_free();
+    pg_query_set_stack_limit(free_bytes > 524288 ? free_bytes - 524288 : 0);
+}
+
 static int validate_input(const char* input) {
     return input != NULL && strlen(input) > 0;
 }
@@ -25,6 +36,7 @@ static void* safe_malloc(size_t size) {
 // Raw struct access functions for parse
 EMSCRIPTEN_KEEPALIVE
 PgQueryParseResult* wasm_parse_query_raw(const char* input) {
+    set_stack_guard();
     if (!validate_input(input)) {
         return NULL;
     }

versions/14/Makefile

@@ -77,7 +77,7 @@ ifdef EMSCRIPTEN
 		-I$(LIBPG_QUERY_DIR)/vendor \
 		-L$(LIBPG_QUERY_DIR) \
 		-sEXPORTED_FUNCTIONS="['_malloc','_free','_wasm_parse_query_raw','_wasm_free_parse_result']" \
-		-sEXPORTED_RUNTIME_METHODS="['lengthBytesUTF8','stringToUTF8','getValue','UTF8ToString','HEAPU8','HEAPU32']" \
+		-sEXPORTED_RUNTIME_METHODS="['lengthBytesUTF8','stringToUTF8','getValue','UTF8ToString','HEAPU8','HEAPU32','stackSave','stackRestore']" \
 		-sEXPORT_NAME="$(WASM_MODULE_NAME)" \
 		-sENVIRONMENT="web,node,worker" \
         -sASSERTIONS=0 \

versions/14/src/index.ts

@@ -172,9 +172,14 @@ instantiate();
 function callWasm(fn: () => number): number {
   const proc = (globalThis as any).process;
   const exitCode = proc?.exitCode;
+  const module = wasmModule;
+  const stackPointer = module.stackSave();
   try {
     return fn();
   } catch (error) {
+    // The throw unwound the JS frames but not the C shadow stack; give it
+    // back so later sync calls on this instance do not start half-way down it.
+    module.stackRestore(stackPointer);
     failures++;
     stale = true;
     if (++staleFailures >= MAX_STALE_FAILURES) wasmModule = undefined;

versions/14/src/wasm_wrapper.c

@@ -6,10 +6,21 @@
  */
 
 #include <emscripten.h>
+#include <emscripten/stack.h>
 #include <pg_query.h>
 #include <stdlib.h>
 #include <string.h>
 
+/* libpg_query's check_stack_depth() needs a reference point and a limit per
+ * entry: the shadow stack left after the previous call, minus a 512 KiB
+ * margin for the frames between here and the recursion. */
+extern void pg_query_set_stack_limit(size_t bytes) __attribute__((weak));
+static void set_stack_guard(void) {
+    if (!pg_query_set_stack_limit) return; /* libpg_query without the guard (upstream pins) */
+    size_t free_bytes = emscripten_stack_get_free();
+    pg_query_set_stack_limit(free_bytes > 524288 ? free_bytes - 524288 : 0);
+}
+
 static int validate_input(const char* input) {
     return input != NULL && strlen(input) > 0;
 }
@@ -25,6 +36,7 @@ static void* safe_malloc(size_t size) {
 // Raw struct access functions for parse
 EMSCRIPTEN_KEEPALIVE
 PgQueryParseResult* wasm_parse_query_raw(const char* input) {
+    set_stack_guard();
     if (!validate_input(input)) {
         return NULL;
     }

versions/15/Makefile

@@ -77,7 +77,7 @@ ifdef EMSCRIPTEN
 		-I$(LIBPG_QUERY_DIR)/vendor \
 		-L$(LIBPG_QUERY_DIR) \
 		-sEXPORTED_FUNCTIONS="['_malloc','_free','_wasm_parse_query_raw','_wasm_free_parse_result']" \
-		-sEXPORTED_RUNTIME_METHODS="['lengthBytesUTF8','stringToUTF8','getValue','UTF8ToString','HEAPU8','HEAPU32']" \
+		-sEXPORTED_RUNTIME_METHODS="['lengthBytesUTF8','stringToUTF8','getValue','UTF8ToString','HEAPU8','HEAPU32','stackSave','stackRestore']" \
 		-sEXPORT_NAME="$(WASM_MODULE_NAME)" \
 		-sENVIRONMENT="web,node,worker" \
         -sASSERTIONS=0 \

versions/15/src/index.ts

@@ -172,9 +172,14 @@ instantiate();
 function callWasm(fn: () => number): number {
   const proc = (globalThis as any).process;
   const exitCode = proc?.exitCode;
+  const module = wasmModule;
+  const stackPointer = module.stackSave();
   try {
     return fn();
   } catch (error) {
+    // The throw unwound the JS frames but not the C shadow stack; give it
+    // back so later sync calls on this instance do not start half-way down it.
+    module.stackRestore(stackPointer);
     failures++;
     stale = true;
     if (++staleFailures >= MAX_STALE_FAILURES) wasmModule = undefined;

versions/15/src/wasm_wrapper.c

@@ -6,10 +6,21 @@
  */
 
 #include <emscripten.h>
+#include <emscripten/stack.h>
 #include <pg_query.h>
 #include <stdlib.h>
 #include <string.h>
 
+/* libpg_query's check_stack_depth() needs a reference point and a limit per
+ * entry: the shadow stack left after the previous call, minus a 512 KiB
+ * margin for the frames between here and the recursion. */
+extern void pg_query_set_stack_limit(size_t bytes) __attribute__((weak));
+static void set_stack_guard(void) {
+    if (!pg_query_set_stack_limit) return; /* libpg_query without the guard (upstream pins) */
+    size_t free_bytes = emscripten_stack_get_free();
+    pg_query_set_stack_limit(free_bytes > 524288 ? free_bytes - 524288 : 0);
+}
+
 static int validate_input(const char* input) {
     return input != NULL && strlen(input) > 0;
 }
@@ -25,6 +36,7 @@ static void* safe_malloc(size_t size) {
 // Raw struct access functions for parse
 EMSCRIPTEN_KEEPALIVE
 PgQueryParseResult* wasm_parse_query_raw(const char* input) {
+    set_stack_guard();
     if (!validate_input(input)) {
         return NULL;
     }

versions/16/Makefile

@@ -77,7 +77,7 @@ ifdef EMSCRIPTEN
 		-I$(LIBPG_QUERY_DIR)/vendor \
 		-L$(LIBPG_QUERY_DIR) \
 		-sEXPORTED_FUNCTIONS="['_malloc','_free','_wasm_parse_query_raw','_wasm_free_parse_result']" \
-		-sEXPORTED_RUNTIME_METHODS="['lengthBytesUTF8','stringToUTF8','getValue','UTF8ToString','HEAPU8','HEAPU32']" \
+		-sEXPORTED_RUNTIME_METHODS="['lengthBytesUTF8','stringToUTF8','getValue','UTF8ToString','HEAPU8','HEAPU32','stackSave','stackRestore']" \
 		-sEXPORT_NAME="$(WASM_MODULE_NAME)" \
 		-sENVIRONMENT="web,node,worker" \
         -sASSERTIONS=0 \

versions/16/src/index.ts

@@ -172,9 +172,14 @@ instantiate();
 function callWasm(fn: () => number): number {
   const proc = (globalThis as any).process;
   const exitCode = proc?.exitCode;
+  const module = wasmModule;
+  const stackPointer = module.stackSave();
   try {
     return fn();
   } catch (error) {
+    // The throw unwound the JS frames but not the C shadow stack; give it
+    // back so later sync calls on this instance do not start half-way down it.
+    module.stackRestore(stackPointer);
     failures++;
     stale = true;
     if (++staleFailures >= MAX_STALE_FAILURES) wasmModule = undefined;

versions/16/src/wasm_wrapper.c

@@ -6,10 +6,21 @@
  */
 
 #include <emscripten.h>
+#include <emscripten/stack.h>
 #include <pg_query.h>
 #include <stdlib.h>
 #include <string.h>
 
+/* libpg_query's check_stack_depth() needs a reference point and a limit per
+ * entry: the shadow stack left after the previous call, minus a 512 KiB
+ * margin for the frames between here and the recursion. */
+extern void pg_query_set_stack_limit(size_t bytes) __attribute__((weak));
+static void set_stack_guard(void) {
+    if (!pg_query_set_stack_limit) return; /* libpg_query without the guard (upstream pins) */
+    size_t free_bytes = emscripten_stack_get_free();
+    pg_query_set_stack_limit(free_bytes > 524288 ? free_bytes - 524288 : 0);
+}
+
 static int validate_input(const char* input) {
     return input != NULL && strlen(input) > 0;
 }
@@ -25,6 +36,7 @@ static void* safe_malloc(size_t size) {
 // Raw struct access functions for parse
 EMSCRIPTEN_KEEPALIVE
 PgQueryParseResult* wasm_parse_query_raw(const char* input) {
+    set_stack_guard();
     if (!validate_input(input)) {
         return NULL;
     }

versions/17/Makefile

@@ -19,7 +19,7 @@ WASM_STACK_SIZE ?= 33554432
 # byte-for-byte. Set in package.json `x-publish` and applied by copy:templates.
 LIBPG_QUERY_REPO := https://github.com/constructive-io/libpg_query.git
 LIBPG_QUERY_TAG := 17-constructive
-LIBPG_QUERY_COMMIT := 1ade880fd1bbc6e89c33ea2424ddc366c41c1408
+LIBPG_QUERY_COMMIT := ab619b807feea69fdaaa4ea1516236b575bc6d1d
 
 CACHE_DIR := .cache
 
@@ -77,7 +77,7 @@ ifdef EMSCRIPTEN
 		-I$(LIBPG_QUERY_DIR)/vendor \
 		-L$(LIBPG_QUERY_DIR) \
 		-sEXPORTED_FUNCTIONS="['_malloc','_free','_wasm_parse_query_raw','_wasm_free_parse_result']" \
-		-sEXPORTED_RUNTIME_METHODS="['lengthBytesUTF8','stringToUTF8','getValue','UTF8ToString','HEAPU8','HEAPU32']" \
+		-sEXPORTED_RUNTIME_METHODS="['lengthBytesUTF8','stringToUTF8','getValue','UTF8ToString','HEAPU8','HEAPU32','stackSave','stackRestore']" \
 		-sEXPORT_NAME="$(WASM_MODULE_NAME)" \
 		-sENVIRONMENT="web,node,worker" \
         -sASSERTIONS=0 \

versions/17/package.json

@@ -15,7 +15,7 @@
     "distTag": "pg17",
     "libpgQueryRepo": "https://github.com/constructive-io/libpg_query.git",
     "libpgQueryTag": "17-constructive",
-    "libpgQueryCommit": "1ade880fd1bbc6e89c33ea2424ddc366c41c1408"
+    "libpgQueryCommit": "ab619b807feea69fdaaa4ea1516236b575bc6d1d"
   },
   "files": [
     "wasm/*"
@@ -30,7 +30,7 @@
     "wasm:rebuild": "pnpm wasm:make rebuild",
     "wasm:clean": "pnpm wasm:make clean",
     "wasm:clean-cache": "pnpm wasm:make clean-cache",
-    "test": "node --test test/parsing.test.js test/errors.test.js test/input-encoding.test.js test/wasm-recovery.test.js"
+    "test": "node --test test/parsing.test.js test/errors.test.js test/input-encoding.test.js test/wasm-recovery.test.js test/wasm-stale-stack.test.js"
   },
   "author": "Constructive <[email protected]>",
   "license": "MIT",