Skip to content

Security: Unsafe Dynamic Module Loading in Patch Auto-Discovery #863

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tomaioo wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Security: Unsafe Dynamic Module Loading in Patch Auto-Discovery #863

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 14 additions & 15 deletions primus/backends/maxtext/patches/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,28 +14,27 @@
This follows the same convention used by ``primus.backends.megatron.patches``.
"""

import importlib
import pkgutil
_ALLOWED_PATCH_MODULES = [
"primus.backends.maxtext.patches.checkpoint_patches",
"primus.backends.maxtext.patches.config_patches",
"primus.backends.maxtext.patches.data_patches",
"primus.backends.maxtext.patches.model_patches",
"primus.backends.maxtext.patches.optimizer_patches",
"primus.backends.maxtext.patches.tokenizer_patches",
]


def _auto_import_patch_modules() -> None:
"""
Automatically import all patch modules under this package.
Import explicitly allowed patch modules under this package.

Any module whose fully-qualified name ends with ``"_patches"`` will be
imported, which in turn triggers its ``@register_patch`` side effects.

This allows adding new ``*_patches.py`` files without having to update
this ``__init__`` module.
Only whitelisted modules are imported, which triggers their
``@register_patch`` side effects. This prevents arbitrary code
execution from untrusted files in the package directory.
"""
package_name = __name__

for module_info in pkgutil.walk_packages(__path__, prefix=package_name + "."):
mod_name = module_info.name

if not (mod_name.endswith("_patches") or mod_name.endswith("_patch")):
continue
import importlib

for mod_name in _ALLOWED_PATCH_MODULES:
importlib.import_module(mod_name)


Expand Down