Skip to content

Fix FicklingContextManager incomplete cleanup and dead code (#241)#243

Open
mldangelo wants to merge 4 commits intotrailofbits:masterfrom
mldangelo:fix/context-manager-241
Open

Fix FicklingContextManager incomplete cleanup and dead code (#241)#243
mldangelo wants to merge 4 commits intotrailofbits:masterfrom
mldangelo:fix/context-manager-241

Conversation

@mldangelo
Copy link
Copy Markdown

@mldangelo mldangelo commented Mar 4, 2026

FicklingContextManager had three issues noted in #241:

  1. __exit__ only restored pickle.load, leaving the other five entry points (_pickle.load, pickle.loads, _pickle.loads, pickle.Unpickler, _pickle.Unpickler) permanently hooked after the context exited.

  2. wrapped_load in __enter__ was dead code — hook.run_hook() immediately overwrote it, so max_acceptable_severity never actually took effect.

  3. Entering a context while a pre-existing hook (from run_hook() or activate_safe_ml_environment()) was active would clobber that hook on exit.

This PR fixes all three by:

  • Snapshotting the current hook state on __enter__ and restoring it on __exit__ (via a stack for re-entrant safety), so pre-existing hooks are preserved.
  • Adding a max_acceptable_severity parameter to run_hook() in hook.py, which creates severity-aware wrappers when the threshold differs from the default. This includes the Unpickler path (via a subclass), following the same pattern as activate_safe_ml_environment().
  • Simplifying context.py to a thin wrapper that delegates entirely to hook.run_hook() / hook.restore_hooks().

Also includes duplicate-keyword protection (kwargs.pop) so callers that pass max_acceptable_severity explicitly don't get a TypeError from the wrappers.

Tests cover hook lifecycle, exception safety, severity enforcement, pre-existing hook preservation, duplicate-keyword handling, and re-entrant/nested usage.

Closes #241

…its#241)

- Snapshot/restore pre-enter hook state on exit (via state stack),
  preserving any pre-existing hooks and supporting re-entrant usage
- Remove dead wrapped_load lambda that was immediately overwritten
- Add max_acceptable_severity parameter to run_hook() with duplicate-
  kwarg protection on all paths including Unpickler
- Simplify context.py to a thin wrapper over hook.run_hook()
Cover hook lifecycle, exception safety, max_acceptable_severity
enforcement, pre-existing hook preservation, duplicate-keyword
safety on all paths, re-entrant/nested usage, and check_safety().
- Replace eval() with getattr-based _get_entry_point() helper
- Move kwargs.pop to SafetyUnpicklerWithSeverity.__init__ to avoid
  per-load() dict copy
@mldangelo mldangelo requested a review from ESultanik as a code owner March 4, 2026 22:24
@mldangelo
Copy link
Copy Markdown
Author

@thomas-chauchefoin-tob can you please take a look at this PR too? Thanks!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

FicklingContextManager improvements

1 participant