@contrast/rewriter 1.4.0 → 1.4.1

package/lib/index.js

@@ -12,6 +12,7 @@
  * engineered, modified, repackaged, sold, redistributed or otherwise used in a
  * way not consistent with the End User License Agreement.
  */
+// @ts-check
 
 'use strict';
 
@@ -21,7 +22,9 @@ const Module = require('module');
 const rewriterPath = require.resolve('@contrast/agent-swc-plugin');
 const unwriterPath = require.resolve('@contrast/agent-swc-plugin-unwrite');
 
+// @ts-expect-error `wrapper` is missing from @types/node
 const prefix = Module.wrapper[0];
+// @ts-expect-error `wrapper` is missing from @types/node
 const suffix = Module.wrapper[1].replace(/;$/, '.apply(this, arguments);');
 
 /** @typedef {'assess' | 'protect'} Mode */
@@ -42,18 +45,26 @@ const rewriter = {
   /**
    * @param {string} content the source code
    * @param {object} opts
-   * @param {string} opts.filename e.g. 'index.js'
-   * @param {boolean} opts.isModule if true, file is parsed as an ES module instead of a CJS script
-   * @param {boolean} opts.inject if true, injects ContrastMethods on the global object
-   * @param {boolean} opts.wrap if true, wraps the content with a modified module wrapper IIFE
+   * @param {string=} opts.filename e.g. 'index.js'
+   * @param {boolean=} opts.isModule if true, file is parsed as an ES module instead of a CJS script
+   * @param {boolean=} opts.inject if true, injects ContrastMethods on the global object
+   * @param {boolean=} opts.wrap if true, wraps the content with a modified module wrapper IIFE
    * @returns {import("@swc/core").Output}
    */
   rewrite(content, opts = {}) {
+    let shebang = '';
+
+    if (content.charAt(0) === '#') {
+      shebang = content.substring(0, content.indexOf('\n') + 1);
+      // see the test output: swc doesn't include the commented shebang in the generated code despite including comments otherwise
+      content = `//${content}`;
+    }
+
     if (opts.wrap) {
-      content = `${prefix}${content}${suffix}`;
+      content = `${shebang}${prefix}${content}${suffix}`;
     }
 
-    return transformSync(content, {
+    const result = transformSync(content, {
       filename: opts.filename,
       isModule: opts.isModule,
       jsc: {
@@ -72,6 +83,25 @@ const rewriter = {
       },
       sourceMaps: true,
     });
+
+    if (!opts.wrap) {
+      let carriageReturn = 0;
+      // swc always adds a newline, so we only need to check the input
+      if (!content.endsWith('\n')) {
+        result.code = result.code.substring(0, result.code.length - 1);
+      } else if (content.endsWith('\r\n')) {
+        // if EOL is \r\n, then we need to account for that when we check the
+        // negative index of the last semicolon below
+        carriageReturn = 1;
+      }
+      const resultSemicolonIdx = result.code.lastIndexOf(';');
+      const contentSemicolonIdx = content.lastIndexOf(';');
+      if (contentSemicolonIdx === -1 || resultSemicolonIdx - result.code.length !== contentSemicolonIdx - content.length + carriageReturn) {
+        result.code = result.code.substring(0, resultSemicolonIdx) + result.code.substring(resultSemicolonIdx + 1, result.code.length);
+      }
+    }
+
+    return result;
   },
 
   /**
@@ -90,11 +120,10 @@ const rewriter = {
   },
 };
 
-/** @typedef {{}} Core */
 /** @typedef {typeof rewriter} Rewriter */
 
 /**
- * @param {Core} core
+ * @param {{ rewriter: Rewriter }} core
  * @returns {Rewriter}
  */
 module.exports = function init(core) {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@contrast/rewriter",
-  "version": "1.4.0",
+  "version": "1.4.1",
   "description": "A transpilation tool mainly used for instrumentation",
   "license": "SEE LICENSE IN LICENSE",
   "author": "Contrast Security <nodejs@contrastsecurity.com> (https://www.contrastsecurity.com)",