LunAST: Misc fixes, patching with AST

packages/core/src/patch.ts

@@ -43,11 +43,39 @@ const moduleCache: Record<string, string> = {};
 const patched: Record<string, Array<string>> = {};
 
 function patchModules(entry: WebpackJsonpEntry[1]) {
- for (const [id, func] of Object.entries(entry)) {
- // `function(e,t,n){}` isn't valid I guess? so make it an IIFE to make ESTree happy
- moonlight.lunast.parseScript(id, `(${func.toString()})()`);
+ function patchModule(id: string, patchId: string, replaced: string) {
+ // Store what extensions patched what modules for easier debugging
+ patched[id] = patched[id] || [];
+ patched[id].push(patchId);
+
+ // Webpack module arguments are minified, so we replace them with consistent names
+ // We have to wrap it so things don't break, though
+ const patchedStr = patched[id].sort().join(", ");
+
+ const wrapped =
+ `(${replaced}).apply(this, arguments)\n` +
+ `// Patched by moonlight: ${patchedStr}\n` +
+ `//# sourceURL=Webpack-Module-${id}`;
+
+ try {
+ const func = new Function(
+ "module",
+ "exports",
+ "require",
+ wrapped
+ ) as WebpackModuleFunc;
+ entry[id] = func;
+ entry[id].__moonlight = true;
+ return true;
+ } catch (e) {
+ logger.warn("Error constructing function for patch", patchId, e);
+ patched[id].pop();
+ return false;
+ }
+ }
 
- let moduleString = Object.prototype.hasOwnProperty.call(moduleCache, id)
+ for (const [id, func] of Object.entries(entry)) {
+ let moduleString = Object.hasOwn(moduleCache, id)
  ? moduleCache[id]
  : func.toString().replace(/\n/g, "");
 
@@ -103,32 +131,8 @@ function patchModules(entry: WebpackJsonpEntry[1]) {
  continue;
  }
 
- // Store what extensions patched what modules for easier debugging
- patched[id] = patched[id] || [];
- patched[id].push(`${patch.ext}#${patch.id}`);
-
- // Webpack module arguments are minified, so we replace them with consistent names
- // We have to wrap it so things don't break, though
- const patchedStr = patched[id].sort().join(", ");
-
- const wrapped =
- `(${replaced}).apply(this, arguments)\n` +
- `// Patched by moonlight: ${patchedStr}\n` +
- `//# sourceURL=Webpack-Module-${id}`;
-
- try {
- const func = new Function(
- "module",
- "exports",
- "require",
- wrapped
- ) as WebpackModuleFunc;
- entry[id] = func;
- entry[id].__moonlight = true;
+ if (patchModule(id, `${patch.ext}#${patch.id}`, replaced)) {
  moduleString = replaced;
- } catch (e) {
- logger.warn("Error constructing function for patch", patch, e);
- patched[id].pop();
  }
  } else if (replace.type === PatchReplaceType.Module) {
  // Directly replace the module with a new one
@@ -146,6 +150,15 @@ function patchModules(entry: WebpackJsonpEntry[1]) {
  }
  }
 
+ let parsed = moonlight.lunast.parseScript(id, `(${moduleString})`);
+ if (parsed != null) {
+ // parseScript adds an extra ; for some reason
+ parsed = parsed.trimEnd().substring(0, parsed.lastIndexOf(";"));
+ if (patchModule(id, "lunast", parsed)) {
+ moduleString = parsed;
+ }
+ }
+
  if (moonlightNode.config.patchAll === true) {
  if (
  (typeof id !== "string" || !id.includes("_")) &&

packages/lunast/TODO.md

@@ -3,7 +3,7 @@
 - [ ] Experiment more! We need to know what's bad with this
 - [ ] Write utility functions for imports, exports, etc.
 - [ ] Map Z/ZP to default
-- [ ] Steal Webpack require and use it in our LunAST instance
+- [x] Steal Webpack require and use it in our LunAST instance
 - [ ] Map `import` statements to LunAST
 - [ ] Support patching in the AST
  - Let user modify the AST, have a function to flag it as modified, if it's modified we serialize it back into a string and put it back into Webpack

packages/lunast/src/index.ts

@@ -3,6 +3,7 @@ import { Remapped } from "./modules";
 import { getProcessors } from "./utils";
 import { parse } from "meriyah";
 import { Processor, ProcessorState } from "./remap";
+import { generate } from "astring";
 
 export default class LunAST {
  private modules: Record<string, RemapModule>;
@@ -33,30 +34,33 @@ export default class LunAST {
  return "dev";
  }
 
- public parseScript(id: string, code: string) {
- const moduleString = code.toString().replace(/\n/g, "");
+ public parseScript(id: string, code: string): string | null {
  const available = [...this.processors]
  .sort((a, b) => (a.priority ?? 0) - (b.priority ?? 0))
  .filter((x) =>
  x.find != null
  ? typeof x.find === "string"
- ? moduleString.indexOf(x.find) !== -1
- : x.find.test(moduleString)
+ ? code.indexOf(x.find) !== -1
+ : x.find.test(code)
  : true
  )
  .filter((x) =>
  x.dependencies != null
  ? x.dependencies.every((dep) => this.successful.has(dep))
  : true
  );
- if (available.length === 0) return;
+ if (available.length === 0) return null;
 
  const module = parse(code);
+ let dirty = false;
  const state: ProcessorState = {
  id,
  // @ts-expect-error The ESTree types are mismatched with estree-toolkit, but ESTree is a standard so this is fine
  ast: module,
- lunast: this
+ lunast: this,
+ markDirty: () => {
+ dirty = true;
+ }
  };
 
  for (const processor of available) {
@@ -65,6 +69,12 @@ export default class LunAST {
  this.successful.add(processor.name);
  }
  }
+
+ if (dirty) {
+ return generate(module);
+ } else {
+ return null;
+ }
  }
 
  public getType(name: string) {

packages/lunast/src/modules/test.ts

@@ -1,5 +1,6 @@
 import { traverse, is } from "estree-toolkit";
-import { getExports, getGetters, register } from "../utils";
+import { getExports, getPropertyGetters, register, magicAST } from "../utils";
+import { BlockStatement } from "estree-toolkit/dist/generated/types";
 
 // These aren't actual modules yet, I'm just using this as a testbed for stuff
 register({
@@ -13,11 +14,46 @@ register({
  }
 });
 
+// Exports example
+/*register({
+ name: "ApplicationStoreDirectoryStore",
+ find: '"displayName","ApplicationStoreDirectoryStore"',
+ process({ ast }) {
+ const exports = getExports(ast);
+ return Object.keys(exports).length > 0;
+ }
+});*/
+
+// Patching example
+register({
+ name: "ImagePreview",
+ find: ".Messages.OPEN_IN_BROWSER",
+ process({ id, ast, lunast, markDirty }) {
+ const getters = getPropertyGetters(ast);
+ const replacement = magicAST(`return require("common_react").createElement(
+ "div",
+ {
+ style: {
+ color: "white",
+ },
+ },
+ "balls"
+)`)!;
+ for (const node of Object.values(getters)) {
+ const body = node.path.get<BlockStatement>("body");
+ body.replaceWith(replacement);
+ }
+ markDirty();
+
+ return true;
+ }
+});
+
 register({
  name: "ClipboardUtils",
  find: 'document.queryCommandEnabled("copy")',
  process({ id, ast, lunast }) {
- const getters = getGetters(ast);
+ const getters = getPropertyGetters(ast);
  const fields = [];
 
  for (const [name, node] of Object.entries(getters)) {

packages/lunast/src/remap.ts

@@ -12,4 +12,5 @@ export type ProcessorState = {
  id: string;
  ast: Program;
  lunast: LunAST;
+ markDirty: () => void;
 };

packages/lunast/src/utils.ts

@@ -2,11 +2,13 @@ import type { Processor } from "./remap";
 import { traverse, is, Binding } from "estree-toolkit";
 // FIXME something's fishy with these types
 import type {
+ ExpressionStatement,
  ObjectExpression,
  Program,
  Property,
  ReturnStatement
 } from "estree-toolkit/dist/generated/types";
+import { parse } from "meriyah";
 
 export const processors: Processor[] = [];
 
@@ -24,56 +26,58 @@ export function getExports(ast: Program) {
 
  traverse(ast, {
  $: { scope: true },
- BlockStatement: {
- enter(path) {
- // Walk up to make sure we are indeed the top level
- let parent = path.parentPath;
- while (!is.program(parent)) {
- parent = parent?.parentPath ?? null;
- if (
- parent == null ||
- parent.node == null ||
- ![
- "FunctionExpression",
- "ExpressionStatement",
- "CallExpression",
- "Program"
- ].includes(parent.node.type)
- ) {
- this.stop();
- return;
- }
+ BlockStatement(path) {
+ // Walk up to make sure we are indeed the top level
+ let parent = path.parentPath;
+ while (!is.program(parent)) {
+ parent = parent?.parentPath ?? null;
+ if (
+ parent == null ||
+ parent.node == null ||
+ ![
+ "FunctionExpression",
+ "ExpressionStatement",
+ "CallExpression",
+ "Program"
+ ].includes(parent.node.type)
+ ) {
+ return;
  }
- },
-
- leave(path) {
- path.scope?.crawl();
- if (!is.functionExpression(path.parent)) return;
-
- for (let i = 0; i < path.parent.params.length; i++) {
- const param = path.parent.params[i];
- if (!is.identifier(param)) continue;
- const binding = path.scope?.getBinding(param.name);
- if (!binding) continue;
-
- // module
- if (i === 0) {
- for (const reference of binding.references) {
- if (!is.identifier(reference.node)) continue;
- if (!is.assignmentExpression(reference.parentPath?.parentPath))
- continue;
-
- const exports = reference.parentPath?.parentPath.node?.right;
- if (!is.objectExpression(exports)) continue;
-
- for (const property of exports.properties) {
- if (!is.property(property)) continue;
- if (!is.identifier(property.key)) continue;
- ret[property.key.name] = property.value;
- }
+ }
+
+ if (!is.functionExpression(path.parent)) return;
+
+ for (let i = 0; i < path.parent.params.length; i++) {
+ const param = path.parent.params[i];
+ if (!is.identifier(param)) continue;
+ const binding = path.scope?.getBinding(param.name);
+ if (!binding) continue;
+
+ // module
+ if (i === 0) {
+ for (const reference of binding.references) {
+ if (!is.identifier(reference.node)) continue;
+ if (!is.assignmentExpression(reference.parentPath?.parentPath))
+ continue;
+
+ const exportsNode = reference.parentPath?.parentPath.node;
+ if (!is.memberExpression(exportsNode?.left)) continue;
+ if (!is.identifier(exportsNode.left.property)) continue;
+ if (exportsNode.left.property.name !== "exports") continue;
+
+ const exports = exportsNode?.right;
+ if (!is.objectExpression(exports)) continue;
+
+ for (const property of exports.properties) {
+ if (!is.property(property)) continue;
+ if (!is.identifier(property.key)) continue;
+ ret[property.key.name] = property.value;
  }
  }
- // TODO: exports
+ }
+ // TODO: exports
+ else if (i === 1) {
+ // console.log("getExports:", path, param, binding);
  }
  }
  }
@@ -82,7 +86,7 @@ export function getExports(ast: Program) {
  return ret;
 }
 
-export function getGetters(ast: Program) {
+export function getPropertyGetters(ast: Program) {
  const ret: Record<string, Binding> = {};
 
  traverse(ast, {
@@ -124,3 +128,18 @@ export function getGetters(ast: Program) {
 
  return ret;
 }
+
+export function magicAST(code: string) {
+ // Wraps code in an IIFE so you can type `return` and all that goodies
+ // Might not work for some other syntax issues but oh well
+ const tree = parse("(()=>{" + code + "})()");
+
+ const expressionStatement = tree.body[0] as ExpressionStatement;
+ if (!is.expressionStatement(expressionStatement)) return null;
+ if (!is.callExpression(expressionStatement.expression)) return null;
+ if (!is.arrowFunctionExpression(expressionStatement.expression.callee))
+ return null;
+ if (!is.blockStatement(expressionStatement.expression.callee.body))
+ return null;
+ return expressionStatement.expression.callee.body;
+}

tsconfig.json

@@ -1,6 +1,6 @@
 {
  "compilerOptions": {
- "target": "es2016",
+ "target": "es2022",
  "module": "es6",
  "esModuleInterop": true,
  "forceConsistentCasingInFileNames": true,