sumantro93

import torch
import helion
import helion.language as hl
import time

# ============================================================
# Example 1: Matrix Multiplication
# ============================================================
@helion.kernel()
def matmul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:

sumantro93

import torch
... import torch.nn as nn
... 
... # The kernel at line 941 computes: buffer_data[t * n_channel + c]
... # If n_channel is extremely large, t * n_channel can overflow int64_t
... 
... # Demonstration of overflow math:
... int64_max = (1 << 63) - 1
... n_channel = 2**40  # Hypothetical huge channel count
... num_threads = 8

sumantro93

#!/usr/bin/env python3
"""
Repro: CWE-122 Heap-based Buffer Overflow via strcpy() in libshm/socket.h

Target: torch/lib/libshm/socket.h:39
Vulnerable Code:
    strcpy(address.sun_path, path);  // No bounds checking!

The sockaddr_un.sun_path buffer is typically 108 bytes on Linux.
If path exceeds this, strcpy() overflows into adjacent memory.

sumantro93

import torch
import gc
def get_memory_mb():
    import resource
    return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024
initial_mem = get_memory_mb()
print(f"Initial memory: {initial_mem:.2f} MB")
# Repeated operations cause memory leak
for i in range(10000):
    t = torch.randn(2, 3, 4, 5, 6)  # 5 dims

sumantro93

#!/usr/bin/env python3
"""
Exploit for SpmmReduceKernel.cpp vulnerabilities

Target File: aten/src/ATen/native/cpu/SpmmReduceKernel.cpp
Vulnerable Code Location: Lines 36, 60-62, 131, 222, 225, 276-277, 280-282, 331-332, 370

=============================================================================
CWE CLASSIFICATION (Reference: https://cwe.mitre.org/data/definitions/699.html)
=============================================================================

sumantro93

#!/usr/bin/env python3
"""
Proof of Concept for Integer Overflow in QuantizedOpKernels.cpp:3385
Target: float* buffer_ptr = buffer_data + t * Bs * 2 * C + n * 2 * C;

This vulnerability has COMPOUND MULTIPLICATIONS making it more severe than previous ones.
"""

import sys

sumantro93

#!/usr/bin/env python3
"""
Proof of Concept for Integer Overflow in kai_kernels.cpp:255 (and surrounding lines)
Target: Multiple overflow points in KleidiAI kernel code
"""

import sys

INT64_MAX = 2**63 - 1
SIZE_MAX = 2**64 - 1  # For size_t on 64-bit systems

sumantro93

#!/usr/bin/env python3
"""
Attempt to exploit batch_norm_kernel.cpp:922 through PyTorch APIs
Target: tid * n_channel overflow in batch normalization
"""

import sys
import os

print("="*70)

sumantro93

#!/usr/bin/env python3
"""
Proof of Concept for Integer Overflow in PyTorch ParallelNative.cpp:137
Demonstrates the overflow behavior in Python (simulating the C++ logic)
"""

import sys

# Python 3 has arbitrary precision integers, so we need to simulate int64_t behavior
INT64_MAX = 2**63 - 1

sumantro93

flowchart LR

A[Static Analysis & Scanning\n(Snyk + Internal Tooling)] --> B[Log Aggregation & Finding Classification]

B --> C[Jira Ticket Creation\n(security: potential vuln)]
C --> D[ProdSec Advisor Review / ACK]

D --> E[Reporter Begins PoC\n(reproduce & validate)]
E --> F[Report to Upstream PyTorch\n(Security Advisory Channels)]