Improve specialisation after call

vimeo · Jun 19, 2020 · b1c836e · b1c836e
1 parent 8f2e28c
commit b1c836e
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 2 deletions.
diff --git a/src/Psalm/Internal/Analyzer/FunctionLikeAnalyzer.php b/src/Psalm/Internal/Analyzer/FunctionLikeAnalyzer.php
@@ -751,7 +751,7 @@ function (FunctionLikeParameter $p) {
         if ($codebase->taint
             && $this->function instanceof ClassMethod
             && $cased_method_id
-            && $this->function->name->name === '__construct'
+            && $storage->specialize_call
             && isset($context->vars_in_scope['$this'])
             && $context->vars_in_scope['$this']->parent_nodes
         ) {

diff --git a/...Psalm/Internal/Analyzer/Statements/Expression/Call/Method/MethodCallReturnTypeFetcher.php b/...Psalm/Internal/Analyzer/Statements/Expression/Call/Method/MethodCallReturnTypeFetcher.php
@@ -4,6 +4,7 @@
 use PhpParser;
 use Psalm\Internal\Analyzer\StatementsAnalyzer;
 use Psalm\Internal\Codebase\InternalCallMapHandler;
+use Psalm\Internal\Analyzer\Statements\Expression\ExpressionIdentifier;
 use Psalm\Codebase;
 use Psalm\CodeLocation;
 use Psalm\Context;
@@ -193,7 +194,7 @@ public static function fetch(
                 $declaring_method_id
             );
 
-            $node_location = new CodeLocation($statements_analyzer, $stmt);
+            $node_location = new CodeLocation($statements_analyzer, $stmt->name);
 
             $method_call_node = TaintNode::getForMethodReturn(
                 (string) $method_id,
@@ -207,6 +208,41 @@ public static function fetch(
             $return_type_candidate->parent_nodes = [
                 $method_call_node
             ];
+
+            if ($method_storage->specialize_call) {
+                $var_id = ExpressionIdentifier::getArrayVarId(
+                    $stmt->var,
+                    null,
+                    $statements_analyzer
+                );
+
+                if ($var_id && isset($context->vars_in_scope[$var_id])) {
+                    $var_node = TaintNode::getForAssignment(
+                        $var_id,
+                        new CodeLocation($statements_analyzer, $stmt->var)
+                    );
+
+                    $codebase->taint->addTaintNode($var_node);
+
+                    $codebase->taint->addPath(
+                        $method_call_node,
+                        $var_node,
+                        'method-call-' . $method_id->method_name
+                    );
+
+                    $stmt_var_type = clone $context->vars_in_scope[$var_id];
+
+                    if ($context->vars_in_scope[$var_id]->parent_nodes) {
+                        foreach ($context->vars_in_scope[$var_id]->parent_nodes as $parent_node) {
+                            $codebase->taint->addPath($parent_node, $var_node, '=');
+                        }
+                    }
+
+                    $stmt_var_type->parent_nodes = [$var_node];
+
+                    $context->vars_in_scope[$var_id] = $stmt_var_type;
+                }
+            }
         }
 
         return $return_type_candidate;

diff --git a/tests/TaintTest.php b/tests/TaintTest.php
@@ -1501,6 +1501,47 @@ public static function deleteUser(PDO $pdo, string $userId) : void {
         $this->analyzeFile('somefile.php', new Context());
     }
 
+    public function testTaintPropertyPassingObjectSettingValueLater() : void
+    {
+        $this->expectException(\Psalm\Exception\CodeException::class);
+        $this->expectExceptionMessage('TaintedInput');
+
+        $this->project_analyzer->trackTaintedInputs();
+
+        $this->addFile(
+            'somefile.php',
+            '<?phps
+                /** @psalm-taint-specialize */
+                class User {
+                    public string $id;
+
+                    public function __construct(string $userId) {
+                        $this->id = $userId;
+                    }
+
+                    public function setId(string $userId) : void {
+                        $this->id = $userId;
+                    }
+                }
+
+                class UserUpdater {
+                    public static function doDelete(PDO $pdo, User $user) : void {
+                        self::deleteUser($pdo, $user->id);
+                    }
+
+                    public static function deleteUser(PDO $pdo, string $userId) : void {
+                        $pdo->exec("delete from users where user_id = " . $userId);
+                    }
+                }
+
+                $userObj = new User("5");
+                $userObj->setId((string) $_GET["user_id"]);
+                UserUpdater::doDelete(new PDO(), $userObj);'
+        );
+
+        $this->analyzeFile('somefile.php', new Context());
+    }
+
     public function testTaintPropertyPassingObjectWithDifferentValue() : void
     {
         $this->project_analyzer->trackTaintedInputs();