perf: use node names instead of enter/exit in visitors

j4k0xb · May 24, 2023 · 28145cc · 28145cc
1 parent 4a058f3
commit 28145cc
Show file tree

Hide file tree

Showing 20 changed files with 174 additions and 106 deletions.
diff --git a/src/deobfuscator/controlFlowObject.ts b/src/deobfuscator/controlFlowObject.ts
@@ -93,7 +93,7 @@ export default {
       return binding.constant || binding.constantViolations[0] === binding.path;
     }
 
-    function transform(path: NodePath) {
+    function transform(path: NodePath<t.VariableDeclarator>) {
       let changes = 0;
       if (varMatcher.match(path.node)) {
         // Verify all references to make sure they match how the obfuscator
@@ -135,7 +135,10 @@ export default {
 
         oldRefs.forEach(ref => {
           const varDeclarator = ref.findParent(p => p.isVariableDeclarator());
-          if (varDeclarator) changes += transform(varDeclarator);
+          if (varDeclarator)
+            changes += transform(
+              varDeclarator as NodePath<t.VariableDeclarator>
+            );
         });
 
         path.remove();
@@ -187,7 +190,7 @@ export default {
     }
 
     return {
-      enter(path) {
+      VariableDeclarator(path) {
         this.changes += transform(path);
       },
     };

diff --git a/src/deobfuscator/controlFlowSwitch.ts b/src/deobfuscator/controlFlowSwitch.ts
@@ -60,7 +60,7 @@ export default {
     );
 
     return {
-      enter(path) {
+      BlockStatement(path) {
         if (!matcher.match(path.node)) return;
 
         const caseStatements = new Map(

diff --git a/src/deobfuscator/debugProtection.ts b/src/deobfuscator/debugProtection.ts
@@ -84,7 +84,7 @@ export default {
     );
 
     return {
-      enter(path) {
+      FunctionDeclaration(path) {
         if (!matcher.match(path.node)) return;
 
         const binding = path.scope.getBinding(

diff --git a/src/deobfuscator/inlineDecodedStrings.ts b/src/deobfuscator/inlineDecodedStrings.ts
@@ -54,6 +54,8 @@ function collectCalls(ast: t.Node, vm: VMDecoder) {
     conditional,
   ]);
 
+  const buildExtractedConditional = expression`TEST ? CALLEE(CONSEQUENT) : CALLEE(ALTERNATE)`;
+
   traverse(ast, {
     CallExpression(path) {
       // decode(test ? 1 : 2) -> test ? decode(1) : decode(2)
@@ -62,7 +64,12 @@ function collectCalls(ast: t.Node, vm: VMDecoder) {
         const { test, consequent, alternate } = conditional.current!;
 
         path.replaceWith(
-          expression`${test} ? ${callee}(${consequent}) : ${callee}(${alternate})`()
+          buildExtractedConditional({
+            TEST: test,
+            CALLEE: callee,
+            CONSEQUENT: consequent,
+            ALTERNATE: alternate,
+          })
         );
       } else if (matcher.match(path.node)) {
         calls.push(path);

diff --git a/src/deobfuscator/objectLiterals.ts b/src/deobfuscator/objectLiterals.ts
@@ -34,7 +34,7 @@ export default {
     );
 
     return {
-      enter(path) {
+      VariableDeclarator(path) {
         if (!varMatcher.match(path.node)) return;
 
         const binding = path.scope.getBinding(varId.current!.name);

diff --git a/src/deobfuscator/selfDefending.ts b/src/deobfuscator/selfDefending.ts
@@ -109,7 +109,7 @@ export default {
     );
 
     return {
-      enter(path) {
+      VariableDeclarator(path) {
         if (!matcher.match(path.node)) return;
         const binding = path.scope.getBinding(callController.current!)!;
         // const callControllerFunctionName = (function() { ... })();

diff --git a/src/extractor/bundle.ts b/src/extractor/bundle.ts
@@ -19,20 +19,23 @@ export class Bundle {
   }
 
   applyMappings(mappings: Record<string, m.Matcher<unknown>>): void {
-    const unusedMappings = new Set(Object.keys(mappings));
+    const mappingPaths = Object.keys(mappings);
+    if (!mappingPaths.length) return;
+
+    const unusedMappings = new Set(mappingPaths);
 
     for (const module of this.modules.values()) {
       traverse(module.ast, {
         enter(path) {
-          for (const [name, matcher] of Object.entries(mappings)) {
-            if (matcher.match(path.node)) {
-              if (unusedMappings.has(name)) {
-                unusedMappings.delete(name);
+          for (const mappingPath of mappingPaths) {
+            if (mappings[mappingPath].match(path.node)) {
+              if (unusedMappings.has(mappingPath)) {
+                unusedMappings.delete(mappingPath);
               } else {
-                console.warn(`Mapping ${name} is already used.`);
+                console.warn(`Mapping ${mappingPath} is already used.`);
                 continue;
               }
-              module.path = name;
+              module.path = mappingPath;
               path.stop();
               break;
             }

diff --git a/src/extractor/webpack/bundle.ts b/src/extractor/webpack/bundle.ts
@@ -36,7 +36,7 @@ export class WebpackBundle extends Bundle {
 
     this.modules.forEach(module => {
       traverse(module.ast, {
-        enter: path => {
+        CallExpression: path => {
           if (requireMatcher.match(path.node)) {
             const requiredModule = this.modules.get(requireId.current!.value);
             if (requiredModule) {
@@ -45,7 +45,10 @@ export class WebpackBundle extends Bundle {
                 t.stringLiteral(relativePath(module.path, requiredModule.path))
               );
             }
-          } else if (importMatcher.match(path.node)) {
+          }
+        },
+        ImportDeclaration: path => {
+          if (importMatcher.match(path.node)) {
             const requiredModule = this.modules.get(
               Number(importId.current!.value)
             );

diff --git a/src/extractor/webpack/esm.ts b/src/extractor/webpack/esm.ts
@@ -85,6 +85,8 @@ export function convertESM(module: WebpackModule): void {
     )
   );
 
+  const buildImport = statement`import * as NAME from "PATH";`;
+
   traverse(module.ast, {
     enter(path) {
       // Only traverse the top-level
@@ -98,9 +100,10 @@ export function convertESM(module: WebpackModule): void {
         requireMatcher.match(path.node)
       ) {
         path.replaceWith(
-          statement`import * as ${requireVariable.current} from "${String(
-            requiredModuleId.current
-          )}";`()
+          buildImport({
+            NAME: requireVariable.current,
+            PATH: String(requiredModuleId.current),
+          })
         );
       } else if (defineExportsMatcher.match(path.node)) {
         const exportsBinding = path.scope.getBinding(exportsName.current!.name);
@@ -164,13 +167,9 @@ function exportVariable(
       renameFast(binding, exportName);
       declaration.replaceWith(t.exportNamedDeclaration(declaration.node));
     }
+  } else if (exportName === 'default') {
+    requireDPath.insertAfter(statement`export default ${value}`());
   } else {
-    if (exportName === 'default') {
-      requireDPath.insertAfter(statement`export default ${value}`());
-    } else {
-      requireDPath.insertAfter(
-        statement`export let ${exportName} = ${value}`()
-      );
-    }
+    requireDPath.insertAfter(statement`export let ${exportName} = ${value}`());
   }
 }
diff --git a/src/extractor/webpack/getDefaultExport.ts b/src/extractor/webpack/getDefaultExport.ts
@@ -66,14 +66,18 @@ export function convertDefaultRequire(bundle: WebpackBundle): void {
     m.callExpression(requireN, [])
   );
 
+  const buildDefaultAccess = expression`OBJECT.default`;
+
   bundle.modules.forEach(module => {
     traverse(module.ast, {
-      enter(path) {
+      ['CallExpression|MemberExpression' as 'Expression'](path: NodePath) {
         if (defaultRequireMatcherAlternative.match(path.node)) {
           // Replace require.n(m).a or require.n(m)() with m or m.default
           const requiredModule = getRequiredModule(path);
           if (requiredModule?.ast.program.sourceType === 'module') {
-            path.replaceWith(expression`${moduleArg.current!}.default`());
+            path.replaceWith(
+              buildDefaultAccess({ OBJECT: moduleArg.current! })
+            );
           } else {
             path.replaceWith(moduleArg.current!);
           }
@@ -85,7 +89,9 @@ export function convertDefaultRequire(bundle: WebpackBundle): void {
           const requiredModule = getRequiredModule(path);
           const init = path.get('init');
           if (requiredModule?.ast.program.sourceType === 'module') {
-            init.replaceWith(expression`${moduleArg.current!}.default`());
+            init.replaceWith(
+              buildDefaultAccess({ OBJECT: moduleArg.current! })
+            );
           } else {
             init.replaceWith(moduleArg.current!);
           }

diff --git a/src/extractor/webpack/varInjection.ts b/src/extractor/webpack/varInjection.ts
@@ -4,6 +4,8 @@ import * as m from '@codemod/matchers';
 import { constMemberExpression } from '../../utils/matcher';
 import { WebpackModule } from './module';
 
+const buildVar = statement`var NAME = INIT;`;
+
 /**
  * ```js
  * (function(global) {
@@ -38,7 +40,7 @@ export function inlineVarInjections(module: WebpackModule): void {
   for (const node of program.body) {
     if (matcher.match(node)) {
       const vars = params.current!.map((param, i) =>
-        statement`var ${param} = ${args.current![i + 1]};`()
+        buildVar({ NAME: param, INIT: args.current![i + 1] })
       );
       newBody.push(...vars);
       newBody.push(...body.current!.body);

diff --git a/src/transforms/booleanIf.ts b/src/transforms/booleanIf.ts
@@ -6,26 +6,37 @@ import { Transform } from '.';
 export default {
   name: 'booleanIf',
   tags: ['safe'],
-  visitor: () => ({
-    ExpressionStatement: {
-      exit(path) {
-        const expression = path.node.expression as t.LogicalExpression;
-        if (andMatcher.match(path.node)) {
-          path.replaceWith(
-            statement`if (${expression.left}) { ${expression.right}; }`()
-          );
-          this.changes++;
-        } else if (orMatcher.match(path.node)) {
-          path.replaceWith(
-            statement`if (!${expression.left}) { ${expression.right}; }`()
-          );
-          this.changes++;
-        }
+  visitor: () => {
+    const andMatcher = m.expressionStatement(m.logicalExpression('&&'));
+    const orMatcher = m.expressionStatement(m.logicalExpression('||'));
+
+    const buildIf = statement`if (TEST) { BODY; }`;
+    const buildIfNot = statement`if (!TEST) { BODY; }`;
+
+    return {
+      ExpressionStatement: {
+        exit(path) {
+          const expression = path.node.expression as t.LogicalExpression;
+          if (andMatcher.match(path.node)) {
+            path.replaceWith(
+              buildIf({
+                TEST: expression.left,
+                BODY: expression.right,
+              })
+            );
+            this.changes++;
+          } else if (orMatcher.match(path.node)) {
+            path.replaceWith(
+              buildIfNot({
+                TEST: expression.left,
+                BODY: expression.right,
+              })
+            );
+            this.changes++;
+          }
+        },
       },
-    },
-    noScope: true,
-  }),
+      noScope: true,
+    };
+  },
 } satisfies Transform;
-
-const andMatcher = m.expressionStatement(m.logicalExpression('&&'));
-const orMatcher = m.expressionStatement(m.logicalExpression('||'));
diff --git a/src/transforms/computedProperties.ts b/src/transforms/computedProperties.ts
@@ -1,7 +1,8 @@
 import { isIdentifierName } from '@babel/helper-validator-identifier';
+import { NodePath } from '@babel/traverse';
 import * as t from '@babel/types';
 import * as m from '@codemod/matchers';
-import { Transform } from '.';
+import { Transform, TransformState } from '.';
 
 export default {
   name: 'computedProperties',
@@ -22,16 +23,25 @@ export default {
     );
 
     return {
-      exit(path) {
-        if (propertyMatcher.match(path.node)) {
-          path.node.computed = false;
-          path.node.property = t.identifier(stringMatcher.current!.value);
-          this.changes++;
-        } else if (keyMatcher.match(path.node)) {
-          path.node.computed = false;
-          path.node.key = t.identifier(stringMatcher.current!.value);
-          this.changes++;
-        }
+      // https://github.com/babel/babel/pull/14862/files
+      // isn't included in the @types/babel__traverse package and can't be augmented
+      ['MemberExpression|OptionalMemberExpression' as 'Expression']: {
+        exit(this: TransformState, path: NodePath) {
+          if (propertyMatcher.match(path.node)) {
+            path.node.computed = false;
+            path.node.property = t.identifier(stringMatcher.current!.value);
+            this.changes++;
+          }
+        },
+      },
+      ['ObjectProperty|ClassProperty|ObjectMethod|ClassMethod' as 'Expression']: {
+        exit(this: TransformState, path: NodePath) {
+          if (keyMatcher.match(path.node)) {
+            path.node.computed = false;
+            path.node.key = t.identifier(stringMatcher.current!.value);
+            this.changes++;
+          }
+        },
       },
       noScope: true,
     };

diff --git a/src/transforms/mergeStrings.ts b/src/transforms/mergeStrings.ts
@@ -17,19 +17,26 @@ export default {
     );
 
     return {
-      exit(path) {
-        if (matcher.match(path.node)) {
-          // "a" + "b" -> "ab"
-          path.replaceWith(
-            t.stringLiteral(left.current!.value + right.current!.value)
-          );
-          this.changes++;
-        } else if (nestedMatcher.match(path.parent) && path.isStringLiteral()) {
-          // a + "b" + "c" -> a + "bc"
-          left.current!.value += right.current!.value;
-          path.remove();
-          this.changes++;
-        }
+      BinaryExpression: {
+        exit(path) {
+          if (matcher.match(path.node)) {
+            // "a" + "b" -> "ab"
+            path.replaceWith(
+              t.stringLiteral(left.current!.value + right.current!.value)
+            );
+            this.changes++;
+          }
+        },
+      },
+      StringLiteral: {
+        exit(path) {
+          if (nestedMatcher.match(path.parent)) {
+            // a + "b" + "c" -> a + "bc"
+            left.current!.value += right.current!.value;
+            path.remove();
+            this.changes++;
+          }
+        },
       },
 
       noScope: true,