#!/usr/bin/env python3
# find_push.py - locate the client's CPU mapping that owns the GR ring page
# (GPA 0x12C600000 as content-located on the host) and print the CPU VA<->
# GPU VA correspondence for the low GPU VAs (0x4ca0000 etc).
import os, struct, sys, subprocess, mmap, time

RING_GPA = 0x12C600000
PUSH_GVA = 0x4ca0000
RING_GVA = 0x4cc0000

def pids():
    out = subprocess.run(["pgrep", "-x", "vk_frame"], capture_output=True).stdout
    return [int(x) for x in out.split()]

def read_pagemap(pid, va):
    try:
        with open("/proc/%d/pagemap" % pid, "rb") as f:
            f.seek((va // 4096) * 8)
            d = f.read(8)
        if len(d) < 8:
            return None
        e = int.from_bytes(d, "little")
        if not (e >> 63) & 1:
            return None
        return (e & ((1 << 55) - 1)) << 12
    except OSError:
        return None

def main():
    for pid in pids():
        print("pid=%d" % pid)
        try:
            maps = open("/proc/%d/maps" % pid).read().splitlines()
        except OSError:
            continue
        for line in maps:
            p = line.split()
            if len(p) < 6 or "nvidia" not in p[-1]:
                continue
            lo, hi = [int(x, 16) for x in p[0].split("-")]
            if hi - lo > 0x1000000:
                continue
            # scan pages; look for RING_GPA
            for va in range(lo, hi, 0x10000):
                for off in range(0, 0x10000, 0x1000):
                    pa = read_pagemap(pid, va + off)
                    if pa == RING_GPA:
                        print("FOUND ring page: cpu_va=0x%x (map 0x%x-0x%x %s)" % (va+off, lo, hi, p[-1]))
                        delta_cpu = (va + off) - RING_GVA
                        push_cpu = PUSH_GVA + delta_cpu
                        print("  delta_cpu = 0x%x" % delta_cpu)
                        for gva in (PUSH_GVA, PUSH_GVA+0x1000, 0x4ca059c, 0x4ca0980, 0x4ca5870, 0x4903000):
                            cva = gva + delta_cpu
                            pa2 = read_pagemap(pid, cva & ~0xfff)
                            print("  gva=0x%x -> cpu=0x%x -> pa=%s" % (gva, cva, hex(pa2) if pa2 else None))
                        return
        print("  ring GPA 0x%x not found in this pid's mappings" % RING_GPA)

main()