js-confuser-vm 0.1.1 → 0.1.2

package/src/transforms/bytecode/controlFlowFlattening.ts

@@ -1,566 +0,0 @@
-// Control Flow Flattening (CFF)
-//
-// Splits each function into basic blocks and routes all execution through a
-// while-loop + switch-style comparison chain that dispatches based on a
-// `state` register.  Original jumps become state transitions.
-//
-// ── How it works ─────────────────────────────────────────────────────────────
-//
-// 1. Each function's instruction stream is split into basic blocks at every
-//    label definition and after every terminator (JUMP, JUMP_IF_*, RETURN,
-//    THROW).
-//
-// 2. Each block is assigned a random u16 state value.  A sentinel endState
-//    (not used by any block) marks loop termination.
-//
-// 3. A dispatch loop is compiled from a Template:
-//
-//      var state = <startState>;
-//      while (state !== <endState>) {
-//        if (state === <s0>) _VM_JUMP_("<block0>");
-//        if (state === <s1>) _VM_JUMP_("<block1>");
-//        ...
-//      }
-//
-//    The Template's `state` register is extracted via compileInline() so that
-//    block bodies can write state transitions to it.
-//
-// 4. Block bodies are emitted with their original instructions.  Terminators
-//    are rewritten:
-//
-//      JUMP target         → LOAD_INT state, targetBlock.stateValue
-//                             JUMP <loopTop>
-//
-//      JUMP_IF_FALSE c, t  → JUMP_IF_TRUE c, <skipLabel>
-//                             LOAD_INT state, targetBlock.stateValue
-//                             JUMP <loopTop>
-//                             <skipLabel>:
-//                             LOAD_INT state, fallthroughBlock.stateValue
-//                             JUMP <loopTop>
-//
-//      RETURN / THROW      → kept in-place (exits the VM frame directly)
-//
-// 5. Block order is shuffled randomly so spatial locality gives no hints.
-//
-// ── Pipeline position ─────────────────────────────────────────────────────────
-// Same slot as Dispatcher: before resolveRegisters and resolveLabels.
-// Can run alongside Dispatcher (they are composable).
-
-import type {
-  Bytecode,
-  Instruction,
-  RegisterOperand,
-} from "../../types.ts";
-import { Compiler } from "../../compiler.ts";
-import { getRandomInt } from "../../utils/random-utils.ts";
-import { U16_MAX } from "../../utils/op-utils.ts";
-import { Template } from "../../template.ts";
-import {
-  ref,
-  buildMaxIdMap,
-  forEachFunction,
-  extractLabel,
-} from "../../utils/pass-utils.ts";
-
-// ── Basic block splitting ────────────────────────────────────────────────────
-
-interface BasicBlock {
-  label: string;
-  body: Bytecode;
-  terminator: Instruction | null;
-  stateValue: number;
-  // Index of the block that originally followed this one (for fallthroughs).
-  // -1 means "no successor" (last block, or ends with RETURN/THROW).
-  originalNextIndex: number;
-}
-
-function isTerminator(op: number, compiler: Compiler): boolean {
-  const OP = compiler.OP;
-  return (
-    op === OP.JUMP ||
-    op === OP.JUMP_IF_FALSE ||
-    op === OP.JUMP_IF_TRUE ||
-    op === OP.RETURN ||
-    op === OP.THROW
-  );
-}
-
-function splitBasicBlocks(
-  instrs: Bytecode,
-  compiler: Compiler,
-): BasicBlock[] {
-  const blocks: BasicBlock[] = [];
-  const usedStates = new Set<number>();
-
-  const assignState = (): number => {
-    let s: number;
-    do {
-      s = getRandomInt(0, U16_MAX);
-    } while (usedStates.has(s));
-    usedStates.add(s);
-    return s;
-  };
-
-  let currentLabel: string | null = null;
-  let currentBody: Bytecode = [];
-
-  const flushBlock = (terminator: Instruction | null) => {
-    if (currentBody.length === 0 && terminator === null && currentLabel === null)
-      return;
-
-    const label = currentLabel ?? compiler._makeLabel("cff_block");
-    blocks.push({
-      label,
-      body: currentBody,
-      terminator,
-      stateValue: assignState(),
-      originalNextIndex: -1, // filled in after all blocks are created
-    });
-    currentBody = [];
-    currentLabel = null;
-  };
-
-  for (const instr of instrs) {
-    const op = instr[0];
-
-    // defineLabel → start a new block boundary
-    if (op === null && (instr[1] as any)?.type === "defineLabel") {
-      flushBlock(null);
-      currentLabel = (instr[1] as any).label;
-      continue;
-    }
-
-    // Terminator → ends the current block
-    if (op !== null && isTerminator(op, compiler)) {
-      flushBlock(instr);
-      continue;
-    }
-
-    currentBody.push(instr);
-  }
-
-  // Flush trailing instructions
-  flushBlock(null);
-
-  // Split large blocks (> MAX_BLOCK_SIZE instructions) into smaller chunks
-  // so that no single block reveals too much sequential code.
-  const MAX_BLOCK_SIZE = 3;
-  const splitBlocks: BasicBlock[] = [];
-  for (const block of blocks) {
-    if (block.body.length <= MAX_BLOCK_SIZE) {
-      splitBlocks.push(block);
-      continue;
-    }
-    // Chunk the body into pieces of MAX_BLOCK_SIZE
-    for (let j = 0; j < block.body.length; j += MAX_BLOCK_SIZE) {
-      const isFirst = j === 0;
-      const isLast = j + MAX_BLOCK_SIZE >= block.body.length;
-      splitBlocks.push({
-        label: isFirst ? block.label : compiler._makeLabel("cff_split"),
-        body: block.body.slice(j, j + MAX_BLOCK_SIZE),
-        terminator: isLast ? block.terminator : null,
-        stateValue: isFirst ? block.stateValue : assignState(),
-        originalNextIndex: -1,
-      });
-    }
-  }
-  // Replace blocks with split result
-  blocks.length = 0;
-  blocks.push(...splitBlocks);
-
-  // Wire up originalNextIndex for fallthrough resolution
-  for (let i = 0; i < blocks.length - 1; i++) {
-    blocks[i].originalNextIndex = i + 1;
-  }
-  // Last block has no successor
-  if (blocks.length > 0) {
-    blocks[blocks.length - 1].originalNextIndex = -1;
-  }
-
-  return blocks;
-}
-
-// ── Cross-block register promotion ───────────────────────────────────────────
-// Scans all blocks (bodies + terminators) and finds register operands that
-// appear in more than one block.  Those registers must not be in the "temp"
-// pool because resolveRegisters' linear scan doesn't understand the CFF
-// dispatch loop and would reuse their slots between blocks.
-//
-// Promotion is done in-place: we delete the `kind` property on the operand
-// objects so they default to the "local::" pool (which never reuses slots).
-
-function promoteMultiBlockRegisters(blocks: BasicBlock[]): void {
-  // (fnId, regId) → index of first block where this register was seen
-  const regFirstBlock = new Map<string, number>();
-  // Set of register keys that appear in 2+ blocks
-  const multiBlockRegs = new Set<string>();
-
-  for (let bi = 0; bi < blocks.length; bi++) {
-    const allInstrs = blocks[bi].terminator
-      ? [...blocks[bi].body, blocks[bi].terminator!]
-      : blocks[bi].body;
-
-    for (const instr of allInstrs) {
-      for (let j = 1; j < instr.length; j++) {
-        const op = instr[j] as any;
-        if (op && typeof op === "object" && op.type === "register") {
-          const key = `${op.fnId}:${op.id}`;
-          const first = regFirstBlock.get(key);
-          if (first === undefined) {
-            regFirstBlock.set(key, bi);
-          } else if (first !== bi) {
-            multiBlockRegs.add(key);
-          }
-        }
-      }
-    }
-  }
-
-  if (multiBlockRegs.size === 0) return;
-
-  // Second pass: pin all operand instances of multi-block registers so that
-  // resolveRegisters assigns them to the "local::" pool (no slot reuse).
-  for (const block of blocks) {
-    const allInstrs = block.terminator
-      ? [...block.body, block.terminator!]
-      : block.body;
-
-    for (const instr of allInstrs) {
-      for (let j = 1; j < instr.length; j++) {
-        const op = instr[j] as any;
-        if (op && typeof op === "object" && op.type === "register") {
-          const key = `${op.fnId}:${op.id}`;
-          if (multiBlockRegs.has(key)) {
-            op.pinned = true;
-          }
-        }
-      }
-    }
-  }
-}
-
-// ── Generate the dispatch loop via Template ──────────────────────────────────
-
-function buildDispatchTemplate(
-  blocks: BasicBlock[],
-  endState: number,
-  startState: number,
-  compiler: Compiler,
-  fnId: number,
-  maxId: Map<number, number>,
-): {
-  bytecode: Bytecode;
-  rState: RegisterOperand;
-  loopTopLabel: string;
-  loopExitLabel: string;
-  innerBytecode: Bytecode;
-} {
-  // Build the if-chain cases
-  const cases = blocks
-    .map(
-      (block) =>
-        `if (state === ${block.stateValue}) _VM_JUMP_("${block.label}");`,
-    )
-    .join("\n    ");
-
-  const source = `
-    var state = ${startState};
-    while (state !== ${endState}) {
-      ${cases}
-    }
-  `;
-
-  const tmpl = new Template(source);
-  const result = tmpl.compileInline({}, compiler, fnId, maxId);
-
-  // Pin ALL dispatch-loop registers so resolveRegisters assigns them to the
-  // "local::" pool (no slot reuse).  The dispatch loop is re-entered on every
-  // state transition (backward JUMP to while_top), but the linear-scan liveness
-  // in resolveRegisters doesn't track loops and would incorrectly treat dispatch
-  // temps as dead after one pass, allowing their slots to be reused by body
-  // registers that are live across blocks.
-  for (const instr of result.bytecode) {
-    for (let j = 1; j < instr.length; j++) {
-      const op = instr[j] as any;
-      if (op && typeof op === "object" && op.type === "register") {
-        op.pinned = true;
-      }
-    }
-  }
-
-  const rState = result.registers.get("state");
-  if (!rState) {
-    throw new Error("CFF: Template did not produce a 'state' register");
-  }
-
-  // Find the while loop labels from the compiled IR
-  let loopTopLabel: string | null = null;
-  let loopExitLabel: string | null = null;
-
-  for (const instr of result.bytecode) {
-    if (instr[0] === null && (instr[1] as any)?.type === "defineLabel") {
-      const label = (instr[1] as any).label as string;
-      if (label.includes("while_top") && !loopTopLabel) {
-        loopTopLabel = label;
-      }
-      if (label.includes("while_exit") && !loopExitLabel) {
-        loopExitLabel = label;
-      }
-    }
-  }
-
-  if (!loopTopLabel || !loopExitLabel) {
-    throw new Error("CFF: Could not find while loop labels in Template output");
-  }
-
-  return {
-    bytecode: result.bytecode,
-    rState,
-    loopTopLabel,
-    loopExitLabel,
-    innerBytecode: result.innerBytecode,
-  };
-}
-
-// ── State transition helpers ─────────────────────────────────────────────────
-
-function emitStateTransition(
-  out: Bytecode,
-  rState: RegisterOperand,
-  targetState: number,
-  loopTopLabel: string,
-  compiler: Compiler,
-): void {
-  out.push([compiler.OP.LOAD_INT!, ref(rState), targetState]);
-  out.push([compiler.OP.JUMP!, { type: "label", label: loopTopLabel }]);
-}
-
-// ── Per-function transformation ──────────────────────────────────────────────
-
-function processFunctionBlock(
-  instrs: Bytecode,
-  fnId: number,
-  compiler: Compiler,
-  maxId: Map<number, number>,
-): { instrs: Bytecode; tail: Bytecode } {
-  const OP = compiler.OP;
-
-  // Only transform functions that contain simple jumps
-  const hasRoutableJump = instrs.some((instr) => {
-    const op = instr[0];
-    return op === OP.JUMP || op === OP.JUMP_IF_FALSE || op === OP.JUMP_IF_TRUE;
-  });
-  if (!hasRoutableJump) return { instrs, tail: [] };
-
-  // ── 1. Split into basic blocks ──────────────────────────────────────────
-  const blocks = splitBasicBlocks(instrs, compiler);
-  if (blocks.length < 2) return { instrs, tail: [] };
-
-  // ── 1b. Promote cross-block registers to "local" pool ──────────────────
-  // resolveRegisters does a linear-scan liveness analysis that doesn't
-  // understand the CFF dispatch loop (backward jumps).  A "temp" register
-  // that's live across two blocks would appear to die within its first
-  // block and get its slot reused, corrupting values read in later blocks.
-  //
-  // Fix: find every register that appears in more than one block and
-  // delete its "temp" kind so it lands in the "local::" pool (no reuse).
-  promoteMultiBlockRegisters(blocks);
-
-  const usedStates = new Set(blocks.map((b) => b.stateValue));
-
-  // Pick endState sentinel
-  let endState: number;
-  do {
-    endState = getRandomInt(0, U16_MAX);
-  } while (usedStates.has(endState));
-
-  const startState = blocks[0].stateValue;
-
-  // ── 2. Build dispatch loop from Template ────────────────────────────────
-  const dispatch = buildDispatchTemplate(
-    blocks,
-    endState,
-    startState,
-    compiler,
-    fnId,
-    maxId,
-  );
-  const { rState, loopTopLabel, loopExitLabel } = dispatch;
-
-  // ── 3. Pre-compute all state mappings BEFORE shuffle ─────────────────
-  // These maps capture the correct stateValues while the blocks array is
-  // still in its original split order.  After the shuffle, indexing into
-  // blocks[] by original index would give the wrong block.
-
-  // label → stateValue (for jump target resolution)
-  const labelToState = new Map<string, number>();
-  for (const block of blocks) {
-    labelToState.set(block.label, block.stateValue);
-  }
-
-  // originalIndex → fallthrough stateValue
-  const fallthroughStateMap = new Map<number, number>();
-  for (let i = 0; i < blocks.length; i++) {
-    const next = blocks[i].originalNextIndex;
-    fallthroughStateMap.set(
-      i,
-      next >= 0 ? blocks[next].stateValue : endState,
-    );
-  }
-
-  // ── 4. Shuffle block order ──────────────────────────────────────────────
-  // Track which original index each shuffled position came from, so we can
-  // look up fallthroughStateMap correctly during emission.
-  const originalIndices = blocks.map((_, i) => i);
-
-  // Fisher-Yates shuffle
-  for (let i = blocks.length - 1; i > 0; i--) {
-    const j = getRandomInt(0, i);
-    [blocks[i], blocks[j]] = [blocks[j], blocks[i]];
-    [originalIndices[i], originalIndices[j]] = [
-      originalIndices[j],
-      originalIndices[i],
-    ];
-  }
-
-  // ── 5. Emit: dispatch loop + block bodies ───────────────────────────────
-  const out: Bytecode = [];
-
-  // Dispatch loop (var state = ...; while(...) { if-chain })
-  out.push(...dispatch.bytecode);
-
-  // Each block: defineLabel → body → state transition → JUMP loopTop
-  for (let i = 0; i < blocks.length; i++) {
-    const block = blocks[i];
-    const origIdx = originalIndices[i];
-
-    // Block label
-    out.push([null, { type: "defineLabel", label: block.label }]);
-
-    // Block body
-    out.push(...block.body);
-
-    // Terminator rewriting
-    const term = block.terminator;
-
-    if (term === null) {
-      // Fallthrough → transition to the original next block's state
-      emitStateTransition(
-        out,
-        rState,
-        fallthroughStateMap.get(origIdx)!,
-        loopTopLabel,
-        compiler,
-      );
-    } else if (term[0] === OP.RETURN || term[0] === OP.THROW) {
-      // Exits the frame — emit as-is
-      out.push(term);
-    } else if (term[0] === OP.JUMP) {
-      const targetLabel = extractLabel(term[1]);
-      if (targetLabel !== null) {
-        const targetState = labelToState.get(targetLabel);
-        if (targetState !== undefined) {
-          emitStateTransition(out, rState, targetState, loopTopLabel, compiler);
-        } else {
-          // Target outside this function's blocks — keep original
-          out.push(term);
-        }
-      } else {
-        out.push(term);
-      }
-    } else if (term[0] === OP.JUMP_IF_FALSE) {
-      // Original: if (!cond) goto target; else fallthrough
-      // → if (cond) goto skipLabel  (inverted)
-      //   state = targetState; goto loopTop
-      //   skipLabel:
-      //   state = fallthroughState; goto loopTop
-      const cond = term[1] as RegisterOperand;
-      const targetLabel = extractLabel(term[2]);
-
-      if (targetLabel !== null) {
-        const targetState = labelToState.get(targetLabel);
-        if (targetState !== undefined) {
-          const skipLabel = compiler._makeLabel("cff_skip");
-
-          out.push([
-            OP.JUMP_IF_TRUE!,
-            cond,
-            { type: "label", label: skipLabel },
-          ]);
-          emitStateTransition(
-            out,
-            rState,
-            targetState,
-            loopTopLabel,
-            compiler,
-          );
-          out.push([null, { type: "defineLabel", label: skipLabel }]);
-          emitStateTransition(
-            out,
-            rState,
-            fallthroughStateMap.get(origIdx)!,
-            loopTopLabel,
-            compiler,
-          );
-        } else {
-          out.push(term);
-        }
-      } else {
-        out.push(term);
-      }
-    } else if (term[0] === OP.JUMP_IF_TRUE) {
-      // Original: if (cond) goto target; else fallthrough
-      // → if (!cond) goto skipLabel  (inverted)
-      //   state = targetState; goto loopTop
-      //   skipLabel:
-      //   state = fallthroughState; goto loopTop
-      const cond = term[1] as RegisterOperand;
-      const targetLabel = extractLabel(term[2]);
-
-      if (targetLabel !== null) {
-        const targetState = labelToState.get(targetLabel);
-        if (targetState !== undefined) {
-          const skipLabel = compiler._makeLabel("cff_skip");
-
-          out.push([
-            OP.JUMP_IF_FALSE!,
-            cond,
-            { type: "label", label: skipLabel },
-          ]);
-          emitStateTransition(
-            out,
-            rState,
-            targetState,
-            loopTopLabel,
-            compiler,
-          );
-          out.push([null, { type: "defineLabel", label: skipLabel }]);
-          emitStateTransition(
-            out,
-            rState,
-            fallthroughStateMap.get(origIdx)!,
-            loopTopLabel,
-            compiler,
-          );
-        } else {
-          out.push(term);
-        }
-      } else {
-        out.push(term);
-      }
-    }
-  }
-
-  return { instrs: out, tail: dispatch.innerBytecode };
-}
-
-// ── Pass entry point ──────────────────────────────────────────────────────────
-export function controlFlowFlattening(
-  bc: Bytecode,
-  compiler: Compiler,
-): { bytecode: Bytecode } {
-  const maxId = buildMaxIdMap(bc);
-  return forEachFunction(bc, compiler, (fnInstrs, fnId) =>
-    processFunctionBlock(fnInstrs, fnId, compiler, maxId),
-  );
-}