Fix run_hook() missing several pickle entry points (GHSA-wccx-j62j-r448)

run_hook() only hooked pickle.load and pickle.Unpickler, leaving
pickle.loads, _pickle.load, and _pickle.loads unprotected.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: thomas-chauchefoin-tob

README.md

@@ -87,7 +87,7 @@ raises an `UnsafeFileError` exception if malicious content is detected in the fi
 #### Option 1 (recommended): check safety of all pickle files loaded
 
 ```python
-# This enforces safety checks every time pickle.load() is used
+# This enforces safety checks every time pickle is used to deserialize
 fickling.always_check_safety()
 
 # Attempt to load an unsafe file now raises an exception

fickling/hook.py

@@ -31,8 +31,11 @@ def load(self):
 
 def run_hook():
     """Replace pickle.load() and pickle.Unpickler by fickling's safe versions"""
-    # Hook the function
+    # Hook functions
     pickle.load = loader.load
+    _pickle.load = loader.load
+    pickle.loads = loader.loads
+    _pickle.loads = loader.loads
 
     # Hook the Unpickler class
     pickle.Unpickler = FicklingSafetyUnpickler

fickling/loader.py

@@ -1,6 +1,11 @@
 import pickle
 from io import BytesIO
 
+# Save the original pickle.loads before any hooks can replace it.
+# loader.load() must use the real pickle.loads for final deserialization,
+# otherwise hooking pickle.loads causes infinite recursion.
+_original_pickle_loads = pickle.loads
+
 from fickling.analysis import Severity, check_safety
 from fickling.exception import UnsafeFileError
 from fickling.fickle import Pickled
@@ -21,7 +26,7 @@ def load(
         # We don't do pickle.load(file) because it could allow for a race
         # condition where the file we check is not the same that gets
         # loaded after the analysis.
-        return pickle.loads(pickled_data.dumps(), *args, **kwargs)
+        return _original_pickle_loads(pickled_data.dumps(), *args, **kwargs)
     if pickled_data.has_invalid_opcode:
         raise UnsafeFileError(
             file,

test/test_hook.py

@@ -1,3 +1,4 @@
+import _pickle
 import io
 import pickle
 import unittest
@@ -25,6 +26,9 @@ def setUp(self):
         # Set up global fickling hook
         hook.run_hook()
 
+    def tearDown(self):
+        hook.remove_hook()
+
     def test_safe_pickle(self):
         # Fickling can check a pickle file for safety prior to running it
         test_list = [1, 2, 3]
@@ -44,3 +48,19 @@ def test_unsafe_pickle(self):
                 pass
             else:
                 self.fail(e)
+
+    # https://github.com/trailofbits/fickling/security/advisories/GHSA-wccx-j62j-r448
+    def test_run_hook_covers_all_entry_points(self):
+        data = pickle.dumps(Payload())
+        cases = {
+            "pickle.load": lambda: pickle.load(io.BytesIO(data)),
+            "pickle.loads": lambda: pickle.loads(data),
+            "pickle.Unpickler": lambda: pickle.Unpickler(io.BytesIO(data)).load(),
+            "_pickle.load": lambda: _pickle.load(io.BytesIO(data)),
+            "_pickle.loads": lambda: _pickle.loads(data),
+            "_pickle.Unpickler": lambda: _pickle.Unpickler(io.BytesIO(data)).load(),
+        }
+        for name, call in cases.items():
+            with self.subTest(entry_point=name):
+                with self.assertRaises(UnsafeFileError, msg=f"{name} was not intercepted"):
+                    call()