feat: support callback argument

lib/index.js

@@ -9,8 +9,8 @@ const generateBefore = (t, id) =>
     ),
   ]);
 
-const generateInside = (t, id, line, timeout) =>
-  t.ifStatement(
+const generateInside = ({ t, id, line, timeout, extra } = {}) => {
+  return t.ifStatement(
     t.binaryExpression(
       '>',
       t.binaryExpression(
@@ -23,13 +23,27 @@ const generateInside = (t, id, line, timeout) =>
       ),
       t.numericLiteral(timeout)
     ),
-    t.breakStatement()
+    extra
+      ? t.blockStatement([
+          t.expressionStatement(
+            t.callExpression(extra, [t.numericLiteral(line)])
+          ),
+          t.breakStatement(),
+        ])
+      : t.breakStatement()
   );
+};
 
-const protect = (t, timeout) => path => {
+const protect = (t, timeout, extra) => path => {
   const id = path.scope.generateUidIdentifier('LP');
   const before = generateBefore(t, id);
-  const inside = generateInside(t, id, path.node.loc.start.line, timeout);
+  const inside = generateInside({
+    t,
+    id,
+    line: path.node.loc.start.line,
+    timeout,
+    extra,
+  });
   const body = path.get('body');
 
   // if we have an expression statement, convert it to a block
@@ -40,10 +54,22 @@ const protect = (t, timeout) => path => {
   body.unshiftContainer('body', inside);
 };
 
-module.exports = (timeout = 100) => ({ types: t }) => ({
-  visitor: {
-    WhileStatement: protect(t, timeout),
-    ForStatement: protect(t, timeout),
-    DoWhileStatement: protect(t, timeout),
-  },
-});
+module.exports = (timeout = 100, extra = null) => {
+  if (typeof extra === 'string') {
+    const string = extra;
+    extra = `() => console.error("${string.replace(/"/g, '\\"')}")`;
+  }
+  return ({ types: t, transform }) => {
+    const callback = extra
+      ? transform(extra).ast.program.body[0].expression
+      : null;
+
+    return {
+      visitor: {
+        WhileStatement: protect(t, timeout, callback),
+        ForStatement: protect(t, timeout, callback),
+        DoWhileStatement: protect(t, timeout, callback),
+      },
+    };
+  };
+};

package.json

@@ -3,7 +3,7 @@
   "name": "loop-protect",
   "description": "Prevent infinite loops in dynamically eval'd JavaScript.",
   "main": "lib/",
-  "version": "2.0.0",
+  "version": "2.1.0",
   "homepage": "https://github.com/jsbin/loop-protect",
   "repository": {
     "type": "git",

test/loop-protect.test.js

@@ -1,5 +1,5 @@
 const Babel = require('babel-standalone');
-Babel.registerPlugin('loopProtection', require('../lib')(200));
+Babel.registerPlugin('loopProtection', require('../lib')(100));
 const assert = e => console.assert(e);
 
 const code = {
@@ -16,7 +16,7 @@ const code = {
   irl1:
     'var nums = [0,1];\n var total = 8;\n for(var i = 0; i <= total; i++){\n var newest = nums[i--]\n nums.push(newest);\n }\n return i;',
   irl2:
-    'var a = 0;\n for(var j=1;j<=2;j++){\n for(var i=1;i<=60000;i++) {\n a += 1;\n }\n }\n return a;',
+    'var a = 0;\n for(var j=1;j<=2;j++){\n for(var i=1;i<=30000;i++) {\n a += 1;\n }\n }\n return a;',
   notloops:
     'console.log("do");\nconsole.log("while");\nconsole.log(" foo do bar ");\nconsole.log(" foo while bar ");\nreturn true;',
   notprops:
@@ -85,6 +85,37 @@ describe('loop', function() {
     assert(run(compiled) === true);
   });
 
+  it('console error when passing string', () => {
+    const code = `var i = 0; while (true) i++; return true`;
+
+    Babel.registerPlugin(
+      'loopProtectionAlt2',
+      require('../lib')(100, 'Loop broken')
+    );
+
+    const compiled = Babel.transform(new Function(code).toString(), {
+      plugins: ['loopProtectionAlt2'],
+    }).code; // eslint-disable-line no-new-func
+    expect(run(compiled)).toBe(true);
+  });
+
+  it('throws when giving a custom function', () => {
+    const code = `var i = 0; while (true) i++; return true`;
+    const callback = line => {
+      throw new Error(`Bad loop on line ${line}`);
+    };
+
+    Babel.registerPlugin('loopProtectionAlt', require('../lib')(100, callback));
+
+    const compiled = Babel.transform(new Function(code).toString(), {
+      plugins: ['loopProtectionAlt'],
+    }).code; // eslint-disable-line no-new-func
+
+    expect(() => {
+      run(compiled);
+    }).toThrowError('Bad loop on line 2');
+  });
+
   // https://github.com/jsbin/loop-protect/issues/5
   it('blank line', () => {
     const code = `const log = () => {};while (1)
@@ -183,7 +214,7 @@ describe('loop', function() {
     var compiled = loopProtect(c);
     var r = run(compiled);
     expect(compiled).not.toBe(c);
-    expect(r).toBe(120000);
+    expect(r).toBe(60000);
   });
 
   it('should rewrite loops when curlies are on the next line', function() {