@contrast/rewriter 1.0.0

package/LICENSE

@@ -0,0 +1,12 @@
+Copyright: 2022 Contrast Security, Inc
+Contact: support@contrastsecurity.com
+License: Commercial
+
+NOTICE: This Software and the patented inventions embodied within may only be
+used as part of Contrast Security’s commercial offerings. Even though it is
+made available through public repositories, use of this Software is subject to
+the applicable End User Licensing Agreement found at
+https://www.contrastsecurity.com/enduser-terms-0317a or as otherwise agreed
+between Contrast Security and the End User. The Software may not be reverse
+engineered, modified, repackaged, sold, redistributed or otherwise used in a
+way not consistent with the End User License Agreement.

package/README.md

@@ -0,0 +1,39 @@
+## `@contrast/rewriter`
+
+Provides a configured service for rewriting code.
+
+Basic idea: Enabled features can register rewrite transforms to support their end goal.
+
+Example: Assess
+
+Assess will register transforms for `+` -> `contrast_add()` so that it can perform propagation
+via instrumentation of `contrast_add()`.
+
+
+#### Example Service Usage
+
+```typescript
+const { Rewriter } = require('.');
+
+const rewriter = new Rewriter({ config, logger });
+
+rewriter.addTransforms({
+  BinaryExpression(path) {
+    const method = methodLookups[path.node.operator];
+    if (method) {
+      path.replaceWith(
+        t.callExpression(
+          expression('ContrastMethods.%%method%%')({ method }), [
+            path.node.left,
+            path.node.right
+          ]
+        )
+      );
+    }
+  }
+});
+
+const result = rewriter.rewriteFile({
+  content: 'function add(x, y) { return x + y; }'
+});
+```

package/lib/index.js

@@ -0,0 +1,157 @@
+'use strict';
+
+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');
+
+/**
+ * factory
+ */
+module.exports = function(core) {
+  const rewriter = new Rewriter(core);
+  core.rewriter = rewriter;
+  return rewriter;
+};
+
+
+/**
+ * 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.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
+            })
+          ];
+          return;
+        }
+
+        if (state.inject) {
+          path.unshiftContainer('body', [
+            statement(
+              'const ContrastMethods = global.ContrastMethods || (() => { throw new SyntaxError(\'ContrastMethods undefined during compilation\'); })();'
+            )(),
+          ]);
+        }
+      },
+      CallExpression(path, state) {
+        if (path.node.callee.name === 'eval') {
+          path.node.arguments = [
+            t.callExpression(expression('global.ContrastMethods.eval')(), path.node.arguments)
+          ];
+        }
+      },
+    };
+    this.unwriteTransforms = {
+      CallExpression(path) {
+        const obj = path.node.callee.object;
+        if (obj && obj.property && obj.property.name === 'ContrastMethods') {
+          path.replaceWith(path.node.arguments[0]);
+        }
+      }
+    };
+  }
+
+  /**
+   * @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}
+   */
+  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
+    };
+    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();
+
+    const result = generate(
+      ast,
+      {
+        jsonCompatibleStrings: true,
+        sourceMaps: true,
+        sourceFileName: state.filename
+      },
+      state.orig
+    );
+
+    result.code = removeAddedSemicolons(content, result.code);
+    result.deps = state.deps;
+
+    return result;
+  }
+
+  /**
+   * @param {string} code
+   * @param {object} opts
+   * @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;

package/lib/index.test.js

@@ -0,0 +1,102 @@
+'use strict';
+
+const { expect } = require('chai');
+const sinon = require('sinon');
+
+describe('rewriter', function() {
+  let core;
+  let visitorSpy;
+  let rewriter;
+
+  beforeEach(function() {
+    const mocks = require('../../test/mocks');
+    core = mocks.core();
+    core.config = mocks.config();
+    core.logger = mocks.logger();
+    visitorSpy = sinon.spy();
+
+    const factory = require('.');
+    rewriter = factory(core);
+    rewriter.visitors.push(visitorSpy);
+  });
+
+
+  // these baseline values don't have trailing semicolons
+  const content = 'eval(\'3\')';
+  const rewrittenContent = 'eval(global.ContrastMethods.eval(\'3\'))';
+
+  describe('rewrite()', function() {
+    [
+      {
+        content,
+        rewrittenContent,
+        desc: 'no spurious trailing semicolons are added',
+      },
+      {
+        content: `${content};`,
+        rewrittenContent: `${rewrittenContent};`,
+        desc: 'trailing semicolons are preserved',
+      },
+      {
+        content: 'foo',
+        rewrittenContent: 'foo',
+        desc: 'code is unaltered if no contrast methods need rewriting',
+        rawMappings: { length: 2 },
+        visitorsCall: 3
+      }
+    ].forEach(({ content, rewrittenContent, desc, rawMappings, visitorsCall }) => {
+      it(`rewrites successfully and ${desc}`, function() {
+        const result = rewriter.rewrite(content);
+        expect(result.code).eql(rewrittenContent);
+        expect(result.map).ok;
+        expect(result.rawMappings).lengthOf(rawMappings ? rawMappings.length : 5);
+        expect(visitorSpy.getCalls().length).eql(visitorsCall || 11);
+      });
+    });
+  });
+
+  describe('unwrite()', function() {
+    [
+      {
+        content,
+        rewrittenContent,
+        desc: 'no spurious trailing semicolons are added'
+      },
+      {
+        content: `${content};`,
+        rewrittenContent: `${rewrittenContent};`,
+        desc: 'trailing semicolons are preserved'
+      },
+      {
+        content,
+        rewrittenContent: content,
+        desc: 'code is unaltered if no contrast methods need unwriting',
+      },
+    ].forEach(({ content, rewrittenContent, desc }) => {
+      it(`unwrites successfully and ${desc}`, function() {
+        const ret = rewriter.unwrite(rewrittenContent);
+        expect(ret).eql(content);
+      });
+    });
+  });
+
+  describe('option: wrap', function() {
+    it('{ wrap: true } will add "Module.wrap" IIFE to', function () {
+      const result = rewriter.rewrite(content, { wrap: true });
+      expect(result.code).eql(`(function (exports, require, module, __filename, __dirname) {
+  ${rewrittenContent};
+}).apply(this, arguments)`);
+    });
+  });
+
+  describe('option: inject', function() {
+    it('{ inject: true } will add contrast method declarations', function() {
+      const result = rewriter.rewrite(content, { inject: true });
+      expect(result.code).eql(`const ContrastMethods = global.ContrastMethods || (() => {
+  throw new SyntaxError('ContrastMethods undefined during compilation');
+})();
+
+${rewrittenContent}`);
+    });
+  });
+});

package/lib/source-maps.js

@@ -0,0 +1,37 @@
+'use strict';
+
+const { readFileSync, existsSync } = require('fs');
+const { transfer } = require('multi-stage-sourcemap');
+const { SourceMapConsumer } = require('@contrast/synchronous-source-maps');
+const path = require('path');
+
+module.exports = function(deps) {
+  const { isAgentPath } = deps;
+  const sourceMaps = deps.rewriter.sourceMaps = {};
+  const consumerCache = sourceMaps.consumerCache = {};
+
+  sourceMaps.cacheConsumerMap = function(filename, map) {
+    consumerCache[filename] = new SourceMapConsumer(map);
+  };
+
+  /**
+   */
+  sourceMaps.chain = function(filename, map) {
+    let ret;
+
+    if (existsSync(`${filename}.map`)) {
+      try {
+        const existingMap = JSON.parse(readFileSync(`${filename}.map`, 'utf8'));
+        ret = transfer({ fromSourceMap: map, toSourceMap: existingMap });
+        deps.logger.trace(`Merged sourcemap from ${filename}.map`);
+      } catch (e) {
+        deps.logger.debug(`Unable to read ${filename}.map.js`);
+        deps.logger.debug(`${e}`);
+      }
+    }
+
+    return ret;
+  };
+
+  return sourceMaps;
+};

package/package.json

@@ -0,0 +1,28 @@
+{
+  "name": "@contrast/rewriter",
+  "version": "1.0.0",
+  "description": "A transpilation tool mainly used for instrumentation",
+  "license": "SEE LICENSE IN LICENSE",
+  "author": "Contrast Security <nodejs@contrastsecurity.com> (https://www.contrastsecurity.com)",
+  "files": [
+    "lib/"
+  ],
+  "main": "lib/index.js",
+  "types": "lib/index.d.ts",
+  "engines": {
+    "npm": ">= 8.4.0",
+    "node": ">= 14.15.0"
+  },
+  "scripts": {
+    "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/synchronous-source-maps": "^1.1.3",
+    "multi-stage-sourcemap": "^0.3.1"
+  }
+}