Fix PyTorch v1.3+ hook bypass by hooking pickle.Unpickler class

[dhalf]

Fix PyTorch v1.3+ Hook Bypass by Hooking pickle.Unpickler Class

Summary

This PR fixes a critical issue where fickling's hook system did not intercept pickle operations in PyTorch v1.3+ model files. The hooks worked correctly on legacy PyTorch files but failed on modern PyTorch files that use the new
zipfile serialization format.

Problem

The fickling hook system (run_hook() and always_check_safety()) only replaced pickle.load() and pickle.loads() functions but did NOT replace the pickle.Unpickler class. PyTorch v1.3+ uses pickle.Unpickler() directly
when loading models from zipfile format, completely bypassing the function hooks.

Reproduction

import fickling.hook as hook
import torch
import torchvision.models as models

hook.run_hook()

model = models.mobilenet_v2()

# Legacy format (works)
torch.save(model, "legacy.pt", _use_new_zipfile_serialization=False)
torch.load("legacy.pt")  # ✅ Triggers fickling hook

# Modern format (broken before this PR)
torch.save(model, "modern.pt")  # Default uses zipfile
torch.load("modern.pt")  # ❌ Silently bypasses fickling hook

Root Cause

Legacy PyTorch (v0.1.10):
torch.load() → pickle.load(file) → ✅ HOOKED

Modern PyTorch (v1.3+):
torch.load() → zipfile.ZipFile() → pickle.Unpickler(file) → ❌ NOT HOOKED

Solution

Hook both the pickle functions AND the pickle.Unpickler class to provide complete coverage.

Changes

Core Implementation

fickling/hook.py:
1. Added FicklingSafetyUnpickler class - Wrapper that delegates to fickling.load() for security analysis
2. Updated run_hook() - Now also hooks pickle.Unpickler and _pickle.Unpickler
3. Updated activate_safe_ml_environment() - Hooks Unpickler classes using a dynamically created SafeMLUnpickler subclass
4. Updated remove_hook() - Restores original Unpickler classes
5. Stored original Unpickler references - _original_pickle_Unpickler and _original__pickle_Unpickler

Design Decisions

- Wrapper Pattern: FicklingSafetyUnpickler doesn't inherit from pickle.Unpickler, it's a simple proxy that delegates to fickling.load()
- Subclass for ML Environment: activate_safe_ml_environment() creates a SafeMLUnpickler subclass of FicklingMLUnpickler to allow PyTorch's code that subclasses pickle.Unpickler to work correctly
- Both Python and C: Hooks both pickle.Unpickler (Python) and _pickle.Unpickler (C implementation)

Testing

Existing Tests

All existing tests pass, confirming backward compatibility:
$ pytest test/test_hook.py test/test_unpickler.py -v
============================== 6 passed in 0.53s ===============================

Manual Verification

The fix can be verified with the PoC script from the issue:
import fickling.hook as hook
import torch
import torchvision.models as models

hook.run_hook()
model = models.mobilenet_v2()

torch.save(model, "model.pt")  # Modern format
torch.save(model, "legacy.pt", _use_new_zipfile_serialization=False)

print("MODERN:")
torch.load("model.pt")  # ✅ Now triggers fickling!

print("LEGACY:")
torch.load("legacy.pt")  # ✅ Still works

Impact

Before Fix

- ❌ PyTorch v1.3+ models silently bypass security checks
- ❌ Any library using pickle.Unpickler directly bypasses hooks
- ❌ Incomplete security coverage

After Fix

- ✅ PyTorch v1.3+ models trigger security analysis
- ✅ All pickle entry points (functions + classes) are hooked
- ✅ Complete security coverage
- ✅ Backward compatible

Benefits

1. Complete Hook Coverage - Intercepts both function calls and class instantiation
2. PyTorch v1.3+ Support - Modern PyTorch models now trigger security checks
3. Universal Fix - Benefits any library that uses pickle.Unpickler directly
4. Backward Compatible - All existing tests pass, no breaking changes
5. Consistent Architecture - Follows the pattern established by FicklingMLUnpickler

Related Issues

Fixes the issue reported in: "Function hook does not work on all PyTorch inputs"

The global function hook did not work on PyTorch v1.3+ files but did work on PyTorch v0.1.10 files due to the different serialization formats.

[dhalf]

#80