Implement Secure Tensor Processing and Parameter Conversion

[backspace-agent]

Summary

This PR implements comprehensive security measures for tensor processing and parameter conversion in flaxvision/utils.py to address critical security vulnerabilities that could lead to buffer overflows, memory corruption, and malicious tensor attacks.

Key Security Enhancements

• Tensor Validation: Added validate_tensor_security() function with comprehensive checks:

  • NaN/Inf detection to prevent malicious tensor attacks
  • Memory size limits to prevent buffer overflows (default 1GB limit)
  • Dimension validation to prevent excessive memory usage (max 6 dimensions)
  • Data type validation to ensure only numeric types are processed

• Parameter Structure Validation: Added validate_parameter_structure() function with:

  • Parameter count limits to prevent resource exhaustion attacks
  • Key type and length validation to prevent malicious parameter injection
  • Recursive validation for nested parameter structures

• Memory Monitoring: Added monitor_memory_usage() function with:

  • Real-time memory usage tracking during conversion
  • Warnings for excessive memory consumption
  • Graceful fallback when psutil is not available

• Secure Conversion Functions: Enhanced both torch_to_flax() and torch_to_linen() with:

  • Input validation before processing
  • Tensor security validation for each parameter
  • Memory monitoring during conversion
  • Output validation before returning results

Security Issues Addressed

1. Buffer Overflows: Prevents processing of oversized tensors that could cause memory corruption
2. Malicious Tensor Attacks: Detects and blocks tensors with NaN/Inf values
3. Resource Exhaustion: Limits parameter count and tensor dimensions to prevent DoS attacks
4. Memory Corruption: Validates data types and tensor structure before processing

Testing

• Created comprehensive test suite (tests/test_tensor_security.py) covering:

  • Malicious tensor inputs (NaN, Inf values)
  • Oversized tensor detection
  • Invalid parameter structures
  • Memory exhaustion scenarios
  • Edge cases and error handling

• Added validation scripts to verify implementation:

  • validate_security_direct.py - Direct testing of security functions
  • test_security_direct.py - Comprehensive security test suite

Implementation Details

All changes are defensive security measures that:

• Do not break existing functionality - All existing valid tensors continue to work
• Provide clear error messages - Security violations are reported with specific details
• Include graceful fallbacks - Memory monitoring works even without psutil
• Are thoroughly tested - Comprehensive test coverage for all security scenarios

Files Modified

• flaxvision/utils.py - Core security implementation
• tests/test_tensor_security.py - Comprehensive test suite
• tests/run_tests.sh - Made executable
• Various validation scripts for testing

Verification

The implementation has been thoroughly tested and verified to:

• ✅ Block malicious tensor inputs (NaN, Inf, oversized)
• ✅ Validate parameter structures and prevent injection attacks
• ✅ Monitor memory usage during conversion
• ✅ Maintain backward compatibility with existing code
• ✅ Provide clear error messages for security violations

This implementation addresses all security requirements from the original issue and provides robust protection against tensor-based attacks while maintaining full compatibility with existing functionality.