77 lines
2.5 KiB
Python
77 lines
2.5 KiB
Python
|
import random
|
||
|
|
from typing import Dict, Optional
|
||
|
|
|
||
|
|
from tinygrad.runtime.support.allocator import TLSFAllocator
|
||
|
|
|
||
|
|
class AllocatorFuzzer:
|
||
|
|
def __init__(self, total_size):
|
||
|
|
self.total_size = total_size
|
||
|
|
self.alloc_payload = 0
|
||
|
|
self.mv = memoryview(bytearray(total_size))
|
||
|
|
self.alloctor = TLSFAllocator(total_size, block_size=16)
|
||
|
|
|
||
|
|
self.allocations: Dict[int, tuple[int, int]] = {} # ptr -> (size, pattern)
|
||
|
|
|
||
|
|
self.min_alloc_size = 16
|
||
|
|
self.max_alloc_size = int(total_size * 0.3)
|
||
|
|
self.alloc_probability = 0.7
|
||
|
|
|
||
|
|
def generate_pattern(self, ptr: int, size: int) -> int: return (ptr * 31 + size * 17) & 0xFF
|
||
|
|
|
||
|
|
def fill_memory(self, ptr: int, size: int, pattern: int):
|
||
|
|
for i in range(min(size, 32)):
|
||
|
|
self.mv[ptr + i] = pattern
|
||
|
|
self.mv[ptr + (size - 1 - i)] = pattern
|
||
|
|
|
||
|
|
def verify_memory(self, ptr: int, size: int, pattern: int) -> bool:
|
||
|
|
for i in range(min(size, 32)):
|
||
|
|
assert self.mv[ptr + i] == pattern
|
||
|
|
assert self.mv[ptr + (size - 1 - i)] == pattern
|
||
|
|
return True
|
||
|
|
|
||
|
|
def random_alloc(self) -> Optional[int]:
|
||
|
|
size = random.randint(self.min_alloc_size, min(self.max_alloc_size, self.total_size - self.alloc_payload))
|
||
|
|
|
||
|
|
try:
|
||
|
|
ptr = self.alloctor.alloc(size)
|
||
|
|
except MemoryError:
|
||
|
|
print(f"Failed to allocate {size} bytes. Payload size is {self.alloc_payload}, so fragmenation is {(size / self.total_size)*100.0:.2f}%")
|
||
|
|
return None
|
||
|
|
|
||
|
|
pattern = self.generate_pattern(ptr, size)
|
||
|
|
self.fill_memory(ptr, size, pattern)
|
||
|
|
self.allocations[ptr] = (size, pattern)
|
||
|
|
self.alloc_payload += size
|
||
|
|
print(f"Allocated {size} bytes at {ptr:x}, pattern: {pattern:02x}")
|
||
|
|
return ptr
|
||
|
|
|
||
|
|
def random_free(self) -> bool:
|
||
|
|
if not self.allocations: return False
|
||
|
|
|
||
|
|
ptr = random.choice(list(self.allocations.keys()))
|
||
|
|
size, pattern = self.allocations[ptr]
|
||
|
|
|
||
|
|
# Verify pattern before freeing
|
||
|
|
if not self.verify_memory(ptr, size, pattern):
|
||
|
|
raise RuntimeError(f"Memory corruption detected at {ptr:x}!")
|
||
|
|
|
||
|
|
print(f"Freeing {size} bytes at {ptr:x}, pattern verified: {pattern:02x}")
|
||
|
|
self.alloc_payload -= size
|
||
|
|
self.alloctor.free(ptr)
|
||
|
|
del self.allocations[ptr]
|
||
|
|
return True
|
||
|
|
|
||
|
|
def run(self):
|
||
|
|
for i in range(10000000):
|
||
|
|
if (random.random() < self.alloc_probability or not self.allocations): self.random_alloc()
|
||
|
|
else: self.random_free()
|
||
|
|
|
||
|
|
print("\nCleaning up remaining allocations...")
|
||
|
|
while self.allocations: self.random_free()
|
||
|
|
|
||
|
|
print("Fuzzing completed successfully!")
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
fuzzer = AllocatorFuzzer(1 << 30)
|
||
|
|
fuzzer.run()
|