feat: Optionally only wrap modules hooked in --import (#146)

Author: timfish

README.md

@@ -34,14 +34,14 @@ console.log(foo) // 1 more than whatever that module exported
 This requires the use of an ESM loader hook, which can be added with the following
 command-line option.
 
-```
---loader=import-in-the-middle/hook.mjs
+```shell
+node --loader=import-in-the-middle/hook.mjs my-app.mjs
 ```
 
-It's also possible to register the loader hook programmatically via the Node
+Since `--loader` has been deprecated you can also register the loader hook programmatically via the Node
 [`module.register()`](https://nodejs.org/api/module.html#moduleregisterspecifier-parenturl-options)
 API. However, for this to be able to hook non-dynamic imports, it needs to be
-loaded before your app code is evaluated via the `--import` command-line option.
+registered before your app code is evaluated via the `--import` command-line option.
 
 `my-loader.mjs`
 ```js
@@ -54,9 +54,12 @@ node --import=./my-loader.mjs ./my-code.mjs
 ```
 
 When registering the loader hook programmatically, it's possible to pass a list
-of modules, file URLs or regular expressions to either exclude or specifically
-include which modules are intercepted. This is useful if a module is not
+of modules, file URLs or regular expressions to either `exclude` or specifically
+`include` which modules are intercepted. This is useful if a module is not
 compatible with the loader hook. 
+
+> **Note:** This feature is incompatible with the `{internals: true}` Hook option
+
 ```js
 import * as module from 'module'
 
@@ -71,6 +74,41 @@ module.register('import-in-the-middle/hook.mjs', import.meta.url, {
 })
 ```
 
+### Only Intercepting Hooked modules 
+> **Note:** This feature is experimental and is incompatible with the `{internals: true}` Hook option
+
+If you are `Hook`'ing all modules before they are imported, for example in a
+module loaded via the Node.js `--import` CLI argument, you can configure the
+loader to intercept only modules that were specifically hooked.
+
+`instrument.mjs`
+```js
+import { register } from 'module'
+import { Hook, createAddHookMessageChannel } from 'import-in-the-middle'
+
+const { registerOptions, waitForAllMessagesAcknowledged } = createAddHookMessageChannel()
+
+register('import-in-the-middle/hook.mjs', import.meta.url, registerOptions)
+
+Hook(['fs'], (exported, name, baseDir) => {
+  // Instrument the fs module
+})
+
+// Ensure that the loader has acknowledged all the modules 
+// before we allow execution to continue
+await waitForAllMessagesAcknowledged()
+```
+`my-app.mjs`
+```js
+import * as fs from 'fs'
+// fs will be instrumented!
+fs.readFileSync('file.txt')
+```
+
+```shell
+node --import=./instrument.mjs ./my-app.mjs
+```
+
 ## Limitations
 
 * You cannot add new exports to a module. You can only modify existing ones.

hook.js

@@ -281,13 +281,31 @@ function createHook (meta) {
     if (data) {
       includeModules = ensureArrayWithBareSpecifiersFileUrlsAndRegex(data.include, 'include')
       excludeModules = ensureArrayWithBareSpecifiersFileUrlsAndRegex(data.exclude, 'exclude')
+
+      if (data.addHookMessagePort) {
+        data.addHookMessagePort.on('message', (modules) => {
+          if (includeModules === undefined) {
+            includeModules = []
+          }
+
+          for (const each of modules) {
+            if (!each.startsWith('node:') && builtinModules.includes(each)) {
+              includeModules.push(`node:${each}`)
+            }
+
+            includeModules.push(each)
+          }
+
+          data.addHookMessagePort.postMessage('ack')
+        }).unref()
+      }
     }
   }
 
   async function resolve (specifier, context, parentResolve) {
     cachedResolve = parentResolve
 
-    // See github.com/nodejs/import-in-the-middle/pull/76.
+    // See https://github.com/nodejs/import-in-the-middle/pull/76.
     if (specifier === iitmURL) {
       return {
         url: specifier,

index.d.ts

@@ -84,3 +84,39 @@ export declare function addHook(hookFn: HookFunction): void
  * @param {HookFunction} hookFn The function to be removed.
  */
 export declare function removeHook(hookFn: HookFunction): void
+
+type CreateAddHookMessageChannelReturn<Data> = {
+  addHookMessagePort: MessagePort,
+  waitForAllMessagesAcknowledged: Promise<void>
+  registerOptions: { data?: Data; transferList?: any[]; }
+}
+
+/**
+ * EXPERIMENTAL
+ * This feature is experimental and may change in minor versions.
+ * **NOTE** This feature is incompatible with the {internals: true} Hook option.
+ *
+ * Creates a message channel with a port that can be used to add hooks to the
+ * list of exclusively included modules.
+ *
+ * This can be used to only wrap modules that are Hook'ed, however modules need
+ * to be hooked before they are imported.
+ *
+ * ```ts
+ * import { register } from 'module'
+ * import { Hook, createAddHookMessageChannel } from 'import-in-the-middle'
+ *
+ * const { registerOptions, waitForAllMessagesAcknowledged } = createAddHookMessageChannel()
+ *
+ * register('import-in-the-middle/hook.mjs', import.meta.url, registerOptions)
+ *
+ * Hook(['fs'], (exported, name, baseDir) => {
+ *   // Instrument the fs module
+ * })
+ *
+ * // Ensure that the loader has acknowledged all the modules
+ * // before we allow execution to continue
+ * await waitForAllMessagesAcknowledged()
+ * ```
+ */
+export declare function createAddHookMessageChannel<Data = any>(): CreateAddHookMessageChannelReturn<Data>;

index.js

@@ -5,6 +5,7 @@
 const path = require('path')
 const parse = require('module-details-from-path')
 const { fileURLToPath } = require('url')
+const { MessageChannel } = require('worker_threads')
 
 const {
   importHooks,
@@ -31,6 +32,75 @@ function callHookFn (hookFn, namespace, name, baseDir) {
   }
 }
 
+let sendModulesToLoader
+
+/**
+ * EXPERIMENTAL
+ * This feature is experimental and may change in minor versions.
+ * **NOTE** This feature is incompatible with the {internals: true} Hook option.
+ *
+ * Creates a message channel with a port that can be used to add hooks to the
+ * list of exclusively included modules.
+ *
+ * This can be used to only wrap modules that are Hook'ed, however modules need
+ * to be hooked before they are imported.
+ *
+ * ```ts
+ * import { register } from 'module'
+ * import { Hook, createAddHookMessageChannel } from 'import-in-the-middle'
+ *
+ * const { registerOptions, waitForAllMessagesAcknowledged } = createAddHookMessageChannel()
+ *
+ * register('import-in-the-middle/hook.mjs', import.meta.url, registerOptions)
+ *
+ * Hook(['fs'], (exported, name, baseDir) => {
+ *   // Instrument the fs module
+ * })
+ *
+ * // Ensure that the loader has acknowledged all the modules
+ * // before we allow execution to continue
+ * await waitForAllMessagesAcknowledged()
+ * ```
+ */
+function createAddHookMessageChannel () {
+  const { port1, port2 } = new MessageChannel()
+  let pendingAckCount = 0
+  let resolveFn
+
+  sendModulesToLoader = (modules) => {
+    pendingAckCount++
+    port1.postMessage(modules)
+  }
+
+  port1.on('message', () => {
+    pendingAckCount--
+
+    if (resolveFn && pendingAckCount <= 0) {
+      resolveFn()
+    }
+  }).unref()
+
+  function waitForAllMessagesAcknowledged () {
+    // This timer is to prevent the process from exiting with code 13:
+    // 13: Unsettled Top-Level Await.
+    const timer = setInterval(() => { }, 1000)
+    const promise = new Promise((resolve) => {
+      resolveFn = resolve
+    }).then(() => { clearInterval(timer) })
+
+    if (pendingAckCount === 0) {
+      resolveFn()
+    }
+
+    return promise
+  }
+
+  const addHookMessagePort = port2
+  const registerOptions = { data: { addHookMessagePort, include: [] }, transferList: [addHookMessagePort] }
+
+  return { registerOptions, addHookMessagePort, waitForAllMessagesAcknowledged }
+}
+
 function Hook (modules, options, hookFn) {
   if ((this instanceof Hook) === false) return new Hook(modules, options, hookFn)
   if (typeof modules === 'function') {
@@ -43,6 +113,10 @@ function Hook (modules, options, hookFn) {
   }
   const internals = options ? options.internals === true : false
 
+  if (sendModulesToLoader && Array.isArray(modules)) {
+    sendModulesToLoader(modules)
+  }
+
   this._iitmHook = (name, namespace) => {
     const filename = name
     const isBuiltin = name.startsWith('node:')
@@ -92,3 +166,4 @@ module.exports = Hook
 module.exports.Hook = Hook
 module.exports.addHook = addHook
 module.exports.removeHook = removeHook
+module.exports.createAddHookMessageChannel = createAddHookMessageChannel

test/fixtures/import-after.mjs

@@ -0,0 +1,16 @@
+import { strictEqual } from 'assert'
+import { sep } from 'path'
+import * as os from 'node:os'
+import { Hook } from '../../index.js'
+
+const hooked = []
+
+Hook((_, name) => {
+  hooked.push(name)
+})
+
+strictEqual(hooked.length, 2)
+strictEqual(hooked[0], 'path')
+strictEqual(hooked[1], 'os')
+strictEqual(sep, '@')
+strictEqual(os.arch(), 'new_crazy_arch')

test/fixtures/import.mjs

@@ -0,0 +1,23 @@
+import { register } from 'module'
+import { Hook, createAddHookMessageChannel } from '../../index.js'
+// We've imported path here to ensure that the hook is still applied later even
+// if the library is used here.
+import * as path from 'path'
+
+const { registerOptions, waitForAllMessagesAcknowledged } = createAddHookMessageChannel()
+
+register('../../hook.mjs', import.meta.url, registerOptions)
+
+Hook(['path'], (exports) => {
+  exports.sep = '@'
+})
+
+Hook(['os'], (exports) => {
+  exports.arch = function () {
+    return 'new_crazy_arch'
+  }
+})
+
+console.assert(path.sep !== '@')
+
+await waitForAllMessagesAcknowledged()

test/register/v18.19-include-message-port.mjs

@@ -0,0 +1,14 @@
+import { spawnSync } from 'child_process'
+
+const out = spawnSync(process.execPath,
+  ['--import', './test/fixtures/import.mjs', './test/fixtures/import-after.mjs'],
+  { stdio: 'inherit', env: {} }
+)
+
+if (out.error) {
+  console.error(out.error)
+}
+if (out.status !== 0) {
+  console.error(`Expected exit code 0, got ${out.status}`)
+}
+process.exit(out.status)