#!/usr/bin/env python3
# wall2_resolver2.py - HOST-SIDE wall#2 resolver v2 (r674+).
#
# Why: the vk_frame client's channel-init push (va=0x4500000 len=1432) is
# kernel-built and kernel-only; the sema service cannot translate it
# ("false translation for push ... rescuing"), the client times out (~11 s)
# per attempt, retries, then VKFAIL -4. Guest RAM is host-readable, so the
# push page can be found and injected via sema_map_add.
#
# v2 recipe (per the r672/r673 structure analysis):
#  1. wait for the channel ring line in dmesg (validated RING_PA + handle)
#  2. scan guest RAM for the push page:
#     a) CONTENT: the 24-byte release group the service's own rescue matches
#        (hdr=0x20038090, hi=1, payload sane, w5/v5 shape). The channel-init
#        entry (va=0x4500000 len=1432) puts that group at page offset 0x580.
#     b) TUPLES: pages of 32-byte records {vmap_ptr; flags; PA; size} (both
#        observed layouts) matched on (PA & ~0xfff) == RING_PA with strides
#        0x20/0x40/0x60; entries below the ring entry in VA order map
#        0x4500000.. (push VA = ring VA 0x4520000 - 0x20000 = 32 pages).
#  3. inject (va -> pa, len) through /sys/module/vgsp/parameters/sema_map_add
#  4. keep scanning/re-injecting for the watch window (retries allocate new
#     pages) and log everything.
#
# Run this FIRST (foreground), then launch the client in the guest.
import subprocess, struct, time, sys, re, os, signal

RAM = "/dev/shm/vgpu-guest-ram"
PARAM = "/sys/module/vgsp/parameters/sema_map_add"
LOG = "/tmp/opencode/wall2_resolver2.log"
RING_VA = 0x4520000
PUSH_VA = 0x4500000
DELTA = (RING_VA - PUSH_VA) // 0x1000          # 32 pages
REL_OFF = 0x580                                # group offset in the push page
MASK = ~0xfff & ((1 << 64) - 1)

WAIT_RING_S = int(os.environ.get("W2_WAIT_RING", "420"))
WATCH_S = int(os.environ.get("W2_WATCH", "150"))

logf = open(LOG, "a", buffering=1)
def log(*a):
    s = " ".join(str(x) for x in a)
    print(s, flush=True)
    logf.write("[%s] %s\n" % (time.strftime("%H:%M:%S"), s))

def dmesg():
    return subprocess.run(["sudo", "dmesg"], capture_output=True, text=True).stdout

def inject(va, pa, ln=0x1000, ap=1):
    try:
        with open(PARAM, "w") as f:
            f.write("0x%x,0x%x,0x%x,%d\n" % (va, pa, ln, ap))
        log("INJECT va=0x%x -> pa=0x%x len=0x%x" % (va, pa, ln))
        return True
    except OSError as e:
        log("inject FAIL", e)
        return False

def prefix_ok(d):
    off = 0
    while off + 4 <= len(d):
        hdr = struct.unpack_from("<I", d, off)[0]
        if hdr:
            break
        off += 4
    if off + 4 > len(d):
        return False
    op = (hdr >> 29) & 7
    if op in (4, 7):
        return True
    if op not in (1, 3, 5):
        return False
    return ((hdr >> 16) & 0xfff) != 0

def rel_shape(g):
    hdr, hi, lo, pay, w5, v5 = struct.unpack("<6I", g)
    return (hdr == 0x20038090 and hi == 1 and 0 < pay < (1 << 24)
            and (w5 >> 29) == 1 and ((w5 >> 13) & 7) == 4
            and ((w5 & 0x1fff) << 2) == 0x300 and (v5 & 3) == 0), lo & 0xfff, pay

img = None
def open_ram():
    global img
    if img is None:
        img = open(RAM, "rb", buffering=0)
    return img

def gpa2off(pa):
    if pa < 0x80000000:
        return pa
    if pa < 0x180000000:
        return 0x80000000 + (pa - 0x100000000)
    return None

def off2gpa(off):
    return off if off < 0x80000000 else 0x100000000 + (off - 0x80000000)

def rd(pa, n):
    off = gpa2off(pa)
    if off is None:
        return None
    f = open_ram()
    try:
        f.seek(off)
        return f.read(n)
    except OSError:
        return None

def find_content_candidates():
    """Pages containing the release group; returns [(page_pa, off, slot, pay)]."""
    f = open_ram()
    f.seek(0, 2)
    fsz = f.tell()
    pat = b"\x90\x80\x03\x20"
    out = []
    CH = 64 << 20
    pos = 0
    while pos < fsz:
        f.seek(pos)
        d = f.read(min(CH, fsz - pos))
        st = 0
        while True:
            i = d.find(pat, st)
            if i < 0:
                break
            st = i + 1
            if i + 24 > len(d):
                continue
            ok, slot, pay = rel_shape(d[i:i + 24])
            if not ok:
                continue
            gpa = off2gpa(pos + i)
            page = gpa & MASK
            off = gpa - page
            out.append((page, off, slot, pay))
        pos += CH
    return out

def find_tuple_arrays(ring_pa):
    """Find record pages containing the ring PA; returns list of
    (page_gpa, stride, ring_index, entries[list of pa])."""
    f = open_ram()
    f.seek(0, 2)
    fsz = f.tell()
    pat = struct.pack("<Q", ring_pa)
    arrays = []
    CH = 64 << 20
    pos = 0
    hits = []
    while pos < fsz:
        f.seek(pos)
        d = f.read(min(CH, fsz - pos))
        st = 0
        while True:
            i = d.find(pat, st)
            if i < 0:
                break
            hits.append(off2gpa(pos + i))
            st = i + 1
        pos += CH
    seen = set()
    for gpa in hits:
        page = gpa & MASK
        if page in seen:
            continue
        seen.add(page)
        d = rd(page, 0x1000)
        if not d:
            continue
        rec_off = gpa - page
        for stride in (0x20, 0x40, 0x60):
            for rel in range(0, min(stride, rec_off + 1), 8):
                start = rec_off - rel
                ents = []
                for o in range(start, 0x1000 - stride + 1, stride):
                    # PA field anywhere in the record: try +0x10 first, then any qword
                    vals = struct.unpack_from("<%dQ" % (stride // 8), d, o)
                    pa = None
                    for cand_off in (0x10, 0x8, 0x18, 0x0, 0x20, 0x28):
                        if cand_off + 8 <= stride:
                            v = vals[cand_off // 8]
                            if (v & MASK) == ring_pa:
                                pa = v & MASK
                                break
                    if pa is None:
                        continue
                    ents.append((o, pa))
                if len(ents) >= 4:
                    idx = [e[1] for e in ents].index(ring_pa)
                    arrays.append((page, stride, start, idx, [e[1] for e in ents]))
                    break
            if arrays and arrays[-1][0] == page:
                break
    return arrays

def main():
    log("=== wall2_resolver2 start (wait_ring=%ss watch=%ss) ===" % (WAIT_RING_S, WATCH_S))
    ring_pa = None
    ring_h = None
    deadline = time.time() + WAIT_RING_S
    while time.time() < deadline:
        for ln in dmesg().splitlines():
            m = re.search(r"h=(0x[0-9a-f]+) .*ring rc=0 pa=0x([0-9a-f]+)", ln)
            if m:
                ring_h = m.group(1)
                ring_pa = int(m.group(2), 16)
                break
        if ring_pa:
            break
        time.sleep(1)
    if not ring_pa:
        log("no ring line seen in %ss; exiting" % WAIT_RING_S)
        return 1
    log("RING %s PA=0x%x" % (ring_h, ring_pa))

    injected = set()
    watch_deadline = time.time() + WATCH_S
    last_pass = 0
    while time.time() < watch_deadline:
        # ---- content pass
        try:
            cands = find_content_candidates()
        except Exception as e:
            log("content scan error:", e)
            cands = []
        # prefer the canonical group offset 0x580, then highest payload
        cands.sort(key=lambda c: (c[1] != REL_OFF, -c[3]))
        if cands:
            page, off, slot, pay = cands[0]
            d = rd(page, 0x1000)
            pok = bool(d) and prefix_ok(d[:0x100])
            log("content: %d candidate(s); best page=0x%x off=0x%x slot=0x%x pay=%u prefix_ok=%s"
                % (len(cands), page, off, slot, pay, pok))
            key = ("C", PUSH_VA, page)
            if key not in injected and pok:
                if inject(PUSH_VA, page, 0x1000):
                    injected.add(key)
                    # also cover the following page (entry2 tail) if non-contiguous unknown; skip
        else:
            # ---- tuple fallback (only every ~10 s to save time)
            if time.time() - last_pass > 10:
                last_pass = time.time()
                try:
                    arrs = find_tuple_arrays(ring_pa)
                except Exception as e:
                    log("tuple scan error:", e)
                    arrs = []
                if arrs:
                    page, stride, start, idx, ents = arrs[0]
                    log("tuple: array page=0x%x stride=0x%x ring_index=%d entries=%d"
                        % (page, stride, idx, len(ents)))
                    for k in range(idx - DELTA, idx):
                        if k < 0 or k >= len(ents):
                            continue
                        va = PUSH_VA + (k - (idx - DELTA)) * 0x1000
                        pa = ents[k]
                        key = ("T", va, pa)
                        if key in injected:
                            continue
                        d = rd(pa, 0x40)
                        if d and any(d):
                            inject(va, pa, 0x1000)
                            injected.add(key)
                else:
                    log("tuple: no array with ring PA found")
        # stop early when the service confirms our mapping
        if len(injected) and "injected mapping va=0x%x" % PUSH_VA in dmesg():
            log("service confirmed the injected mapping; done")
            break
        time.sleep(1)
    log("=== done: %d injection(s) ===" % len(injected))
    return 0

if __name__ == "__main__":
    sys.exit(main())