Skip to content

Mitigate pickle deserialization RCE risk in replay buffer checkpoints#237

Open
vijesh-cloud wants to merge 2 commits into
google:masterfrom
vijesh-cloud:fix/pickle-rce-vulnerability
Open

Mitigate pickle deserialization RCE risk in replay buffer checkpoints#237
vijesh-cloud wants to merge 2 commits into
google:masterfrom
vijesh-cloud:fix/pickle-rce-vulnerability

Conversation

@vijesh-cloud

Copy link
Copy Markdown

Summary

This PR mitigates a potential remote code execution (RCE) risk arising from the use of Python's pickle module for replay buffer checkpoint serialization and deserialization.

Changes

  • Replaced pickle-based checkpoint serialization/deserialization with msgpack in replay buffer implementations.
  • Added local path validation to prevent loading checkpoints from remote or URI-based locations.
  • Preserved support for complex data structures through explicit type markers during serialization.
  • Added regression tests to verify that unsafe checkpoint paths are rejected.

Files Updated

  • circular_replay_buffer.py
  • subsequence_replay_buffer.py
  • pickle_collector.py
  • checkpointer.py
  • logger.py
  • Colab utility modules
  • Associated test files

Validation

  • Existing checkpoint save/load tests continue to pass.
  • Local filesystem checkpoint operations remain functional.
  • Remote paths such as gs://, s3://, UNC paths, and other URI-based locations are rejected before file I/O.
  • Serialization round-trip tests confirm correct restoration of supported data types, including tuples and sets.

Security Impact

Python's pickle deserialization can execute arbitrary code when processing untrusted data. Replacing it with a safer serialization format and validating checkpoint paths reduces the risk of arbitrary code execution during checkpoint loading while preserving existing functionality for trusted local checkpoints.

…dd remote path validation

Replace pickle serialization with msgpack in replay buffers to eliminate pickle-based RCE vulnerabilities:
- circular_replay_buffer.py: Added _pack/_unpack helpers, _validate_local_path guard, replaced pickle.dump/load
- subsequence_replay_buffer.py: Same security fixes for labs-based buffer
- Added type preservation for tuples, sets via __dopamine_type__ markers
- Remote path validation rejects gs://, s3://, UNC, and other URI schemes
- Tests: Added testLoadRejectsRemoteCheckpointPaths regression test
- Also patched colab utils, pickle_collector, checkpointer, logger for consistency

Validated:
- All data types roundtrip correctly through msgpack serializer
- Remote paths (gs://, s3://) are rejected before file I/O
- Local filesystem paths continue to work normally
- Existing checkpoint save/load tests pass
@google-cla

google-cla Bot commented Jun 12, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

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.

1 participant