restore commented code

DataDog · Sep 20, 2024 · ab30619 · ab30619
1 parent b546ee8
commit ab30619
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 19 deletions.
diff --git a/ddtrace/appsec/_common_module_patches.py b/ddtrace/appsec/_common_module_patches.py
@@ -31,8 +31,8 @@
 def patch_common_modules():
     try_wrap_function_wrapper("builtins", "open", wrapped_open_CFDDB7ABBA9081B6)
     try_wrap_function_wrapper("urllib.request", "OpenerDirector.open", wrapped_open_ED4CF71136E15EBF)
-    # try_wrap_function_wrapper("_io", "BytesIO.read", wrapped_read_F3E51D71B4EC16EF)
-    # try_wrap_function_wrapper("_io", "StringIO.read", wrapped_read_F3E51D71B4EC16EF)
+    try_wrap_function_wrapper("_io", "BytesIO.read", wrapped_read_F3E51D71B4EC16EF)
+    try_wrap_function_wrapper("_io", "StringIO.read", wrapped_read_F3E51D71B4EC16EF)
     try_wrap_function_wrapper("os", "system", wrapped_system_5542593D237084A7)
     core.on("asm.block.dbapi.execute", execute_4C9BAC8E228EB347)
     if asm_config._iast_enabled:
@@ -42,8 +42,8 @@ def patch_common_modules():
 def unpatch_common_modules():
     try_unwrap("builtins", "open")
     try_unwrap("urllib.request", "OpenerDirector.open")
-    # try_unwrap("_io", "BytesIO.read")
-    # try_unwrap("_io", "StringIO.read")
+    try_unwrap("_io", "BytesIO.read")
+    try_unwrap("_io", "StringIO.read")
 
 
 def wrapped_read_F3E51D71B4EC16EF(original_read_callable, instance, args, kwargs):

diff --git a/ddtrace/appsec/_constants.py b/ddtrace/appsec/_constants.py
@@ -2,8 +2,8 @@
 from re import Match
 import sys
 
-# from _io import BytesIO
-# from _io import StringIO
+from _io import BytesIO
+from _io import StringIO
 
 
 if sys.version_info >= (3, 8):
@@ -125,8 +125,7 @@ class IAST(metaclass=Constant_Class):
     SEP_MODULES: Literal[","] = ","
     REQUEST_IAST_ENABLED: Literal["_dd.iast.request_enabled"] = "_dd.iast.request_enabled"
     TEXT_TYPES = (str, bytes, bytearray)
-    # TAINTEABLE_TYPES = (str, bytes, bytearray, Match, BytesIO, StringIO)
-    TAINTEABLE_TYPES = (str, bytes, bytearray, Match)
+    TAINTEABLE_TYPES = (str, bytes, bytearray, Match, BytesIO, StringIO)
 
 
 class IAST_SPAN_TAGS(metaclass=Constant_Class):

diff --git a/ddtrace/appsec/_iast/_ast/visitor.py b/ddtrace/appsec/_iast/_ast/visitor.py
@@ -40,16 +40,16 @@ def _mark_avoid_convert_recursively(node):
     "definitions_module": "ddtrace.appsec._iast._taint_tracking.aspects",
     "alias_module": "ddtrace_aspects",
     "functions": {
-        # "StringIO": "ddtrace_aspects.stringio_aspect",
-        # "BytesIO": "ddtrace_aspects.bytesio_aspect",
+        "StringIO": "ddtrace_aspects.stringio_aspect",
+        "BytesIO": "ddtrace_aspects.bytesio_aspect",
         "str": "ddtrace_aspects.str_aspect",
         "bytes": "ddtrace_aspects.bytes_aspect",
         "bytearray": "ddtrace_aspects.bytearray_aspect",
         "ddtrace_iast_flask_patch": "ddtrace_aspects.empty_func",  # To avoid recursion
     },
     "stringalike_methods": {
-        # "StringIO": "ddtrace_aspects.stringio_aspect",
-        # "BytesIO": "ddtrace_aspects.bytesio_aspect",
+        "StringIO": "ddtrace_aspects.stringio_aspect",
+        "BytesIO": "ddtrace_aspects.bytesio_aspect",
         "decode": "ddtrace_aspects.decode_aspect",
         "join": "ddtrace_aspects.join_aspect",
         "encode": "ddtrace_aspects.encode_aspect",

diff --git a/ddtrace/appsec/_iast/_taint_tracking/Utils/StringUtils.h b/ddtrace/appsec/_iast/_taint_tracking/Utils/StringUtils.h
@@ -50,7 +50,7 @@ is_text(const PyObject* pyptr)
 inline bool
 is_tainteable(const PyObject* pyptr)
 {
-    return pyptr != nullptr and (is_text(pyptr) or PyReMatch_Check(pyptr));
+    return pyptr != nullptr; // and (is_text(pyptr) or PyReMatch_Check(pyptr));
 }
 
 // Base function for the variadic template

diff --git a/ddtrace/appsec/_iast/_taint_tracking/__init__.py b/ddtrace/appsec/_iast/_taint_tracking/__init__.py
@@ -1,5 +1,5 @@
-# from io import BytesIO
-# from io import StringIO
+from io import BytesIO
+from io import StringIO
 import os
 from typing import Any
 from typing import Tuple
@@ -222,11 +222,9 @@ def trace_calls_and_returns(frame, event, arg):
             if frame in TAINTED_FRAMES:
                 TAINTED_FRAMES.remove(frame)
                 log.debug("Return from %s on line %d of %s, return value: %s", func_name, line_no, filename, arg)
-                # if isinstance(arg, (str, bytes, bytearray, BytesIO, StringIO, list, tuple, dict)):
-                if isinstance(arg, (str, bytes, bytearray, list, tuple, dict)):
+                if isinstance(arg, (str, bytes, bytearray, BytesIO, StringIO, list, tuple, dict)):
                     if (
-                        # (isinstance(arg, (str, bytes, bytearray, BytesIO, StringIO)) and is_pyobject_tainted(arg))
-                        (isinstance(arg, (str, bytes, bytearray)) and is_pyobject_tainted(arg))
+                        (isinstance(arg, (str, bytes, bytearray, BytesIO, StringIO)) and is_pyobject_tainted(arg))
                         or (isinstance(arg, (list, tuple)) and any([is_pyobject_tainted(x) for x in arg]))
                         or (isinstance(arg, dict) and any([is_pyobject_tainted(x) for x in arg.values()]))
                     ):