chore(iast): iast code injection error with locals and globals (#11996)

## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

---------

Co-authored-by: Federico Mon <[email protected]>

ddtrace/appsec/_iast/taint_sinks/code_injection.py

@@ -13,11 +13,10 @@
 from ddtrace.appsec._iast._patch import try_wrap_function_wrapper
 from ddtrace.appsec._iast._taint_tracking._taint_objects import is_pyobject_tainted
 from ddtrace.appsec._iast.constants import VULN_CODE_INJECTION
+from ddtrace.appsec._iast.taint_sinks._base import VulnerabilityBase
 from ddtrace.internal.logger import get_logger
 from ddtrace.settings.asm import config as asm_config
 
-from ._base import VulnerabilityBase
-
 
 log = get_logger(__name__)
 
@@ -52,29 +51,27 @@ def _iast_coi(wrapped, instance, args, kwargs):
     if asm_config._iast_enabled and len(args) >= 1:
         _iast_report_code_injection(args[0])
 
-    return wrapped(*args, **kwargs)
-
-
-def _iast_coi_exec(wrapped, instance, args, kwargs):
-    if asm_config._iast_enabled and len(args) >= 1:
-        _iast_report_code_injection(args[0])
-
-    caller_frame = inspect.currentframe().f_back.f_back
-    if caller_frame is None:
-        return wrapped(*args, **kwargs)
-
-    caller_globals = caller_frame.f_globals
-    caller_locals = caller_frame.f_locals
-
-    original_globals = {}
+    caller_frame = None
     if len(args) > 1:
-        original_globals = args[1]
+        func_globals = args[1]
+    elif kwargs.get("globals"):
+        func_globals = kwargs.get("globals")
+    else:
+        frames = inspect.currentframe()
+        caller_frame = frames.f_back
+        func_globals = caller_frame.f_globals
 
-    original_locals = {}
     if len(args) > 2:
-        original_locals = args[2]
-
-    return wrapped(args[0], {**caller_globals, **original_globals}, {**caller_locals, **original_locals})
+        func_locals = args[2]
+    elif kwargs.get("locals"):
+        func_locals = kwargs.get("locals")
+    else:
+        if caller_frame is None:
+            frames = inspect.currentframe()
+            caller_frame = frames.f_back
+        func_locals = caller_frame.f_locals
+
+    return wrapped(args[0], func_globals, func_locals)
 
 
 @oce.register

tests/appsec/iast/fixtures/taint_sinks/code_injection.py

@@ -10,6 +10,34 @@ def pt_eval(origin_string):
     return r
 
 
+def pt_eval_globals(origin_string):
+    context = {"x": 5, "y": 10}
+    r = eval(origin_string, context)
+    return r
+
+
+def pt_eval_globals_locals(origin_string):
+    z = 15  # noqa: F841
+    globals_dict = {"x": 10}
+    locals_dict = {"y": 20}
+    r = eval(origin_string, globals_dict, locals_dict)
+    return r
+
+
+def pt_eval_lambda(fun):
+    return eval("lambda v,fun=fun:not fun(v)")
+
+
+def is_true(value):
+    return value is True
+
+
+def pt_eval_lambda_globals(origin_string):
+    globals_dict = {"square": lambda x: x * x}
+    r = eval(origin_string, globals=globals_dict)
+    return r
+
+
 def pt_literal_eval(origin_string):
     r = literal_eval(origin_string)
     return r

tests/appsec/iast/taint_sinks/test_code_injection.py

@@ -1,5 +1,7 @@
 import os
 
+import pytest
+
 from ddtrace.appsec._iast._taint_tracking import OriginType
 from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
 from ddtrace.appsec._iast.constants import VULN_CODE_INJECTION
@@ -9,10 +11,10 @@
 
 
 ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+mod = _iast_patched_module("tests.appsec.iast.fixtures.taint_sinks.code_injection")
 
 
 def test_code_injection_eval(iast_context_defaults):
-    mod = _iast_patched_module("tests.appsec.iast.fixtures.taint_sinks.code_injection")
     code_string = '"abc" + "def"'
 
     tainted_string = taint_pyobject(
@@ -35,53 +37,120 @@ def test_code_injection_eval(iast_context_defaults):
     assert vulnerability["evidence"].get("redacted") is None
 
 
-# TODO: wrap exec functions is very dangerous because it needs and modifies locals and globals from the original func
-# def test_code_injection_exec(iast_context_defaults):
-#     mod = _iast_patched_module("tests.appsec.iast.fixtures.taint_sinks.code_injection")
-#     code_string = '"abc" + "def"'
-#
-#     tainted_string = taint_pyobject(
-#         code_string, source_name="path", source_value=code_string, source_origin=OriginType.PATH
-#     )
-#     mod.pt_exec(tainted_string)
-#
-#     data = _get_iast_data()
-#
-#     assert len(data["vulnerabilities"]) == 1
-#     vulnerability = data["vulnerabilities"][0]
-#     source = data["sources"][0]
-#     assert vulnerability["type"] == VULN_CODE_INJECTION
-#     assert source["name"] == "path"
-#     assert source["origin"] == OriginType.PATH
-#     assert source["value"] == '"abc" + "def"'
-#     assert vulnerability["evidence"]["valueParts"] == [{"source": 0, "value": '"abc" + "def"'}]
-#     assert "value" not in vulnerability["evidence"].keys()
-#     assert vulnerability["evidence"].get("pattern") is None
-#     assert vulnerability["evidence"].get("redacted") is None
-#
-#
-# def test_code_injection_exec_with_globals(iast_context_defaults):
-#     mod = _iast_patched_module("tests.appsec.iast.fixtures.taint_sinks.code_injection")
-#     code_string = 'my_var_in_pt_exec_with_globals + "-" + my_var_in_pt_exec_with_globals + "-"'
-#
-#     tainted_string = taint_pyobject(
-#         code_string, source_name="path", source_value=code_string, source_origin=OriginType.PATH
-#     )
-#     mod.pt_exec_with_globals(tainted_string)
-#
-#     data = _get_iast_data()
-#
-#     assert len(data["vulnerabilities"]) == 1
-#     vulnerability = data["vulnerabilities"][0]
-#     source = data["sources"][0]
-#     assert vulnerability["type"] == VULN_CODE_INJECTION
-#     assert source["name"] == "path"
-#     assert source["origin"] == OriginType.PATH
-#     assert source["value"] == '"abc" + "def"'
-#     assert vulnerability["evidence"]["valueParts"] == [{"source": 0, "value": '"abc" + "def"'}]
-#     assert "value" not in vulnerability["evidence"].keys()
-#     assert vulnerability["evidence"].get("pattern") is None
-#     assert vulnerability["evidence"].get("redacted") is None
+def test_code_injection_eval_globals(iast_context_defaults):
+    """Validate globals and locals of the function"""
+
+    code_string = "x + y"
+
+    tainted_string = taint_pyobject(
+        code_string, source_name="path", source_value=code_string, source_origin=OriginType.PATH
+    )
+    mod.pt_eval_globals(tainted_string)
+
+    data = _get_iast_data()
+
+    assert len(data["vulnerabilities"]) == 1
+    vulnerability = data["vulnerabilities"][0]
+    source = data["sources"][0]
+    assert vulnerability["type"] == VULN_CODE_INJECTION
+    assert source["name"] == "path"
+    assert source["origin"] == OriginType.PATH
+    assert source["value"] == "x + y"
+    assert vulnerability["evidence"]["valueParts"] == [{"source": 0, "value": "x + y"}]
+    assert "value" not in vulnerability["evidence"].keys()
+    assert vulnerability["evidence"].get("pattern") is None
+    assert vulnerability["evidence"].get("redacted") is None
+
+
+def test_code_injection_eval_globals_locals(iast_context_defaults):
+    """Validate globals and locals of the function"""
+
+    code_string = "x + y"
+
+    tainted_string = taint_pyobject(
+        code_string, source_name="path", source_value=code_string, source_origin=OriginType.PATH
+    )
+
+    mod.pt_eval_globals_locals(tainted_string)
+
+    data = _get_iast_data()
+
+    assert len(data["vulnerabilities"]) == 1
+    vulnerability = data["vulnerabilities"][0]
+    source = data["sources"][0]
+    assert vulnerability["type"] == VULN_CODE_INJECTION
+    assert source["name"] == "path"
+    assert source["origin"] == OriginType.PATH
+    assert source["value"] == "x + y"
+    assert vulnerability["evidence"]["valueParts"] == [{"source": 0, "value": "x + y"}]
+    assert "value" not in vulnerability["evidence"].keys()
+    assert vulnerability["evidence"].get("pattern") is None
+    assert vulnerability["evidence"].get("redacted") is None
+
+
+def test_code_injection_eval_globals_locals_override(iast_context_defaults):
+    """Validate globals and locals of the function"""
+
+    code_string = "x + y + z"
+
+    tainted_string = taint_pyobject(
+        code_string, source_name="path", source_value=code_string, source_origin=OriginType.PATH
+    )
+    with pytest.raises(NameError):
+        mod.pt_eval_globals_locals(tainted_string)
+
+    data = _get_iast_data()
+
+    assert len(data["vulnerabilities"]) == 1
+    vulnerability = data["vulnerabilities"][0]
+    source = data["sources"][0]
+    assert vulnerability["type"] == VULN_CODE_INJECTION
+    assert source["name"] == "path"
+    assert source["origin"] == OriginType.PATH
+    assert source["value"] == "x + y + z"
+    assert vulnerability["evidence"]["valueParts"] == [{"source": 0, "value": "x + y + z"}]
+    assert "value" not in vulnerability["evidence"].keys()
+    assert vulnerability["evidence"].get("pattern") is None
+    assert vulnerability["evidence"].get("redacted") is None
+
+
+def test_code_injection_eval_lambda(iast_context_defaults):
+    """Validate globals and locals of the function"""
+    mod = _iast_patched_module("tests.appsec.iast.fixtures.taint_sinks.code_injection")
+
+    def pt_eval_lambda_no_tainted(fun):
+        return eval("lambda v,fun=fun:not fun(v)")
+
+    def is_true_no_tainted(value):
+        return value is True
+
+    assert mod.pt_eval_lambda(mod.is_true)(True) is pt_eval_lambda_no_tainted(is_true_no_tainted)(True)
+
+
+def test_code_injection_eval_globals_kwargs_lambda(iast_context_defaults):
+    """Validate globals and locals of the function"""
+
+    code_string = "square(5)"
+
+    tainted_string = taint_pyobject(
+        code_string, source_name="path", source_value=code_string, source_origin=OriginType.PATH
+    )
+
+    mod.pt_eval_lambda_globals(tainted_string)
+
+    data = _get_iast_data()
+
+    assert len(data["vulnerabilities"]) == 1
+    vulnerability = data["vulnerabilities"][0]
+    source = data["sources"][0]
+    assert vulnerability["type"] == VULN_CODE_INJECTION
+    assert source["name"] == "path"
+    assert source["origin"] == OriginType.PATH
+    assert source["value"] == "square(5)"
+    assert vulnerability["evidence"]["valueParts"] == [{"source": 0, "value": "square(5)"}]
+    assert "value" not in vulnerability["evidence"].keys()
+    assert vulnerability["evidence"].get("pattern") is None
+    assert vulnerability["evidence"].get("redacted") is None
 
 
 def test_code_injection_literal_eval(iast_context_defaults):