fix(iast): patch json.encoder to encode LazyTaintDict as dict [backport 2.1] (#7350)

Backport fee5ed0 from #7348 to 2.1.

IAST: Fix an issue where JSON encoder would fail at encoding a tainted
dict or list.

## Checklist
- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
are followed. If no release note is required, add label
`changelog/no-changelog`.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.
- [x] 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)
- [x] If this PR touches code that signs or publishes builds or
packages, or handles credentials of any kind, I've requested a review
from `@DataDog/security-design-and-guidance`.
- [x] This PR doesn't touch any of that.

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

Author: github-actions[bot]

ddtrace/appsec/_iast/_patches/json_tainting.py

@@ -32,6 +32,7 @@ def patch():
     if not set_and_check_module_is_patched("json", default_attr=_DEFAULT_ATTR):
         return
     try_wrap_function_wrapper("json", "loads", wrapped_loads)
+    try_wrap_function_wrapper("json.encoder", "JSONEncoder.default", patched_json_encoder_default)
 
 
 def wrapped_loads(wrapped, instance, args, kwargs):
@@ -60,3 +61,10 @@ def wrapped_loads(wrapped, instance, args, kwargs):
             log.debug("Unexpected exception while reporting vulnerability", exc_info=True)
             raise
     return obj
+
+
+def patched_json_encoder_default(original_func, instance, args, kwargs):
+    if isinstance(args[0], (LazyTaintList, LazyTaintDict)):
+        return args[0]._obj
+
+    return original_func(*args, **kwargs)

releasenotes/notes/iast-fix-lazy-taint-json-encoder-89cdd9aba0351d96.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    IAST: This fix resolves an issue where JSON encoder would throw an exception while encoding a tainted dict or list.

tests/appsec/iast/test_taint_utils.py

@@ -4,17 +4,20 @@
 
 try:
     from ddtrace.appsec._iast import oce
+    from ddtrace.appsec._iast._patch_modules import patch_iast
     from ddtrace.appsec._iast._taint_tracking import OriginType
     from ddtrace.appsec._iast._taint_tracking import create_context
     from ddtrace.appsec._iast._taint_tracking import is_pyobject_tainted
     from ddtrace.appsec._iast._taint_tracking import taint_pyobject
     from ddtrace.appsec._iast._taint_utils import LazyTaintDict
+    from ddtrace.appsec._iast._taint_utils import LazyTaintList
     from ddtrace.appsec._iast._taint_utils import check_tainted_args
 except (ImportError, AttributeError):
     pytest.skip("IAST not supported for this Python version", allow_module_level=True)
 
 
 def setup():
+    patch_iast()
     create_context()
     oce._enabled = True
 
@@ -234,3 +237,31 @@ def test_checked_tainted_args():
     assert check_tainted_args(
         args=(tainted_arg, untainted_arg), kwargs=None, tracer=None, integration_name="psycopg", method=cursor.execute
     )
+
+
+def test_json_encode_dict():
+    import json
+
+    tainted_dict = LazyTaintDict(
+        {
+            "tr_key_001": ["tr_val_001", "tr_val_002", "tr_val_003", {"tr_key_005": "tr_val_004"}],
+            "tr_key_002": {"tr_key_003": {"tr_key_004": "tr_val_005"}},
+        },
+        origins=(OriginType.PARAMETER, OriginType.PARAMETER),
+    )
+
+    assert json.dumps(tainted_dict) == (
+        '{"tr_key_001": ["tr_val_001", "tr_val_002", "tr_val_003", '
+        '{"tr_key_005": "tr_val_004"}], "tr_key_002": {"tr_key_003": {"tr_key_004": "tr_val_005"}}}'
+    )
+
+
+def test_json_encode_list():
+    import json
+
+    tainted_list = LazyTaintList(
+        ["tr_val_001", "tr_val_002", "tr_val_003", {"tr_key_005": "tr_val_004"}],
+        origins=(OriginType.PARAMETER, OriginType.PARAMETER),
+    )
+
+    assert json.dumps(tainted_list) == '["tr_val_001", "tr_val_002", "tr_val_003", {"tr_key_005": "tr_val_004"}]'