diff --git a/lib/index.js b/lib/index.js index b3ead1b6..149e47d1 100644 --- a/lib/index.js +++ b/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), + }, + }; + }; +}; diff --git a/package.json b/package.json index 58d1baca..f3917c70 100644 --- a/package.json +++ b/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", diff --git a/test/loop-protect.test.js b/test/loop-protect.test.js index 9249222e..e6f0cf32 100644 --- a/test/loop-protect.test.js +++ b/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() {