fix: add taint source on plugin-added taints

docs/security_analysis/custom_taint_sources.md

@@ -23,52 +23,32 @@ For example this plugin treats all variables named `$bad_data` as taint sources.
 ```php
 <?php
 
-namespace Some\Ns;
+namespace Psalm\Example\Plugin;
 
-use PhpParser;
-use Psalm\CodeLocation;
-use Psalm\Context;
-use Psalm\FileManipulation;
-use Psalm\Plugin\EventHandler\AfterExpressionAnalysisInterface;
-use Psalm\Plugin\EventHandler\Event\AfterExpressionAnalysisEvent;
+use PhpParser\Node\Expr\Variable;
+use Psalm\Plugin\EventHandler\AddTaintsInterface;
+use Psalm\Plugin\EventHandler\Event\AddRemoveTaintsEvent;
 use Psalm\Type\TaintKindGroup;
 
-class BadSqlTainter implements AfterExpressionAnalysisInterface
+/**
+ * Add input taints to all variables named 'bad_data'
+ */
+class TaintBadDataPlugin implements AddTaintsInterface
 {
     /**
-     * Called after an expression has been checked
-     *
-     * @param  PhpParser\Node\Expr  $expr
-     * @param  Context              $context
-     * @param  FileManipulation[]   $file_replacements
+     * Called to see what taints should be added
      *
-     * @return void
+     * @return list<string>
      */
-    public static function afterExpressionAnalysis(AfterExpressionAnalysisEvent $event): ?bool {
+    public static function addTaints(AddRemoveTaintsEvent $event): array
+    {
         $expr = $event->getExpr();
-        $statements_source = $event->getStatementsSource();
-        $codebase = $event->getCodebase();
-        if ($expr instanceof PhpParser\Node\Expr\Variable
-            && $expr->name === 'bad_data'
-        ) {
-            $expr_type = $statements_source->getNodeTypeProvider()->getType($expr);
 
-            // should be a globally unique id
-            // you can use its line number/start offset
-            $expr_identifier = '$bad_data'
-                . '-' . $statements_source->getFileName()
-                . ':' . $expr->getAttribute('startFilePos');
-
-            if ($expr_type) {
-                $codebase->addTaintSource(
-                    $expr_type,
-                    $expr_identifier,
-                    TaintKindGroup::ALL_INPUT,
-                    new CodeLocation($statements_source, $expr)
-                );
-            }
+        if ($expr instanceof Variable && $expr->name === 'bad_data') {
+            return TaintKindGroup::ALL_INPUT;
         }
-        return null;
+
+        return [];
     }
 }
 ```

examples/plugins/TaintActiveRecords.php

@@ -0,0 +1,94 @@
+<?php
+
+namespace Psalm\Example\Plugin;
+
+use PhpParser\Node\ArrayItem;
+use PhpParser\Node\Expr\PropertyFetch;
+use PhpParser\Node\Expr;
+use Psalm\Plugin\EventHandler\AddTaintsInterface;
+use Psalm\Plugin\EventHandler\Event\AddRemoveTaintsEvent;
+use Psalm\Type\Atomic;
+use Psalm\Type\Atomic\TNamedObject;
+use Psalm\Type\TaintKindGroup;
+use Psalm\Type\Union;
+
+/**
+ * Marks all property fetches of models inside namespace \app\models as tainted.
+ * ActiveRecords are model-representation of database entries, which can always
+ * contain user-input and therefor should be tainted.
+ */
+class TaintActiveRecords implements AddTaintsInterface
+{
+    /**
+     * Called to see what taints should be added
+     *
+     * @return list<string>
+     */
+    public static function addTaints(AddRemoveTaintsEvent $event): array
+    {
+        $expr = $event->getExpr();
+
+        // Model properties are accessed by property fetch, so abort here
+        if ($expr instanceof ArrayItem) {
+            return [];
+        }
+
+        $statements_source = $event->getStatementsSource();
+
+        // For all property fetch expressions, walk through the full fetch path
+        // (e.g. `$model->property->subproperty`) and check if it contains
+        // any class of namespace \app\models\
+        do {
+            $expr_type = $statements_source->getNodeTypeProvider()->getType($expr);
+            if (!$expr_type) {
+                continue;
+            }
+
+            if (self::containsActiveRecord($expr_type)) {
+                return TaintKindGroup::ALL_INPUT;
+            }
+        } while ($expr = self::getParentNode($expr));
+
+        return [];
+    }
+
+    /**
+     * @return bool `true` if union contains a type of model
+     */
+    private static function containsActiveRecord(Union $union_type): bool
+    {
+        foreach ($union_type->getAtomicTypes() as $type) {
+            if (self::isActiveRecord($type)) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
+    /**
+     * @return bool `true` if namespace of type is in namespace `app\models`
+     */
+    private static function isActiveRecord(Atomic $type): bool
+    {
+        if (!$type instanceof TNamedObject) {
+            return false;
+        }
+
+        return strpos($type->value, 'app\models\\') === 0;
+    }
+
+
+    /**
+     * Return next node that should be followed for active record search
+     */
+    private static function getParentNode(ArrayItem|Expr $expr): ?Expr
+    {
+        // Model properties are always accessed by a property fetch
+        if ($expr instanceof PropertyFetch) {
+            return $expr->var;
+        }
+
+        return null;
+    }
+}

src/Psalm/Internal/Analyzer/FunctionLikeAnalyzer.php

@@ -17,6 +17,7 @@
 use Psalm\FileManipulation;
 use Psalm\Internal\Analyzer\FunctionLike\ReturnTypeAnalyzer;
 use Psalm\Internal\Analyzer\FunctionLike\ReturnTypeCollector;
+use Psalm\Internal\Analyzer\Statements\Expression\Call\FunctionCallReturnTypeFetcher;
 use Psalm\Internal\Analyzer\Statements\ExpressionAnalyzer;
 use Psalm\Internal\Codebase\TaintFlowGraph;
 use Psalm\Internal\Codebase\VariableUseGraph;
@@ -830,6 +831,24 @@ public function analyze(
             }
         }
 
+        // Class methods are analyzed deferred, therefor it's required to
+        // add taint sources additionally on analyze not only on call
+        if ($codebase->taint_flow_graph
+            && $this->function instanceof ClassMethod
+            && $cased_method_id) {
+            $method_source = DataFlowNode::getForMethodReturn(
+                (string) $method_id,
+                $cased_method_id,
+                $storage->location,
+            );
+
+            FunctionCallReturnTypeFetcher::taintUsingStorage(
+                $storage,
+                $codebase->taint_flow_graph,
+                $method_source,
+            );
+        }
+
         if ($add_mutations) {
             if ($this->return_vars_in_scope !== null) {
                 $context->vars_in_scope = TypeAnalyzer::combineKeyedTypes(
@@ -1067,7 +1086,7 @@ private function processParams(
 
                     $statements_analyzer->data_flow_graph->addNode($param_assignment);
 
-                    if ($cased_method_id) {
+                    if ($cased_method_id !== null) {
                         $type_source = DataFlowNode::getForMethodArgument(
                             $cased_method_id,
                             $cased_method_id,

src/Psalm/Internal/Analyzer/Statements/Expression/ArrayAnalyzer.php

@@ -14,6 +14,7 @@
 use Psalm\Internal\Codebase\TaintFlowGraph;
 use Psalm\Internal\Codebase\VariableUseGraph;
 use Psalm\Internal\DataFlow\DataFlowNode;
+use Psalm\Internal\DataFlow\TaintSource;
 use Psalm\Internal\Type\Comparator\UnionTypeComparator;
 use Psalm\Internal\Type\TypeCombiner;
 use Psalm\Issue\DuplicateArrayKey;
@@ -42,6 +43,7 @@
 use Psalm\Type\Atomic\TTrue;
 use Psalm\Type\Union;
 
+use function array_diff;
 use function array_merge;
 use function array_values;
 use function count;
@@ -441,6 +443,12 @@ private static function analyzeArrayItem(
                     $added_taints = $codebase->config->eventDispatcher->dispatchAddTaints($event);
                     $removed_taints = $codebase->config->eventDispatcher->dispatchRemoveTaints($event);
 
+                    $taints = array_diff($added_taints, $removed_taints);
+                    if ($taints !== [] && $statements_analyzer->data_flow_graph instanceof TaintFlowGraph) {
+                        $taint_source = TaintSource::fromNode($new_parent_node);
+                        $statements_analyzer->data_flow_graph->addSource($taint_source);
+                    }
+
                     foreach ($item_value_type->parent_nodes as $parent_node) {
                         $data_flow_graph->addPath(
                             $parent_node,
@@ -476,6 +484,13 @@ private static function analyzeArrayItem(
                     $added_taints = $codebase->config->eventDispatcher->dispatchAddTaints($event);
                     $removed_taints = $codebase->config->eventDispatcher->dispatchRemoveTaints($event);
 
+                    $taints = array_diff($added_taints, $removed_taints);
+                    if ($taints !== [] && $statements_analyzer->data_flow_graph instanceof TaintFlowGraph) {
+                        $taint_source = TaintSource::fromNode($new_parent_node);
+                        $taint_source->taints = $taints;
+                        $statements_analyzer->data_flow_graph->addSource($taint_source);
+                    }
+
                     foreach ($item_key_type->parent_nodes as $parent_node) {
                         $data_flow_graph->addPath(
                             $parent_node,

src/Psalm/Internal/Analyzer/Statements/Expression/Assignment/InstancePropertyAssignmentAnalyzer.php

@@ -28,6 +28,7 @@
 use Psalm\Internal\Codebase\TaintFlowGraph;
 use Psalm\Internal\Codebase\VariableUseGraph;
 use Psalm\Internal\DataFlow\DataFlowNode;
+use Psalm\Internal\DataFlow\TaintSource;
 use Psalm\Internal\FileManipulation\FileManipulationBuffer;
 use Psalm\Internal\MethodIdentifier;
 use Psalm\Internal\Type\Comparator\TypeComparisonResult;
@@ -78,6 +79,7 @@
 use Psalm\Type\Union;
 use UnexpectedValueException;
 
+use function array_diff;
 use function array_merge;
 use function array_pop;
 use function count;
@@ -505,6 +507,13 @@ private static function taintProperty(
                 $added_taints = $codebase->config->eventDispatcher->dispatchAddTaints($event);
                 $removed_taints = $codebase->config->eventDispatcher->dispatchRemoveTaints($event);
 
+                $taints = array_diff($added_taints, $removed_taints);
+                if ($taints !== [] && $statements_analyzer->data_flow_graph instanceof TaintFlowGraph) {
+                    $taint_source = TaintSource::fromNode($property_node);
+                    $taint_source->taints = $taints;
+                    $statements_analyzer->data_flow_graph->addSource($taint_source);
+                }
+
                 $data_flow_graph->addPath(
                     $property_node,
                     $var_node,
@@ -600,6 +609,13 @@ public static function taintUnspecializedProperty(
         $added_taints = $codebase->config->eventDispatcher->dispatchAddTaints($event);
         $removed_taints = $codebase->config->eventDispatcher->dispatchRemoveTaints($event);
 
+        $taints = array_diff($added_taints, $removed_taints);
+        if ($taints !== [] && $statements_analyzer->data_flow_graph instanceof TaintFlowGraph) {
+            $taint_source = TaintSource::fromNode($property_node);
+            $taint_source->taints = $taints;
+            $statements_analyzer->data_flow_graph->addSource($taint_source);
+        }
+
         $data_flow_graph->addPath(
             $localized_property_node,
             $property_node,