leakhunt/tools/patch_v8_test.py

"""patch_v8_test.py <pid> [--revert]

EXPERIMENTAL v8: comprehensive fix for UIElement_UIItem leak.

Combines v7's two outer NOPs (so UpdateEmptySlots runs always)
with a 1-byte change to the inner trim loop's break-3 (so the loop
skips non-WAITING items instead of exiting at the first one).

Sites:
  Site 1 — 0x004e439d (v6/v7): NOP visibility guard
    Original: 0F 84 F3 01 00 00   jz +0x1f3
    Patched:  90 90 90 90 90 90

  Site 2 — 0x004e43c0 (v7): NOP auto-fit guard
    Original: 0F 85 CD 01 00 00   jne +0x1cd
    Patched:  90 90 90 90 90 90

  Site 3 — 0x004e4496 (v8 new): turn "break" into "continue"
    Original: 75 0D                jne +0x0d → exit_loop
    Patched:  75 08                jne +0x08 → inc_counter (continue)

After site 3, when an item's state != WAITING, the loop skips the
InternalDeleteItem call and advances the counter instead of exiting.
Combined with sites 1 and 2, UpdateEmptySlots now runs always and
iterates the entire array, calling InternalDeleteItem on every WAITING
UIItem.

Safety considerations:
  - Break 1 (NULL item) and break 2 (not UIItem type) are KEPT — they
    are real safety guards.
  - The counter (ebx) increments for skipped items too, so the loop
    may exit earlier than ideal if many items are non-WAITING. But
    with width=0 (invisible) the delete_count is ~count+1 so most
    items should still be visited.
  - InternalDeleteItem defers actual deletion to a queue, so even if
    we mistakenly delete an active item, the engine should handle it
    on the next frame.
"""
import argparse
import ctypes
import ctypes.wintypes as wt
import sys


SITES = [
    ("visibility guard (jz) ", 0x004e439d,
     bytes([0x0f, 0x84, 0xf3, 0x01, 0x00, 0x00]),
     bytes([0x90] * 6)),
    ("auto-fit guard (jne)  ", 0x004e43c0,
     bytes([0x0f, 0x85, 0xcd, 0x01, 0x00, 0x00]),
     bytes([0x90] * 6)),
    ("trim break-3 (jne)    ", 0x004e4496,
     bytes([0x75, 0x0d]),
     bytes([0x75, 0x08])),
]


PROCESS_VM_READ           = 0x0010
PROCESS_VM_WRITE          = 0x0020
PROCESS_VM_OPERATION      = 0x0008
PROCESS_QUERY_INFORMATION = 0x0400
PAGE_EXECUTE_READWRITE    = 0x40


k32 = ctypes.windll.kernel32
OpenProcess = k32.OpenProcess
OpenProcess.argtypes = [wt.DWORD, wt.BOOL, wt.DWORD]; OpenProcess.restype = wt.HANDLE
CloseHandle = k32.CloseHandle
CloseHandle.argtypes = [wt.HANDLE]; CloseHandle.restype = wt.BOOL
WriteProcessMemory = k32.WriteProcessMemory
WriteProcessMemory.argtypes = [wt.HANDLE, wt.LPVOID, wt.LPCVOID, ctypes.c_size_t,
                               ctypes.POINTER(ctypes.c_size_t)]
WriteProcessMemory.restype = wt.BOOL
ReadProcessMemory = k32.ReadProcessMemory
ReadProcessMemory.argtypes = [wt.HANDLE, wt.LPCVOID, wt.LPVOID, ctypes.c_size_t,
                              ctypes.POINTER(ctypes.c_size_t)]
ReadProcessMemory.restype = wt.BOOL
VirtualProtectEx = k32.VirtualProtectEx
VirtualProtectEx.argtypes = [wt.HANDLE, wt.LPVOID, ctypes.c_size_t, wt.DWORD,
                             ctypes.POINTER(wt.DWORD)]
VirtualProtectEx.restype = wt.BOOL


def read_bytes(h, addr, n):
    buf = (ctypes.c_ubyte * n)()
    sz = ctypes.c_size_t(0)
    if not ReadProcessMemory(h, addr, buf, n, ctypes.byref(sz)):
        raise OSError(f"read 0x{addr:08x} err={ctypes.get_last_error()}")
    return bytes(buf[:sz.value])


def write_bytes(h, addr, data):
    old_prot = wt.DWORD(0)
    if not VirtualProtectEx(h, addr, len(data), PAGE_EXECUTE_READWRITE, ctypes.byref(old_prot)):
        raise OSError(f"VirtualProtectEx 0x{addr:08x} err={ctypes.get_last_error()}")
    sz = ctypes.c_size_t(0)
    ok = WriteProcessMemory(h, addr, data, len(data), ctypes.byref(sz))
    err = ctypes.get_last_error() if not ok else 0
    restored = wt.DWORD(0)
    VirtualProtectEx(h, addr, len(data), old_prot.value, ctypes.byref(restored))
    if not ok:
        raise OSError(f"write 0x{addr:08x} err={err}")


def apply_or_revert(h, label, site_va, orig, patched, revert):
    cur = read_bytes(h, site_va, len(orig))
    print(f"  {label} @ 0x{site_va:08x}: current {cur.hex()}")
    if revert:
        if cur == orig:
            print(f"    already original")
            return
        if cur != patched:
            print(f"    UNEXPECTED — refusing")
            sys.exit(3)
        write_bytes(h, site_va, orig)
        after = read_bytes(h, site_va, len(orig))
        print(f"    reverted; now: {after.hex()}")
        return

    if cur == patched:
        print(f"    already patched")
        return
    if cur != orig:
        print(f"    UNEXPECTED — bytes don't match expected original. refusing.")
        sys.exit(4)
    write_bytes(h, site_va, patched)
    after = read_bytes(h, site_va, len(orig))
    print(f"    patched; now: {after.hex()}")
    if after != patched:
        print(f"    MISMATCH — write didn't take")
        sys.exit(5)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("pid", type=int)
    ap.add_argument("--revert", action="store_true")
    args = ap.parse_args()

    h = OpenProcess(
        PROCESS_VM_READ | PROCESS_VM_WRITE | PROCESS_VM_OPERATION | PROCESS_QUERY_INFORMATION,
        False, args.pid,
    )
    if not h:
        print(f"OpenProcess({args.pid}) err={ctypes.get_last_error()}")
        sys.exit(2)

    print(f"PID {args.pid}")
    for label, va, orig, patched in SITES:
        apply_or_revert(h, label, va, orig, patched, args.revert)
    print("  OK")
    CloseHandle(h)


if __name__ == "__main__":
    main()