diff --git a/scripts/roblox_abi_check.py b/scripts/roblox_abi_check.py new file mode 100644 index 0000000..bcb43ad --- /dev/null +++ b/scripts/roblox_abi_check.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python3 +# Copyright 2026 Mocktail Project Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run the HostAbi deriver and the solver on one update and compare them. + +Usage: + python3 roblox_abi_check.py --reference-lib OLD.so \\ + --reference-profile OLD_PROFILE.json \\ + --reference-compatibility config/roblox_compatibility.json \\ + --candidate-lib NEW.so + +Prints every field where the two disagree. Exit status: 0 agree, 1 differ or +one side failed, 2 both failed. +""" + +from __future__ import annotations + +import argparse +import importlib.util +from pathlib import Path +import sys +from typing import Any, Sequence + + +def _load_solver(): + name = "mocktail_roblox_abi_solver" + if name in sys.modules: + return sys.modules[name] + path = Path(__file__).resolve().with_name("roblox_abi_solver.py") + specification = importlib.util.spec_from_file_location(name, path) + if specification is None or specification.loader is None: + raise RuntimeError("cannot load roblox_abi_solver.py") + module = importlib.util.module_from_spec(specification) + sys.modules[name] = module + specification.loader.exec_module(module) + return module + + +SOLVER = _load_solver() +DERIVER = SOLVER.DERIVER +AnalyzerError = SOLVER.AnalyzerError +format_rva = SOLVER.format_rva + + +def compare_results(derived: Any, solved: Any, path: str = "") -> list[str]: + """Returns one line for every leaf value where the two results differ.""" + if isinstance(derived, dict) and isinstance(solved, dict): + lines = [] + for key in sorted(derived.keys() | solved.keys(), key=str): + child = f"{path}.{key}" if path else str(key) + if key not in solved: + lines.append(f"{child}: deriver {derived[key]}, solver missing") + elif key not in derived: + lines.append(f"{child}: deriver missing, solver {solved[key]}") + else: + lines.extend(compare_results(derived[key], solved[key], child)) + return lines + if isinstance(derived, list) and isinstance(solved, list) and len(derived) == len(solved): + return [ + line + for index, (left, right) in enumerate(zip(derived, solved)) + for line in compare_results(left, right, f"{path}[{index}]") + ] + return [] if derived == solved else [f"{path}: deriver {derived}, solver {solved}"] + + +def check_status( + derived: Any, solved: Any, deriver_error: str | None, solver_error: str | None +) -> tuple[int, list[str]]: + """Exit status and report lines: 0 agree, 1 differ or one failed, 2 both failed.""" + lines = [] + if deriver_error is not None: + lines.append(f"deriver failed: {deriver_error}") + if solver_error is not None: + lines.append(f"solver failed: {solver_error}") + if deriver_error is not None and solver_error is not None: + return 2, lines + if deriver_error is not None or solver_error is not None: + return 1, lines + differences = compare_results(derived, solved) + if differences: + return 1, differences + return 0, ["deriver and solver agree on every field"] + + +def _solver_runtime_source(source: dict[str, Any]) -> dict[str, Any]: + # The deriver's loader returns integers; the solver reads hex strings. + result: dict[str, Any] = {} + for key, value in source.items(): + if isinstance(value, dict): + result[key] = { + field: item if field == "vtable_layout_version" else format_rva(item) + for field, item in value.items() + } + else: + result[key] = format_rva(value) + return result + + +def _describe(error: Exception) -> str: + if isinstance(error, AnalyzerError): + return str(error) + return f"{type(error).__name__}: {error}" + + +def check(arguments: argparse.Namespace) -> int: + """Runs the deriver and the solver on one update and compares them.""" + sidecar = DERIVER.validated_sidecar(Path(arguments.reference_profile)) + manifests = [Path(path) for path in arguments.reference_compatibility] + cache_dir = Path(arguments.cache_dir) if arguments.cache_dir else None + derived = solved = deriver_error = solver_error = None + with DERIVER.ElfImage(Path(arguments.reference_lib)) as reference, DERIVER.ElfImage( + Path(arguments.candidate_lib) + ) as candidate: + try: + profile, anchors = DERIVER.derive_profile(reference, candidate, sidecar) + runtime = DERIVER.derive_runtime_compatibility(reference, candidate, manifests) + derived = {"profile": profile, "derivation_anchors": anchors, "runtime": runtime} + except Exception as error: + deriver_error = _describe(error) + try: + # The input checks are cheap; indexing both libraries is not. + SOLVER.check_inputs(reference, candidate, sidecar) + solver = SOLVER.Solver(reference, candidate, cache_dir) + profile, anchors = solver.derive_profile(sidecar) + source = DERIVER.load_reference_runtime_profile(manifests, reference.build_id) + runtime = solver.derive_runtime_compatibility(_solver_runtime_source(source)) + solved = {"profile": profile, "derivation_anchors": anchors, "runtime": runtime} + except Exception as error: + solver_error = _describe(error) + status, lines = check_status(derived, solved, deriver_error, solver_error) + for line in lines: + print(line) + return status + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser( + description="Run the HostAbi deriver and the solver on one update and compare them" + ) + parser.add_argument("--reference-lib", required=True, metavar="REF") + parser.add_argument("--reference-profile", required=True, metavar="REFJSON") + parser.add_argument( + "--reference-compatibility", action="append", default=[], metavar="MANIFEST" + ) + parser.add_argument("--candidate-lib", required=True, metavar="LIB") + parser.add_argument( + "--cache-dir", + default=str(Path.home() / ".cache" / "mocktail" / "function_index"), + help="where function indexes are kept; empty to disable", + ) + arguments = parser.parse_args(sys.argv[1:] if argv is None else argv) + try: + return check(arguments) + except AnalyzerError as error: + print(f"roblox-abi-check: error: {error}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/roblox_abi_solver.py b/scripts/roblox_abi_solver.py new file mode 100644 index 0000000..a29f4fd --- /dev/null +++ b/scripts/roblox_abi_solver.py @@ -0,0 +1,401 @@ +#!/usr/bin/env python3 +# Copyright 2026 Mocktail Project Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Evidence-scoring solver that finds HostAbi addresses in a new libroblox. + +Each function in both builds is indexed by `roblox_function_index.py`. A +reference function is matched through independent clues: + +- the same instruction shape, unique in both builds +- callers and callees that are already matched +- the strings it references +- its data references, shifted by the distance mapped for those addresses +- its place between matched neighbours + +A match is accepted only with a score of at least ACCEPT_SCORE and a lead of +at least ACCEPT_LEAD over the runner-up. `roblox_abi_targets.py` turns the +matches into a profile with the same shape as the one +`derive_roblox_host_abi_profile.py` writes, and `roblox_abi_check.py` +compares the two on one update. +""" + +from __future__ import annotations + +import bisect +import collections +from dataclasses import dataclass +import difflib +import importlib.util +from pathlib import Path +import sys +from typing import Any, Callable, Sequence + + +def _load_sibling(name: str, file_name: str): + if name in sys.modules: + return sys.modules[name] + path = Path(__file__).resolve().with_name(file_name) + specification = importlib.util.spec_from_file_location(name, path) + if specification is None or specification.loader is None: + raise RuntimeError(f"cannot load {file_name}") + module = importlib.util.module_from_spec(specification) + sys.modules[name] = module + specification.loader.exec_module(module) + return module + + +INDEX = _load_sibling("mocktail_roblox_function_index", "roblox_function_index.py") +TARGETS = _load_sibling("mocktail_roblox_abi_targets", "roblox_abi_targets.py") +DERIVER = INDEX.DERIVER +FunctionIndex = INDEX.FunctionIndex +Instruction = INDEX.Instruction +AnalyzerError = DERIVER.AnalyzerError +format_rva = DERIVER.format_rva +check_inputs = TARGETS.check_inputs +validate_fmod_source = TARGETS.validate_fmod_source + +ACCEPT_SCORE = 0.6 +ACCEPT_LEAD = 0.2 +# Candidates that get the full, disassembling score after the cheap one. +FULL_SCORE_LIMIT = 64 +# A data address used by more functions than this is too common to narrow +# the candidates by itself. +XREF_FANOUT_LIMIT = 64 +# Scores are float sums; a lead that equals ACCEPT_LEAD on paper may be +# a few ulps short. +_TOLERANCE = 1e-9 + + +def choose_candidate( + scored: Sequence[tuple[float, int]], tie_breaker: int | None = None +) -> int | None: + """Returns the leading candidate, or None when no lead is clear. + + `tie_breaker` picks among the candidates that clear ACCEPT_SCORE and sit + within ACCEPT_LEAD of the best; it cannot lift or overrule any other. + """ + if not scored: + return None + ranked = sorted(scored, key=lambda item: item[0], reverse=True) + best_score, best = ranked[0] + runner_up = ranked[1][0] if len(ranked) > 1 else 0.0 + if best_score >= ACCEPT_SCORE and best_score - runner_up >= ACCEPT_LEAD - _TOLERANCE: + return best + tied = { + rva for value, rva in ranked + if value >= ACCEPT_SCORE and best_score - value < ACCEPT_LEAD - _TOLERANCE + } + return tie_breaker if tie_breaker in tied else None + + +def align_instructions( + reference_shapes: Sequence[str], candidate_shapes: Sequence[str] +) -> dict[int, int]: + """Maps reference instruction indexes to candidate indexes of equal shape.""" + matcher = difflib.SequenceMatcher( + None, reference_shapes, candidate_shapes, autojunk=False + ) + mapping = {} + for block in matcher.get_matching_blocks(): + for offset in range(block.size): + mapping[block.a + offset] = block.b + offset + return mapping + + +@dataclass(frozen=True) +class Match: + rva: int + score: float + lead: float + # "score" when the score alone decided, "vtable" when a slot broke a tie. + decided_by: str + + +def _share(expected: Sequence[Any], present: set[Any]) -> float | None: + if not expected: + return None + return sum(1 for item in expected if item in present) / len(expected) + + +def _weighted(parts: Sequence[tuple[float | None, float]]) -> float: + present = [(value, weight) for value, weight in parts if value is not None] + return sum(value * weight for value, weight in present) / sum( + weight for _, weight in present + ) + + +class Solver(TARGETS.TargetDerivation): + def __init__(self, reference: Any, candidate: Any, cache_dir: Path | None = None): + self._start(FunctionIndex(reference, cache_dir), FunctionIndex(candidate, cache_dir)) + + @classmethod + def from_indexes(cls, reference: Any, candidate: Any) -> "Solver": + solver = cls.__new__(cls) + solver._start(reference, candidate) + return solver + + def _start(self, reference: Any, candidate: Any) -> None: + self.reference = reference + self.candidate = candidate + self.pairs: dict[int, int] = {} + self._owners: dict[int, int] = {} + for start, other in self._shape_anchors().items(): + self._pair(start, other) + self._propagate_calls() + self.data_map = self._map_data_references() + self.matches: dict[str, Match] = {} + self._claimed: dict[int, int] = {} + + def _pair(self, start: int, other: int) -> None: + """Pairs two functions, dropping any pair that held either one.""" + owner = self._owners.get(other) + if owner is not None and owner != start: + del self.pairs[owner] + previous = self.pairs.get(start) + if previous is not None and previous != other: + del self._owners[previous] + self.pairs[start] = other + self._owners[other] = start + + def _shape_anchors(self) -> dict[int, int]: + reference_hashes = collections.defaultdict(list) + candidate_hashes = collections.defaultdict(list) + for start, info in self.reference.functions.items(): + reference_hashes[info.shape_hash].append(start) + for start, info in self.candidate.functions.items(): + candidate_hashes[info.shape_hash].append(start) + pairs = {} + for shape_hash, starts in reference_hashes.items(): + others = candidate_hashes.get(shape_hash, ()) + if len(starts) == 1 and len(others) == 1: + pairs[starts[0]] = others[0] + return pairs + + def _propagate_calls(self) -> None: + # Matched functions with as many calls pair their callees in order. + reference, candidate = self.reference.functions, self.candidate.functions + pending = list(self.pairs.items()) + while pending: + discovered = [] + for reference_start, candidate_start in pending: + reference_calls = reference[reference_start].calls + candidate_calls = candidate[candidate_start].calls + if len(reference_calls) != len(candidate_calls): + continue + for callee, other in zip(reference_calls, candidate_calls): + if ( + callee in self.pairs + or other in self._owners + or callee not in reference + or other not in candidate + ): + continue + size = reference[callee].size + if abs(size - candidate[other].size) <= max(16, size // 4): + self._pair(callee, other) + discovered.append((callee, other)) + pending = discovered + + def _map_data_references(self) -> dict[int, int]: + votes: dict[int, collections.Counter] = collections.defaultdict(collections.Counter) + for reference_start, candidate_start in self.pairs.items(): + reference_refs = self.reference.functions[reference_start].refs + candidate_refs = self.candidate.functions[candidate_start].refs + if len(reference_refs) == len(candidate_refs): + for reference_target, candidate_target in zip(reference_refs, candidate_refs): + votes[reference_target][candidate_target] += 1 + return {target: counter.most_common(1)[0][0] for target, counter in votes.items()} + + def _gap_candidates(self, start: int) -> list[int]: + starts = self.reference.starts + index = bisect.bisect_left(starts, start) + below = next( + (self.pairs[starts[i]] for i in range(index - 1, -1, -1) if starts[i] in self.pairs), + self.candidate.text.address, + ) + above = next( + (self.pairs[starts[i]] for i in range(index + 1, len(starts)) if starts[i] in self.pairs), + self.candidate.text.address + self.candidate.text.size, + ) + if above <= below: + return [] + low = bisect.bisect_left(self.candidate.starts, below) + high = bisect.bisect_right(self.candidate.starts, above) + return self.candidate.starts[low:high] + + def _neighbour_candidates(self, start: int) -> set[int]: + found = set() + for callee in self.reference.functions[start].calls: + if callee in self.pairs: + found.update(self.candidate.callers.get(self.pairs[callee], ())) + for caller in self.reference.callers.get(start, ()): + if caller in self.pairs: + found.update( + target + for target in self.candidate.functions[self.pairs[caller]].calls + if target in self.candidate.functions + ) + return found + + def _pattern_candidates(self, start: int) -> set[int]: + # Data a function touches usually moves as one block, so its + # references shifted by a mapped distance find it in the candidate. + refs = set(self.reference.functions[start].refs) + distances = collections.Counter( + self.data_map[target] - target for target in refs if target in self.data_map + ) + # A common address adds a vote but cannot bring in a candidate alone. + votes = collections.Counter() + found_by_rare = set() + for target in refs: + users = set() + for distance, _ in distances.most_common(3): + found = self.candidate.data_users.get(target + distance, ()) + users.update(found) + if len(found) <= XREF_FANOUT_LIMIT: + found_by_rare.update(found) + votes.update(users) + needed = max(1, len(refs) // 2) + return {user for user in found_by_rare if votes[user] >= needed} + + def _graph_score(self, start: int, other: int) -> float: + reference = self.reference.functions[start] + candidate = self.candidate.functions[other] + size = min(reference.size, candidate.size) / max(reference.size, candidate.size) + calls = [self.pairs[c] for c in reference.calls if c in self.pairs] + callers = [self.pairs[c] for c in self.reference.callers.get(start, ()) if c in self.pairs] + data = [self.data_map[t] for t in reference.refs if t in self.data_map] + return _weighted( + ( + (size, 1.0), + (_share(calls, set(candidate.calls)), 2.0), + (_share(callers, set(self.candidate.callers.get(other, ()))), 2.0), + (_share(data, set(candidate.refs)), 2.0), + ) + ) + + def _full_score(self, start: int, other: int, graph: float) -> float: + similarity = difflib.SequenceMatcher( + None, self.reference.shapes(start), self.candidate.shapes(other), autojunk=False + ).ratio() + strings = [ + text + for text in map(self.reference.string_at, self.reference.functions[start].refs) + if text is not None + ] + candidate_strings = set( + map(self.candidate.string_at, self.candidate.functions[other].refs) + ) + return _weighted( + ( + (similarity, 3.0), + (_share(strings, candidate_strings), 2.0), + (graph, 5.0), + ) + ) + + def match_function( + self, + name: str, + start: int, + contract: Callable[[Any, int], bool] | None = None, + vtable_slot: int | None = None, + ) -> int: + """Matches one reference function. + + A contract drops candidates outright. `vtable_slot` is the candidate + that sits in the matching slot of an already located vtable; it only + breaks a tie between candidates that pass on their own evidence. + """ + if start not in self.reference.functions: + raise AnalyzerError(f"reference {name} is not a function start") + if name in self.matches: + return self.matches[name].rva + candidates = set(self._gap_candidates(start)) + candidates |= self._neighbour_candidates(start) + candidates |= self._pattern_candidates(start) + if start in self.pairs: + candidates.add(self.pairs[start]) + if vtable_slot in self.candidate.functions: + candidates.add(vtable_slot) + size = self.reference.functions[start].size + candidates = { + other + for other in candidates + if self._claimed.get(other, start) == start + and self.candidate.functions[other].size <= 8 * size + 64 + and (contract is None or contract(self.candidate.image, other)) + } + graph = {other: self._graph_score(start, other) for other in candidates} + shortlist = sorted(candidates, key=graph.__getitem__, reverse=True)[:FULL_SCORE_LIMIT] + if vtable_slot in candidates and vtable_slot not in shortlist: + shortlist.append(vtable_slot) + scored = [(self._full_score(start, other, graph[other]), other) for other in shortlist] + chosen = choose_candidate(scored) + decided_by = "score" + if chosen is None and vtable_slot is not None: + chosen = choose_candidate(scored, tie_breaker=vtable_slot) + decided_by = "vtable" + if chosen is None: + ranked = sorted(scored, reverse=True)[:2] + raise AnalyzerError( + f"{name} has no clear candidate: " + + ", ".join(f"{format_rva(rva)}={value:.2f}" for value, rva in ranked) + ) + chosen_score = next(value for value, other in scored if other == chosen) + runner_up = max((value for value, other in scored if other != chosen), default=0.0) + self.matches[name] = Match(chosen, chosen_score, chosen_score - runner_up, decided_by) + self._claimed[chosen] = start + self._pair(start, chosen) + return chosen + + def _aligned(self, name: str, start: int) -> tuple[list[Instruction], list[Instruction], dict[int, int]]: + other = self.match_function(name, start) + reference = self.reference.instructions(start) + candidate = self.candidate.instructions(other) + mapping = align_instructions( + [item.shape for item in reference], [item.shape for item in candidate] + ) + return reference, candidate, mapping + + def map_instruction(self, name: str, rva: int) -> int: + start = self.reference.containing(rva) + if start is None: + raise AnalyzerError(f"reference {name} is outside every function") + reference, candidate, mapping = self._aligned(f"{name} function", start) + index = next((i for i, item in enumerate(reference) if item.address == rva), None) + if index is None or index not in mapping: + raise AnalyzerError(f"{name} instruction has no counterpart") + return candidate[mapping[index]].address + + def map_data(self, name: str, function_name: str, start: int, target: int) -> int: + reference, candidate, mapping = self._aligned(function_name, start) + found = set() + for index, item in enumerate(reference): + if target not in item.rip_targets: + continue + if index not in mapping: + raise AnalyzerError(f"{name} reference instruction has no counterpart") + other = candidate[mapping[index]] + if len(other.rip_targets) != len(item.rip_targets): + raise AnalyzerError(f"{name} RIP-relative operand count changed") + found.update( + other.rip_targets[position] + for position, value in enumerate(item.rip_targets) + if value == target + ) + if len(found) != 1: + raise AnalyzerError(f"{name} has no unique mapped data address") + return found.pop() diff --git a/scripts/roblox_abi_targets.py b/scripts/roblox_abi_targets.py new file mode 100644 index 0000000..d5dac61 --- /dev/null +++ b/scripts/roblox_abi_targets.py @@ -0,0 +1,375 @@ +#!/usr/bin/env python3 +# Copyright 2026 Mocktail Project Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""HostAbi targets that `roblox_abi_solver.py` derives from its matches. + +`TargetDerivation` turns matched functions into the HostAbi profile, the +derivation anchors and the runtime compatibility block, in the shape that +`derive_roblox_host_abi_profile.py` writes. It reuses the deriver's input +checks, contracts and post-processing. +""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path +import sys +from typing import Any, Callable + + +def _load_index(): + name = "mocktail_roblox_function_index" + if name in sys.modules: + return sys.modules[name] + path = Path(__file__).resolve().with_name("roblox_function_index.py") + specification = importlib.util.spec_from_file_location(name, path) + if specification is None or specification.loader is None: + raise RuntimeError("cannot load roblox_function_index.py") + module = importlib.util.module_from_spec(specification) + sys.modules[name] = module + specification.loader.exec_module(module) + return module + + +DERIVER = _load_index().DERIVER +AnalyzerError = DERIVER.AnalyzerError +format_rva = DERIVER.format_rva +FMOD_OUTPUT_FIELDS = DERIVER.FMOD_OUTPUT_FIELDS +FMOD_INPUT_FIELDS = DERIVER.FMOD_INPUT_FIELDS +fmod_slots = DERIVER.fmod_vtable_slots + + +def check_inputs(reference: Any, candidate: Any, reference_sidecar: dict[str, Any]) -> None: + """Applies the deriver's identity and preflight checks to both builds.""" + DERIVER.check_derivation_inputs(reference, candidate, reference_sidecar) + + +def validate_fmod_source(source: Any) -> int: + """Checks a reference FMOD bridge the way the deriver does; returns its layout.""" + if ( + not isinstance(source, dict) + or not FMOD_OUTPUT_FIELDS <= source.keys() + or source.keys() - FMOD_OUTPUT_FIELDS - set(FMOD_INPUT_FIELDS) - {"vtable_layout_version"} + ): + raise AnalyzerError("reference FMOD output-device profile is incomplete") + present = [field for field in FMOD_INPUT_FIELDS if field in source] + if present and len(present) != len(FMOD_INPUT_FIELDS): + raise AnalyzerError("reference FMOD input profile is incomplete") + layout = source.get("vtable_layout_version", 1) + if type(layout) is not int or layout not in (1, 2): + raise AnalyzerError("reference FMOD vtable layout is unsupported") + return layout + + +class TargetDerivation: + """Mixin for `Solver`: needs `reference`, `candidate`, `matches`, + `match_function`, `map_data` and `map_instruction`.""" + + def _check_registry_allocator( + self, + reference_registry: int, + candidate_registry: int, + reference_allocate: int, + candidate_allocate: int, + ) -> None: + # The registry initializer's tenth instruction calls the allocator bridge. + for index, start, allocate, side in ( + (self.reference, reference_registry, reference_allocate, "reference"), + (self.candidate, candidate_registry, candidate_allocate, "candidate"), + ): + code = index.instructions(start) + if len(code) <= 9 or code[9].shape != "call" or code[9].branch_target != allocate: + raise AnalyzerError( + f"{side} registry initializer does not call its allocator bridge" + ) + + def derive_profile( + self, reference_sidecar: dict[str, Any] + ) -> tuple[dict[str, Any], dict[str, Any]]: + reference_profile = reference_sidecar["profile"] + anchors = reference_sidecar["derivation_anchors"] + DERIVER.validate_profile_shape(reference_profile, anchors) + reference_image = self.reference.image + candidate_image = self.candidate.image + check_inputs(reference_image, candidate_image, reference_sidecar) + parse = DERIVER.parse_rva + bridges = [] + found = {} + for entry in reference_profile["bridge_entries"]: + rva = self.match_function(entry["label"], parse(entry["rva"], entry["label"])) + found[entry["label"]] = rva + bridges.append({"rva": format_rva(rva), "kind": entry["kind"], "label": entry["label"]}) + + seeds = reference_profile["data_seeds"] + arena_initializer = parse(seeds["arena_initializer"], "arena initializer") + thread_initializer = parse(seeds["allocator_thread_initializer"], "thread initializer") + self.match_function("arena-initializer", arena_initializer) + self.match_function("allocator-thread-initializer", thread_initializer) + + bootstrap = reference_profile["native_pre_jni_bootstrap"] + DERIVER.validate_reference_registry_slot(reference_image, bootstrap) + registry = self.match_function( + "registry-initializer", parse(bootstrap["registry_initializer"], "registry") + ) + reference_allocate = parse( + DERIVER.bridge_map(reference_profile)["allocate"]["rva"], "reference allocate" + ) + self._check_registry_allocator( + parse(bootstrap["registry_initializer"], "registry"), + registry, reference_allocate, found["allocate"], + ) + registry_slot = DERIVER.derive_registry_slot(candidate_image, registry) + + reference_init = reference_image.init_array_relocations() + candidate_init = candidate_image.init_array_relocations() + begin = DERIVER.LEADING_CONSTRUCTOR_BEGIN + count = DERIVER.LEADING_CONSTRUCTOR_COUNT + leading = candidate_init[begin : begin + count] + offsets = [] + for index in range(begin, begin + count): + expected = parse(anchors["constructor_rvas"][str(index)], f"constructor {index}") + if reference_init[index] != expected: + raise AnalyzerError(f"reference constructor {index} differs from its anchor") + rva = self.match_function(f"constructor-{index}", expected) + if leading.count(rva) != 1: + raise AnalyzerError(f"constructor {index} left the leading constructors") + offsets.append(leading.index(rva)) + if sorted(offsets) != list(range(count)): + raise AnalyzerError("leading constructors do not map one to one") + + singleton = parse(anchors["jni_singleton_initializer_rva"], "JNI singleton") + jni_singleton = self.match_function("jni-singleton-initializer", singleton) + allocator_anchor = parse(anchors["allocator_object_initializer_rva"], "allocator anchor") + empty_anchor = parse(anchors["empty_string_initializer_rva"], "empty string anchor") + allocator_function = self.reference.containing(allocator_anchor) + empty_function = self.reference.containing(empty_anchor) + if allocator_function is None or empty_function is None: + raise AnalyzerError("reference anchor is outside every function") + + slots = { + "allocator_object_slot": self.map_data( + "allocator object slot", "allocator-object-initializer function", + allocator_function, parse(seeds["allocator_object_slot"], "allocator slot"), + ), + "empty_string_slot": self.map_data( + "empty string slot", "empty-string-initializer function", + empty_function, parse(seeds["empty_string_slot"], "empty string slot"), + ), + "jni_singleton_slot": self.map_data( + "JNI singleton slot", "jni-singleton-initializer", + singleton, parse(seeds["jni_singleton_slot"], "JNI singleton slot"), + ), + "arena_guard_slot": self.map_data( + "arena guard slot", "arena-initializer", + arena_initializer, parse(seeds["arena_guard_slot"], "arena guard slot"), + ), + "arena_table_slot": self.map_data( + "arena table slot", "usable-size", + parse(DERIVER.bridge_map(reference_profile)["usable-size"]["rva"], "usable size"), + parse(seeds["arena_table_slot"], "arena table slot"), + ), + } + for slot in slots.values(): + candidate_image.require_writable_rva(slot, 8) + + candidate_count = len(candidate_init) + reference_count = reference_profile["init_array_count"] + boundary = reference_profile["native_mimalloc_thread_initializer_after_constructor"] + if not begin <= boundary < begin + count: + raise AnalyzerError("native allocator TLS boundary is no longer verified") + thread_boundary = begin + offsets[boundary - begin] + native_ranges = DERIVER.remapped_constructor_ranges( + reference_profile["native_mimalloc_constructor_run_ranges"], + offsets, reference_count, candidate_count, + ) + if not any(item["begin"] <= thread_boundary < item["end_exclusive"] for item in native_ranges): + raise AnalyzerError("native allocator TLS boundary is no longer verified") + profile = { + "elf_build_id": candidate_image.build_id, + "bridge_entries": bridges, + "data_seeds": { + "allocator_object_slot": format_rva(slots["allocator_object_slot"]), + "empty_string_slot": format_rva(slots["empty_string_slot"]), + "jni_singleton_slot": format_rva(slots["jni_singleton_slot"]), + "jni_singleton_bytes": seeds["jni_singleton_bytes"], + "arena_initializer": format_rva(self.matches["arena-initializer"].rva), + "allocator_thread_initializer": format_rva( + self.matches["allocator-thread-initializer"].rva + ), + "arena_guard_slot": format_rva(slots["arena_guard_slot"]), + "arena_table_slot": format_rva(slots["arena_table_slot"]), + "arena_table_slot_count": seeds["arena_table_slot_count"], + }, + "native_allocator": { + "allocate": format_rva(found["allocate"]), + "deallocate": format_rva(found["free"]), + }, + "init_array_offset": format_rva(candidate_image.sections[".init_array"].address), + "init_array_count": candidate_count, + "constructor_run_ranges": DERIVER.remapped_constructor_ranges( + reference_profile["constructor_run_ranges"], + offsets, reference_count, candidate_count, + ), + "native_mimalloc_constructor_run_ranges": native_ranges, + "native_mimalloc_thread_initializer_after_constructor": thread_boundary, + "native_pre_jni_bootstrap": { + "registry_initializer": format_rva(registry), + "registry_slot": format_rva(registry_slot), + }, + "default_allocator_strategy": "native_mimalloc", + } + derived_anchors = { + "signature_version": 1, + "allocator_object_initializer_rva": format_rva( + self.map_instruction("allocator-object-initializer", allocator_anchor) + ), + "empty_string_initializer_rva": format_rva( + self.map_instruction("empty-string-initializer", empty_anchor) + ), + "jni_singleton_initializer_rva": format_rva(jni_singleton), + "constructor_rvas": { + str(index): format_rva(candidate_init[index]) + for index in range(begin, begin + count) + }, + } + DERIVER.validate_profile_shape(profile, derived_anchors) + return profile, derived_anchors + + def derive_runtime_compatibility(self, source: dict[str, Any]) -> dict[str, Any]: + result: dict[str, Any] = {} + fullscreen = source.get("user_game_settings_fullscreen_setter_rva") + if fullscreen is not None: + reference_setter = DERIVER.parse_rva(fullscreen, "fullscreen setter") + if not DERIVER.has_fullscreen_setter_contract(self.reference.image, reference_setter): + raise AnalyzerError("reference fullscreen setter contract changed") + setter = self.match_function( + "fullscreen-setter", reference_setter, DERIVER.has_fullscreen_setter_contract + ) + result["user_game_settings_fullscreen_setter_rva"] = format_rva(setter) + fmod = source.get("fmod_output_device_bridge") + if fmod is not None: + result["fmod_output_device_bridge"] = self._derive_fmod_bridge(fmod) + return result + + def _derive_fmod_bridge(self, source: dict[str, Any]) -> dict[str, Any]: + layout = validate_fmod_source(source) + has_input = FMOD_INPUT_FIELDS[0] in source + slots = fmod_slots(layout, has_input) + parse = DERIVER.parse_rva + vtable = parse(source["vtable_rva"], "FMOD vtable") + reference_relocations = DERIVER.relative_relocation_map(self.reference.image) + methods = {} + unresolved = [] + for field, slot in slots.items(): + method = parse(source[field], field) + if reference_relocations.get(vtable + slot * 8) != method: + raise AnalyzerError("reference FMOD vtable does not match its methods") + try: + methods[field] = self.match_function( + f"fmod-{field}", method, self._fmod_contract(layout, field) + ) + except AnalyzerError: + unresolved.append(field) + string_constructor = parse(source["string_constructor_rva"], "FMOD string constructor") + if not DERIVER.has_fmod_string_constructor_contract(self.reference.image, string_constructor): + raise AnalyzerError("reference FMOD string constructor contract changed") + candidate_string_constructor = self.match_function( + "fmod-string-constructor", + string_constructor, + DERIVER.has_fmod_string_constructor_contract, + ) + + relocations = DERIVER.relative_relocation_map(self.candidate.image) + located = self._locate_fmod_vtables(relocations, methods) + vtables = sorted({vtable for vtable, _ in located}) + if len(vtables) != 1: + raise AnalyzerError(f"FMOD output vtable matched {len(vtables)} candidate locations") + # The located vtable breaks ties for the methods that did not resolve. + # Count and info share slots 5 and 6 in both layouts, so the vtable can + # fit both; exactly one layout must then hold every method. + solutions = [] + for located_vtable, target_layout in located: + target_slots = fmod_slots(target_layout, has_input) + resolved = dict(methods) + try: + for field in unresolved: + resolved[field] = self.match_function( + f"fmod-{field}", + parse(source[field], field), + self._fmod_contract(target_layout, field), + relocations.get(located_vtable + target_slots[field] * 8), + ) + except AnalyzerError: + continue + if self._locate_fmod_vtables(relocations, resolved) != [ + (located_vtable, target_layout) + ]: + continue + if target_layout == 2 and not DERIVER.has_fmod_device_lists_select_contract( + self.candidate.image, resolved["select_method_rva"] + ): + continue + solutions.append((located_vtable, target_layout, target_slots, resolved)) + if len(solutions) != 1: + raise AnalyzerError(f"FMOD output vtable holds every method in {len(solutions)} layouts") + located_vtable, target_layout, target_slots, methods = solutions[0] + self.candidate.image.require_relro_rva( + located_vtable, (max(target_slots.values()) + 1) * 8 + ) + result = {"vtable_rva": format_rva(located_vtable)} + result["string_constructor_rva"] = format_rva(candidate_string_constructor) + result.update((field, format_rva(rva)) for field, rva in methods.items()) + if target_layout != 1: + result["vtable_layout_version"] = target_layout + return result + + @staticmethod + def _fmod_contract(layout: int, field: str) -> Callable[[Any, int], bool] | None: + if layout == 2 and field == "select_method_rva": + return DERIVER.has_fmod_device_lists_select_contract + return None + + @staticmethod + def _locate_fmod_vtables( + relocations: dict[int, int], methods: dict[str, int] + ) -> list[tuple[int, int]]: + """Returns (vtable, layout) for each vtable holding every method.""" + located = [] + for layout in (1, 2): + slots = fmod_slots(layout, True) + held = {slots[field]: rva for field, rva in methods.items()} + located.extend( + (vtable, layout) + for vtable in TargetDerivation._vtables_holding(relocations, held) + ) + return sorted(located) + + @staticmethod + def _vtables_holding(relocations: dict[int, int], methods: dict[int, int]) -> list[int]: + """Returns every vtable whose slots hold all of `methods`.""" + if not methods: + raise AnalyzerError("no vtable method was matched") + slot, method = next(iter(methods.items())) + return sorted( + offset - slot * 8 + for offset, value in relocations.items() + if value == method + and offset >= slot * 8 + and (offset - slot * 8) % 8 == 0 + and all( + relocations.get(offset - slot * 8 + index * 8) == rva + for index, rva in methods.items() + ) + ) diff --git a/scripts/roblox_function_index.py b/scripts/roblox_function_index.py new file mode 100644 index 0000000..40e2cbf --- /dev/null +++ b/scripts/roblox_function_index.py @@ -0,0 +1,291 @@ +#!/usr/bin/env python3 +# Copyright 2026 Mocktail Project Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Function index of one libroblox image for `roblox_abi_solver.py`. + +Every function listed in `.eh_frame_hdr` is decoded once and reduced to a +normalized instruction shape, its direct calls and its RIP-relative data +references. Shapes drop branch targets, RIP displacements and immediates of +0x1000 and above, so a function moved by the linker keeps its shape. +""" + +from __future__ import annotations + +import bisect +import collections +from dataclasses import dataclass +import gzip +import hashlib +import importlib.util +import json +import os +from pathlib import Path +import re +import struct +import sys +import tempfile +from typing import Any + + +def _load_deriver(): + name = "mocktail_host_abi_deriver_for_solver" + if name in sys.modules: + return sys.modules[name] + path = Path(__file__).resolve().with_name("derive_roblox_host_abi_profile.py") + specification = importlib.util.spec_from_file_location(name, path) + if specification is None or specification.loader is None: + raise RuntimeError("cannot load the HostAbi deriver") + module = importlib.util.module_from_spec(specification) + sys.modules[name] = module + specification.loader.exec_module(module) + return module + + +DERIVER = _load_deriver() +AnalyzerError = DERIVER.AnalyzerError + +_RIP_OPERAND = re.compile(r"\[rip ([+-]) (0x[0-9a-f]+|\d+)\]") +_IMMEDIATE = re.compile(r"0x[0-9a-f]+|\b\d+\b") +# Smaller immediates are offsets and flags that stay put between builds. +# Larger ones are addresses and sizes that move. +_STABLE_IMMEDIATE_LIMIT = 0x1000 + + +def _stable_immediate(match: re.Match) -> str: + return "I" if int(match.group(0), 0) >= _STABLE_IMMEDIATE_LIMIT else match.group(0) + + +@dataclass(frozen=True) +class Instruction: + address: int + shape: str + rip_targets: tuple[int, ...] + branch_target: int | None + + +def decode_function( + decoder: Any, code: bytes, start: int +) -> list[Instruction]: + end = start + len(code) + instructions = [] + for address, size, mnemonic, operands in decoder.disasm_lite(code, start): + if (mnemonic == "call" or mnemonic.startswith("j")) and operands.startswith("0x"): + target = int(operands, 16) + if mnemonic == "call": + shape = "call" + else: + shape = mnemonic + (" in" if start <= target < end else " out") + instructions.append(Instruction(address, shape, (), target)) + continue + rip_targets = [] + for match in _RIP_OPERAND.finditer(operands): + displacement = int(match.group(2), 0) + rip_targets.append( + address + size + (-displacement if match.group(1) == "-" else displacement) + ) + operands = _IMMEDIATE.sub(_stable_immediate, _RIP_OPERAND.sub("[rip]", operands)) + instructions.append( + Instruction(address, f"{mnemonic} {operands}", tuple(rip_targets), None) + ) + return instructions + + +@dataclass(frozen=True) +class FunctionInfo: + size: int + shape_hash: bytes + calls: tuple[int, ...] + refs: tuple[int, ...] + + +def function_table(image: Any) -> list[tuple[int, int]]: + """Returns (start, size) for every FDE indexed by `.eh_frame_hdr`.""" + header = image.sections.get(".eh_frame_hdr") + if header is None: + raise AnalyzerError("libroblox has no .eh_frame_hdr") + raw = image.section_bytes(header) + # Version 1, a 4-byte eh_frame_ptr, udata4 count, datarel sdata4 table. + if len(raw) < 12 or raw[0] != 1 or raw[2] != 0x03 or raw[3] != 0x3B: + raise AnalyzerError(".eh_frame_hdr uses an unsupported encoding") + if raw[1] == 0xFF or raw[1] & 0x0F not in (0x03, 0x0B): + raise AnalyzerError(".eh_frame_hdr eh_frame_ptr encoding is not 4 bytes") + count = struct.unpack_from(" len(raw): + raise AnalyzerError(".eh_frame_hdr table exceeds its section") + table = [] + for index in range(count): + location, fde = struct.unpack_from(" Path: + """Cache file for one library; shapes depend on this code and capstone.""" + version = getattr(DERIVER.capstone, "__version__", "unknown") + return cache_dir / f"{library_sha256}-{_SOURCE_DIGEST}-capstone{version}.json.gz" + + +def link_functions( + functions: dict[int, FunctionInfo], +) -> tuple[dict[int, list[int]], dict[int, list[int]]]: + """Returns each function's callers and each data address's users, once each.""" + callers: dict[int, list[int]] = collections.defaultdict(list) + data_users: dict[int, list[int]] = collections.defaultdict(list) + for start, info in functions.items(): + for target in dict.fromkeys(info.calls): + if target in functions: + callers[target].append(start) + for target in dict.fromkeys(info.refs): + data_users[target].append(start) + return callers, data_users + + +def _int_list(value: Any) -> tuple[int, ...]: + if not isinstance(value, list) or not all(type(item) is int for item in value): + raise ValueError("cached index entry is malformed") + return tuple(value) + + +def _read_cache(path: Path) -> dict[int, FunctionInfo] | None: + """Returns the cached functions, or None when the file is missing or bad.""" + try: + with gzip.open(path, "rt", encoding="utf-8") as stream: + entries = json.load(stream) + if not isinstance(entries, list) or not entries: + raise ValueError("cached index is empty") + functions = {} + for entry in entries: + start, size, shape_hash, calls, refs = entry + if type(start) is not int or type(size) is not int: + raise ValueError("cached index entry is malformed") + functions[start] = FunctionInfo( + size, bytes.fromhex(shape_hash), _int_list(calls), _int_list(refs) + ) + return functions + except (OSError, EOFError, ValueError, TypeError): + return None + + +def _write_cache(path: Path, functions: dict[int, FunctionInfo]) -> None: + # A cache that cannot be written only costs the next run a rebuild. + entries = [ + [start, info.size, info.shape_hash.hex(), list(info.calls), list(info.refs)] + for start, info in functions.items() + ] + try: + path.parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary = tempfile.mkstemp(dir=path.parent, suffix=".tmp") + except OSError: + return + try: + with os.fdopen(descriptor, "wb") as raw, gzip.open(raw, "wt", encoding="utf-8") as stream: + json.dump(entries, stream, separators=(",", ":")) + os.replace(temporary, path) + except OSError: + Path(temporary).unlink(missing_ok=True) + + +class FunctionIndex: + """Normalized shapes, calls and data references of every function. + + With `cache_dir`, the per-function data is kept in a file named after the + library's SHA-256, this module's source and the capstone version, and read + back instead of decoding the library again. + """ + + def __init__(self, image: Any, cache_dir: Path | None = None): + DERIVER.require_capstone() + self.image = image + self.decoder = DERIVER.capstone.Cs( + DERIVER.capstone.CS_ARCH_X86, DERIVER.capstone.CS_MODE_64 + ) + self.text = image.sections[".text"] + self._decoded: dict[int, list[Instruction]] = {} + path = None if cache_dir is None else cache_path(Path(cache_dir), image.sha256) + functions = None if path is None else _read_cache(path) + if functions is None: + functions = self._decode_all() + if path is not None: + _write_cache(path, functions) + self.functions = functions + self.starts = sorted(self.functions) + self.callers, self.data_users = link_functions(self.functions) + + def _decode_all(self) -> dict[int, FunctionInfo]: + text_end = self.text.address + self.text.size + functions = {} + for start, size in function_table(self.image): + if size == 0 or not self.text.address <= start < text_end: + continue + instructions = self.instructions(start, size, cache=False) + shape_hash = hashlib.blake2b( + "\n".join(item.shape for item in instructions).encode(), digest_size=8 + ).digest() + calls = tuple( + item.branch_target + for item in instructions + if item.branch_target is not None + and (item.shape == "call" or item.shape.endswith(" out")) + ) + refs = tuple(target for item in instructions for target in item.rip_targets) + functions[start] = FunctionInfo(size, shape_hash, calls, refs) + return functions + + def instructions( + self, start: int, size: int | None = None, cache: bool = True + ) -> list[Instruction]: + if cache and start in self._decoded: + return self._decoded[start] + if size is None: + size = self.functions[start].size + offset = self.text.offset + start - self.text.address + decoded = decode_function(self.decoder, self.image.data[offset : offset + size], start) + if cache: + self._decoded[start] = decoded + return decoded + + def shapes(self, start: int) -> list[str]: + # Scoring reads each candidate once; only aligned functions stay cached. + return [item.shape for item in self.instructions(start, cache=start in self._decoded)] + + def containing(self, rva: int) -> int | None: + index = bisect.bisect_right(self.starts, rva) - 1 + if index < 0: + return None + start = self.starts[index] + return start if rva < start + self.functions[start].size else None + + def string_at(self, rva: int) -> bytes | None: + segment = self.image.segment_for_rva(rva) + if segment is None or not segment.contains_file_data(rva): + return None + offset = self.image.rva_to_offset(rva) + raw = self.image.data[offset : offset + 96] + end = raw.find(b"\0") + if end < 4 or not all(32 <= byte < 127 for byte in raw[:end]): + return None + return raw[:end] diff --git a/tests/roblox_abi_solver_test.py b/tests/roblox_abi_solver_test.py new file mode 100644 index 0000000..85588be --- /dev/null +++ b/tests/roblox_abi_solver_test.py @@ -0,0 +1,528 @@ +#!/usr/bin/env python3 +# Copyright 2026 Mocktail Project Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import contextlib +import importlib.util +import io +import json +from pathlib import Path +import sys +import tempfile +import types +import unittest +from unittest import mock + +from host_abi_regression import ( + REGRESSION_CACHE, + expected_runtime, + load_regression_build, + regression_library, + regression_manifest, + regression_sidecar, +) + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +# Each pair solves the second build from the first. +REGRESSION_PAIRS = ( + ("2908", "2998"), ("2998", "3092"), ("3092", "3222"), ("3222", "3092") +) + + +def load_script(name: str, file_name: str): + specification = importlib.util.spec_from_file_location( + name, PROJECT_ROOT / "scripts" / file_name + ) + if specification is None or specification.loader is None: + raise RuntimeError(f"cannot load {file_name}") + module = importlib.util.module_from_spec(specification) + sys.modules[specification.name] = module + specification.loader.exec_module(module) + return module + + +# The check script reuses the solver loaded under this module name. +SOLVER = load_script("mocktail_roblox_abi_solver", "roblox_abi_solver.py") +CHECK = load_script("mocktail_roblox_abi_check", "roblox_abi_check.py") +DERIVER = SOLVER.DERIVER + + +class CandidateChoiceTest(unittest.TestCase): + def test_accepts_a_clear_leader(self) -> None: + self.assertEqual(SOLVER.choose_candidate([(0.80, 0x10), (0.41, 0x20)]), 0x10) + + def test_rejects_a_near_tie(self) -> None: + self.assertIsNone(SOLVER.choose_candidate([(0.80, 0x10), (0.70, 0x20)])) + + def test_order_of_the_scores_does_not_matter(self) -> None: + self.assertEqual(SOLVER.choose_candidate([(0.41, 0x20), (0.80, 0x10)]), 0x10) + + def test_a_lead_of_exactly_the_threshold_is_accepted(self) -> None: + self.assertEqual(SOLVER.choose_candidate([(0.70, 0x10), (0.50, 0x20)]), 0x10) + + def test_a_single_candidate_still_needs_the_floor(self) -> None: + self.assertEqual(SOLVER.choose_candidate([(0.80, 0x10)]), 0x10) + self.assertIsNone(SOLVER.choose_candidate([(0.50, 0x10)])) + self.assertIsNone(SOLVER.choose_candidate([])) + + +class InstructionAlignmentTest(unittest.TestCase): + def test_maps_matching_instructions_across_an_insertion(self) -> None: + self.assertEqual( + SOLVER.align_instructions(["push", "mov", "ret"], ["push", "nop", "mov", "ret"]), + {0: 0, 1: 2, 2: 3}, + ) + + def test_leaves_replaced_instructions_unmapped(self) -> None: + self.assertEqual( + SOLVER.align_instructions(["push", "add", "ret"], ["push", "sub", "ret"]), + {0: 0, 2: 2}, + ) + + +class VtableTieBreakTest(unittest.TestCase): + def test_breaks_a_tie_between_candidates_that_pass_on_their_own(self) -> None: + scored = [(0.90, 0x10), (0.85, 0x20)] + self.assertIsNone(SOLVER.choose_candidate(scored)) + self.assertEqual(SOLVER.choose_candidate(scored, tie_breaker=0x20), 0x20) + + def test_cannot_lift_a_candidate_below_the_floor(self) -> None: + scored = [(0.90, 0x10), (0.85, 0x20), (0.50, 0x30)] + self.assertIsNone(SOLVER.choose_candidate(scored, tie_breaker=0x30)) + + def test_cannot_overrule_a_clear_leader(self) -> None: + scored = [(0.95, 0x10), (0.65, 0x20)] + self.assertEqual(SOLVER.choose_candidate(scored, tie_breaker=0x20), 0x10) + + def test_a_missing_slot_candidate_costs_nothing(self) -> None: + scored = [(0.90, 0x10), (0.30, 0x20)] + self.assertEqual(SOLVER.choose_candidate(scored, tie_breaker=0x99), 0x10) + + +def function(size: int = 32, shape: bytes = b"", calls=(), refs=()) -> object: + return SOLVER.INDEX.FunctionInfo(size, shape, tuple(calls), tuple(refs)) + + +class FakeIndex: + """The parts of FunctionIndex the solver reads, from hand-made data.""" + + def __init__(self, functions, data_users=None, instructions=None, image=None): + self.functions = functions + self.starts = sorted(functions) + self.callers = SOLVER.INDEX.link_functions(functions)[0] + self.data_users = data_users or {} + self._instructions = instructions or {} + self.image = image + + def instructions(self, start): + return self._instructions[start] + + +def call(address: int, target: int): + return SOLVER.Instruction(address, "call", (), target) + + +def plain(address: int): + return SOLVER.Instruction(address, "nop ", (), None) + + +class FunctionLinkTest(unittest.TestCase): + def test_lists_each_caller_and_data_user_once(self) -> None: + callers, data_users = SOLVER.INDEX.link_functions({ + 0x100: function(calls=(0x300, 0x300, 0x300), refs=(0x900, 0x900)), + 0x200: function(calls=(0x300,), refs=(0x900,)), + 0x300: function(), + }) + self.assertEqual(callers[0x300], [0x100, 0x200]) + self.assertEqual(data_users[0x900], [0x100, 0x200]) + + +class DataPatternTest(unittest.TestCase): + def solver(self, refs, data_users): + solver = SOLVER.Solver.from_indexes( + FakeIndex({0x100: function(refs=refs)}), FakeIndex({}, data_users) + ) + solver.data_map = {target: target + 0x1000 for target in refs} + return solver + + def test_an_address_used_by_too_many_functions_casts_no_votes(self) -> None: + hot_users = list(range(0x2000, 0x2001 + SOLVER.XREF_FANOUT_LIMIT)) + solver = self.solver((0x9000, 0x9100), {0xA000: hot_users, 0xA100: [0x3000]}) + self.assertEqual(solver._pattern_candidates(0x100), {0x3000}) + + def test_a_common_address_still_backs_a_candidate_found_by_a_rare_one(self) -> None: + hot_users = [0x3000, *range(0x2000, 0x2001 + SOLVER.XREF_FANOUT_LIMIT)] + refs = (0x9000, 0x9100, 0x9200, 0x9300) + solver = self.solver(refs, {0xA000: hot_users, 0xA100: [0x3000]}) + self.assertEqual(solver._pattern_candidates(0x100), {0x3000}) + + def test_a_repeated_reference_casts_one_vote(self) -> None: + refs = (0x9000, 0x9100, 0x9200, 0x9300) + solver = self.solver( + refs, {0xA000: [0x4000, 0x4000, 0x4000], 0xA100: [0x5000], 0xA200: [0x5000]} + ) + self.assertEqual(solver._pattern_candidates(0x100), {0x5000}) + + +class OneToOnePairTest(unittest.TestCase): + def test_call_propagation_never_pairs_two_functions_with_one(self) -> None: + reference = FakeIndex({ + 0x100: function(shape=b"A", calls=(0x200, 0x300)), + 0x200: function(shape=b"x1"), + 0x300: function(shape=b"x2"), + }) + candidate = FakeIndex({ + 0x1100: function(shape=b"A", calls=(0x1200, 0x1200)), + 0x1200: function(shape=b"c"), + }) + solver = SOLVER.Solver.from_indexes(reference, candidate) + self.assertEqual(solver.pairs, {0x100: 0x1100, 0x200: 0x1200}) + + def test_a_new_pair_replaces_the_one_holding_its_candidate(self) -> None: + solver = SOLVER.Solver.from_indexes(FakeIndex({}), FakeIndex({})) + solver._pair(0x200, 0x1200) + solver._pair(0x300, 0x1200) + self.assertEqual(solver.pairs, {0x300: 0x1200}) + + +class RegistryAllocatorTest(unittest.TestCase): + def solver(self, candidate_call_index: int): + reference_code = [plain(0x100 + i) for i in range(12)] + reference_code[9] = call(0x109, 0x500) + candidate_code = [plain(0x1100 + i) for i in range(12)] + candidate_code[candidate_call_index] = call(0x1100 + candidate_call_index, 0x1500) + return SOLVER.Solver.from_indexes( + FakeIndex({0x100: function()}, instructions={0x100: reference_code}), + FakeIndex({0x1100: function()}, instructions={0x1100: candidate_code}), + ) + + def test_accepts_the_allocator_call_at_the_tenth_instruction(self) -> None: + self.solver(9)._check_registry_allocator(0x100, 0x1100, 0x500, 0x1500) + + def test_rejects_an_allocator_call_anywhere_else(self) -> None: + with self.assertRaisesRegex(SOLVER.AnalyzerError, "allocator"): + self.solver(3)._check_registry_allocator(0x100, 0x1100, 0x500, 0x1500) + + def test_a_reference_profile_without_its_free_bridge_is_rejected(self) -> None: + build = load_regression_build("3092") + sidecar = regression_sidecar(build) + sidecar["profile"]["bridge_entries"] = [ + entry for entry in sidecar["profile"]["bridge_entries"] if entry["label"] != "free" + ] + solver = SOLVER.Solver.from_indexes(FakeIndex({}), FakeIndex({})) + with self.assertRaisesRegex(SOLVER.AnalyzerError, "allocator bridges"): + solver.derive_profile(sidecar) + + +FMOD_FIELDS = { + "vtable_rva": "0x1000", + "string_constructor_rva": "0x10", + "count_method_rva": "0x20", + "info_method_rva": "0x30", + "current_method_rva": "0x40", + "select_method_rva": "0x50", +} +FMOD_INPUT_FIELDS = DERIVER.FMOD_INPUT_FIELDS + + +class FmodSourceTest(unittest.TestCase): + def test_accepts_both_layouts(self) -> None: + self.assertEqual(SOLVER.validate_fmod_source(FMOD_FIELDS), 1) + self.assertEqual( + SOLVER.validate_fmod_source({**FMOD_FIELDS, "vtable_layout_version": 2}), 2 + ) + + def test_rejects_a_partial_input_profile(self) -> None: + source = {**FMOD_FIELDS, FMOD_INPUT_FIELDS[0]: "0x60"} + with self.assertRaisesRegex(SOLVER.AnalyzerError, "incomplete"): + SOLVER.validate_fmod_source(source) + + def test_rejects_a_missing_output_field(self) -> None: + source = dict(FMOD_FIELDS) + del source["select_method_rva"] + with self.assertRaisesRegex(SOLVER.AnalyzerError, "incomplete"): + SOLVER.validate_fmod_source(source) + + def test_rejects_an_unknown_layout(self) -> None: + for layout in (3, "2", True): + with self.subTest(layout=layout), self.assertRaisesRegex( + SOLVER.AnalyzerError, "layout" + ): + SOLVER.validate_fmod_source({**FMOD_FIELDS, "vtable_layout_version": layout}) + + def test_finds_a_vtable_that_moved_to_the_other_layout(self) -> None: + vtable = 0x8000 + relocations = {vtable + 5 * 8: 0xA, vtable + 6 * 8: 0xB, vtable + 8 * 8: 0xC} + methods = {"count_method_rva": 0xA, "info_method_rva": 0xB, "current_method_rva": 0xC} + self.assertEqual( + SOLVER.Solver._locate_fmod_vtables(relocations, methods), [(vtable, 2)] + ) + + +class EhFrameHeaderTest(unittest.TestCase): + def test_rejects_an_eight_byte_frame_pointer(self) -> None: + # Version 1, eh_frame_ptr udata8, count udata4, table datarel sdata4. + raw = bytes((1, 0x04, 0x03, 0x3B)) + bytes(16) + header = types.SimpleNamespace(address=0x1000) + image = types.SimpleNamespace( + sections={".eh_frame_hdr": header}, section_bytes=lambda section: raw + ) + with self.assertRaisesRegex(SOLVER.AnalyzerError, "eh_frame_ptr"): + SOLVER.INDEX.function_table(image) + + +class FmodLayoutTieTest(unittest.TestCase): + def test_resolves_current_and_select_from_the_located_vtable(self) -> None: + reference_vtable, vtable = 0x1000, 0x8000 + source = { + "vtable_rva": hex(reference_vtable), "string_constructor_rva": "0x10", + "count_method_rva": "0x20", "info_method_rva": "0x30", + "current_method_rva": "0x40", "select_method_rva": "0x50", + } + reference_relocations = { + reference_vtable + slot * 8: method + for slot, method in ((5, 0x20), (6, 0x30), (7, 0x40), (17, 0x50)) + } + candidate_relocations = { + vtable + slot * 8: method + for slot, method in ((5, 0xA), (6, 0xB), (7, 0xC), (17, 0xD)) + } + reference_image = types.SimpleNamespace(require_relro_rva=lambda *_: None) + candidate_image = types.SimpleNamespace(require_relro_rva=lambda *_: None) + solver = SOLVER.Solver.from_indexes( + FakeIndex({}, image=reference_image), FakeIndex({}, image=candidate_image) + ) + first_pass = { + "fmod-count_method_rva": 0xA, + "fmod-info_method_rva": 0xB, + "fmod-string-constructor": 0x99, + } + + def match(name, start, contract=None, vtable_slot=None): + if name in first_pass: + return first_pass[name] + if vtable_slot is not None: + return vtable_slot + raise SOLVER.AnalyzerError(f"{name} has no clear candidate") + + with mock.patch.object(solver, "match_function", side_effect=match), mock.patch.object( + DERIVER, + "relative_relocation_map", + side_effect=lambda image: ( + reference_relocations if image is reference_image else candidate_relocations + ), + ), mock.patch.object( + DERIVER, "has_fmod_string_constructor_contract", return_value=True + ), mock.patch.object( + DERIVER, "has_fmod_device_lists_select_contract", return_value=False + ): + result = solver._derive_fmod_bridge(source) + self.assertEqual(result["vtable_rva"], DERIVER.format_rva(vtable)) + self.assertEqual(result["current_method_rva"], DERIVER.format_rva(0xC)) + self.assertEqual(result["select_method_rva"], DERIVER.format_rva(0xD)) + self.assertNotIn("vtable_layout_version", result) + + +def regression_inputs(reference_name: str, candidate_name: str): + """Both fixtures and libraries, or None when either library is missing.""" + reference = load_regression_build(reference_name) + candidate = load_regression_build(candidate_name) + reference_library = regression_library(reference) + candidate_library = regression_library(candidate) + if reference_library is None or candidate_library is None: + return None + return reference, candidate, reference_library, candidate_library + + +def regression_images(test: unittest.TestCase, reference_name: str, candidate_name: str): + inputs = regression_inputs(reference_name, candidate_name) + if inputs is None: + test.skipTest(f"exact {reference_name} and {candidate_name} payloads are unavailable") + return inputs + + +@unittest.skipUnless(DERIVER.capstone is not None, "Python capstone is unavailable") +class InputCheckTest(unittest.TestCase): + def test_rejects_inputs_the_deriver_rejects(self) -> None: + reference, _, reference_library, candidate_library = regression_images( + self, "3092", "3222" + ) + sidecar = regression_sidecar(reference) + with DERIVER.ElfImage(reference_library) as reference_image, DERIVER.ElfImage( + candidate_library + ) as candidate_image: + SOLVER.check_inputs(reference_image, candidate_image, sidecar) + with self.assertRaisesRegex(SOLVER.AnalyzerError, "distinct"): + SOLVER.check_inputs(reference_image, reference_image, sidecar) + wrong = {**sidecar, "elf_build_id": "00" * 20} + with self.assertRaisesRegex(SOLVER.AnalyzerError, "identity"): + SOLVER.check_inputs(reference_image, candidate_image, wrong) + solver = SOLVER.Solver.from_indexes( + FakeIndex({}, image=reference_image), FakeIndex({}, image=candidate_image) + ) + with self.assertRaisesRegex(SOLVER.AnalyzerError, "identity"): + solver.derive_profile(wrong) + + +@unittest.skipUnless(DERIVER.capstone is not None, "Python capstone is unavailable") +class FunctionIndexCacheTest(unittest.TestCase): + def library(self) -> Path: + library = regression_library(load_regression_build("3222")) + if library is None: + self.skipTest("exact 3222 payload is unavailable") + return library + + def test_reuses_a_cached_index_without_decoding(self) -> None: + with tempfile.TemporaryDirectory() as cache, DERIVER.ElfImage(self.library()) as image: + built = SOLVER.FunctionIndex(image, Path(cache)) + self.assertEqual(len(list(Path(cache).iterdir())), 1) + with mock.patch.object( + SOLVER.INDEX, "decode_function", side_effect=AssertionError("decoded") + ): + loaded = SOLVER.FunctionIndex(image, Path(cache)) + self.assertEqual(loaded.functions, built.functions) + self.assertEqual(loaded.callers, built.callers) + self.assertEqual(loaded.data_users, built.data_users) + + def test_rebuilds_a_damaged_cache(self) -> None: + with tempfile.TemporaryDirectory() as cache, DERIVER.ElfImage(self.library()) as image: + built = SOLVER.FunctionIndex(image, Path(cache)) + (entry,) = Path(cache).iterdir() + entry.write_bytes(b"not an index") + self.assertEqual(SOLVER.FunctionIndex(image, Path(cache)).functions, built.functions) + + def test_a_new_disassembler_version_gets_a_new_entry(self) -> None: + before = SOLVER.INDEX.cache_path(Path("/cache"), "ab" * 32) + with mock.patch.object(DERIVER.capstone, "__version__", "0.0.0"): + after = SOLVER.INDEX.cache_path(Path("/cache"), "ab" * 32) + self.assertNotEqual(before, after) + + +class CheckReportTest(unittest.TestCase): + def test_lists_each_differing_field_by_path(self) -> None: + derived = {"profile": {"a": "0x1", "b": ["0x2", "0x3"]}, "runtime": {"c": "0x4"}} + solved = {"profile": {"a": "0x1", "b": ["0x2", "0x9"]}, "runtime": {}} + self.assertEqual( + CHECK.compare_results(derived, solved), + ["profile.b[1]: deriver 0x3, solver 0x9", "runtime.c: deriver 0x4, solver missing"], + ) + + def test_exit_status_follows_agreement(self) -> None: + same = {"profile": {"a": "0x1"}} + self.assertEqual(CHECK.check_status(same, same, None, None)[0], 0) + self.assertEqual(CHECK.check_status(same, {"profile": {}}, None, None)[0], 1) + status, lines = CHECK.check_status(None, same, "no unique match", None) + self.assertEqual(status, 1) + self.assertIn("deriver failed: no unique match", lines) + self.assertEqual(CHECK.check_status(None, None, "a", "b")[0], 2) + + +@unittest.skipUnless(DERIVER.capstone is not None, "Python capstone is unavailable") +class CheckCommandTest(unittest.TestCase): + def run_check(self, reference: dict, reference_library, candidate_library): + with tempfile.TemporaryDirectory() as temporary: + sidecar = Path(temporary) / "reference.json" + sidecar.write_text(json.dumps(regression_sidecar(reference)), encoding="utf-8") + manifest = Path(temporary) / "compatibility.json" + manifest.write_text(json.dumps(regression_manifest(reference)), encoding="utf-8") + output = io.StringIO() + with contextlib.redirect_stdout(output): + status = CHECK.main([ + "--reference-lib", str(reference_library), + "--reference-profile", str(sidecar), + "--reference-compatibility", str(manifest), + "--candidate-lib", str(candidate_library), + "--cache-dir", str(REGRESSION_CACHE / "function_index"), + ]) + return status, output.getvalue() + + def test_reports_agreement_on_a_real_update(self) -> None: + reference, _, reference_library, candidate_library = regression_images( + self, "2998", "3092" + ) + status, output = self.run_check(reference, reference_library, candidate_library) + self.assertEqual(status, 0, output) + self.assertIn("deriver and solver agree", output) + + def test_rejects_a_wrong_reference_before_indexing(self) -> None: + reference, _, _, candidate_library = regression_images(self, "2998", "3092") + third = regression_library(load_regression_build("3222")) + if third is None: + self.skipTest("exact 3222 payload is unavailable") + with mock.patch.object( + SOLVER, "FunctionIndex", side_effect=AssertionError("indexed") + ): + status, output = self.run_check(reference, third, candidate_library) + self.assertEqual(status, 2, output) + self.assertIn("solver failed: reference ELF bytes do not match", output) + + def test_reports_any_error_as_that_side_failing(self) -> None: + reference, _, reference_library, candidate_library = regression_images( + self, "2998", "3092" + ) + with mock.patch.object( + DERIVER, "derive_profile", side_effect=ValueError("bad slot") + ), mock.patch.object( + SOLVER.Solver, "derive_profile", side_effect=KeyError("free") + ), mock.patch.object(SOLVER, "FunctionIndex", return_value=FakeIndex({})): + status, output = self.run_check(reference, reference_library, candidate_library) + self.assertEqual(status, 2, output) + self.assertIn("deriver failed: ValueError: bad slot", output) + self.assertIn("solver failed: KeyError: 'free'", output) + + +# The smallest lead of any match the score decided alone, per pair. A +# change that brings matches closer to a tie fails here first. +WEAKEST_SCORE_LEAD = { + ("2908", "2998"): 0.22, + ("2998", "3092"): 0.22, + ("3092", "3222"): 0.22, + ("3222", "3092"): 0.21, +} + + +@unittest.skipUnless(DERIVER.capstone is not None, "Python capstone is unavailable") +class SolverRegressionTest(unittest.TestCase): + def test_solves_each_build_from_its_predecessor(self) -> None: + for reference_name, candidate_name in REGRESSION_PAIRS: + with self.subTest(reference=reference_name, candidate=candidate_name): + self.solve_pair(reference_name, candidate_name) + + def solve_pair(self, reference_name: str, candidate_name: str) -> None: + reference, candidate, reference_library, candidate_library = regression_images( + self, reference_name, candidate_name + ) + with DERIVER.ElfImage(reference_library) as reference_image, DERIVER.ElfImage( + candidate_library + ) as candidate_image: + solver = SOLVER.Solver( + reference_image, candidate_image, REGRESSION_CACHE / "function_index" + ) + profile, anchors = solver.derive_profile(regression_sidecar(reference)) + runtime = solver.derive_runtime_compatibility(reference["runtime_compatibility"]) + self.assertEqual(profile, candidate["profile"]) + self.assertEqual(anchors, candidate["derivation_anchors"]) + self.assertEqual(runtime, expected_runtime(reference, candidate)) + weakest = min( + match.lead for match in solver.matches.values() if match.decided_by == "score" + ) + self.assertGreaterEqual( + round(weakest, 2), WEAKEST_SCORE_LEAD[(reference_name, candidate_name)] + ) + + +if __name__ == "__main__": + unittest.main()