#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# scrub_pin_qmp.py - host-side channel-buffer (PMA scrubber / X DDX CE)
# semaphore+finish page scanner.
#
# The guest RM polls the channel-buffer host semaphore (channelWaitForFreeEntry)
# and finish payload (channelWaitForFinishPayload) through the *vmap* page of
# the buffer memdesc. After the guest RM re-creates its kernel VAS the host's
# GPU-VAS walk can keep resolving a stale page (r1712: pass authored
# 0x126b60000 while the guest polled 0x124fec000), so the waits never
# completed. This tool resolves the guest's own vmap pages and registers them
# through sema_scrubbuf_add; the service then authors the proven all-ones pair
# into the pages the guest actually reads (CONT.409).
#
# Usage (host, root for the QMP socket):
#   sudo python3 tools/guest-diag/scrub_pin_qmp.py [--duration 0] [--interval 2]
#
# Every registered page must look like a channel buffer: the ring page
# (base+0x64000) is all-zero or 906F entries whose GET names a push inside
# the channel's own GPU window. A page that fails the shape test is skipped.
import argparse
import importlib.util
import os
import re
import subprocess
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
OIV = os.path.normpath(os.path.join(HERE, "..", ".."))
VMEXEC = os.path.join(OIV, "scripts", "vgpu_vm.py")
PARAM = "/sys/module/vgsp/parameters/sema_scrubbuf_add"

# nv_vm_map_pages areas seen for channel buffers (pb 0x64000 + gpfifo + sema
# + finish, with per-size variants). The fwimage boots observed 0x6e000 and
# neighbours; keep a small band and rely on the ring shape test.
SCRUB_SIZES = (0x6b000, 0x6c000, 0x6d000, 0x6e000, 0x6f000, 0x70000,
               0x71000, 0x72000)
PAGE = 0x1000


def log(msg):
    print("[%s] %s" % (time.strftime("%H:%M:%S"), msg), flush=True)


QMP_PATH = "/tmp/qemu-qmp.sock"


def load_qmp(wait=0.0):
    """Connect to the QEMU monitor, retrying until it exists.

    The systemd unit starts at host boot, before the VM exists, and cycle.sh
    removes and re-creates the socket on every VM cycle. Wait for the socket
    instead of exiting, so one long-lived scanner covers the whole session
    (it must already be scanning when Xorg creates its DDX channel buffer -
    that is the point of running it from boot)."""
    spec = importlib.util.spec_from_file_location(
        "semfix_watch", os.path.join(HERE, "semfix_watch.py"))
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    deadline = time.time() + wait if wait else None
    last = "socket missing"
    while True:
        if os.path.exists(QMP_PATH):
            try:
                return mod.Qmp(QMP_PATH)
            except Exception as e:
                last = e
        if deadline is not None and time.time() >= deadline:
            raise SystemExit("no QMP socket: %s" % last)
        time.sleep(2)


def vmexec(cmd, timeout=25):
    try:
        out = subprocess.run(
            ["python3", VMEXEC, "exec", "--timeout", str(timeout), cmd],
            capture_output=True, text=True, timeout=timeout + 15).stdout
    except Exception:
        return ""
    return out


def guest_vmaps():
    out = vmexec(
        "sudo -n grep -E 'nv_vm_map_pages' /proc/vmallocinfo 2>/dev/null | "
        "awk '{split($1,a,\"-\"); "
        "sz=strtonum(\"0x\" substr(a[2],3))-strtonum(\"0x\" substr(a[1],3)); "
        "print a[1], sz}'", timeout=20)
    found = []
    for line in out.splitlines():
        m = re.match(r"^(0x[0-9a-fA-F]+)\s+(\d+)$", line.strip())
        if m:
            base, size = int(m.group(1), 16), int(m.group(2))
            if size in SCRUB_SIZES:
                found.append((base, size))
    return found


def phys_read64(q, pa, n=2):
    """Read n qwords of guest physical memory through the QEMU monitor.

    HMP prints each output line as "<addr>: <values>"; only the part after
    the colon is data (the address tokens are not).
    """
    out = q.hmp("xp/%dgx 0x%x" % (n, pa))
    vals = []
    for line in out.splitlines():
        if ":" not in line:
            continue
        vals += [int(x, 16) for x in re.findall(r"0x([0-9a-fA-F]+)",
                                               line.split(":", 1)[1])]
    return vals[:n]


def ring_ok(q, ring_pa):
    """906F GPFIFO ring page: all-zero or entries whose GET names a push
    inside the channel's own GPU window (the guest scanner's shape test,
    guest_pin_stream.py _scrub_ring_ok)."""
    vals = phys_read64(q, ring_pa, n=16)
    if not vals:
        return False
    for i in range(0, len(vals), 2):
        e0 = vals[i] & 0xffffffff
        if e0 == 0:
            continue
        # Low 32 bits of the push GPU VA: the CeUtils scrubber window is
        # 0x120000000..0x120080000, the DDX channel buffer pushes reach
        # 0x1200a0000+, so accept through 0x120200000.
        if not (0x20000000 <= e0 < 0x20200000) or (e0 & 3):
            return False
    return True


def wait_shape(q, hsem):
    """A channel-buffer wait signature: the host semaphore reads a
    plausible consumed count and the finish word is still zero.

    r1749: the DDX's own channel buffer does not expose the CeUtils ring
    layout at +0x64000 (its ring reads the semaphore word 0xfff), so the
    ring test rejected exactly the page X was blocked on
    (_nv015102rm timed wait on vmap+0x6c004). The consumed value with a
    zero finish word is the wait the release pair exists to end."""
    hv = phys_read64(q, hsem, n=4)
    if len(hv) < 2:
        return False
    return 0 < hv[0] <= 0x20000 and hv[1] == 0 and hv[2] == 0 \
        and hv[3] == 0


def register(fin_pa, hsem_pa):
    req = "0x%x,0x%x" % (fin_pa, hsem_pa)
    r = subprocess.run(["sudo", "tee", PARAM], input=req + "\n",
                       capture_output=True, text=True)
    return r.returncode == 0


def clear():
    register(0, 0)


def scan(q):
    pages = []
    for base, size in guest_vmaps():
        try:
            hsem = q.gva2gpa(base + 0x6c000)
            ring = q.gva2gpa(base + 0x64000)
        except Exception:
            continue
        if not hsem or not ring:
            continue
        # r1749/r1752: the ring-shape test rejected exactly the buffers
        # the guest was blocked on, and the QMP physical reads are not
        # stable enough to shape-test a moving page. The candidate set is
        # already narrow (nv_vm_map_pages area of a channel-buffer size),
        # so register every stable candidate; the service only authors a
        # record the guest itself polls (and only while !submitted).
        #
        # r1760: in the single green cycle the strict "two identical reads"
        # gate skipped the live DDX wait page (hsem_pa=0x103789000) for the
        # WHOLE cycle -> no registration -> the Xorg DDX wedge the release
        # exists to end. Retry the read a few times (the flakiness is often
        # a single bad sample), and accept a page that reads wait-shaped
        # even when the samples disagree: consumed count in (0, 0x20000]
        # with a zero finish word is the exact signature of the wait this
        # registration releases (wait_shape below).
        stable = None
        prev = None
        last = []
        for _ in range(6):
            hv = phys_read64(q, hsem, n=2)
            if len(hv) < 2:
                continue
            last = hv
            if prev is not None and hv == prev:
                stable = hv
                break
            prev = hv
        if stable is None:
            # r1760: the live page read hsem=0x20000064/fin=0x6001 while the
            # DDX spun on it and every sample pair disagreed. The page the
            # guest waits on IS a live channel buffer; a nonzero sema word
            # at the channel-buffer sema offset is that signature. Register
            # with a loud log; the service only authors records the guest
            # polls and (for all-ones) only while no USERD submissions are
            # observed, so a false registration is bounded.
            if wait_shape(q, hsem) or (last and last[0]):
                log("register vmap %#x (unstable but live: last=%s) "
                    "hsem_pa=%#x" % (base, [hex(v) for v in last], hsem))
            else:
                log("skip vmap %#x (unstable read) hsem_pa=%#x" %
                    (base, hsem))
                continue
        pages.append((hsem + 4, hsem))
    wp = walk_page_cached(SCRUB_SEMA_VA)
    if wp:
        pages.append((wp + 4, wp))
    return pages


# r1762: the DDX scrub wait polls the CeUtils scrubber channel's sema/fin
# page at the GPU VA 0x12006c000 (+4 = fin, the mem_scrub finish word), which
# the service's kernel-VAS walk resolves (walk_probe) - NOT the guest-vmap
# candidate pages this scanner used to register (0x128013000 was already
# all-ones and the guest still spun; registering the walk-resolved
# 0x12801b000 released the wait in ~23 ms).  Register both: the vmap path
# stays the fallback for boots where the walk is stale (r1712).
SCRUB_SEMA_VA = 0x12006c000
WALK_PARAM = "/sys/module/vgsp/parameters/walk_probe"
_walk_cache = {"va": 0, "pa": 0, "at": 0.0}


def host_walk_page(va):
    """Resolve a GPU VA to a guest PA through the service's kernel-VAS walk.

    walk_probe is a module parameter; the result lands in dmesg as the last
    'sema-walk-probe: ... discovered rc=0 pa=0x...' block."""
    try:
        r = subprocess.run(["sudo", "tee", WALK_PARAM],
                           input="0x%x\n" % va,
                           capture_output=True, text=True)
        if r.returncode != 0:
            return 0
        out = subprocess.run(["sudo", "dmesg"], capture_output=True,
                             text=True).stdout
    except Exception:
        return 0
    pa = 0
    kern_pa = 0
    vas_h = None
    for line in out.splitlines():
        if "sema-walk-probe: va=0x" in line:
            pa = 0          # a newer block begins; only its result counts
            kern_pa = 0
            vas_h = None
        elif "sema-walk-probe: vas[" in line:
            try:
                vas_h = int(line.split("h=0x", 1)[1].split()[0], 16)
            except Exception:
                vas_h = None
        elif "sema-walk-probe:" in line and "rc=0 pa=0x" in line:
            try:
                p = int(line.rsplit("pa=0x", 1)[1].split()[0], 16)
            except Exception:
                p = 0
            if "discovered" in line:
                pa = p
            elif vas_h == 0xc and p:
                # r1764: the summary 'discovered' path can fail (rc=-2)
                # while the per-VAS kernel row (h=0xc) still resolves the
                # polled scrubber sema page (r1667-authoritative).
                kern_pa = p
    return pa or kern_pa


def walk_page_cached(va):
    now = time.time()
    if _walk_cache["va"] == va and now - _walk_cache["at"] < 15:
        return _walk_cache["pa"]
    pa = host_walk_page(va)
    if pa != _walk_cache["pa"] or not _walk_cache["at"]:
        log("scrubber walk page va=%#x -> pa=%#x (fin=%#x)" %
            (va, pa, pa + 4))
    _walk_cache.update({"va": va, "pa": pa, "at": now})
    return pa


def module_generation():
    """Identity of the current vgsp module + QEMU instance.

    scrub_pin_qmp is long-lived: cycle.sh rmmod/insmods vgsp and relaunches
    QEMU between rounds. Both reset the registration state the scanner
    tracks implicitly, so a changed generation must force a full re-register
    even when the guest's vmap addresses happen to repeat."""
    mod = "?"
    try:
        with open("/proc/modules") as f:
            for line in f:
                if line.startswith("vgsp "):
                    f = line.split()
                    mod = f[0] + "@" + (f[5] if len(f) > 5 else "?")
                    break
    except Exception:
        pass
    pid = "?"
    try:
        with open("/tmp/qemu-harness.pid") as f:
            pid = f.read().strip()
    except Exception:
        pass
    return (mod, pid)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--duration", type=float, default=0, help="seconds (0 = forever)")
    ap.add_argument("--interval", type=float, default=2.0)
    args = ap.parse_args()

    q = load_qmp()
    cur = None
    gen = None
    last_refresh = 0.0
    start = time.time()
    while True:
        if args.duration and time.time() - start >= args.duration:
            break
        g = module_generation()
        if g != gen:
            log("generation %s -> %s, re-registering" % (gen, g))
            gen = g
            cur = None
        try:
            pages = scan(q)
        except Exception as e:
            log("scan failed: %s" % e)
            pages = []
        key = tuple(sorted(pages))
        now = time.time()
        if key != cur:
            if cur is not None and not key:
                # r1743: a transiently-empty scan (the live buffer was only
                # visible through a rejected alias for a pass) used to
                # clear() and the release stopped, wedging the DDX wait the
                # registration exists to release. Keep the live records;
                # only a *different non-empty* set replaces them.
                log("scan empty; keeping %d registration(s)" % len(cur))
            else:
                if cur is not None:
                    clear()
                for fin, hsem in pages:
                    if register(fin, hsem):
                        log("registered fin=%#x hsem=%#x" % (fin, hsem))
                    else:
                        log("register failed fin=%#x hsem=%#x" % (fin, hsem))
                cur = key
                last_refresh = now
        elif key and now - last_refresh > 30:
            # Additive re-assert: the service reaps records on its own
            # (bounded registry); re-registering is idempotent because
            # records are keyed by hsem_pa.
            for fin, hsem in pages:
                register(fin, hsem)
            last_refresh = now
        time.sleep(args.interval)
    return 0


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