feat: Optionally only wrap modules hooked in --import

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
@@ -71,6 +71,42 @@ module.register('import-in-the-middle/hook.mjs', import.meta.url, {
 })
 ```
 
+### Only Intercepting Hooked modules
+
+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 { addHookMessagePort, waitForAllMessagesAcknowledged } = createAddHookMessageChannel()
+
+const options = { data: { addHookMessagePort }, transferList: [addHookMessagePort] }
+
+register('import-in-the-middle/hook.mjs', import.meta.url, options)
+
+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

@@ -255,13 +255,24 @@ function createHook (meta) {
     if (data) {
       includeModules = ensureArrayWithBareSpecifiersAndFileUrls(data.include, 'include')
       excludeModules = ensureArrayWithBareSpecifiersAndFileUrls(data.exclude, 'exclude')
+
+      if (data.addHookMessagePort) {
+        data.addHookMessagePort.on('message', (modules) => {
+          if (includeModules === undefined) {
+            includeModules = []
+          }
+
+          includeModules.push(...modules)
+          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,21 @@ export declare function addHook(hookFn: HookFunction): void
  * @param {HookFunction} hookFn The function to be removed.
  */
 export declare function removeHook(hookFn: HookFunction): void
+
+/**
+ * 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 { createAddHookMessageChannel } from 'import-in-the-middle'
+ *
+ * const addHookMessagePort = createAddHookMessageChannel()
+ *
+ * register('import-in-the-middle/hook.mjs', import.meta.url, { data: { addHookMessagePort }, transferList: [addHookMessagePort] })
+ * ```
+ */
+export declare function createAddHookMessageChannel(): MessagePort;

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,44 @@ function callHookFn (hookFn, namespace, name, baseDir) {
   }
 }
 
+let sendModulesToLoader
+
+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
+  }
+
+  return { addHookMessagePort: port2, waitForAllMessagesAcknowledged }
+}
+
 function Hook (modules, options, hookFn) {
   if ((this instanceof Hook) === false) return new Hook(modules, options, hookFn)
   if (typeof modules === 'function') {
@@ -43,6 +82,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 +135,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,13 @@
+import { strictEqual } from 'assert'
+import { sep } from 'path'
+import { Hook } from '../../index.js'
+
+const hooked = []
+
+Hook((_, name) => {
+  hooked.push(name)
+})
+
+strictEqual(hooked.length, 1)
+strictEqual(hooked[0], 'path')
+strictEqual(sep, '@')

test/fixtures/import.mjs

@@ -0,0 +1,17 @@
+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 { addHookMessagePort, waitForAllMessagesAcknowledged } = createAddHookMessageChannel()
+
+register('../../hook.mjs', import.meta.url, { data: { addHookMessagePort }, transferList: [addHookMessagePort] })
+
+Hook(['path'], (exports) => {
+  exports.sep = '@'
+})
+
+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)