From e5ae6cf106ee575fa6ad03c1b82885f4c4fc0111 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sa=C3=BAl=20Ibarra=20Corretg=C3=A9?= <s@saghul.net>
Date: Tue, 16 Apr 2024 16:42:11 +0200
Subject: [PATCH] Fix handling of memory limit

Default to 0, which is "disabled", just like the stack limit.
---
 qjs.c     | 22 ++++++++++++----------
 quickjs.c |  8 +++++---
 quickjs.h |  1 +
 3 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/qjs.c b/qjs.c
index 344fc76..4e1b44a 100644
--- a/qjs.c
+++ b/qjs.c
@@ -199,7 +199,8 @@ static void *js_trace_malloc(JSMallocState *s, size_t size)
     /* Do not allocate zero bytes: behavior is platform dependent */
     assert(size != 0);
 
-    if (unlikely(s->malloc_size + size > s->malloc_limit))
+    /* When malloc_limit is 0 (unlimited), malloc_limit - 1 will be SIZE_MAX. */
+    if (unlikely(s->malloc_size + size > s->malloc_limit - 1))
         return NULL;
     ptr = malloc(size);
     js_trace_malloc_printf(s, "A %zd -> %p\n", size, ptr);
@@ -238,7 +239,8 @@ static void *js_trace_realloc(JSMallocState *s, void *ptr, size_t size)
         free(ptr);
         return NULL;
     }
-    if (s->malloc_size + size - old_size > s->malloc_limit)
+    /* When malloc_limit is 0 (unlimited), malloc_limit - 1 will be SIZE_MAX. */
+    if (s->malloc_size + size - old_size > s->malloc_limit - 1)
         return NULL;
 
     js_trace_malloc_printf(s, "R %zd %p", size, ptr);
@@ -295,10 +297,10 @@ int main(int argc, char **argv)
     int module = -1;
     int load_std = 0;
     int dump_unhandled_promise_rejection = 0;
-    size_t memory_limit = 0;
     char *include_list[32];
     int i, include_count = 0;
-    size_t stack_size = 0;
+    int64_t memory_limit = -1;
+    int64_t stack_size = -1;
 
     argv0 = (JSCFunctionListEntry)JS_PROP_STRING_DEF("argv0", argv[0],
                                                      JS_PROP_C_W_E);
@@ -403,7 +405,7 @@ int main(int argc, char **argv)
                     opt_arg = argv[optind++];
                 }
                 // TODO(chqrlie): accept kmg suffixes
-                memory_limit = (size_t)strtod(opt_arg, NULL);
+                memory_limit = strtoull(opt_arg, NULL, 0);
                 break;
             }
             if (!strcmp(longopt, "stack-size")) {
@@ -415,7 +417,7 @@ int main(int argc, char **argv)
                     opt_arg = argv[optind++];
                 }
                 // TODO(chqrlie): accept kmg suffixes
-                stack_size = (size_t)strtod(opt_arg, NULL);
+                stack_size = strtoull(opt_arg, NULL, 0);
                 break;
             }
             if (opt) {
@@ -437,10 +439,10 @@ int main(int argc, char **argv)
         fprintf(stderr, "qjs: cannot allocate JS runtime\n");
         exit(2);
     }
-    if (memory_limit != 0)
-        JS_SetMemoryLimit(rt, memory_limit);
-    if (stack_size != 0)
-        JS_SetMaxStackSize(rt, stack_size);
+    if (memory_limit >= 0)
+        JS_SetMemoryLimit(rt, (size_t)memory_limit);
+    if (stack_size >= 0)
+        JS_SetMaxStackSize(rt, (size_t)stack_size);
     if (dump_flags != 0)
         JS_SetDumpFlags(rt, dump_flags);
     js_std_set_worker_new_context_func(JS_NewCustomContext);
diff --git a/quickjs.c b/quickjs.c
index 4dc9627..844361b 100644
--- a/quickjs.c
+++ b/quickjs.c
@@ -1608,7 +1608,7 @@ JSRuntime *JS_NewRuntime2(const JSMallocFunctions *mf, void *opaque)
 
     memset(&ms, 0, sizeof(ms));
     ms.opaque = opaque;
-    ms.malloc_limit = -1;
+    ms.malloc_limit = 0;
 
     rt = mf->js_malloc(&ms, sizeof(JSRuntime));
     if (!rt)
@@ -1685,7 +1685,8 @@ static void *js_def_malloc(JSMallocState *s, size_t size)
     /* Do not allocate zero bytes: behavior is platform dependent */
     assert(size != 0);
 
-    if (unlikely(s->malloc_size + size > s->malloc_limit))
+    /* When malloc_limit is 0 (unlimited), malloc_limit - 1 will be SIZE_MAX. */
+    if (unlikely(s->malloc_size + size > s->malloc_limit - 1))
         return NULL;
 
     ptr = malloc(size);
@@ -1723,7 +1724,8 @@ static void *js_def_realloc(JSMallocState *s, void *ptr, size_t size)
         free(ptr);
         return NULL;
     }
-    if (s->malloc_size + size - old_size > s->malloc_limit)
+    /* When malloc_limit is 0 (unlimited), malloc_limit - 1 will be SIZE_MAX. */
+    if (s->malloc_size + size - old_size > s->malloc_limit - 1)
         return NULL;
 
     ptr = realloc(ptr, size);
diff --git a/quickjs.h b/quickjs.h
index 179528c..73f12b0 100644
--- a/quickjs.h
+++ b/quickjs.h
@@ -291,6 +291,7 @@ typedef struct JSGCObjectHeader JSGCObjectHeader;
 JS_EXTERN JSRuntime *JS_NewRuntime(void);
 /* info lifetime must exceed that of rt */
 JS_EXTERN void JS_SetRuntimeInfo(JSRuntime *rt, const char *info);
+/* use 0 to disable memory limit */
 JS_EXTERN void JS_SetMemoryLimit(JSRuntime *rt, size_t limit);
 JS_EXTERN void JS_SetDumpFlags(JSRuntime *rt, uint64_t flags);
 JS_EXTERN void JS_SetGCThreshold(JSRuntime *rt, size_t gc_threshold);