Fix #4661 - support conditional escaping for functions

src/Psalm/Internal/Analyzer/Statements/Expression/Call/FunctionCallAnalyzer.php

@@ -19,6 +19,7 @@ use Psalm\Internal\FileManipulation\FileManipulationBuffer;
 use Psalm\Internal\DataFlow\TaintSource;
 use Psalm\Internal\DataFlow\DataFlowNode;
 use Psalm\Internal\Codebase\TaintFlowGraph;
+use Psalm\Internal\Type\TypeExpander;
 use Psalm\Issue\DeprecatedFunction;
 use Psalm\Issue\ForbiddenCode;
 use Psalm\Issue\MixedFunctionCall;
@@ -1000,7 +1001,7 @@ class FunctionCallAnalyzer extends CallAnalyzer
                         $return_type = clone $function_storage->return_type;
 
                         if ($template_result->upper_bounds && $function_storage->template_types) {
-                            $return_type = \Psalm\Internal\Type\TypeExpander::expandUnion(
+                            $return_type = TypeExpander::expandUnion(
                                 $codebase,
                                 $return_type,
                                 null,
@@ -1014,7 +1015,7 @@ class FunctionCallAnalyzer extends CallAnalyzer
                             );
                         }
 
-                        $return_type = \Psalm\Internal\Type\TypeExpander::expandUnion(
+                        $return_type = TypeExpander::expandUnion(
                             $codebase,
                             $return_type,
                             null,
@@ -1099,7 +1100,8 @@ class FunctionCallAnalyzer extends CallAnalyzer
             $stmt,
             $function_id,
             $function_storage,
-            $stmt_type
+            $stmt_type,
+            $template_result
         );
 
         if ($function_storage->proxy_calls !== null) {
@@ -1370,7 +1372,8 @@ class FunctionCallAnalyzer extends CallAnalyzer
         PhpParser\Node\Expr\FuncCall $stmt,
         string $function_id,
         FunctionLikeStorage $function_storage,
-        Type\Union $stmt_type
+        Type\Union $stmt_type,
+        TemplateResult $template_result
     ) : ?DataFlowNode {
         if (!$statements_analyzer->data_flow_graph instanceof TaintFlowGraph
             || \in_array('TaintedInput', $statements_analyzer->getSuppressedIssues())
@@ -1389,7 +1392,52 @@ class FunctionCallAnalyzer extends CallAnalyzer
 
         $statements_analyzer->data_flow_graph->addNode($function_call_node);
 
+        $codebase = $statements_analyzer->getCodebase();
+
+        $conditionally_removed_taints = [];
+
+        foreach ($function_storage->conditionally_removed_taints as $conditionally_removed_taint) {
+            $conditionally_removed_taint = clone $conditionally_removed_taint;
+
+            $conditionally_removed_taint->replaceTemplateTypesWithArgTypes(
+                $template_result,
+                $codebase
+            );
+
+            $expanded_type = TypeExpander::expandUnion(
+                $statements_analyzer->getCodebase(),
+                $conditionally_removed_taint,
+                null,
+                null,
+                null,
+                true,
+                true
+            );
+
+            foreach ($expanded_type->getLiteralStrings() as $literal_string) {
+                $conditionally_removed_taints[] = $literal_string->value;
+            }
+        }
+
+        if ($conditionally_removed_taints && $function_storage->location) {
+            $assignment_node = DataFlowNode::getForAssignment(
+                $function_id . '-escaped',
+                $function_storage->signature_return_type_location ?: $function_storage->location,
+                $function_call_node->specialization_key
+            );
+
+            $statements_analyzer->data_flow_graph->addPath(
+                $function_call_node,
+                $assignment_node,
+                'conditionally-escaped',
+                [],
+                $conditionally_removed_taints
+            );
+
+            $stmt_type->parent_nodes[$assignment_node->id] = $assignment_node;
+        } else {
             $stmt_type->parent_nodes[$function_call_node->id] = $function_call_node;
+        }
 
         if ($function_storage->return_source_params) {
             $removed_taints = $function_storage->removed_taints;

src/Psalm/Internal/PhpVisitor/Reflector/FunctionLikeDocblockParser.php

@@ -226,7 +226,13 @@ class FunctionLikeDocblockParser
         if (isset($parsed_docblock->tags['psalm-taint-escape'])) {
             foreach ($parsed_docblock->tags['psalm-taint-escape'] as $param) {
                 $param = trim($param);
-                $info->removed_taints[] = $param;
+                if ($param[0] === '(') {
+                    $line_parts = CommentAnalyzer::splitDocLine($param);
+
+                    $info->removed_taints[] = CommentAnalyzer::sanitizeDocblockType($line_parts[0]);
+                } else {
+                    $info->removed_taints[] = explode(' ', $param)[0];
+                }
             }
         }
 

src/Psalm/Internal/PhpVisitor/Reflector/FunctionLikeDocblockScanner.php

@@ -471,7 +471,47 @@ class FunctionLikeDocblockScanner
         }
 
         $storage->added_taints = $docblock_info->added_taints;
-        $storage->removed_taints = $docblock_info->removed_taints;
+
+        foreach ($docblock_info->removed_taints as $removed_taint) {
+            if ($removed_taint[0] === '(') {
+                try {
+                    [$fixed_type_tokens, $function_template_types] = self::getConditionalSanitizedTypeTokens(
+                        $removed_taint,
+                        $aliases,
+                        $function_template_types + $class_template_types,
+                        $type_aliases,
+                        $storage,
+                        $classlike_storage,
+                        $cased_function_id,
+                        $function_template_types
+                    );
+
+                    $removed_taint = TypeParser::parseTokens(
+                        \array_values($fixed_type_tokens),
+                        null,
+                        $function_template_types + $class_template_types,
+                        $type_aliases
+                    );
+
+                    $removed_taint->queueClassLikesForScanning($codebase, $file_storage);
+
+                    $removed_taint_single = \array_values($removed_taint->getAtomicTypes())[0];
+
+                    if (!$removed_taint_single instanceof Type\Atomic\TConditional) {
+                        throw new TypeParseTreeException('Escaped taint must be a conditional');
+                    }
+
+                    $storage->conditionally_removed_taints[] = $removed_taint;
+                } catch (TypeParseTreeException $e) {
+                    $storage->docblock_issues[] = new InvalidDocblock(
+                        $e->getMessage() . ' in docblock for ' . $cased_function_id,
+                        new CodeLocation($file_scanner, $stmt, null, true)
+                    );
+                }
+            } else {
+                $storage->removed_taints[] = $removed_taint;
+            }
+        }
 
         if ($docblock_info->flows) {
             foreach ($docblock_info->flows as $flow) {

src/Psalm/Storage/FunctionLikeStorage.php

@@ -194,6 +194,11 @@ abstract class FunctionLikeStorage
      */
     public $removed_taints = [];
 
+    /**
+     * @var array<Type\Union>
+     */
+    public $conditionally_removed_taints = [];
+
     /**
      * @var array<int, string>
      */

tests/TaintTest.php

@@ -562,6 +562,22 @@ class TaintTest extends TestCase
                 $input = strtr(\'data\', \'data\', \'data\');
                 setcookie($input, \'value\');',
             ],
+            'conditionallyEscapedTaintPassedTrue' => [
+                '<?php
+                    /**
+                     * @psalm-taint-escape ($escape is true ? "html" : null)
+                     */
+                    function foo(string $string, bool $escape = true): string {
+                        if ($escape) {
+                            $string = htmlspecialchars($string);
+                        }
+
+                        return $string;
+                    }
+
+                    echo foo($_GET["foo"], true);
+                    echo foo($_GET["foo"]);'
+            ],
         ];
     }
 
@@ -1889,6 +1905,22 @@ class TaintTest extends TestCase
                     setcookie($input, \'value\');',
                 'error_message' => 'TaintedCookie',
             ],
+            'conditionallyEscapedTaintPassedFalse' => [
+                '<?php
+                    /**
+                     * @psalm-taint-escape ($escape is true ? "html" : null)
+                     */
+                    function foo(string $string, bool $escape = true): string {
+                        if ($escape) {
+                            $string = htmlspecialchars($string);
+                        }
+
+                        return $string;
+                    }
+
+                    echo foo($_GET["foo"], false);',
+                'error_message' => 'TaintedHtml',
+            ],
             /*
             // TODO: Stubs do not support this type of inference even with $this->message = $message.
             // Most uses of getMessage() would be with caught exceptions, so this is not representative of real code.