gistfile1.txt

"""
Demonstrate the NEON overread in F.interpolate by placing uint8 tensor data
right before an unmapped guard page.

vld3_u8 in the block-of-4 loop reads 24 bytes (8 pixels × 3 channels)
but only needs 12 bytes (4 pixels × 3 channels). If the extra 12 bytes
cross into an unmapped page, we get SIGBUS.
"""

import ctypes
import ctypes.util
import os
import signal
import sys

import numpy as np
import torch
import torch.nn.functional as F


def run_interpolate_on_guard_page(w_in, w_out, h_in=1):
    """
    Allocates tensor data ending exactly at a page boundary with a guard
    page after it. Returns True if interpolate crashes (overread detected).
    """
    page_size = os.sysconf("SC_PAGE_SIZE")
    num_channels = 3
    tensor_bytes = h_in * w_in * num_channels

    libc = ctypes.CDLL(ctypes.util.find_library("c"))
    libc.mmap.restype = ctypes.c_void_p
    libc.mmap.argtypes = [
        ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int,
        ctypes.c_int, ctypes.c_int, ctypes.c_longlong,
    ]
    libc.mprotect.restype = ctypes.c_int
    libc.mprotect.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
    libc.munmap.restype = ctypes.c_int
    libc.munmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t]

    num_data_pages = (tensor_bytes + page_size - 1) // page_size
    total_pages = num_data_pages + 1

    addr = libc.mmap(0, total_pages * page_size, 0x03, 0x1002, -1, 0)
    if addr in (0, 2**64 - 1, -1):
        return None

    guard_start = addr + num_data_pages * page_size
    libc.mprotect(guard_start, page_size, 0x00)  # PROT_NONE

    # Data ends exactly at guard page
    data_start = guard_start - tensor_bytes

    # Create numpy array pointing at mmap'd memory, zero-copy
    ArrayType = ctypes.c_uint8 * tensor_bytes
    c_arr = ArrayType.from_address(data_start)
    np_arr = np.frombuffer(c_arr, dtype=np.uint8)
    np_arr[:] = 128  # fill with data

    # torch.from_numpy shares memory (zero-copy)
    flat = torch.from_numpy(np_arr)

    # Build channels-last tensor: (1, 3, h_in, w_in)
    # channels-last strides: (H*W*C, 1, W*C, C)
    t = flat.as_strided(
        size=[1, num_channels, h_in, w_in],
        stride=[h_in * w_in * num_channels, 1, w_in * num_channels, num_channels],
    )

    pid = os.fork()
    if pid == 0:
        # Child: run interpolate, crash if overread hits guard page
        try:
            out = F.interpolate(
                t, size=(h_in, w_out),
                mode='bilinear', antialias=True, align_corners=False,
            )
            os._exit(0)
        except Exception as e:
            print(f"    exception: {e}", file=sys.stderr)
            os._exit(1)
    else:
        _, status = os.waitpid(pid, 0)
        libc.munmap(addr, total_pages * page_size)

        if os.WIFSIGNALED(status):
            return os.WTERMSIG(status)
        return 0


def main():
    print(f"Page size: {os.sysconf('SC_PAGE_SIZE')}")
    print("Testing NEON overread with guard pages...\n")

    crashes = 0
    tests = 0

    for w_in in range(4, 60):
        for w_out in [1, 2, 3]:
            tests += 1
            result = run_interpolate_on_guard_page(w_in, w_out)
            if result is None:
                print(f"  SKIP: w_in={w_in} w_out={w_out} (mmap failed)")
            elif result > 0:
                sig_name = signal.Signals(result).name
                print(f"  CRASH: w_in={w_in}, w_out={w_out} → {sig_name}")
                crashes += 1

    print(f"\n{crashes}/{tests} configurations crashed.")
    if crashes > 0:
        print("Confirmed: NEON vld3_u8 reads past the end of the buffer!")


if __name__ == "__main__":
    main()