Fixed wrong arrow function block scope detection

sanex3339 · sanex3339 · commit 7170a2b476eb · 2018-07-24T23:53:34.000+03:00
diff --git a/dist/index.cli.js b/dist/index.cli.js
diff --git a/dist/index.js b/dist/index.js
diff --git a/src/node/NodeGuards.ts b/src/node/NodeGuards.ts
@@ -253,10 +253,12 @@ export class NodeGuards {
      * @returns {boolean}
      */
     public static isNodeHasBlockScope (node: ESTree.Node, parentNode: ESTree.Node): node is TNodeWithBlockScope {
-        return NodeGuards.isProgramNode(node) || (
-            NodeGuards.isBlockStatementNode(node)
-            && NodeGuards.nodesWithBlockScope.includes(parentNode.type)
-        );
+        return NodeGuards.isProgramNode(node)
+            /**
+             * Should correctly check arrow functions with expression body
+             */
+            || (NodeGuards.isArrowFunctionExpressionNode(node) && !NodeGuards.isBlockStatementNode(node.body))
+            || (NodeGuards.isBlockStatementNode(node) && NodeGuards.nodesWithBlockScope.includes(parentNode.type));
     }
 
     /**
@@ -265,6 +267,10 @@ export class NodeGuards {
      */
     public static isNodeHasScope (node: ESTree.Node): node is TNodeWithScope {
         return NodeGuards.isProgramNode(node)
+            /**
+             * Should correctly check arrow functions with expression body
+             */
+            || (NodeGuards.isArrowFunctionExpressionNode(node) && !NodeGuards.isBlockStatementNode(node.body))
             || NodeGuards.isBlockStatementNode(node)
             || NodeGuards.isSwitchCaseNode(node);
     }
diff --git a/test/dev/dev.ts b/test/dev/dev.ts
@@ -6,15 +6,8 @@ import { NO_ADDITIONAL_NODES_PRESET } from '../../src/options/presets/NoCustomNo
 
     let obfuscatedCode: string = JavaScriptObfuscator.obfuscate(
         `
-        function foo (data, options) {
-            function bar ({data, ...rest}) {
-                function baz ({options}) {
-                    return data + options + rest;
-                }
-            }
-    
-            return data;
-        }
+        const foo = 1;
+        [].map(foo=>1).map(bar=>[foo]);
         `,
         {
             ...NO_ADDITIONAL_NODES_PRESET,
diff --git a/test/functional-tests/node-transformers/obfuscating-transformers/function-transformer/FunctionTransformer.spec.ts b/test/functional-tests/node-transformers/obfuscating-transformers/function-transformer/FunctionTransformer.spec.ts
@@ -436,4 +436,60 @@ describe('FunctionTransformer', () => {
             });
         });
     });
+
+    describe('correct block scope detection of arrow function expression', () => {
+        describe('Variant #1: block statement body', () => {
+            const regExpMatch: string = `` +
+                `\\[]` +
+                `\\['map']\\(_0x[a-f0-9]{4,6} *=> *\\{ *return 0x1; *\\}\\)` +
+                `\\['map']\\(_0x[a-f0-9]{4,6} *=> *\\[foo]\\);` +
+            ``;
+            const regExp: RegExp = new RegExp(regExpMatch);
+
+            let obfuscatedCode: string;
+
+            before(() => {
+                const code: string = readFileAsString(__dirname + '/fixtures/arrow-function-with-expression-body-block-scope-detection-1.js');
+                const obfuscationResult: IObfuscationResult = JavaScriptObfuscator.obfuscate(
+                    code,
+                    {
+                        ...NO_ADDITIONAL_NODES_PRESET
+                    }
+                );
+
+                obfuscatedCode = obfuscationResult.getObfuscatedCode();
+            });
+
+            it('should transform identifiers in arrow function expression body', () => {
+                assert.match(obfuscatedCode, regExp);
+            });
+        });
+
+        describe('Variant #1: expression statement body', () => {
+            const regExpMatch: string = `` +
+                `\\[]` +
+                `\\['map']\\(_0x[a-f0-9]{4,6} *=> *0x1\\)` +
+                `\\['map']\\(_0x[a-f0-9]{4,6} *=> *\\[foo]\\);` +
+            ``;
+            const regExp: RegExp = new RegExp(regExpMatch);
+
+            let obfuscatedCode: string;
+
+            before(() => {
+                const code: string = readFileAsString(__dirname + '/fixtures/arrow-function-with-expression-body-block-scope-detection-2.js');
+                const obfuscationResult: IObfuscationResult = JavaScriptObfuscator.obfuscate(
+                    code,
+                    {
+                        ...NO_ADDITIONAL_NODES_PRESET
+                    }
+                );
+
+                obfuscatedCode = obfuscationResult.getObfuscatedCode();
+            });
+
+            it('should transform identifiers in arrow function expression body', () => {
+                assert.match(obfuscatedCode, regExp);
+            });
+        });
+    });
 });
diff --git a/test/functional-tests/node-transformers/obfuscating-transformers/function-transformer/fixtures/arrow-function-with-expression-body-block-scope-detection-1.js b/test/functional-tests/node-transformers/obfuscating-transformers/function-transformer/fixtures/arrow-function-with-expression-body-block-scope-detection-1.js
@@ -0,0 +1,6 @@
+const foo = 1;
+[]
+    .map(foo => {
+        return 1;
+    })
+    .map(bar => [foo]);
diff --git a/test/functional-tests/node-transformers/obfuscating-transformers/function-transformer/fixtures/arrow-function-with-expression-body-block-scope-detection-2.js b/test/functional-tests/node-transformers/obfuscating-transformers/function-transformer/fixtures/arrow-function-with-expression-body-block-scope-detection-2.js
@@ -0,0 +1,4 @@
+const foo = 1;
+[]
+    .map(foo => 1)
+    .map(bar => [foo]);
