From 583bfd706f69c31504db0738a9f0de3fc5743b8e Mon Sep 17 00:00:00 2001
From: Adrien LUCAS <adrien.lucas@smile.fr>
Date: Tue, 26 Jan 2021 13:29:33 +0100
Subject: [PATCH] [tainting] fix taint flow (#137)

---
 src/Twig/AnalyzedTemplatesTainter.php | 33 +++++++++++++++------------
 1 file changed, 18 insertions(+), 15 deletions(-)

diff --git a/src/Twig/AnalyzedTemplatesTainter.php b/src/Twig/AnalyzedTemplatesTainter.php
index a08936f..d6a3492 100644
--- a/src/Twig/AnalyzedTemplatesTainter.php
+++ b/src/Twig/AnalyzedTemplatesTainter.php
@@ -9,6 +9,7 @@ use PhpParser\Node\Expr\Array_;
 use PhpParser\Node\Expr\MethodCall;
 use PhpParser\Node\Expr\Variable;
 use Psalm\Codebase;
+use Psalm\CodeLocation;
 use Psalm\Context;
 use Psalm\Internal\DataFlow\DataFlowNode;
 use Psalm\Plugin\Hook\AfterMethodCallAnalysisInterface;
@@ -16,7 +17,6 @@ use Psalm\StatementsSource;
 use Psalm\Type\Atomic\TKeyedArray;
 use Psalm\Type\Union;
 use RuntimeException;
-use SplObjectStorage;
 use Twig\Environment;
 
 /**
@@ -39,13 +39,20 @@ class AnalyzedTemplatesTainter implements AfterMethodCallAnalysisInterface
         $templateName = TwigUtils::extractTemplateNameFromExpression($expr->args[0]->value, $statements_source);
 
         // Taints going _in_ the template
+        $methodNode = DataFlowNode::getForMethodArgument(
+            $method_id,
+            $method_id,
+            1,
+            new CodeLocation($statements_source, $expr->args[1]),
+            new CodeLocation($statements_source, $expr->name)
+        );
+
         $templateParameters = self::generateTemplateParameters($expr->args[1]->value, $statements_source);
-        foreach ($templateParameters as $sourceNode) {
-            $parameterName = $templateParameters[$sourceNode];
+        foreach ($templateParameters as $parameterName) {
             $label = $argumentId = strtolower($templateName).'#'.strtolower($parameterName);
             $destinationNode = new DataFlowNode($argumentId, $label, null, null);
 
-            $codebase->taint_flow_graph->addPath($sourceNode, $destinationNode, 'arg');
+            $codebase->taint_flow_graph->addPath($methodNode, $destinationNode, 'arg');
         }
 
         // Taints going _out_ of the template
@@ -58,9 +65,9 @@ class AnalyzedTemplatesTainter implements AfterMethodCallAnalysisInterface
     }
 
     /**
-     * @return SplObjectStorage<DataFlowNode, string>
+     * @return list<string>
      */
-    private static function generateTemplateParameters(Expr $templateParameters, StatementsSource $source): SplObjectStorage
+    private static function generateTemplateParameters(Expr $templateParameters, StatementsSource $source): array
     {
         $type = $source->getNodeTypeProvider()->getType($templateParameters);
         if (null === $type) {
@@ -68,11 +75,10 @@ class AnalyzedTemplatesTainter implements AfterMethodCallAnalysisInterface
         }
 
         if ($templateParameters instanceof Array_) {
-            /** @var SplObjectStorage<DataFlowNode, string> $parameters */
-            $parameters = new SplObjectStorage();
+            $parameters = [];
             foreach ($type->parent_nodes as $node) {
                 if (preg_match('/array\[\'([a-zA-Z]+)\'\]/', $node->label, $matches)) {
-                    $parameters[$node] = $matches[1];
+                    $parameters[] = $matches[1];
                 }
             }
 
@@ -83,12 +89,9 @@ class AnalyzedTemplatesTainter implements AfterMethodCallAnalysisInterface
             /** @var TKeyedArray $arrayValues */
             $arrayValues = $type->getAtomicTypes()['array'];
 
-            /** @var SplObjectStorage<DataFlowNode, string> $parameters */
-            $parameters = new SplObjectStorage();
-            foreach ($arrayValues->properties as $parameterName => $parameterType) {
-                foreach ($parameterType->parent_nodes as $node) {
-                    $parameters[$node] = (string) $parameterName;
-                }
+            $parameters = [];
+            foreach (array_keys($arrayValues->properties) as $parameterName) {
+                $parameters[] = (string) $parameterName;
             }
 
             return $parameters;