@contrast/rewriter 1.3.1 → 1.4.1

package/lib/index.js

@@ -12,226 +12,121 @@
  * engineered, modified, repackaged, sold, redistributed or otherwise used in a
  * way not consistent with the End User License Agreement.
  */
+// @ts-check
 
 'use strict';
 
+const { transformSync } = require('@swc/core');
 const Module = require('module');
-const { default: traverse } = require('@babel/traverse');
-const parser = require('@babel/parser');
-const { default: generate } = require('@babel/generator');
-const t = require('@babel/types');
-const { expression, statement } = require('@babel/template');
 
+const rewriterPath = require.resolve('@contrast/agent-swc-plugin');
+const unwriterPath = require.resolve('@contrast/agent-swc-plugin-unwrite');
 
-/**
- * factory
- */
-module.exports = function(core) {
-  const rewriter = new Rewriter(core);
-  return core.rewriter = rewriter;
-};
+// @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 */
 
-/**
- * Babel will add a trailing semicolon in some cases. This will remove it if it
- * wasn't there to begin with.
- * @param {string} rewritten rewritten content
- * @param {string} orig original content before rewriting
- * @returns {string}
- */
-function removeAddedSemicolons(orig, rewritten) {
-  if (
-    rewritten.charCodeAt(rewritten.length - 1) == 59 &&
-    orig.charCodeAt(orig.length - 1) != 59
-  ) {
-    rewritten = rewritten.substr(0, rewritten.length - 1);
-  }
-  return rewritten;
-}
-
-class Rewriter {
-  constructor(deps) {
-    const self = this;
-    this.logger = deps.logger;
-    this.visitors = [];
-    this.installedModes = [];
-    this.tokens = [];
-    this.injections = [];
-    this.methodLookups = {};
-    this.rewriteTransforms = {
-      enter(...args) {
-        for (const v of self.visitors) {
-          v(...args);
-        }
-      },
-      Program: function Program(path, state) {
-        if (state.wrap) {
-          let [prefix, suffix] = Module.wrapper;
-          prefix = prefix.trim();
-          suffix = suffix.trim().replace(/;$/, '.apply(this, arguments);');
-
-          path.node.body = [
-            statement(`${prefix} %%body%% ${suffix}`)({
-              body: path.node.body
-            })
-          ];
-        }
-
-        if (state.inject) {
-          path.unshiftContainer('body', self.injections);
-        }
-      },
-      CallExpression(path) {
-        if (path.node.callee.name === 'eval') {
-          path.node.arguments = [
-            t.callExpression(expression('global.ContrastMethods.eval')(), path.node.arguments)
-          ];
-        }
-      },
-      BinaryExpression: function BinaryExpression(path) {
-        const method = self.methodLookups[path.node.operator];
-        if (method) {
-          path.replaceWith(
-            t.callExpression(
-              expression('ContrastMethods.%%method%%')({ method }), [
-                path.node.left,
-                path.node.right
-              ]
-            )
-          );
-        }
-      }
-    };
-    this.unwriteTransforms = {
-      CallExpression(path) {
-        const obj = path.node.callee.object;
-        if (obj && obj.property && obj.property.name === 'ContrastMethods') {
-          path.replaceWith(path.node.arguments[0]);
-        }
-      }
-    };
-    this.install = function(mode) {
-      self.installedModes.push(mode);
-      !self.methodLookups.eval && (self.methodLookups = {
-        eval: 'eval'
-      });
-
-      !(self.injections.length === 2) && self.injections.push(
-        statement(
-          'const %%id%% = global.%%id%% || (() => { throw new SyntaxError(%%errMessage%%); })();'
-        )({
-          id: 'ContrastMethods',
-          errMessage: t.stringLiteral(
-            'ContrastMethods undefined during compilation'
-          )
-        }),
-        statement(
-          'var %%name%% = global.ContrastMethods.%%id%% || %%name%%;'
-        )({ name: 'Function', id: 'Function' }),
-      );
-
-      if (self.installedModes.includes('protect')) {
-      // Protect doesn't have anything specific
-      // for rewriting that should not be rewritten
-      // in Assess (at least for now)
-      }
-      if (self.installedModes.includes('assess')) {
-        Object.assign(self.methodLookups, {
-          '+': 'plus',
-          '===': 'tripleEqual',
-          '!==': 'notTripleEqual',
-          '==': 'doubleEqual',
-          '!=': 'notDoubleEqual'
-        });
-        Object.assign(self.rewriteTransforms, {
-          AssignmentExpression(path) {
-            if (path.node.operator !== '+=') return;
-            path.replaceWith(
-              t.assignmentExpression(
-                '=',
-                path.node.left,
-                t.callExpression(expression('global.ContrastMethods.plus')(), [path.node.left, path.node.right])
-              )
-            );
-          }
-        });
-      }
+const rewriter = {
+  /** @type {Set<Mode>} */
+  modes: new Set(),
 
-      self.tokens = Object.keys(self.methodLookups);
-    };
-  }
+  /**
+   * Sets the rewriter to 'assess' or 'protect' mode, enabling different
+   * transforms.
+   * @param {Mode} mode
+   */
+  install(mode) {
+    this.modes.add(mode);
+  },
 
   /**
    * @param {string} content the source code
    * @param {object} opts
-   * @param {string} opts.filename e.g. 'index.js'
-   * @param {boolean} opts.inject whether to inject contrast methods
-   * @param {string} opts.sourceType script or module
-   * @param {boolean} opts.wrap whether to wrap code in module wrap IIFE
-   * @returns {object}
+   * @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 = {
-    inject: false,
-    wrap: false,
-    sourceType: 'script',
-  }) {
-    opts.filename = opts.filename || 'no filename';
-    opts.sourceType = opts.sourceType || 'script';
-
-    const state = {
-      orig: String(content),
-      deps: [],
-      filename: opts.filename,
-      ...opts
-    };
+  rewrite(content, opts = {}) {
+    let shebang = '';
 
-    if (this.tokens.every((token) => state.orig.indexOf(token) === -1)) {
-      return { code: state.orig };
+    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}`;
     }
 
-    const ast = parser.parse(state.orig, {
-      plugins: [
-        'classPrivateMethods',
-        'classPrivateProperties',
-        'classProperties'
-      ],
-      ranges: true,
-      sourceType: state.sourceType,
-      sourceFilename: state.filename,
-      tokens: true
-    });
-
-    traverse(ast, this.rewriteTransforms, null, state);
-    // TODO: Look into how effective this is
-    traverse.cache.clear();
+    if (opts.wrap) {
+      content = `${shebang}${prefix}${content}${suffix}`;
+    }
 
-    const result = generate(
-      ast,
-      {
-        jsonCompatibleStrings: true,
-        sourceMaps: true,
-        sourceFileName: state.filename
+    const result = transformSync(content, {
+      filename: opts.filename,
+      isModule: opts.isModule,
+      jsc: {
+        target: 'es2019', // should work for node >14
+        experimental: {
+          plugins: [
+            [
+              rewriterPath,
+              {
+                assess: this.modes.has('assess'),
+                inject: this.modes.has('assess') && opts.inject,
+              },
+            ],
+          ],
+        },
       },
-      state.orig
-    );
+      sourceMaps: true,
+    });
 
-    result.code = removeAddedSemicolons(content, result.code);
-    result.deps = state.deps;
+    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;
-  }
+  },
 
   /**
-   * @param {string} code
-   * @param {object} opts
+   * @param {string} content
    * @returns {string}
    */
-  unwrite(code, opts) {
-    const ast = parser.parse(code);
-    traverse(ast, this.unwriteTransforms);
-    const unwritten = generate(ast, { jsonCompatibleStrings: true }).code;
-    return removeAddedSemicolons(code, unwritten);
-  }
-}
-
-module.exports.Rewriter = Rewriter;
+  unwrite(content) {
+    return transformSync(content, {
+      jsc: {
+        target: 'es2019', // should work for node >14
+        experimental: {
+          plugins: [[unwriterPath, {}]],
+        },
+      },
+    }).code;
+  },
+};
+
+/** @typedef {typeof rewriter} Rewriter */
+
+/**
+ * @param {{ rewriter: Rewriter }} core
+ * @returns {Rewriter}
+ */
+module.exports = function init(core) {
+  core.rewriter = rewriter;
+  return rewriter;
+};

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@contrast/rewriter",
-  "version": "1.3.1",
+  "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)",
@@ -17,12 +17,10 @@
     "test": "../scripts/test.sh"
   },
   "dependencies": {
-    "@babel/generator": "^7.17.9",
-    "@babel/parser": "^7.17.9",
-    "@babel/template": "^7.16.7",
-    "@babel/traverse": "^7.17.9",
-    "@babel/types": "^7.16.7",
+    "@contrast/agent-swc-plugin": "^1.1.0",
+    "@contrast/agent-swc-plugin-unwrite": "^1.1.0",
     "@contrast/synchronous-source-maps": "^1.1.3",
+    "@swc/core": "1.3.39",
     "multi-stage-sourcemap": "^0.3.1"
   }
-}
+}