From 5c8d409a59d42b404349f4da9854cbbee1f6bf11 Mon Sep 17 00:00:00 2001 From: tomaioo Date: Tue, 7 Jul 2026 23:10:32 -0700 Subject: [PATCH] fix(security): unsafe dynamic module loading in patch auto-discov The `_auto_import_patch_modules` function in `primus/backends/maxtext/patches/__init__.py` automatically imports all modules ending with `_patches` or `_patch` from the package directory. This is a form of dynamic module loading based on file system contents. If an attacker can write a file matching this pattern to the package directory, they can achieve arbitrary code execution when the package is imported. The function uses `importlib.import_module()` which will execute any code in the imported module's top-level scope. Signed-off-by: tomaioo <203048277+tomaioo@users.noreply.github.com> --- primus/backends/maxtext/patches/__init__.py | 29 ++++++++++----------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/primus/backends/maxtext/patches/__init__.py b/primus/backends/maxtext/patches/__init__.py index 3df160f11..99a007a9e 100644 --- a/primus/backends/maxtext/patches/__init__.py +++ b/primus/backends/maxtext/patches/__init__.py @@ -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)