fix: track assignment aliases in component code scanner (#14041)

fix: bind loop and comprehension aliases

Author: erichare

src/backend/base/langflow/agentic/helpers/code_security.py

@@ -306,6 +306,27 @@ def _bind_assignment_target(self, target: ast.AST, value: ast.AST) -> None:
             for target_element in target.elts:
                 self._bind_assignment_target(target_element, value)
 
+    def _bind_iterated_target(self, target: ast.AST, iterable: ast.AST) -> None:
+        """Bind a loop target to every statically visible iterable value."""
+        if isinstance(iterable, (ast.List, ast.Tuple, ast.Set)):
+            values = iterable.elts
+        elif isinstance(iterable, ast.Dict):
+            values = [key for key in iterable.keys if key is not None]
+        else:
+            values = [iterable]
+
+        before_iteration = self._snapshot_alias_state()
+        iteration_states: list[_AliasState] = []
+        for value in values:
+            self._restore_alias_state(before_iteration)
+            self._bind_assignment_target(target, value)
+            iteration_states.append(self._snapshot_alias_state())
+
+        if iteration_states:
+            self._merge_alias_states(iteration_states)
+        else:
+            self._bind_assignment_target(target, iterable)
+
     def _snapshot_alias_state(self) -> _AliasState:
         return self.module_aliases.copy(), self.shadowed_aliases.copy()
 
@@ -314,6 +335,28 @@ def _restore_alias_state(self, state: _AliasState) -> None:
         self.module_aliases = aliases.copy()
         self.shadowed_aliases = shadowed.copy()
 
+    def _restore_target_names(self, target_names: set[str], state: _AliasState) -> None:
+        aliases, shadowed = state
+        for name in target_names:
+            if name in aliases:
+                self.module_aliases[name] = aliases[name]
+                self.shadowed_aliases.discard(name)
+            elif name in shadowed:
+                self.module_aliases.pop(name, None)
+                self.shadowed_aliases.add(name)
+            else:
+                self.module_aliases.pop(name, None)
+                self.shadowed_aliases.discard(name)
+
+    def _assignment_target_names(self, target: ast.AST) -> set[str]:
+        if isinstance(target, ast.Name):
+            return {target.id}
+        if isinstance(target, ast.Starred):
+            return self._assignment_target_names(target.value)
+        if isinstance(target, (ast.Tuple, ast.List)):
+            return set().union(*(self._assignment_target_names(element) for element in target.elts))
+        return set()
+
     def _merge_alias_states(self, states: list[_AliasState]) -> None:
         """Conservatively retain every module value reachable from a branch."""
         names = set().union(*(set(aliases) | shadowed for aliases, shadowed in states))
@@ -414,6 +457,7 @@ def _visit_loop(self, node: ast.For | ast.AsyncFor | ast.While) -> None:
         else:
             self.visit(node.iter)
             self.visit(node.target)
+            self._bind_iterated_target(node.target, node.iter)
 
         for statement in node.body:
             self.visit(statement)
@@ -437,6 +481,35 @@ def visit_AsyncFor(self, node: ast.AsyncFor):
     def visit_While(self, node: ast.While):
         self._visit_loop(node)
 
+    def _visit_comprehension_expression(
+        self, generators: list[ast.comprehension], expressions: tuple[ast.AST, ...]
+    ) -> None:
+        """Visit comprehensions in evaluation order with isolated target bindings."""
+        enclosing_state = self._snapshot_alias_state()
+        target_names: set[str] = set()
+        for generator in generators:
+            self.visit(generator.iter)
+            self.visit(generator.target)
+            self._bind_iterated_target(generator.target, generator.iter)
+            target_names.update(self._assignment_target_names(generator.target))
+            for condition in generator.ifs:
+                self.visit(condition)
+        for expression in expressions:
+            self.visit(expression)
+        self._restore_target_names(target_names, enclosing_state)
+
+    def visit_ListComp(self, node: ast.ListComp):
+        self._visit_comprehension_expression(node.generators, (node.elt,))
+
+    def visit_SetComp(self, node: ast.SetComp):
+        self._visit_comprehension_expression(node.generators, (node.elt,))
+
+    def visit_DictComp(self, node: ast.DictComp):
+        self._visit_comprehension_expression(node.generators, (node.key, node.value))
+
+    def visit_GeneratorExp(self, node: ast.GeneratorExp):
+        self._visit_comprehension_expression(node.generators, (node.elt,))
+
     def _shadow_arguments(self, arguments: ast.arguments) -> None:
         positional = [*arguments.posonlyargs, *arguments.args, *arguments.kwonlyargs]
         for argument in positional:

src/backend/tests/unit/agentic/helpers/test_code_security.py

@@ -620,6 +620,38 @@ def test_should_allow_parameter_shadowing_module_name(self):
         result = scan_code_security("import os\ndef use_safe_object(os):\n    os.system('not the os module')")
         assert result.is_safe is True
 
+    @pytest.mark.parametrize(
+        "code",
+        [
+            "import os\nfor module in (os,):\n    module.system('id')",
+            "import os\nasync def run():\n    async for module in (os,):\n        module.system('id')",
+            "import os\n[module.system('id') for module in (os,)]",
+            "import os\n{module.system('id') for module in (os,)}",
+            "import os\n{module: module.system('id') for module in (os,)}",
+            "import os\n(module.system('id') for module in (os,))",
+        ],
+    )
+    def test_should_detect_iterated_module_alias_call(self, code):
+        result = scan_code_security(code)
+        assert result.is_safe is False
+        assert any("os.system()" in violation for violation in result.violations)
+
+    def test_should_detect_destructured_loop_target_alias_call(self):
+        result = scan_code_security("import os\nfor (module,) in ((os,),):\n    module.system('id')")
+        assert result.is_safe is False
+        assert any("os.system()" in violation for violation in result.violations)
+
+    def test_should_not_leak_comprehension_target_alias(self):
+        code = "import os\nmodule = object()\n[module for module in (os,)]\nmodule.system('not the os module')"
+        result = scan_code_security(code)
+        assert result.is_safe is True
+
+    def test_should_preserve_named_expression_alias_from_comprehension(self):
+        code = "import os\n[(module := os) for _ in (None,)]\nmodule.system('id')"
+        result = scan_code_security(code)
+        assert result.is_safe is False
+        assert any("os.system()" in violation for violation in result.violations)
+
     def test_should_detect_alias_after_zero_iteration_for_loop(self):
         code = """
 import os