diff --git a/bandit/plugins/pytorch_load.py b/bandit/plugins/pytorch_load.py index 667cbb0d1..9946c3979 100644 --- a/bandit/plugins/pytorch_load.py +++ b/bandit/plugins/pytorch_load.py @@ -19,6 +19,13 @@ With `weights_only=True`, PyTorch enforces a strict type check, ensuring that only torch.Tensor objects are loaded. +Note: ``torch.jit.load`` uses TorchScript serialization (not pickle) and +is **not** flagged by this check regardless of arguments. + +Note: When ``weights_only`` is supplied as a variable or non-literal +expression, the check cannot statically resolve its value and will not +raise a finding (benefit of the doubt, consistent with B615 behaviour). + :Example: .. code-block:: none @@ -42,6 +49,8 @@ .. versionadded:: 1.7.10 """ +import ast + import bandit from bandit.core import issue from bandit.core import test_properties as test @@ -56,6 +65,16 @@ def pytorch_load(context): `torch.serialization.load` with untrusted data can lead to arbitrary code execution. The safe alternative is to use `weights_only=True` or the safetensors library. + + ``torch.jit.load`` is excluded because it uses TorchScript + serialization, not pickle, and therefore does not deserialize + arbitrary Python objects. + + When ``weights_only`` is a non-literal expression (variable, + attribute, subscript, etc.) the value cannot be resolved at + static-analysis time; the check is suppressed to avoid false + positives, consistent with the treatment of similar arguments in + other B6xx checks. """ imported = context.is_module_imported_exact("torch") qualname = context.call_function_name_qual @@ -63,6 +82,18 @@ def pytorch_load(context): return if qualname in {"torch.load", "torch.serialization.load"}: + # If weights_only is a non-literal expression (Name, Attribute, + # Subscript, Call, …) we cannot resolve its value statically. + # Give the caller the benefit of the doubt rather than raising a + # false positive. + call_node = context._context.get("call") + if call_node is not None: + for kw in getattr(call_node, "keywords", []): + if kw.arg == "weights_only" and not isinstance( + kw.value, ast.Constant + ): + return + # For torch.load, check if weights_only=True is specified weights_only = context.get_call_arg_value("weights_only") if weights_only == "True" or weights_only is True: diff --git a/examples/pytorch_load.py b/examples/pytorch_load.py index 39b05fab5..dbafbafa6 100644 --- a/examples/pytorch_load.py +++ b/examples/pytorch_load.py @@ -28,3 +28,16 @@ # Example of a torch.*.load call that should NOT trigger B614 # Only pickle deserializers should trigger B614 torch.utils.cpp_extension.load(name="example_ext", sources=[]) + +# torch.jit.load uses TorchScript serialization, not pickle. +# It has no weights_only parameter and should NOT trigger B614. +jit_model = torch.jit.load('script_model.pt') + +# torch.jit.load with map_location is still TorchScript; should NOT +# trigger B614. +jit_model_cpu = torch.jit.load('script_model.pt', map_location='cpu') + +# weights_only passed as a variable (non-literal). +# B614 cannot resolve the value statically; should NOT trigger B614. +_weights_flag = True +model_var = torch.load('model_weights.pth', weights_only=_weights_flag)