#!/usr/bin/env python3
"""
DF-5229 arithmetic harness: radeon CIK UVD v4.2 VCPU-cache window chain vs
vcpu_bo size.

Kernel sources modeled (DragonFly sys/dev/drm/radeon/):

  bo allocation -- radeon_uvd.c:184-186 (shared by every uvd variant):
      bo_size = RADEON_GPU_PAGE_ALIGN(uvd_fw->datasize + 8)
              + RADEON_UVD_STACK_SIZE + RADEON_UVD_HEAP_SIZE
              + RADEON_UVD_SESSION_SIZE * uvd.max_handles

  window programming -- uvd_v4_2.c:38-64, reached from cik.c:8175
  (cik_uvd_start) and kv_dpm.c:1628 (kv_dpm_powergate_uvd ungate):
      region0: offset = gpu_addr + 0x200 (iff fw_header_present), size =
               RADEON_GPU_PAGE_ALIGN(datasize + 4)          uvd_v4_2.c:46-53
      region1: size = RADEON_UVD_HEAP_SIZE                  uvd_v4_2.c:56-58
      region2: size = STACK + SESSION_SIZE*max_handles      uvd_v4_2.c:61-64
      (regions contiguous: each offset = previous offset + previous size)
      span = (0x200 if header) + ALIGN(datasize+4) + HEAP + STACK + SESSION*N

  fw_header_present is set only for CIK ASICs when the new-format fw
  "radeonkmsfw_bonaire_uvd" loads (radeon_uvd.c:124-131,153); that fw
  advertises ucode_version 1.64 >= 1.37 -> max_handles = RADEON_MAX_UVD_HANDLES
  = 30 (radeon_uvd.c:165-166, radeon.h:1673-1674).

  The shipping blob is decoded from the in-tree
  sys/contrib/dev/drm/radeonfw/bonaire_uvd.bin.uu (see fw_header.txt):
      datasize = 232752 (0x38d30), datasize % 4096 = 3376
      -> ALIGN(ds+4) == ALIGN(ds+8) == 233472 -> overhang exactly +512.

Constants from radeon.h:1675-1677 (STACK 200KiB, HEAP 256KiB, SESSION 50KiB),
GPU page 4096 (radeon.h:643-646).

Control: headerless path (uvd_v4_2.c:49, size0 = ALIGN(ds+4) at +0) fits for
every residue -- identical to the accepted-clean uvd_v2_2 shape
(uvd_v2_2.c:113-126), and to DF-4832's control.
"""
import subprocess, sys

GPU_PAGE = 4096
def align(x):
    return (x + GPU_PAGE - 1) & ~(GPU_PAGE - 1)

STACK   = 200 * 1024          # RADEON_UVD_STACK_SIZE      radeon.h:1675
HEAP    = 256 * 1024          # RADEON_UVD_HEAP_SIZE       radeon.h:1676
SESSION = 50 * 1024           # RADEON_UVD_SESSION_SIZE    radeon.h:1677
N_DEFAULT = 10                # RADEON_DEFAULT_UVD_HANDLES radeon.h:1673
N_MAX     = 30                # RADEON_MAX_UVD_HANDLES     radeon.h:1674

def bo_size(datasize, n):
    """radeon_uvd.c:184-186"""
    return align(datasize + 8) + STACK + HEAP + SESSION * n

def v4_2_span(datasize, header, n):
    """uvd_v4_2.c:46-64"""
    skip = 0x200 if header else 0
    return skip + align(datasize + 4) + HEAP + STACK + SESSION * n

def v4_2_span_fixed(datasize, header, n):
    """fix.diff: shrink region0 so offset+size == the aligned fw region"""
    skip = 0x200 if header else 0
    size0 = align(datasize + 8) - skip if header else align(datasize + 4)
    return skip + size0 + HEAP + STACK + SESSION * n

def main():
    print("== exhaustive ds in [1, 65536] (full mod-4096 cycle, header fw, N=30) ==")
    vuln = 0; safe = []
    mn = mx = None
    for ds in range(1, 65537):
        over = v4_2_span(ds, True, N_MAX) - bo_size(ds, N_MAX)
        if over > 0:
            vuln += 1
            mn = over if mn is None else min(mn, over)
            mx = over if mx is None else max(mx, over)
        else:
            safe.append(ds)
    print(f"  header path: overrun bytes: min={mn} max={mx} ; "
          f"vulnerable sizes: {vuln}/65536 ({100.0*vuln/65536:.2f}%)")
    rsafe = sorted({s % 4096 for s in safe})
    print(f"  safe residues (ds%4096): {rsafe}   "
          f"(expect [4089..4092]: {rsafe == [4089,4090,4091,4092]})")

    print("== control: headerless path (uvd_v4_2.c:49 +0 / ALIGN(ds+4)) ==")
    c = sum(1 for ds in range(1, 65537)
            if v4_2_span(ds, False, N_MAX) > bo_size(ds, N_MAX))
    print(f"  overrun sizes: {c}/65536 (expect 0)")

    print("== shipping firmware (decoded from in-tree bonaire_uvd.bin.uu) ==")
    ds = 232752
    for n, tag in ((N_MAX, "v1.64 fw -> N=30"), (N_DEFAULT, "N=10")):
        bo = bo_size(ds, n); sp = v4_2_span(ds, True, n)
        print(f"  datasize={ds} ({tag}): bo_size={bo} span={sp} "
              f"overrun={sp - bo} bytes")
    print(f"  last window tail = gpu_addr + {hex(v4_2_span(ds, True, N_MAX))}"
          f" vs bo end gpu_addr + {hex(bo_size(ds, N_MAX))}")

    print("== N-independence (SESSION*N cancels) ==")
    allsame = all((v4_2_span(ds, True, a) - bo_size(ds, a)) ==
                  (v4_2_span(ds, True, b) - bo_size(ds, b))
                  for ds in (232752, 232240, 1, 4096, 12345)
                  for a, b in ((1, 30), (10, 40)))
    print(f"  overhang independent of max_handles: {allsame}")

    print("== fix validation: span == bo_size for every ds in [0, 0x2FFFF] ==")
    ok = all(v4_2_span_fixed(ds, True, n) == bo_size(ds, n)
             for ds in range(0, 0x30000, 7)
             for n in (N_DEFAULT, N_MAX))
    ok_ctrl = all(v4_2_span_fixed(ds, False, n) <= bo_size(ds, n)
                  for ds in range(0, 0x30000, 7)
                  for n in (N_DEFAULT, N_MAX))
    print(f"  post-fix span == bo_size (header path): {ok}")
    print(f"  post-fix span <= bo_size (headerless path): {ok_ctrl}")
    ds = 232752
    print(f"  shipping fw post-fix: span={v4_2_span_fixed(ds, True, N_MAX)} "
          f"== bo_size={bo_size(ds, N_MAX)}: "
          f"{v4_2_span_fixed(ds, True, N_MAX) == bo_size(ds, N_MAX)}")

if __name__ == "__main__":
    main()