SONARPY-1914 Usage of "torch.load" can lead to untrusted code execution (#1956)

* SONARPY-1914 update metadata for S6985

* SONARPY-1914 implement S6985: Usage of "torch.load" can lead to untrusted code execution

* SONARPY-1914 add S6985 to checklist

* SONARPY-1914 add expected ruling issues

* SONARPY-1914 implement to not to raise when weights_only is not false

* SONARPY-1914 fix small issues of the PR

* SONARPY-1914 Find variable definition and check if False

When a variable is passed as the weights_only parameter, the `ReachingDefinitionsAnalysis`
class is used to check if the variable is set to true

* SONARPY-1914 add unit tests for TreeUtils.toStreamInstanceOfMapper

Author: Seppli11

its/ruling/src/test/resources/expected/python-S6985.json

@@ -0,0 +1,25 @@
+{
+"project:pecos/examples/MACLR/dataset.py": [
+116
+],
+"project:pecos/examples/MACLR/evaluate.py": [
+92
+],
+"project:pecos/examples/MACLR/main.py": [
+280
+],
+"project:pecos/examples/giant-xrt/OGB_baselines/ogbn-arxiv/mlp.py": [
+109
+],
+"project:pecos/examples/giant-xrt/OGB_baselines/ogbn-papers100M/mlp_sgc.py": [
+109
+],
+"project:pecos/pecos/xmc/xtransformer/matcher.py": [
+411,
+1323,
+1338
+],
+"project:pecos/pecos/xmc/xtransformer/module.py": [
+460
+]
+}

python-checks/src/main/java/org/sonar/python/checks/CheckList.java

@@ -366,6 +366,7 @@ public static Iterable<Class> getChecks() {
       TooManyParametersCheck.class,
       TooManyReturnsCheck.class,
       TorchAutogradVariableShouldNotBeUsedCheck.class,
+      TorchLoadLeadsToUntrustedCodeExecutionCheck.class,
       TrailingCommentCheck.class,
       TrailingWhitespaceCheck.class,
       TypeAliasAnnotationCheck.class,

python-checks/src/main/java/org/sonar/python/checks/TorchLoadLeadsToUntrustedCodeExecutionCheck.java

@@ -0,0 +1,77 @@
+/*
+ * SonarQube Python Plugin
+ * Copyright (C) 2011-2024 SonarSource SA
+ * mailto:info AT sonarsource DOT com
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 3 of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+package org.sonar.python.checks;
+
+import java.util.List;
+import java.util.Set;
+import org.sonar.check.Rule;
+import org.sonar.plugins.python.api.PythonSubscriptionCheck;
+import org.sonar.plugins.python.api.symbols.Symbol;
+import org.sonar.plugins.python.api.tree.Argument;
+import org.sonar.plugins.python.api.tree.CallExpression;
+import org.sonar.plugins.python.api.tree.Expression;
+import org.sonar.plugins.python.api.tree.Name;
+import org.sonar.plugins.python.api.tree.RegularArgument;
+import org.sonar.plugins.python.api.tree.Tree;
+import org.sonar.python.cfg.fixpoint.ReachingDefinitionsAnalysis;
+import org.sonar.python.tree.TreeUtils;
+
+@Rule(key = "S6985")
+public class TorchLoadLeadsToUntrustedCodeExecutionCheck extends PythonSubscriptionCheck {
+
+  public static final String TORCH_LOAD = "torch.load";
+  public static final String MESSAGE = "Replace this call with a safe alternative.";
+  public static final String PYTHON_FALSE = "False";
+  public static final String WEIGHTS_ONLY = "weights_only";
+
+  private ReachingDefinitionsAnalysis reachingDefinitionsAnalysis;
+
+  @Override
+  public void initialize(Context context) {
+    context.registerSyntaxNodeConsumer(Tree.Kind.FILE_INPUT, ctx -> reachingDefinitionsAnalysis =
+      new ReachingDefinitionsAnalysis(ctx.pythonFile()));
+
+    context.registerSyntaxNodeConsumer(Tree.Kind.CALL_EXPR, ctx -> {
+      CallExpression callExpression = (CallExpression) ctx.syntaxNode();
+      Symbol calleeSymbol = callExpression.calleeSymbol();
+      if (calleeSymbol != null && TORCH_LOAD.equals(calleeSymbol.fullyQualifiedName()) && isWeightsOnlyNotFoundOrSetToFalse(callExpression.arguments())) {
+        ctx.addIssue(callExpression.callee(), MESSAGE);
+      }
+    });
+  }
+
+  private boolean isWeightsOnlyNotFoundOrSetToFalse(List<Argument> arguments) {
+    RegularArgument weightsOnlyArg = TreeUtils.argumentByKeyword(WEIGHTS_ONLY, arguments);
+    if (weightsOnlyArg == null) return true;
+    if (weightsOnlyArg.expression() instanceof Name name) {
+      return PYTHON_FALSE.equals(name.name()) || isNameSetToFalse(name);
+    }
+    return false;
+  }
+
+  private boolean isNameSetToFalse(Name name) {
+    Set<Expression> values = reachingDefinitionsAnalysis.valuesAtLocation(name);
+    return values.size() == 1 && values.stream()
+      .flatMap(TreeUtils.toStreamInstanceOfMapper(Name.class))
+      .map(Name::name).allMatch(PYTHON_FALSE::equals);
+  }
+
+
+}

python-checks/src/main/resources/org/sonar/l10n/py/rules/python/S6985.html

@@ -0,0 +1,29 @@
+<p>This rule raises an issue when <code>pytorch.load</code> is used to load a model.</p>
+<h2>Why is this an issue?</h2>
+<p>In PyTorch, it is common to load serialized models using the <code>torch.load</code> function. Under the hood, <code>torch.load</code> uses the
+<code>pickle</code> library to load the model and the weights. If the model comes from an untrusted source, an attacker could inject a malicious
+payload which would be executed during the deserialization.</p>
+<h2>How to fix it</h2>
+<p>Use a safer alternative to load the model, such as <code>safetensors.torch.load_model</code>.</p>
+<h3>Code examples</h3>
+<h4>Noncompliant code example</h4>
+<pre data-diff-id="1" data-diff-type="noncompliant">
+import torch
+
+model = torch.load('model.pth') # Noncompliant: torch.load is used to load the model
+</pre>
+<h4>Compliant solution</h4>
+<pre data-diff-id="1" data-diff-type="compliant">
+import torch
+import safetensors
+
+model = MyModel()
+safetensors.torch.load_model(model, 'model.pth')
+</pre>
+<h2>Resources</h2>
+<h3>Documentation</h3>
+<ul>
+  <li> Pytorch documentation: <a href="https://pytorch.org/tutorials/beginner/saving_loading_models.html#save-load-entire-model">Save/Load Entire
+  Model</a> </li>
+</ul>
+

python-checks/src/main/resources/org/sonar/l10n/py/rules/python/S6985.json

@@ -0,0 +1,24 @@
+{
+  "title": "Usage of \"torch.load\" can lead to untrusted code execution",
+  "type": "SECURITY_HOTSPOT",
+  "status": "ready",
+  "remediation": {
+    "func": "Constant\/Issue",
+    "constantCost": "15min"
+  },
+  "tags": [
+    "pytorch",
+    "machine-learning"
+  ],
+  "defaultSeverity": "Major",
+  "ruleSpecification": "RSPEC-6985",
+  "sqKey": "S6985",
+  "scope": "All",
+  "quickfix": "infeasible",
+  "code": {
+    "impacts": {
+      "SECURITY": "HIGH"
+    },
+    "attribute": "CONVENTIONAL"
+  }
+}

python-checks/src/main/resources/org/sonar/l10n/py/rules/python/Sonar_way_profile.json

@@ -247,6 +247,7 @@
     "S6973",
     "S6974",
     "S6979",
-    "S6983"
+    "S6983",
+    "S6985"
   ]
 }

python-checks/src/test/java/org/sonar/python/checks/TorchLoadLeadsToUntrustedCodeExecutionCheckTest.java

@@ -0,0 +1,31 @@
+/*
+ * SonarQube Python Plugin
+ * Copyright (C) 2011-2024 SonarSource SA
+ * mailto:info AT sonarsource DOT com
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 3 of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+package org.sonar.python.checks;
+
+import org.junit.jupiter.api.Test;
+import org.sonar.python.checks.utils.PythonCheckVerifier;
+
+class TorchLoadLeadsToUntrustedCodeExecutionCheckTest {
+
+  @Test
+  void test() {
+    PythonCheckVerifier.verify("src/test/resources/checks/torchLoadLeadsToUntrustedCodeExecution.py", new TorchLoadLeadsToUntrustedCodeExecutionCheck());
+  }
+}

python-checks/src/test/resources/checks/torchLoadLeadsToUntrustedCodeExecution.py

@@ -0,0 +1,37 @@
+import torch
+from torch import load
+import safetensors
+
+model = torch.load('model.pth') # Noncompliant {{Replace this call with a safe alternative.}}
+       #^^^^^^^^^^
+
+some_path = ...
+model = load(some_path) # Noncompliant
+
+torch.load(model2, 'model.pth', weights_only=False) #Noncompliant
+
+model2 = ...
+safetensors.torch.load_model(model2, 'model.pth')
+
+torch.load(model2, 'model.pth', weights_only=True)
+
+def unknown_weights_only_value(some_value, some_func):
+    torch.load(model2, 'model.pth', weights_only=some_value)
+    torch.load(model2, 'model.pth', weights_only=some_func())
+
+def conditional_weights_only(cond):
+    weights_only = True
+    if cond:
+        weights_only = False
+
+    torch.load(model2, 'model.pth', weights_only=weights_only)
+
+def only_one_definition():
+    weights_only = True
+    torch.load(model2, 'model.pth', weights_only=weights_only)
+
+    weights_only = False
+    torch.load(model2, 'model.pth', weights_only=weights_only) #Noncompliant
+
+# test if no issue is raised if there is no symbol for the callee
+something[42]()

python-frontend/src/main/java/org/sonar/python/tree/TreeUtils.java

@@ -461,6 +461,10 @@ public static <T extends Tree> Optional<T> toOptionalInstanceOf(Class<T> castToC
     return Optional.ofNullable(tree).filter(castToClass::isInstance).map(castToClass::cast);
   }
 
+  public static <T extends Tree> Function<Tree, Stream<T>> toStreamInstanceOfMapper(Class<T> castToClass) {
+    return tree -> toOptionalInstanceOf(castToClass, tree).map(Stream::of).orElse(Stream.empty());
+  }
+
   public static Optional<Tree> firstChild(Tree tree, Predicate<Tree> filter) {
     if (filter.test(tree)) {
       return Optional.of(tree);

python-frontend/src/test/java/org/sonar/python/tree/TreeUtilsTest.java

@@ -25,8 +25,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
-import java.util.function.Function;
-import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import org.junit.jupiter.api.Test;
 import org.sonar.plugins.python.api.symbols.Symbol;
 import org.sonar.plugins.python.api.tree.AnyParameter;
@@ -544,6 +543,31 @@ void test_toInstanceOfMapper() {
     assertThat(funcDefPresent).isFalse();
   }
 
+  @Test
+  void test_toStreamInstanceOfMapper() {
+    var fileInput = PythonTestUtils.parse(
+      "class A:",
+      "    x = True",
+      "    def foo(self):",
+      "        def foo2(x, y): return x + y",
+      "        return foo2(1, 1)",
+      "    class B:",
+      "        def bar(self): pass"
+    );
+    Tree tree = PythonTestUtils.getFirstChild(fileInput, t -> t.is(Kind.CLASSDEF));
+
+    boolean classPresent = Stream.of(tree)
+      .flatMap(TreeUtils.toStreamInstanceOfMapper(ClassDef.class))
+      .count() > 0;
+
+    assertThat(classPresent).isTrue();
+
+    boolean funcDefPresent = Stream.of(tree)
+      .flatMap(TreeUtils.toStreamInstanceOfMapper(FunctionDef.class))
+      .count() > 0;
+
+    assertThat(funcDefPresent).isFalse();
+  }
   @Test
   void test_findIndentationSize() {
     var fileInput = PythonTestUtils.parse("def foo():\n" +