#!/usr/bin/env python3
"""
DF-4832 arithmetic harness: radeon UVD v1.0 VCPU-cache window vs vcpu_bo size.

Kernel sources modeled:
  bo allocation (radeon_uvd.c:184-186, shared by every uvd_vX_0 consumer):
      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_v1_0.c:121-135, dispatched from r600.c:3080 for the
  whole R600 class and from uvd_v2_2.c:106 for RV770):
      region0: offset = gpu_addr + 16*8 (+128B), size = ALIGN(datasize)
      region1:                    size = RADEON_UVD_HEAP_SIZE
      region2:                    size = STACK + SESSION_SIZE*max_handles
      span     = 128 + ALIGN(datasize) + HEAP + STACK + SESSION*N

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

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

STACK  = 200 * 1024          # RADEON_UVD_STACK_SIZE
HEAP   = 256 * 1024          # RADEON_UVD_HEAP_SIZE
SESSION = 50 * 1024          # RADEON_UVD_SESSION_SIZE
NDEF = 10                    # RADEON_DEFAULT_UVD_HANDLES (legacy fw: header
                            # absent -> max_handles stays default, radeon_uvd.c:138)

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

def v1_span(datasize, n=NDEF):
    # uvd_v1_0.c:121-135 -- +16 window-offset units of 8 bytes = +128 bytes
    return 128 + align(datasize) + HEAP + STACK + SESSION * n

def v2_2_span(datasize, n=NDEF):
    # uvd_v2_2.c:113-126 (non-RV770 path) -- starts at +0, sizes ALIGN(ds+4)
    return align(datasize + 4) + HEAP + STACK + SESSION * n

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

def sweep():
    print("datasize ds%%4096 | bo_size      v1_span       overrun | v2_2 (ctrl)")
    overruns = 0
    for ds in list(range(1, 0x2000)) + [0x1000 * k + r
                                        for k in (1, 7, 0x2A) for r in (0, 1, 128, 4088, 4089, 4095)]:
        bo, sp = bo_size(ds), v1_span(ds)
        if sp > bo:
            overruns += 1
    # exhaustive class check over one full modulo cycle + page boundary set
    bad = [r for r in range(4096) if v1_span(r) > bo_size(r)]
    print("exhaustive ds in [1,4095] (one full mod-4096 cycle):")
    print("  overrun bytes: min=%d max=%d ; vulnerable residues: %d/4096 (%.1f%%)"
          % (min(v1_span(r) - bo_size(r) for r in bad),
             max(v1_span(r) - bo_size(r) for r in bad),
             len(bad), 100.0 * len(bad) / 4096))
    print("  vulnerable residues are exactly ds %% 4096 in [1,4088]: %s"
          % (bad == list(range(1, 4089))))
    ok = [r for r in range(4096) if v1_span(r) <= bo_size(r)]
    print("  safe residues (ds %% 4096 == 0 or in [4089,4095]): %s"
          % (ok == [0] + list(range(4089, 4096))))

    print()
    print("representative datasizes (real legacy R600/RV770 UVD blobs are")
    print("arbitrary-sized; linux-firmware R600_uvd.bin-class sizes shown):")
    for ds in (0x7DC9, 0x9A44, 0xC000, 0xC080, 0x1B040):
        bo, sp = bo_size(ds), v1_span(ds)
        print("  ds=0x%-6x (%%4096=%4d): bo=%7d span=%7d -> %+4d B %s"
              % (ds, ds % 4096, bo, sp, sp - bo,
                 "OVERRUN" if sp > bo else "fits"))

    print()
    print("control (uvd_v2_2 non-RV770): span<=bo for ALL residues:",
          all(v2_2_span(r) <= bo_size(r) for r in range(4096)))
    print("cross-note (uvd_v4_2 w/ header, out of this file's scope): "
          "overrun residues exist:",
          any(v4_2_span(r, True) > bo_size(r) for r in range(1, 4085)))

    # page-boundary demonstration of region2 tail poking out of the bo
    ds = 0x9A44  # 39492, 39492 % 4096 = 2956 -> vulnerable class
    bo = bo_size(ds)
    span = v1_span(ds)
    print()
    print("worked example ds=0x9A44 (ds%%4096=%d in [1,4088]):" % (ds % 4096))
    print("  vcpu_bo            = [gpu,             gpu+%-7d)" % bo)
    print("  VCPU CACHE windows = [gpu+128,         gpu+%-7d)" % span)
    print("  -> last %d bytes of the STACK/SESSION window lie past the end of"
          % (span - bo))
    print("     vcpu_bo, i.e. inside the next VRAM object allocated after it.")

if __name__ == "__main__":
    sweep()