From 59a6ca5d09b123b9849377c8ed4f10bada3849c5 Mon Sep 17 00:00:00 2001
From: Joyee Cheung <joyeec9h3@gmail.com>
Date: Fri, 23 Jun 2023 18:32:57 +0200
Subject: [PATCH] module: fix the leak in SourceTextModule and ContextifySript

This is a POC that shows how we can replace the persistent
handles to v8::Module and v8::UnboundScript with an internal
reference and fix the leaks.
---
 src/module_wrap.cc                            |  9 ++++---
 src/module_wrap.h                             |  3 ++-
 src/node_contextify.cc                        |  4 ++++
 src/node_contextify.h                         |  5 ++++
 .../test-vm-contextified-script-leak.js       | 17 +++++++++++++
 .../test-vm-source-text-module-leak.js        | 24 +++++++++++++++++++
 6 files changed, 58 insertions(+), 4 deletions(-)
 create mode 100644 test/es-module/test-vm-contextified-script-leak.js
 create mode 100644 test/es-module/test-vm-source-text-module-leak.js

diff --git a/src/module_wrap.cc b/src/module_wrap.cc
index 6f9fc3286deacc..94c773b63e9528 100644
--- a/src/module_wrap.cc
+++ b/src/module_wrap.cc
@@ -53,18 +53,21 @@ ModuleWrap::ModuleWrap(Environment* env,
                        Local<Object> object,
                        Local<Module> module,
                        Local<String> url)
-    : BaseObject(env, object), module_(env->isolate(), module) {
+    : BaseObject(env, object),
+      module_(env->isolate(), module),
+      module_hash_(module->GetIdentityHash()) {
   Local<Value> undefined = Undefined(env->isolate());
+  object->SetDataInInternalField(kModuleSlot, module);
   object->SetInternalField(kURLSlot, url);
   object->SetInternalField(kSyntheticEvaluationStepsSlot, undefined);
   object->SetInternalField(kContextObjectSlot, undefined);
   MakeWeak();
+  module_.SetWeak();
 }
 
 ModuleWrap::~ModuleWrap() {
   HandleScope scope(env()->isolate());
-  Local<Module> module = module_.Get(env()->isolate());
-  auto range = env()->hash_to_module_map.equal_range(module->GetIdentityHash());
+  auto range = env()->hash_to_module_map.equal_range(module_hash_);
   for (auto it = range.first; it != range.second; ++it) {
     if (it->second == this) {
       env()->hash_to_module_map.erase(it);
diff --git a/src/module_wrap.h b/src/module_wrap.h
index 0a800150a90479..245f8f47b40ad6 100644
--- a/src/module_wrap.h
+++ b/src/module_wrap.h
@@ -33,7 +33,7 @@ enum HostDefinedOptions : int {
 class ModuleWrap : public BaseObject {
  public:
   enum InternalFields {
-    kModuleWrapBaseField = BaseObject::kInternalFieldCount,
+    kModuleSlot = BaseObject::kInternalFieldCount,
     kURLSlot,
     kSyntheticEvaluationStepsSlot,
     kContextObjectSlot,  // Object whose creation context is the target Context
@@ -104,6 +104,7 @@ class ModuleWrap : public BaseObject {
   contextify::ContextifyContext* contextify_context_ = nullptr;
   bool synthetic_ = false;
   bool linked_ = false;
+  int module_hash_;
 };
 
 }  // namespace loader
diff --git a/src/node_contextify.cc b/src/node_contextify.cc
index f17848fc97e9eb..0873e482a602b3 100644
--- a/src/node_contextify.cc
+++ b/src/node_contextify.cc
@@ -862,7 +862,11 @@ void ContextifyScript::New(const FunctionCallbackInfo<Value>& args) {
                      "ContextifyScript::New");
     return;
   }
+
   contextify_script->script_.Reset(isolate, v8_script);
+  contextify_script->script_.SetWeak();
+  contextify_script->object()->SetDataInInternalField(kUnboundScriptSlot,
+                                                      v8_script);
 
   std::unique_ptr<ScriptCompiler::CachedData> new_cached_data;
   if (produce_cached_data) {
diff --git a/src/node_contextify.h b/src/node_contextify.h
index af646a69763b5f..b81843ec81ced9 100644
--- a/src/node_contextify.h
+++ b/src/node_contextify.h
@@ -151,6 +151,11 @@ class ContextifyContext : public BaseObject {
 
 class ContextifyScript : public BaseObject {
  public:
+  enum InternalFields {
+    kUnboundScriptSlot = BaseObject::kInternalFieldCount,
+    kInternalFieldCount
+  };
+
   SET_NO_MEMORY_INFO()
   SET_MEMORY_INFO_NAME(ContextifyScript)
   SET_SELF_SIZE(ContextifyScript)
diff --git a/test/es-module/test-vm-contextified-script-leak.js b/test/es-module/test-vm-contextified-script-leak.js
new file mode 100644
index 00000000000000..1e0aa680ad2ed2
--- /dev/null
+++ b/test/es-module/test-vm-contextified-script-leak.js
@@ -0,0 +1,17 @@
+'use strict';
+
+// Flags: --max-old-space-size=16 --trace-gc
+
+require('../common');
+const vm = require('vm');
+let count = 0;
+
+function main() {
+  new vm.Script(`"${Math.random().toString().repeat(512)}";`, {
+    async importModuleDynamically() {},
+  });
+  if (count++ < 2 * 1024) {
+    setTimeout(main, 1);
+  }
+}
+main();
diff --git a/test/es-module/test-vm-source-text-module-leak.js b/test/es-module/test-vm-source-text-module-leak.js
new file mode 100644
index 00000000000000..8394d09ebf2640
--- /dev/null
+++ b/test/es-module/test-vm-source-text-module-leak.js
@@ -0,0 +1,24 @@
+'use strict';
+
+// This tests that vm.SourceTextModule() does not leak.
+// See: https://github.com/nodejs/node/issues/33439
+// Flags: --experimental-vm-modules --max-old-space-size=16 --trace-gc
+
+require('../common');
+
+const vm = require('vm');
+let count = 0;
+async function createModule() {
+  const m = new vm.SourceTextModule(`
+  const bar = new Array(512).fill("----");
+  export { bar };
+`);
+  await m.link(() => {});
+  await m.evaluate();
+  count++;
+  if (count < 4096) {
+    setTimeout(createModule, 1);
+  }
+  return m;
+}
+createModule();