Finalize refactoring

Author: danog

src/Psalm/Codebase.php

@@ -4,6 +4,7 @@
 
 namespace Psalm;
 
+use AssertionError;
 use Exception;
 use InvalidArgumentException;
 use LanguageServerProtocol\Command;
@@ -65,10 +66,11 @@
 use Psalm\Type\Atomic\TLiteralInt;
 use Psalm\Type\Atomic\TLiteralString;
 use Psalm\Type\Atomic\TNamedObject;
-use Psalm\Type\TaintKindGroup;
+use Psalm\Type\TaintKind;
 use Psalm\Type\Union;
 use ReflectionProperty;
 use ReflectionType;
+use RuntimeException;
 use UnexpectedValueException;
 
 use function array_combine;
@@ -99,6 +101,7 @@
 use function substr;
 use function substr_count;
 
+use const PHP_INT_SIZE;
 use const PHP_VERSION_ID;
 
 /**
@@ -254,6 +257,42 @@ final class Codebase
 
     public bool $literal_array_key_check = false;
 
+    /** @internal */
+    public int $taint_count = TaintKind::BUILTIN_TAINT_COUNT;
+
+    /**
+     * @var array<string, int>
+     */
+    private array $taint_map = [
+        'callable' => TaintKind::INPUT_CALLABLE,
+        'unserialize' => TaintKind::INPUT_UNSERIALIZE,
+        'include' => TaintKind::INPUT_INCLUDE,
+        'eval' => TaintKind::INPUT_EVAL,
+        'ldap' => TaintKind::INPUT_LDAP,
+        'sql' => TaintKind::INPUT_SQL,
+        'html' => TaintKind::INPUT_HTML,
+        'has_quotes' => TaintKind::INPUT_HAS_QUOTES,
+        'shell' => TaintKind::INPUT_SHELL,
+        'ssrf' => TaintKind::INPUT_SSRF,
+        'file' => TaintKind::INPUT_FILE,
+        'cookie' => TaintKind::INPUT_COOKIE,
+        'header' => TaintKind::INPUT_HEADER,
+        'xpath' => TaintKind::INPUT_XPATH,
+        'sleep' => TaintKind::INPUT_SLEEP,
+        'extract' => TaintKind::INPUT_EXTRACT,
+        'user_secret' => TaintKind::USER_SECRET,
+        'system_secret' => TaintKind::SYSTEM_SECRET,
+
+        'input' => TaintKind::ALL_INPUT,
+        'tainted' => TaintKind::ALL_INPUT,
+    ];
+
+    /**
+     * @internal
+     * @var array<int, string>
+     */
+    public array $custom_taints = [];
+
     /** @internal */
     public function __construct(
         public Config $config,
@@ -319,6 +358,48 @@ public function __construct(
         $this->loadAnalyzer();
     }
 
+    /**
+     * Used to register a taint, or to fetch the ID of an already registered taint by its alias.
+     *
+     * @throws AssertionError if no more taint slots are left
+     * @throws RuntimeException if the passed alias uses some unregistered taint slots
+     * @param ?int $alias Used to register an alias of one or more pre-existing taints.
+     */
+    public function getOrRegisterTaint(string $taint_type, ?int $alias = null): int
+    {
+        if (isset($this->taint_map[$taint_type])) {
+            return $this->taint_map[$taint_type];
+        }
+        if ($alias === null) {
+            if ($this->taint_count+1 === (PHP_INT_SIZE * 8)) {
+                if (PHP_INT_SIZE === 8) {
+                    throw new RuntimeException("No more taint slots left, please register fewer taints or use some of the built-in taints!");
+                }
+                throw new RuntimeException("No more taint slots left, please switch to a 64-bit build of PHP to get 32 more taint slots, or register fewer taints or use some of the built-in taints!");
+            }
+            $id = 1 << ($this->taint_count++);
+            $this->custom_taints[$id] + $taint_type;
+        } else {
+            if ($this->taint_count+1 !== (PHP_INT_SIZE * 8)) {
+                $mask = (1 << $this->taint_count) - 1;
+                if ($alias & ~$mask) {
+                    throw new AssertionError("The passed alias $alias uses some not yet registered taint slots!");
+                }
+            }
+            $id = $alias;
+        }
+        $this->taint_map[$taint_type] = $id;
+        return $id;
+    }
+
+    /**
+     * Used to to fetch the ID of an already registered taint by its alias, or null if no taint is registered for the alias.
+     */
+    public function getTaint(string $taint_type): ?int
+    {
+        return $this->taint_map[$taint_type] ?? null;
+    }
+
     private function loadAnalyzer(): void
     {
         $this->analyzer = new Analyzer(
@@ -2145,7 +2226,7 @@ public function queueClassLikeForScanning(
     public function addTaintSource(
         Union $expr_type,
         string $taint_id,
-        int $taints = TaintKindGroup::ALL_INPUT,
+        int $taints = TaintKind::ALL_INPUT,
         ?CodeLocation $code_location = null,
     ): Union {
         if (!$this->taint_flow_graph) {
@@ -2167,7 +2248,7 @@ public function addTaintSource(
 
     public function addTaintSink(
         string $taint_id,
-        int $taints = TaintKindGroup::ALL_INPUT,
+        int $taints = TaintKind::ALL_INPUT,
         ?CodeLocation $code_location = null,
     ): void {
         if (!$this->taint_flow_graph) {

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

@@ -21,6 +21,8 @@
 use Psalm\Internal\Type\TemplateInferredTypeReplacer;
 use Psalm\Internal\Type\TemplateResult;
 use Psalm\Internal\Type\TypeExpander;
+use Psalm\Issue\InvalidDocblock;
+use Psalm\IssueBuffer;
 use Psalm\Plugin\EventHandler\Event\AddRemoveTaintsEvent;
 use Psalm\Plugin\EventHandler\Event\AfterFunctionCallAnalysisEvent;
 use Psalm\Storage\FunctionLikeStorage;
@@ -39,7 +41,6 @@
 use Psalm\Type\Atomic\TNull;
 use Psalm\Type\Atomic\TString;
 use Psalm\Type\TaintKind;
-use Psalm\Type\TaintKindGroup;
 use Psalm\Type\Union;
 use UnexpectedValueException;
 
@@ -591,7 +592,15 @@ private static function taintReturnType(
 
             if (!$expanded_type->isNullable()) {
                 foreach ($expanded_type->getLiteralStrings() as $literal_string) {
-                    $conditionally_removed_taints |= TaintKindGroup::NAME_TO_TAINT[$literal_string->value];
+                    $taint = $codebase->getTaint($literal_string->value);
+                    if ($taint === null) {
+                        IssueBuffer::maybeAdd(new InvalidDocblock(
+                            "Invalid taint name {$literal_string->value} provided",
+                            $function_storage->location,
+                        ));
+                        continue;
+                    }
+                    $conditionally_removed_taints |= $taint;
                 }
             }
         }

src/Psalm/Internal/Analyzer/Statements/Expression/Call/Method/MethodCallReturnTypeFetcher.php

@@ -34,7 +34,6 @@
 use Throwable;
 use UnexpectedValueException;
 
-use function array_filter;
 use function count;
 use function in_array;
 use function strtolower;

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

@@ -20,14 +20,14 @@
 use Psalm\Internal\Type\TemplateInferredTypeReplacer;
 use Psalm\Internal\Type\TemplateResult;
 use Psalm\Internal\Type\TypeExpander;
+use Psalm\Issue\InvalidDocblock;
 use Psalm\Issue\NonStaticSelfCall;
 use Psalm\Issue\ParentNotFound;
 use Psalm\IssueBuffer;
 use Psalm\Plugin\EventHandler\Event\AddRemoveTaintsEvent;
 use Psalm\Storage\MethodStorage;
 use Psalm\Type;
 use Psalm\Type\Atomic\TNamedObject;
-use Psalm\Type\TaintKindGroup;
 use Psalm\Type\Union;
 
 use function count;
@@ -330,7 +330,15 @@ public static function taintReturnType(
                 );
 
                 foreach ($expanded_type->getLiteralStrings() as $literal_string) {
-                    $conditionally_removed_taints |= TaintKindGroup::NAME_TO_TAINT[$literal_string->value];
+                    $taint = $codebase->getTaint($literal_string->value);
+                    if ($taint === null) {
+                        IssueBuffer::maybeAdd(new InvalidDocblock(
+                            "Invalid taint name {$literal_string->value} provided",
+                            $method_storage->location,
+                        ));
+                        continue;
+                    }
+                    $conditionally_removed_taints |= $taint;
                 }
             }
         }

src/Psalm/Internal/Analyzer/Statements/Expression/Fetch/VariableFetchAnalyzer.php

@@ -34,7 +34,7 @@
 use Psalm\Type\Atomic\TNonEmptyString;
 use Psalm\Type\Atomic\TNull;
 use Psalm\Type\Atomic\TString;
-use Psalm\Type\TaintKindGroup;
+use Psalm\Type\TaintKind;
 use Psalm\Type\Union;
 
 use function in_array;
@@ -509,7 +509,7 @@ private static function taintVariable(
             || $var_name === '$_COOKIE'
             || $var_name === '$_REQUEST'
         ) {
-            $taints = TaintKindGroup::ALL_INPUT;
+            $taints = TaintKind::ALL_INPUT;
         } else {
             $taints = 0;
         }

src/Psalm/Internal/Codebase/TaintFlowGraph.php

@@ -6,6 +6,7 @@
 
 use Override;
 use Psalm\CodeLocation;
+use Psalm\Codebase;
 use Psalm\Config;
 use Psalm\Internal\Analyzer\ProjectAnalyzer;
 use Psalm\Internal\DataFlow\DataFlowNode;
@@ -34,7 +35,6 @@
 use Psalm\Progress\Phase;
 use Psalm\Progress\Progress;
 use Psalm\Type\TaintKind;
-use Psalm\Type\TaintKindGroup;
 
 use function count;
 use function end;
@@ -215,6 +215,8 @@ public function connectSinksAndSources(Progress $progress): void
 
         $project_analyzer = ProjectAnalyzer::getInstance();
 
+        $codebase = $project_analyzer->getCodebase();
+
         // Remove all specializations without an outgoing edge
         foreach ($this->specializations as $k => &$map) {
             foreach ($map as $kk => $specialized_id) {
@@ -250,6 +252,7 @@ public function connectSinksAndSources(Progress $progress): void
                         $visited_source_ids,
                         $config,
                         $project_analyzer,
+                        $codebase,
                     );
                     continue;
                 }
@@ -275,6 +278,7 @@ public function connectSinksAndSources(Progress $progress): void
                         $visited_source_ids,
                         $config,
                         $project_analyzer,
+                        $codebase,
                     );
 
                     // If this node has first level specializations (=> is first-level & unspecialized),
@@ -299,6 +303,7 @@ public function connectSinksAndSources(Progress $progress): void
                                 $visited_source_ids,
                                 $config,
                                 $project_analyzer,
+                                $codebase,
                             );
                         }
                     } else {
@@ -316,11 +321,12 @@ public function connectSinksAndSources(Progress $progress): void
                                 $visited_source_ids,
                                 $config,
                                 $project_analyzer,
+                                $codebase,
                             );
                         }
                     }
                 } else {
-                    // Process all descendants 
+                    // Process all descendants
                     foreach ($source->processing_specialized_descendants_of as $map) {
                         if (isset($map[$source->id])) {
                             $specialized_id = $map[$source->id];
@@ -339,6 +345,7 @@ public function connectSinksAndSources(Progress $progress): void
                                 $visited_source_ids,
                                 $config,
                                 $project_analyzer,
+                                $codebase,
                             );
                         }
                     }
@@ -365,6 +372,7 @@ private function getChildNodes(
         array $visited_source_ids,
         Config $config,
         ProjectAnalyzer $project_analyzer,
+        Codebase $codebase,
     ): void {
         foreach ($this->forward_edges[$generated_source->id] as $to_id => $path) {
             if (!isset($this->nodes[$to_id])) {
@@ -415,11 +423,12 @@ private function getChildNodes(
                     $path = $this->getPredecessorPath($generated_source)
                     . ' -> ' . $this->getSuccessorPath($sink);
 
-                    foreach (TaintKindGroup::TAINT_TO_NAME as $matching_taint => $_) {
-                        if (!($matching_taints & $matching_taint)) {
+                    for ($x = $codebase->taint_count-1; $x >= 0; $x--) {
+                        $t = 1 << $x;
+                        if (!($matching_taints & $t)) {
                             continue;
                         }
-                        $issue = match ($matching_taint) {
+                        $issue = match ($t) {
                             TaintKind::INPUT_CALLABLE => new TaintedCallable(
                                 'Detected tainted text',
                                 $issue_location,
@@ -529,7 +538,7 @@ private function getChildNodes(
                                 $path,
                             ),
                             default => new TaintedCustom(
-                                'Detected tainted ' . $matching_taint,
+                                'Detected tainted ' . $codebase->custom_taints[$t],
                                 $issue_location,
                                 $issue_trace,
                                 $path,

src/Psalm/Internal/PreloaderList.php

@@ -1755,7 +1755,6 @@ final class PreloaderList {
         \Psalm\Type\MutableUnion::class,
         \Psalm\Type\Reconciler::class,
         \Psalm\Type\TaintKind::class,
-        \Psalm\Type\TaintKindGroup::class,
         \Psalm\Type\TypeNode::class,
         \Psalm\Type\TypeVisitor::class,
         \Psalm\Type\Union::class,

src/Psalm/Type/TaintKind.php

@@ -30,4 +30,9 @@ final class TaintKind
     public const INPUT_EXTRACT = (1 << 16);
     public const USER_SECRET = (1 << 17);
     public const SYSTEM_SECRET = (1 << 18);
+
+    public const ALL_INPUT = (1 << 17) - 1;
+
+    /** @internal Keep this synced with the above */
+    public const BUILTIN_TAINT_COUNT = 19;
 }