control flow flattening improvements: less storage nodes, faster speed

src/custom-nodes/AbstractCustomNode.ts

@@ -31,7 +31,9 @@ export abstract class AbstractCustomNode implements ICustomNode {
     /**
      * @returns {string}
      */
-    public abstract getCode (): string;
+    public getCode (): string {
+        return NodeUtils.convertStructureToCode(this.getNode());
+    }
 
     /**
      * @returns {TStatement[]}
@@ -44,6 +46,11 @@ export abstract class AbstractCustomNode implements ICustomNode {
      * @returns {TStatement[]}
      */
     protected getNodeStructure (): TStatement[] {
-        return NodeUtils.convertCodeToStructure(this.getCode());
+        return NodeUtils.convertCodeToStructure(this.getTemplate());
     }
+
+    /**
+     * @returns {string}
+     */
+    protected abstract getTemplate (): string;
 }

src/custom-nodes/console-output-nodes/ConsoleOutputDisableExpressionNode.ts

@@ -39,7 +39,7 @@ export class ConsoleOutputDisableExpressionNode extends AbstractCustomNode {
     /**
      * @returns {string}
      */
-    public getCode (): string {
+    protected getTemplate (): string {
         return format(ConsoleOutputDisableExpressionTemplate(), {
             consoleLogDisableFunctionName: RandomGeneratorUtils.getRandomVariableName(),
             singleNodeCallControllerFunctionName: this.callsControllerFunctionName

src/custom-nodes/control-flow-replacers-nodes/binary-expression-control-flow-replacer-nodes/BinaryExpressionFunctionNode.ts

@@ -32,14 +32,14 @@ export class BinaryExpressionFunctionNode extends AbstractCustomNode {
     /**
      * @param operator
      */
-    initialize (operator: string): void {
+    public initialize (operator: string): void {
         this.operator = operator;
     }
 
     /**
      * @returns {string}
      */
-    public getCode (): string {
+    protected getTemplate (): string {
         return format(BinaryExpressionFunctionTemplate(), {
             functionName: RandomGeneratorUtils.getRandomVariableName(1),
             operator: this.operator

src/custom-nodes/control-flow-replacers-nodes/binary-expression-control-flow-replacer-nodes/ControlFlowStorageCallNode.ts

@@ -67,7 +67,7 @@ export class ControlFlowStorageCallNode extends AbstractCustomNode {
     /**
      * @returns {string}
      */
-    public getCode (): string {
+    protected getTemplate (): string {
         return format(ControlFlowStorageCallTemplate(), {
             controlFlowStorageKey: this.controlFlowStorageKey,
             controlFlowStorageName: this.controlFlowStorageName,

src/custom-nodes/control-flow-storage-nodes/ControlFlowStorageNode.ts

@@ -48,7 +48,7 @@ export class ControlFlowStorageNode extends AbstractCustomNode {
     /**
      * @returns {string}
      */
-    public getCode (): string {
+    protected getTemplate (): string {
         return format(ControlFlowStorageTemplate(), {
             controlFlowStorage: this.controlFlowStorage.toString(),
             controlFlowStorageName: this.controlFlowStorageName

src/custom-nodes/debug-protection-nodes/DebugProtectionFunctionCallNode.ts

@@ -38,7 +38,7 @@ export class DebugProtectionFunctionCallNode extends AbstractCustomNode {
     /**
      * @returns {string}
      */
-    public getCode (): string {
+    protected getTemplate (): string {
         return format(DebugProtectionFunctionCallTemplate(), {
             debugProtectionFunctionName: this.debugProtectionFunctionName
         });

src/custom-nodes/debug-protection-nodes/DebugProtectionFunctionIntervalNode.ts

@@ -38,7 +38,7 @@ export class DebugProtectionFunctionIntervalNode extends AbstractCustomNode {
     /**
      * @returns {string}
      */
-    public getCode (): string {
+    protected getTemplate (): string {
         return format(DebugProtectionFunctionIntervalTemplate(), {
             debugProtectionFunctionName: this.debugProtectionFunctionName
         });

src/custom-nodes/debug-protection-nodes/DebugProtectionFunctionNode.ts

@@ -38,7 +38,7 @@ export class DebugProtectionFunctionNode extends AbstractCustomNode {
     /**
      * @returns {string}
      */
-    public getCode (): string {
+    protected getTemplate (): string {
         return format(DebugProtectionFunctionTemplate(), {
             debugProtectionFunctionName: this.debugProtectionFunctionName
         });

src/custom-nodes/domain-lock-nodes/DomainLockNode.ts

@@ -40,7 +40,7 @@ export class DomainLockNode extends AbstractCustomNode {
     /**
      * @returns {string}
      */
-    public getCode (): string {
+    protected getTemplate (): string {
         let domainsString: string = this.options.domainLock.join(';'),
             [hiddenDomainsString, diff]: string[] = CryptUtils.hideString(domainsString, domainsString.length * 3);
 

src/custom-nodes/node-calls-controller-nodes/NodeCallsControllerFunctionNode.ts

@@ -53,7 +53,7 @@ export class NodeCallsControllerFunctionNode extends AbstractCustomNode {
     /**
      * @returns {string}
      */
-    public getCode (): string {
+    protected getTemplate (): string {
         if (this.appendEvent === ObfuscationEvents.AfterObfuscation) {
             return JavaScriptObfuscator.obfuscate(
                 format(SingleNodeCallControllerTemplate(), {

src/custom-nodes/self-defending-nodes/SelfDefendingUnicodeNode.ts

@@ -42,7 +42,7 @@ export class SelfDefendingUnicodeNode extends AbstractCustomNode {
     /**
      * @returns {string}
      */
-    public getCode (): string {
+    protected getTemplate (): string {
         return JavaScriptObfuscator.obfuscate(
             format(SelfDefendingTemplate(), {
                 selfDefendingFunctionName: RandomGeneratorUtils.getRandomVariableName(),

src/custom-nodes/string-array-nodes/StringArrayCallsWrapper.ts

@@ -67,25 +67,6 @@ export class StringArrayCallsWrapper extends AbstractCustomNode implements ICust
         this.stringArrayCallsWrapperName = stringArrayCallsWrapperName;
     }
 
-    /**
-     * @returns {string}
-     */
-    public getCode (): string {
-        const decodeNodeTemplate: string = this.getDecodeStringArrayTemplate();
-
-        return JavaScriptObfuscator.obfuscate(
-            format(StringArrayCallsWrapperTemplate(), {
-                decodeNodeTemplate,
-                stringArrayCallsWrapperName: this.stringArrayCallsWrapperName,
-                stringArrayName: this.stringArrayName
-            }),
-            {
-                ...NO_CUSTOM_NODES_PRESET,
-                seed: this.options.seed
-            }
-        ).getObfuscatedCode();
-    }
-
     /**
      * @returns {string}
      */
@@ -130,4 +111,23 @@ export class StringArrayCallsWrapper extends AbstractCustomNode implements ICust
 
         return decodeStringArrayTemplate;
     }
+
+    /**
+     * @returns {string}
+     */
+    protected getTemplate (): string {
+        const decodeNodeTemplate: string = this.getDecodeStringArrayTemplate();
+
+        return JavaScriptObfuscator.obfuscate(
+            format(StringArrayCallsWrapperTemplate(), {
+                decodeNodeTemplate,
+                stringArrayCallsWrapperName: this.stringArrayCallsWrapperName,
+                stringArrayName: this.stringArrayName
+            }),
+            {
+                ...NO_CUSTOM_NODES_PRESET,
+                seed: this.options.seed
+            }
+        ).getObfuscatedCode();
+    }
 }

src/custom-nodes/string-array-nodes/StringArrayNode.ts

@@ -64,16 +64,6 @@ export class StringArrayNode extends AbstractCustomNode {
         this.stringArrayRotateValue = stringArrayRotateValue;
     }
 
-    /**
-     * @returns {string}
-     */
-    public getCode (): string {
-        return format(StringArrayTemplate(), {
-            stringArrayName: this.stringArrayName,
-            stringArray: this.stringArray.toString()
-        });
-    }
-
     /**
      * @returns {TStatement[]}
      */
@@ -82,4 +72,14 @@ export class StringArrayNode extends AbstractCustomNode {
 
         return super.getNode();
     }
+
+    /**
+     * @returns {string}
+     */
+    protected getTemplate (): string {
+        return format(StringArrayTemplate(), {
+            stringArrayName: this.stringArrayName,
+            stringArray: this.stringArray.toString()
+        });
+    }
 }

src/custom-nodes/string-array-nodes/StringArrayRotateFunctionNode.ts

@@ -65,7 +65,7 @@ export class StringArrayRotateFunctionNode extends AbstractCustomNode {
     /**
      * @returns {string}
      */
-    public getCode (): string {
+    protected getTemplate (): string {
         let code: string = '',
             timesName: string = RandomGeneratorUtils.getRandomVariableName(),
             whileFunctionName: string = RandomGeneratorUtils.getRandomVariableName();

src/declarations/ESTree.d.ts

@@ -6,6 +6,7 @@ declare module 'estree' {
     interface BaseNode {
         parentNode?: ESTree.Node;
         obfuscated?: boolean;
+        skipByControlFlow?: boolean;
     }
 
     interface SimpleLiteral extends ESTree.BaseNode, ESTree.BaseExpression {

src/enums/NodeType.ts

@@ -25,6 +25,8 @@ export const NodeType: any = Utils.strEnumify({
     Program: 'Program',
     Property: 'Property',
     ReturnStatement: 'ReturnStatement',
+    SwitchCase: 'SwitchCase',
+    SwitchStatement: 'SwitchStatement',
     TryStatement: 'TryStatement',
     UnaryExpression: 'UnaryExpression',
     UpdateExpression: 'UpdateExpression',

src/interfaces/node-transformers/IControlFlowReplacer.d.ts

@@ -9,12 +9,12 @@ export interface IControlFlowReplacer {
      * @param parentNode
      * @param controlFlowStorage
      * @param controlFlowStorageCustomNodeName
-     * @returns ICustomNode | undefined
+     * @returns ESTree.Node
      */
     replace (
         node: ESTree.Node,
         parentNode: ESTree.Node,
         controlFlowStorage: IStorage <ICustomNode>,
         controlFlowStorageCustomNodeName: string
-    ): ICustomNode | undefined;
+    ): ESTree.Node;
 }

src/node-transformers/node-control-flow-transformers/FunctionControlFlowTransformer.ts

@@ -7,7 +7,6 @@ import * as ESTree from 'estree';
 import { TControlFlowReplacerFactory } from '../../types/container/TControlFlowReplacerFactory';
 import { TControlFlowStorageFactory } from '../../types/container/TControlFlowStorageFactory';
 import { TCustomNodeFactory } from '../../types/container/TCustomNodeFactory';
-import { TStatement } from '../../types/node/TStatement';
 
 import { ICustomNode } from '../../interfaces/custom-nodes/ICustomNode';
 import { IOptions } from '../../interfaces/options/IOptions';
@@ -20,6 +19,7 @@ import { AbstractNodeTransformer } from '../AbstractNodeTransformer';
 import { Node } from '../../node/Node';
 import { NodeAppender } from '../../node/NodeAppender';
 import { NodeControlFlowReplacers } from '../../enums/container/NodeControlFlowReplacers';
+import { NodeUtils } from '../../node/NodeUtils';
 import { RandomGeneratorUtils } from '../../utils/RandomGeneratorUtils';
 
 @injectable()
@@ -92,24 +92,11 @@ export class FunctionControlFlowTransformer extends AbstractNodeTransformer {
                     return;
                 }
 
-                const controlFlowStorageCallCustomNode: ICustomNode | undefined = this.controlFlowReplacerFactory(controlFlowReplacerName)
-                    .replace(node, parentNode, controlFlowStorage, controlFlowStorageCustomNodeName);
-
-                if (!controlFlowStorageCallCustomNode) {
-                    return;
-                }
-
-                // controlFlowStorageCallCustomNode will always have only one TStatement node,
-                // so we can get it by index `0`
-                // also we need to return `expression` property of `ExpressionStatement` node because bug:
-                // https://github.com/estools/escodegen/issues/289
-                const statementNode: TStatement = controlFlowStorageCallCustomNode.getNode()[0];
-
-                if (!statementNode || !Node.isExpressionStatementNode(statementNode)) {
-                    throw new Error(`\`controlFlowStorageCallCustomNode.getNode()\` should returns array with \`ExpressionStatement\` node`);
-                }
-
-                return statementNode.expression;
+                return {
+                    ...this.controlFlowReplacerFactory(controlFlowReplacerName)
+                        .replace(node, parentNode, controlFlowStorage, controlFlowStorageCustomNodeName),
+                    parentNode
+                };
             }
         });
 
@@ -121,6 +108,12 @@ export class FunctionControlFlowTransformer extends AbstractNodeTransformer {
 
         controlFlowStorageCustomNode.initialize(controlFlowStorage, controlFlowStorageCustomNodeName);
 
-        NodeAppender.prependNode(functionNode.body, controlFlowStorageCustomNode.getNode());
+        NodeAppender.prependNode(
+            NodeUtils.getBlockScopeOfNode(
+                functionNode.body,
+                RandomGeneratorUtils.getRandomInteger(0, 5)
+            ),
+            controlFlowStorageCustomNode.getNode()
+        );
     }
 }

src/node-transformers/node-control-flow-transformers/control-flow-replacers/AbstractControlFlowReplacer.ts

@@ -41,12 +41,12 @@ export abstract class AbstractControlFlowReplacer implements IControlFlowReplace
      * @param parentNode
      * @param controlFlowStorage
      * @param controlFlowStorageCustomNodeName
-     * @returns {ICustomNode | undefined}
+     * @returns {ESTree.Node}
      */
     public abstract replace (
         node: ESTree.Node,
         parentNode: ESTree.Node,
         controlFlowStorage: IStorage <ICustomNode>,
         controlFlowStorageCustomNodeName: string
-    ): ICustomNode | undefined;
+    ): ESTree.Node;
 }

src/node-transformers/node-control-flow-transformers/control-flow-replacers/BinaryExpressionControlFlowReplacer.ts

@@ -5,6 +5,7 @@ import * as escodegen from 'escodegen';
 import * as ESTree from 'estree';
 
 import { TCustomNodeFactory } from '../../../types/container/TCustomNodeFactory';
+import { TStatement } from '../../../types/node/TStatement';
 
 import { ICustomNode } from '../../../interfaces/custom-nodes/ICustomNode';
 import { IOptions } from '../../../interfaces/options/IOptions';
@@ -13,6 +14,7 @@ import { IStorage } from '../../../interfaces/storages/IStorage';
 import { CustomNodes } from '../../../enums/container/CustomNodes';
 
 import { AbstractControlFlowReplacer } from './AbstractControlFlowReplacer';
+import { Node } from '../../../node/Node';
 
 @injectable()
 export class BinaryExpressionControlFlowReplacer extends AbstractControlFlowReplacer {
@@ -49,14 +51,14 @@ export class BinaryExpressionControlFlowReplacer extends AbstractControlFlowRepl
      * @param parentNode
      * @param controlFlowStorage
      * @param controlFlowStorageCustomNodeName
-     * @returns {ICustomNode}
+     * @returns {ESTree.Node}
      */
     public replace (
         binaryExpressionNode: ESTree.BinaryExpression,
         parentNode: ESTree.Node,
         controlFlowStorage: IStorage <ICustomNode>,
         controlFlowStorageCustomNodeName: string
-    ): ICustomNode {
+    ): ESTree.Node {
         const key: string = AbstractControlFlowReplacer.getStorageKey();
         const binaryExpressionFunctionNode: ICustomNode = this.customNodeFactory(CustomNodes.BinaryExpressionFunctionNode);
         const controlFlowStorageCallNode: ICustomNode = this.customNodeFactory(CustomNodes.ControlFlowStorageCallNode);
@@ -71,6 +73,12 @@ export class BinaryExpressionControlFlowReplacer extends AbstractControlFlowRepl
 
         controlFlowStorage.set(key, binaryExpressionFunctionNode);
 
-        return controlFlowStorageCallNode;
+        const statementNode: TStatement = controlFlowStorageCallNode.getNode()[0];
+
+        if (!statementNode || !Node.isExpressionStatementNode(statementNode)) {
+            throw new Error(`\`controlFlowStorageCallCustomNode.getNode()\` should returns array with \`ExpressionStatement\` node`);
+        }
+
+        return statementNode.expression;
     }
 }

src/node/Node.ts

@@ -91,6 +91,14 @@ export class Node {
         return node.type === NodeType.Identifier;
     }
 
+    /**
+     * @param node
+     * @returns {boolean}
+     */
+    public static isIfStatementNode (node: ESTree.Node): node is ESTree.IfStatement {
+        return node.type === NodeType.IfStatement;
+    }
+
     /**
      * @param node
      * @param parentNode

src/node/NodeUtils.ts

@@ -50,6 +50,22 @@ export class NodeUtils {
         return <TStatement[]>structure.body;
     }
 
+    /**
+     * @param structure
+     * @returns {string}
+     */
+    public static convertStructureToCode (structure: ESTree.Node[]): string {
+        let code: string = '';
+
+        for (let node of structure) {
+            code += escodegen.generate(node, {
+                sourceMapWithCode: true
+            }).code;
+        }
+
+        return code;
+    }
+
     /**
      * @param node
      * @param index
@@ -97,7 +113,7 @@ export class NodeUtils {
             return parentNode;
         }
 
-        return NodeUtils.getBlockScopeOfNode(parentNode);
+        return NodeUtils.getBlockScopeOfNode(parentNode, depth);
     }
 
     /**
@@ -180,12 +196,12 @@ export class NodeUtils {
         (<any>estraverse)[traverseType](node, {
             enter: (node: ESTree.Node, parentNode: ESTree.Node): any => {
                 if (node.type === nodeType && visitor.enter) {
-                    visitor.enter(node, parentNode);
+                    return visitor.enter(node, parentNode);
                 }
             },
             leave: (node: ESTree.Node, parentNode: ESTree.Node): any => {
                 if (node.type === nodeType && visitor.leave) {
-                    visitor.leave(node, parentNode);
+                    return visitor.leave(node, parentNode);
                 }
             }
         });