From c70b2f7a7608272e08ecdb93ceeda5415471b0a5 Mon Sep 17 00:00:00 2001
From: Yagiz Nizipli <yagiz@nizipli.com>
Date: Sat, 18 May 2024 21:37:25 -0400
Subject: [PATCH] cli: add `NODE_RUN_SCRIPT_NAME` env to `node --run`

PR-URL: https://github.com/nodejs/node/pull/53032
Refs: https://github.com/nodejs/node/issues/52673
Reviewed-By: Daniel Lemire <daniel@lemire.me>
Reviewed-By: Benjamin Gruenbaum <benjamingr@gmail.com>
---
 doc/api/cli.md                                | 11 +++
 src/node_task_runner.cc                       | 76 ++++++++++---------
 src/node_task_runner.h                        |  3 +
 .../node_modules/.bin/special-env-variables   |  2 +
 .../.bin/special-env-variables.bat            |  2 +
 test/fixtures/run-script/package.json         |  4 +-
 test/message/node_run_non_existent.out        |  2 +
 test/parallel/test-node-run.js                | 12 +++
 8 files changed, 77 insertions(+), 35 deletions(-)
 create mode 100755 test/fixtures/run-script/node_modules/.bin/special-env-variables
 create mode 100755 test/fixtures/run-script/node_modules/.bin/special-env-variables.bat

diff --git a/doc/api/cli.md b/doc/api/cli.md
index b7ade9ffbe09b1..ad080e559987e1 100644
--- a/doc/api/cli.md
+++ b/doc/api/cli.md
@@ -1857,6 +1857,10 @@ Modules preloaded with `--require` will run before modules preloaded with `--imp
 
 <!-- YAML
 added: v22.0.0
+changes:
+  - version: REPLACEME
+    pr-url: https://github.com/nodejs/node/pull/53032
+    description: NODE_RUN_SCRIPT_NAME environment variable is added.
 -->
 
 > Stability: 1.1 - Active development
@@ -1897,6 +1901,13 @@ are:
 * Running `pre` or `post` scripts in addition to the specified script.
 * Defining package manager-specific environment variables.
 
+#### Environment variables
+
+The following environment variables are set when running a script with `--run`:
+
+* `NODE_RUN_SCRIPT_NAME`: The name of the script being run. For example, if
+  `--run` is used to run `test`, the value of this variable will be `test`.
+
 ### `--secure-heap=n`
 
 <!-- YAML
diff --git a/src/node_task_runner.cc b/src/node_task_runner.cc
index 4b7d541340a79d..45f6e0f307e8ec 100644
--- a/src/node_task_runner.cc
+++ b/src/node_task_runner.cc
@@ -12,6 +12,7 @@ static constexpr const char* bin_path = "/node_modules/.bin";
 #endif  // _WIN32
 
 ProcessRunner::ProcessRunner(std::shared_ptr<InitializationResultImpl> result,
+                             const std::string& script_name,
                              std::string_view command,
                              const PositionalArgs& positional_args) {
   memset(&options_, 0, sizeof(uv_process_options_t));
@@ -51,39 +52,9 @@ ProcessRunner::ProcessRunner(std::shared_ptr<InitializationResultImpl> result,
   // callback.
   process_.data = this;
 
-  std::string command_str(command);
+  SetEnvironmentVariables(current_bin_path, script_name);
 
-  // Set environment variables
-  uv_env_item_t* env_items;
-  int env_count;
-  CHECK_EQ(0, uv_os_environ(&env_items, &env_count));
-  env = std::unique_ptr<char*[]>(new char*[env_count + 1]);
-  options_.env = env.get();
-
-  // Iterate over environment variables once to store them in the current
-  // ProcessRunner instance.
-  for (int i = 0; i < env_count; i++) {
-    std::string name = env_items[i].name;
-    auto value = env_items[i].value;
-
-#ifdef _WIN32
-    // We use comspec environment variable to find cmd.exe path on Windows
-    // Example: 'C:\\Windows\\system32\\cmd.exe'
-    // If we don't find it, we fallback to 'cmd.exe' for Windows
-    if (StringEqualNoCase(name.c_str(), "comspec")) {
-      file_ = value;
-    }
-#endif  // _WIN32
-
-    // Check if environment variable key is matching case-insensitive "path"
-    if (StringEqualNoCase(name.c_str(), "path")) {
-      env_vars_.push_back(name + "=" + current_bin_path + value);
-    } else {
-      // Environment variables should be in "KEY=value" format
-      env_vars_.push_back(name + "=" + value);
-    }
-  }
-  uv_os_free_environ(env_items, env_count);
+  std::string command_str(command);
 
   // Use the stored reference on the instance.
   options_.file = file_.c_str();
@@ -128,11 +99,47 @@ ProcessRunner::ProcessRunner(std::shared_ptr<InitializationResultImpl> result,
     options_.args[i] = const_cast<char*>(command_args_[i].c_str());
   }
   options_.args[argc] = nullptr;
+}
 
+void ProcessRunner::SetEnvironmentVariables(const std::string& current_bin_path,
+                                            const std::string& script_name) {
+  uv_env_item_t* env_items;
+  int env_count;
+  CHECK_EQ(0, uv_os_environ(&env_items, &env_count));
+
+  // Iterate over environment variables once to store them in the current
+  // ProcessRunner instance.
   for (int i = 0; i < env_count; i++) {
+    std::string name = env_items[i].name;
+    std::string value = env_items[i].value;
+
+#ifdef _WIN32
+    // We use comspec environment variable to find cmd.exe path on Windows
+    // Example: 'C:\\Windows\\system32\\cmd.exe'
+    // If we don't find it, we fallback to 'cmd.exe' for Windows
+    if (StringEqualNoCase(name.c_str(), "comspec")) {
+      file_ = value;
+    }
+#endif  // _WIN32
+
+    if (StringEqualNoCase(name.c_str(), "path")) {
+      // Add bin_path to the beginning of the PATH
+      value = current_bin_path + value;
+    }
+    env_vars_.push_back(name + "=" + value);
+  }
+  uv_os_free_environ(env_items, env_count);
+
+  // Add NODE_RUN_SCRIPT_NAME environment variable to the environment
+  // to indicate which script is being run.
+  env_vars_.push_back("NODE_RUN_SCRIPT_NAME=" + script_name);
+
+  env = std::unique_ptr<char*[]>(new char*[env_vars_.size() + 1]);
+  options_.env = env.get();
+  for (size_t i = 0; i < env_vars_.size(); i++) {
     options_.env[i] = const_cast<char*>(env_vars_[i].c_str());
   }
-  options_.env[env_count] = nullptr;
+  options_.env[env_vars_.size()] = nullptr;
 }
 
 // EscapeShell escapes a string to be used as a command line argument.
@@ -276,7 +283,8 @@ void RunTask(std::shared_ptr<InitializationResultImpl> result,
     return;
   }
 
-  auto runner = ProcessRunner(result, command, positional_args);
+  auto runner =
+      ProcessRunner(result, std::string(command_id), command, positional_args);
   runner.Run();
 }
 
diff --git a/src/node_task_runner.h b/src/node_task_runner.h
index 3d8973db98ab42..65ddef71d6f487 100644
--- a/src/node_task_runner.h
+++ b/src/node_task_runner.h
@@ -23,6 +23,7 @@ using PositionalArgs = std::vector<std::string_view>;
 class ProcessRunner {
  public:
   ProcessRunner(std::shared_ptr<InitializationResultImpl> result,
+                const std::string& script_name,
                 std::string_view command_id,
                 const PositionalArgs& positional_args);
   void Run();
@@ -43,6 +44,8 @@ class ProcessRunner {
 
   // OnExit is the callback function that is called when the process exits.
   void OnExit(int64_t exit_status, int term_signal);
+  void SetEnvironmentVariables(const std::string& bin_path,
+                               const std::string& script_name);
 
 #ifdef _WIN32
   std::string file_ = "cmd.exe";
diff --git a/test/fixtures/run-script/node_modules/.bin/special-env-variables b/test/fixtures/run-script/node_modules/.bin/special-env-variables
new file mode 100755
index 00000000000000..df687e823224f6
--- /dev/null
+++ b/test/fixtures/run-script/node_modules/.bin/special-env-variables
@@ -0,0 +1,2 @@
+#!/bin/sh
+echo "$NODE_RUN_SCRIPT_NAME"
diff --git a/test/fixtures/run-script/node_modules/.bin/special-env-variables.bat b/test/fixtures/run-script/node_modules/.bin/special-env-variables.bat
new file mode 100755
index 00000000000000..ecfcf9262f68c0
--- /dev/null
+++ b/test/fixtures/run-script/node_modules/.bin/special-env-variables.bat
@@ -0,0 +1,2 @@
+@shift
+echo %NODE_RUN_SCRIPT_NAME%
diff --git a/test/fixtures/run-script/package.json b/test/fixtures/run-script/package.json
index 822b753e3db61e..c103b48840fa2d 100644
--- a/test/fixtures/run-script/package.json
+++ b/test/fixtures/run-script/package.json
@@ -8,6 +8,8 @@
     "custom-env": "custom-env",
     "custom-env-windows": "custom-env.bat",
     "path-env": "path-env",
-    "path-env-windows": "path-env.bat"
+    "path-env-windows": "path-env.bat",
+    "special-env-variables": "special-env-variables",
+    "special-env-variables-windows": "special-env-variables.bat"
   }
 }
diff --git a/test/message/node_run_non_existent.out b/test/message/node_run_non_existent.out
index a248a010a9b365..3397d837b0491c 100644
--- a/test/message/node_run_non_existent.out
+++ b/test/message/node_run_non_existent.out
@@ -10,3 +10,5 @@ Available scripts are:
   custom-env-windows: custom-env.bat
   path-env: path-env
   path-env-windows: path-env.bat
+  special-env-variables: special-env-variables
+  special-env-variables-windows: special-env-variables.bat
diff --git a/test/parallel/test-node-run.js b/test/parallel/test-node-run.js
index 11a5b413148f4e..ffbf227049d763 100644
--- a/test/parallel/test-node-run.js
+++ b/test/parallel/test-node-run.js
@@ -80,4 +80,16 @@ describe('node run [command]', () => {
     assert.strictEqual(child.stderr, '');
     assert.strictEqual(child.code, 0);
   });
+
+  it('should set special environment variables', async () => {
+    const scriptName = `special-env-variables${envSuffix}`;
+    const child = await common.spawnPromisified(
+      process.execPath,
+      [ '--no-warnings', '--run', scriptName],
+      { cwd: fixtures.path('run-script') },
+    );
+    assert.ok(child.stdout.includes(scriptName));
+    assert.strictEqual(child.stderr, '');
+    assert.strictEqual(child.code, 0);
+  });
 });