@contrast/rewriter 1.1.0 → 1.3.0

package/README.md

@@ -1,12 +1,8 @@
 ## `@contrast/rewriter`
 
-Provides a configured service for rewriting code.
+Rewrite javascript code with custom rewrite transforms.
 
-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
+For example, Assess will register transforms for `+` -> `contrast_add()` so that it can perform propagation
 via instrumentation of `contrast_add()`.
 
 
@@ -15,7 +11,7 @@ via instrumentation of `contrast_add()`.
 ```typescript
 const { Rewriter } = require('.');
 
-const rewriter = new Rewriter({ config, logger });
+const rewriter = new Rewriter({ logger });
 
 rewriter.addTransforms({
   BinaryExpression(path) {
@@ -33,7 +29,5 @@ rewriter.addTransforms({
   }
 });
 
-const result = rewriter.rewriteFile({
-  content: 'function add(x, y) { return x + y; }'
-});
+const result = rewriter.rewrite('function add(x, y) { return x + y; }');
 ```

package/lib/index.js

@@ -22,13 +22,13 @@ 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;
+  return core.rewriter = rewriter;
 };
 
 
@@ -54,6 +54,10 @@ class Rewriter {
     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) {
@@ -62,7 +66,6 @@ class Rewriter {
       },
       Program: function Program(path, state) {
         if (state.wrap) {
-
           let [prefix, suffix] = Module.wrapper;
           prefix = prefix.trim();
           suffix = suffix.trim().replace(/;$/, '.apply(this, arguments);');
@@ -75,25 +78,29 @@ class Rewriter {
         }
 
         if (state.inject) {
-          path.unshiftContainer('body', [
-            statement(
-              'const ContrastMethods = global.ContrastMethods || (() => { throw new SyntaxError(\'ContrastMethods undefined during compilation\'); })();'
-            )(),
-          ]);
+          path.unshiftContainer('body', self.injections);
         }
       },
-      CallExpression(path, state) {
+      CallExpression(path) {
         if (path.node.callee.name === 'eval') {
           path.node.arguments = [
             t.callExpression(expression('global.ContrastMethods.eval')(), path.node.arguments)
           ];
         }
-        if (path.node.callee.name === 'Function') {
-          path.node.arguments = [
-            t.callExpression(expression('global.ContrastMethods.Function')(), 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) {
@@ -103,6 +110,66 @@ class Rewriter {
         }
       }
     };
+    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'
+        });
+        self.injections.push(
+          statement(
+            'const %%name%% = global.%%id%% || %%name%%;'
+          )({ name: 'JSON', id: 'ContrastJSON' }),
+          statement(
+            'const %%name%% = global.%%id%% || %%name%%;'
+          )({ name: 'Object', id: 'ContrastObject' }),
+          statement(
+            'const %%name%% = global.%%id%% || %%name%%;'
+          )({ name: 'Number', id: 'ContrastNumber' })
+        );
+        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])
+              )
+            );
+          }
+        });
+      }
+
+      self.tokens = Object.keys(self.methodLookups);
+    };
   }
 
   /**
@@ -129,7 +196,7 @@ class Rewriter {
       ...opts
     };
 
-    if (state.orig.indexOf('eval') === -1 && state.orig.indexOf('Function') === -1) {
+    if (this.tokens.every((token) => state.orig.indexOf(token) === -1)) {
       return { code: state.orig };
     }
 

package/package.json

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