fix(taint): resolve direct lambda calls at call time

src/wardline/scanner/taint/variable_level.py

@@ -92,6 +92,10 @@
     "_CURRENT_VAR_TYPES", default=None
 )
 
+_CURRENT_LAMBDA_BINDINGS: contextvars.ContextVar[dict[str, ast.Lambda] | None] = contextvars.ContextVar(
+    "_CURRENT_LAMBDA_BINDINGS", default=None
+)
+
 _CURRENT_MODULE_PREFIX: contextvars.ContextVar[str | None] = contextvars.ContextVar(
     "_CURRENT_MODULE_PREFIX", default=None
 )
@@ -209,6 +213,47 @@ def _resolve_lambda_bodies(
         _resolve_expr(lam.body, function_taint, taint_map, scope)
 
 
+def _resolve_lambda_body_at_call(
+    lam: ast.Lambda,
+    function_taint: TaintState,
+    taint_map: dict[str, TaintState],
+    var_taints: dict[str, TaintState],
+    pos_taints: list[TaintState],
+    kw_taints: dict[str | None, TaintState],
+) -> None:
+    """Record a bound lambda body's call-site arg taints at a direct call site.
+
+    The final lambda-body pass is deliberately conservative for escaping/deferred
+    lambdas, but direct ``cb()`` calls have stronger ordering information: the
+    lambda executes with the current enclosing variable taints at this statement.
+    Recording the body here prevents a later clean assignment from laundering a
+    raw value that was captured when the direct call actually ran.
+    """
+    scope = dict(var_taints)
+    args = lam.args
+    positional_params = [*args.posonlyargs, *args.args]
+    for param, taint in zip(positional_params, pos_taints, strict=False):
+        scope[param.arg] = taint
+    for param in positional_params[len(pos_taints) :]:
+        scope[param.arg] = function_taint
+    if args.vararg is not None:
+        extra = pos_taints[len(positional_params) :]
+        scope[args.vararg.arg] = extra[0] if extra else function_taint
+        for taint in extra[1:]:
+            scope[args.vararg.arg] = combine(scope[args.vararg.arg], taint)
+    for param in args.kwonlyargs:
+        scope[param.arg] = kw_taints.get(param.arg, function_taint)
+    for param in positional_params:
+        if param.arg in kw_taints:
+            scope[param.arg] = combine(scope.get(param.arg, function_taint), kw_taints[param.arg])
+    if args.kwarg is not None:
+        keyword_values = [taint for name, taint in kw_taints.items() if name is not None]
+        scope[args.kwarg.arg] = keyword_values[0] if keyword_values else function_taint
+        for taint in keyword_values[1:]:
+            scope[args.kwarg.arg] = combine(scope[args.kwarg.arg], taint)
+    _resolve_expr(lam.body, function_taint, taint_map, scope)
+
+
 def compute_variable_taints(
     func_node: ast.FunctionDef | ast.AsyncFunctionDef,
     function_taint: TaintState,
@@ -255,6 +300,7 @@ def compute_variable_taints(
     if provenance_clash is not None:
         token_clash = _PROVENANCE_CLASH.set(provenance_clash)
     token_types = _CURRENT_VAR_TYPES.set({})
+    token_lambdas = _CURRENT_LAMBDA_BINDINGS.set({})
     if alias_map is not None:
         token = _CURRENT_ALIAS_MAP.set(alias_map)
     if call_site_arg_taints is not None:
@@ -272,6 +318,7 @@ def compute_variable_taints(
         if token_clash is not None:
             _PROVENANCE_CLASH.reset(token_clash)
         _CURRENT_VAR_TYPES.reset(token_types)
+        _CURRENT_LAMBDA_BINDINGS.reset(token_lambdas)
         if token is not None:
             _CURRENT_ALIAS_MAP.reset(token)
         if token_args is not None:
@@ -554,9 +601,14 @@ def _resolve_call(
         if isinstance(arg, ast.Starred):
             resolved_args[f"*{i}"] = t
     kw_taints = []
+    kw_arg_taints: dict[str | None, TaintState] = {}
     for kw in node.keywords:
         t = _resolve_expr(kw.value, function_taint, taint_map, var_taints)
         kw_taints.append(t)
+        if kw.arg in kw_arg_taints:
+            kw_arg_taints[kw.arg] = combine(kw_arg_taints[kw.arg], t)
+        else:
+            kw_arg_taints[kw.arg] = t
         if kw.arg in resolved_args:
             resolved_args[kw.arg] = combine(resolved_args[kw.arg], t)
         else:
@@ -565,7 +617,32 @@ def _resolve_call(
 
     call_site_arg_taints = _CURRENT_CALL_SITE_ARG_TAINTS.get()
     if call_site_arg_taints is not None:
-        call_site_arg_taints[id(node)] = resolved_args
+        existing = call_site_arg_taints.get(id(node))
+        if existing is None:
+            call_site_arg_taints[id(node)] = resolved_args
+        else:
+            for key, taint in resolved_args.items():
+                existing[key] = combine(existing[key], taint) if key in existing else taint
+
+    lambda_bindings = _CURRENT_LAMBDA_BINDINGS.get()
+    if isinstance(node.func, ast.Name) and lambda_bindings is not None and node.func.id in lambda_bindings:
+        _resolve_lambda_body_at_call(
+            lambda_bindings[node.func.id],
+            function_taint,
+            taint_map,
+            var_taints,
+            pos_taints,
+            kw_arg_taints,
+        )
+    elif isinstance(node.func, ast.Lambda):
+        _resolve_lambda_body_at_call(
+            node.func,
+            function_taint,
+            taint_map,
+            var_taints,
+            pos_taints,
+            kw_arg_taints,
+        )
 
     alias_map = _CURRENT_ALIAS_MAP.get()
     imported_fqn: str | None = None
@@ -708,6 +785,13 @@ def _process_stmt(
         if isinstance(stmt.target, ast.Name):
             taint = _resolve_expr(value, function_taint, taint_map, var_taints)
             var_taints[stmt.target.id] = taint
+
+            lambda_bindings = _CURRENT_LAMBDA_BINDINGS.get()
+            if lambda_bindings is not None:
+                if isinstance(value, ast.Lambda):
+                    lambda_bindings[stmt.target.id] = value
+                else:
+                    lambda_bindings.pop(stmt.target.id, None)
         elif isinstance(stmt.target, (ast.Subscript, ast.Attribute)):
             # Annotated container/attribute write (``self.x: str = expr``,
             # ``d['k']: T = expr``) — same fail-open shape as the plain-Assign Part
@@ -808,6 +892,13 @@ def _handle_assign(
             )
             var_taints[target.id] = taint
 
+            lambda_bindings = _CURRENT_LAMBDA_BINDINGS.get()
+            if lambda_bindings is not None:
+                if isinstance(stmt.value, ast.Lambda):
+                    lambda_bindings[target.id] = stmt.value
+                else:
+                    lambda_bindings.pop(target.id, None)
+
             var_types = _CURRENT_VAR_TYPES.get()
             if var_types is not None:
                 alias_map = _CURRENT_ALIAS_MAP.get()

tests/unit/scanner/rules/test_sink_rules.py

@@ -235,6 +235,64 @@ def f(p):
     assert not any(str(w.message).startswith("WLN-ENGINE-FLOW-INSENSITIVE-FALLBACK") for w in caught)
 
 
+def test_107_lambda_called_with_raw_arg_fires(tmp_path) -> None:
+    # The load-bearing case this change fixes: a raw value is passed as a CALL
+    # ARGUMENT bound to the lambda's parameter, and the sink consumes the param.
+    # The worst-ever pass resets lambda params to neutral, so it cannot see the
+    # argument taint — only call-time resolution can. Silent false-negative on
+    # the prior engine (no fallback warning), so this test genuinely guards it.
+    ctx = _analyze(
+        tmp_path,
+        """
+        @trusted(level='ASSURED')
+        def f(p):
+            raw = read_raw(p)
+            cb = lambda x: eval(x)
+            cb(raw)
+        """,
+    )
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        findings = UntrustedToExec().check(ctx)
+    assert [(x.rule_id, x.qualname) for x in findings] == [("PY-WL-107", "m.f")]
+    assert not any(str(w.message).startswith("WLN-ENGINE-FLOW-INSENSITIVE-FALLBACK") for w in caught)
+
+
+def test_107_inline_lambda_call_with_raw_arg_fires(tmp_path) -> None:
+    # Same defect via an inline (un-named) lambda invoked immediately with a raw
+    # argument: ``(lambda x: eval(x))(raw)``. Also a silent FN on the prior engine.
+    ctx = _analyze(
+        tmp_path,
+        """
+        @trusted(level='ASSURED')
+        def f(p):
+            raw = read_raw(p)
+            (lambda x: eval(x))(raw)
+        """,
+    )
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        findings = UntrustedToExec().check(ctx)
+    assert [(x.rule_id, x.qualname) for x in findings] == [("PY-WL-107", "m.f")]
+    assert not any(str(w.message).startswith("WLN-ENGINE-FLOW-INSENSITIVE-FALLBACK") for w in caught)
+
+
+def test_107_lambda_called_with_clean_arg_does_not_fire(tmp_path) -> None:
+    # Negative guard: call-time argument binding must NOT manufacture a finding
+    # when the argument is a clean literal. Protects against the call-time path
+    # over-firing (false positive).
+    ctx = _analyze(
+        tmp_path,
+        """
+        @trusted(level='ASSURED')
+        def f(p):
+            cb = lambda x: eval(x)
+            cb('ls -la')
+        """,
+    )
+    assert UntrustedToExec().check(ctx) == []
+
+
 def test_107_raw_assigned_after_lambda_def_still_fires(tmp_path) -> None:
     # Closure-by-reference soundness: the lambda is DEFINED while ``src`` is clean,
     # but ``src`` is reassigned raw before the lambda is called — at runtime eval()