src: preload function for Environment

zcbenz · marco-ippolito · commit b3a11b574bc6 · 2024-05-02T10:58:01.000+02:00
This PR adds a |preload| arg to the node::LoadEnvironment to allow embedders to set a preload function for the environment, which will run after the environment is loaded and before the main script runs. This is similiar to the --require CLI option, but runs a C++ function, and can only be set by embedders. The preload function can be used by embedders to inject scripts before running the main script, for example: 1. In Electron it is used to initialize the ASAR virtual filesystem, inject custom process properties, etc. 2. In VS Code it can be used to reset the module search paths for extensions. PR-URL: #51539 Reviewed-By: Joyee Cheung <joyeec9h3@gmail.com> Reviewed-By: Yagiz Nizipli <yagiz.nizipli@sentry.io>
diff --git a/lib/internal/process/pre_execution.js b/lib/internal/process/pre_execution.js
@@ -196,6 +196,9 @@ function setupUserModules(forceDefaultLoader = false) {
   } = require('internal/modules/helpers');
   assert(!hasStartedUserCJSExecution());
   assert(!hasStartedUserESMExecution());
+  if (getEmbedderOptions().hasEmbedderPreload) {
+    runEmbedderPreload();
+  }
   // Do not enable preload modules if custom loaders are disabled.
   // For example, loader workers are responsible for doing this themselves.
   // And preload modules are not supported in ShadowRealm as well.
@@ -711,6 +714,10 @@ function initializeFrozenIntrinsics() {
   }
 }
 
+function runEmbedderPreload() {
+  internalBinding('mksnapshot').runEmbedderPreload(process, require);
+}
+
 function loadPreloadModules() {
   // For user code, we preload modules if `-r` is passed
   const preloadModules = getOptionValue('--require');
diff --git a/src/api/environment.cc b/src/api/environment.cc
@@ -538,25 +538,31 @@ NODE_EXTERN std::unique_ptr<InspectorParentHandle> GetInspectorParentHandle(
 #endif
 }
 
-MaybeLocal<Value> LoadEnvironment(
-    Environment* env,
-    StartExecutionCallback cb) {
+MaybeLocal<Value> LoadEnvironment(Environment* env,
+                                  StartExecutionCallback cb,
+                                  EmbedderPreloadCallback preload) {
   env->InitializeLibuv();
   env->InitializeDiagnostics();
+  if (preload) {
+    env->set_embedder_preload(std::move(preload));
+  }
 
   return StartExecution(env, cb);
 }
 
 MaybeLocal<Value> LoadEnvironment(Environment* env,
-                                  std::string_view main_script_source_utf8) {
+                                  std::string_view main_script_source_utf8,
+                                  EmbedderPreloadCallback preload) {
   CHECK_NOT_NULL(main_script_source_utf8.data());
   return LoadEnvironment(
-      env, [&](const StartExecutionCallbackInfo& info) -> MaybeLocal<Value> {
+      env,
+      [&](const StartExecutionCallbackInfo& info) -> MaybeLocal<Value> {
         Local<Value> main_script =
             ToV8Value(env->context(), main_script_source_utf8).ToLocalChecked();
         return info.run_cjs->Call(
             env->context(), Null(env->isolate()), 1, &main_script);
-      });
+      },
+      std::move(preload));
 }
 
 Environment* GetCurrentEnvironment(Local<Context> context) {
diff --git a/src/env-inl.h b/src/env-inl.h
@@ -430,6 +430,14 @@ inline builtins::BuiltinLoader* Environment::builtin_loader() {
   return &builtin_loader_;
 }
 
+inline const EmbedderPreloadCallback& Environment::embedder_preload() const {
+  return embedder_preload_;
+}
+
+inline void Environment::set_embedder_preload(EmbedderPreloadCallback fn) {
+  embedder_preload_ = std::move(fn);
+}
+
 inline double Environment::new_async_id() {
   async_hooks()->async_id_fields()[AsyncHooks::kAsyncIdCounter] += 1;
   return async_hooks()->async_id_fields()[AsyncHooks::kAsyncIdCounter];
diff --git a/src/env.h b/src/env.h
@@ -1002,6 +1002,9 @@ class Environment : public MemoryRetainer {
 
 #endif  // HAVE_INSPECTOR
 
+  inline const EmbedderPreloadCallback& embedder_preload() const;
+  inline void set_embedder_preload(EmbedderPreloadCallback fn);
+
   inline void set_process_exit_handler(
       std::function<void(Environment*, ExitCode)>&& handler);
 
@@ -1212,6 +1215,7 @@ class Environment : public MemoryRetainer {
   std::unique_ptr<PrincipalRealm> principal_realm_ = nullptr;
 
   builtins::BuiltinLoader builtin_loader_;
+  EmbedderPreloadCallback embedder_preload_;
 
   // Used by allocate_managed_buffer() and release_managed_buffer() to keep
   // track of the BackingStore for a given pointer.
diff --git a/src/node.h b/src/node.h
@@ -731,12 +731,33 @@ struct StartExecutionCallbackInfo {
 
 using StartExecutionCallback =
     std::function<v8::MaybeLocal<v8::Value>(const StartExecutionCallbackInfo&)>;
+using EmbedderPreloadCallback =
+    std::function<void(Environment* env,
+                       v8::Local<v8::Value> process,
+                       v8::Local<v8::Value> require)>;
 
+// Run initialization for the environment.
+//
+// The |preload| function, usually used by embedders to inject scripts,
+// will be run by Node.js before Node.js executes the entry point.
+// The function is guaranteed to run before the user land module loader running
+// any user code, so it is safe to assume that at this point, no user code has
+// been run yet.
+// The function will be executed with preload(process, require), and the passed
+// require function has access to internal Node.js modules. There is no
+// stability guarantee about the internals exposed to the internal require
+// function. Expect breakages when updating Node.js versions if the embedder
+// imports internal modules with the internal require function.
+// Worker threads created in the environment will also respect The |preload|
+// function, so make sure the function is thread-safe.
 NODE_EXTERN v8::MaybeLocal<v8::Value> LoadEnvironment(
     Environment* env,
-    StartExecutionCallback cb);
+    StartExecutionCallback cb,
+    EmbedderPreloadCallback preload = nullptr);
 NODE_EXTERN v8::MaybeLocal<v8::Value> LoadEnvironment(
-    Environment* env, std::string_view main_script_source_utf8);
+    Environment* env,
+    std::string_view main_script_source_utf8,
+    EmbedderPreloadCallback preload = nullptr);
 NODE_EXTERN void FreeEnvironment(Environment* env);
 
 // Set a callback that is called when process.exit() is called from JS,
diff --git a/src/node_options.cc b/src/node_options.cc
@@ -1305,6 +1305,12 @@ void GetEmbedderOptions(const FunctionCallbackInfo<Value>& args) {
           .IsNothing())
     return;
 
+  if (ret->Set(context,
+               FIXED_ONE_BYTE_STRING(env->isolate(), "hasEmbedderPreload"),
+               Boolean::New(isolate, env->embedder_preload() != nullptr))
+          .IsNothing())
+    return;
+
   args.GetReturnValue().Set(ret);
 }
 
diff --git a/src/node_snapshotable.cc b/src/node_snapshotable.cc
@@ -1417,6 +1417,17 @@ void SerializeSnapshotableObjects(Realm* realm,
   });
 }
 
+void RunEmbedderPreload(const FunctionCallbackInfo<Value>& args) {
+  Environment* env = Environment::GetCurrent(args);
+  CHECK(env->embedder_preload());
+  CHECK_EQ(args.Length(), 2);
+  Local<Value> process_obj = args[0];
+  Local<Value> require_fn = args[1];
+  CHECK(process_obj->IsObject());
+  CHECK(require_fn->IsFunction());
+  env->embedder_preload()(env, process_obj, require_fn);
+}
+
 void CompileSerializeMain(const FunctionCallbackInfo<Value>& args) {
   CHECK(args[0]->IsString());
   Local<String> filename = args[0].As<String>();
@@ -1540,6 +1551,7 @@ void CreatePerContextProperties(Local<Object> target,
 void CreatePerIsolateProperties(IsolateData* isolate_data,
                                 Local<ObjectTemplate> target) {
   Isolate* isolate = isolate_data->isolate();
+  SetMethod(isolate, target, "runEmbedderPreload", RunEmbedderPreload);
   SetMethod(isolate, target, "compileSerializeMain", CompileSerializeMain);
   SetMethod(isolate, target, "setSerializeCallback", SetSerializeCallback);
   SetMethod(isolate, target, "setDeserializeCallback", SetDeserializeCallback);
@@ -1552,6 +1564,7 @@ void CreatePerIsolateProperties(IsolateData* isolate_data,
 }
 
 void RegisterExternalReferences(ExternalReferenceRegistry* registry) {
+  registry->Register(RunEmbedderPreload);
   registry->Register(CompileSerializeMain);
   registry->Register(SetSerializeCallback);
   registry->Register(SetDeserializeCallback);
diff --git a/src/node_worker.cc b/src/node_worker.cc
@@ -62,6 +62,7 @@ Worker::Worker(Environment* env,
       thread_id_(AllocateEnvironmentThreadId()),
       name_(name),
       env_vars_(env_vars),
+      embedder_preload_(env->embedder_preload()),
       snapshot_data_(snapshot_data) {
   Debug(this, "Creating new worker instance with thread id %llu",
         thread_id_.id);
@@ -386,8 +387,12 @@ void Worker::Run() {
         }
 
         Debug(this, "Created message port for worker %llu", thread_id_.id);
-        if (LoadEnvironment(env_.get(), StartExecutionCallback{}).IsEmpty())
+        if (LoadEnvironment(env_.get(),
+                            StartExecutionCallback{},
+                            std::move(embedder_preload_))
+                .IsEmpty()) {
           return;
+        }
 
         Debug(this, "Loaded environment for worker %llu", thread_id_.id);
       }
diff --git a/src/node_worker.h b/src/node_worker.h
@@ -114,6 +114,7 @@ class Worker : public AsyncWrap {
 
   std::unique_ptr<MessagePortData> child_port_data_;
   std::shared_ptr<KVStore> env_vars_;
+  EmbedderPreloadCallback embedder_preload_;
 
   // A raw flag that is used by creator and worker threads to
   // sync up on pre-mature termination of worker  - while in the
diff --git a/test/cctest/test_environment.cc b/test/cctest/test_environment.cc
@@ -778,3 +778,31 @@ TEST_F(EnvironmentTest, RequestInterruptAtExit) {
 
   context->Exit();
 }
+
+TEST_F(EnvironmentTest, EmbedderPreload) {
+  v8::HandleScope handle_scope(isolate_);
+  v8::Local<v8::Context> context = node::NewContext(isolate_);
+  v8::Context::Scope context_scope(context);
+
+  node::EmbedderPreloadCallback preload = [](node::Environment* env,
+                                             v8::Local<v8::Value> process,
+                                             v8::Local<v8::Value> require) {
+    CHECK(process->IsObject());
+    CHECK(require->IsFunction());
+    process.As<v8::Object>()
+        ->Set(env->context(),
+              v8::String::NewFromUtf8Literal(env->isolate(), "prop"),
+              v8::String::NewFromUtf8Literal(env->isolate(), "preload"))
+        .Check();
+  };
+
+  std::unique_ptr<node::Environment, decltype(&node::FreeEnvironment)> env(
+      node::CreateEnvironment(isolate_data_, context, {}, {}),
+      node::FreeEnvironment);
+
+  v8::Local<v8::Value> main_ret =
+      node::LoadEnvironment(env.get(), "return process.prop;", preload)
+          .ToLocalChecked();
+  node::Utf8Value main_ret_str(isolate_, main_ret);
+  EXPECT_EQ(std::string(*main_ret_str), "preload");
+}
