chore(iast): add extra pyobject checks on is_text and new_pyobject_id

ddtrace/appsec/_iast/_taint_tracking/Aspects/AspectIndex.cpp

@@ -67,6 +67,10 @@ api_index_aspect(PyObject* self, PyObject* const* args, const Py_ssize_t nargs)
         return nullptr;
     }
 
+    if (args == nullptr) {
+        return nullptr;
+    }
+
     PyObject* candidate_text = args[0];
     PyObject* idx = args[1];
     const auto result_o = PyObject_GetItem(candidate_text, idx);

ddtrace/appsec/_iast/_taint_tracking/Utils/StringUtils.cpp

@@ -69,6 +69,14 @@ new_pyobject_id(PyObject* tainted_object)
     if (!tainted_object)
         return nullptr;
 
+    // Check that it's aligned correctly
+    if (reinterpret_cast<uintptr_t>(tainted_object) % alignof(PyObject) != 0) return tainted_object;
+
+    // Try to safely access ob_type
+    if (const PyObject* temp = tainted_object;!temp->ob_type) return tainted_object;
+
+    py::gil_scoped_acquire acquire;
+
     if (PyUnicode_Check(tainted_object)) {
         PyObject* empty_unicode = PyUnicode_New(0, 127);
         if (!empty_unicode)
@@ -86,6 +94,7 @@ new_pyobject_id(PyObject* tainted_object)
         Py_XDECREF(val);
         return result;
     }
+
     if (PyBytes_Check(tainted_object)) {
         PyObject* empty_bytes = PyBytes_FromString("");
         if (!empty_bytes)
@@ -102,7 +111,9 @@ new_pyobject_id(PyObject* tainted_object)
         Py_XDECREF(val);
         Py_XDECREF(empty_bytes);
         return res;
-    } else if (PyByteArray_Check(tainted_object)) {
+    }
+
+    if (PyByteArray_Check(tainted_object)) {
         PyObject* empty_bytes = PyBytes_FromString("");
         if (!empty_bytes)
             return tainted_object;

ddtrace/appsec/_iast/_taint_tracking/Utils/StringUtils.h

@@ -47,9 +47,20 @@ set_fast_tainted_if_notinterned_unicode(PyObject* objptr);
 inline bool
 is_text(const PyObject* pyptr)
 {
-    return (pyptr != nullptr) and (PyUnicode_Check(pyptr) or PyBytes_Check(pyptr) or PyByteArray_Check(pyptr));
+    if (pyptr == nullptr) {
+        return false;
+    };
+
+    // Check that it's aligned correctly
+    if (reinterpret_cast<uintptr_t>(pyptr) % alignof(PyObject) != 0) return false;;
+
+    // Try to safely access ob_type
+    if (const PyObject* temp = pyptr;!temp->ob_type) return false;
+
+    return PyUnicode_Check(pyptr) or PyBytes_Check(pyptr) or PyByteArray_Check(pyptr);
 }
 
+
 inline bool
 is_tainteable(const PyObject* pyptr)
 {

ddtrace/appsec/_iast/_taint_tracking/clean.sh

@@ -3,7 +3,7 @@ set -exu
 #cd -- "$(dirname -- "${BASH_SOURCE[0]}")" || exit
 
 rm -rf CMakeFiles/ CMakeCache.txt Makefile cmake_install.cmake __pycache__/ .cmake *.cbp Testing
-rm -rf tests/CMakeFiles/ tests/CMakeCache.txt tests/Makefile tests/cmake_install.cmake tests/__pycache__/ tests/.cmake *.cbp 
+rm -rf tests/CMakeFiles/ tests/CMakeCache.txt tests/Makefile tests/cmake_install.cmake tests/__pycache__/ tests/*.cmake *.cbp 
 rm -rf cmake-build-debug cmake-build-default cmake-build-tests
 rm -rf tests/cmake-build-debug tests/cmake-build-default tests/cmake-build-tests
 rm -rf tests/CMakeFiles/native_tests.dir
@@ -12,3 +12,4 @@ rm -rf coverage_html
 yes|rm -f *.so
 yes|rm -f tests/*.so
 yes|rm -f tests/native_tests
+yes|rm -f gcov_reports/*.gcov

ddtrace/appsec/_iast/_taint_tracking/tests/test_index_aspect.cpp

@@ -0,0 +1,43 @@
+#include "Aspects/Helpers.h"
+
+#include <Aspects/AspectIndex.h>
+#include <Initializer/Initializer.h>
+#include <tests/test_common.hpp>
+
+using AspectIndexCheck = PyEnvWithContext;
+
+TEST_F(AspectIndexCheck, check_index_internal_all_nullptr)
+{
+    const TaintRangeRefs refs;
+    index_aspect(nullptr, nullptr, nullptr, refs, Initializer::get_tainting_map());
+}
+
+TEST_F(AspectIndexCheck, check_index_internal_all_nullptr_negative_index)
+{
+    PyObject* idx = PyLong_FromLong(-1);
+    const TaintRangeRefs refs;
+    auto ret = index_aspect(nullptr, nullptr, idx, refs, Initializer::get_tainting_map());
+    EXPECT_EQ(ret, nullptr);
+    Py_DecRef(idx);
+}
+
+TEST_F(AspectIndexCheck, check_api_index_aspect_all_nullptr)
+{
+    auto ret = api_index_aspect(nullptr, nullptr, 2);
+    EXPECT_EQ(ret, nullptr);
+}
+
+TEST_F(AspectIndexCheck, check_api_index_aspect_wrong_index)
+{
+    PyObject* py_str = PyUnicode_FromString("abc");
+    PyObject* idx = PyLong_FromLong(4);
+    PyObject* args_array[2];
+    args_array[0] = py_str;
+    args_array[1] = idx;
+    auto res = api_index_aspect(nullptr, args_array, 2);
+    ASSERT_EQ(res, nullptr);
+    EXPECT_EQ(has_pyerr_as_string(), std::string("string index out of range"));
+    PyErr_Clear();
+    Py_DecRef(py_str);
+    Py_DecRef(idx);
+}

ddtrace/appsec/_iast/_taint_tracking/tests/test_other.cpp

@@ -0,0 +1,100 @@
+#include <gtest/gtest.h>
+#include <Python.h>
+#include <tests/test_common.hpp>
+#include <iostream>  // jjj
+
+using PyByteArray_Check = PyEnvCheck;
+
+TEST_F(PyByteArray_Check, WithOtherTypes)
+{
+    PyObject* non_text_obj = PyLong_FromLong(42);
+    EXPECT_FALSE(PyByteArray_Check(non_text_obj));
+    // Test with None
+    PyObject* none_obj = Py_None;
+    Py_INCREF(none_obj);
+    EXPECT_FALSE(PyByteArray_Check(none_obj)) << "Failed for None type";
+    Py_DECREF(none_obj);
+
+    // Test with Bool
+    PyObject* bool_obj = Py_True;
+    Py_INCREF(bool_obj);
+    EXPECT_FALSE(PyByteArray_Check(bool_obj)) << "Failed for Bool type";
+    Py_DECREF(bool_obj);
+
+    // Test with Integer
+    PyObject* int_obj = PyLong_FromLong(42);
+    EXPECT_FALSE(PyByteArray_Check(int_obj)) << "Failed for Integer type";
+    Py_DECREF(int_obj);
+
+    // Test with Float
+    PyObject* float_obj = PyFloat_FromDouble(3.14);
+    EXPECT_FALSE(PyByteArray_Check(float_obj)) << "Failed for Float type";
+    Py_DECREF(float_obj);
+
+    // Test with String
+    PyObject* str_obj = PyUnicode_FromString("test");
+    EXPECT_FALSE(PyByteArray_Check(str_obj)) << "Failed for String type";
+    Py_DECREF(str_obj);
+
+    // Test with Bytes
+    PyObject* bytes_obj = PyBytes_FromString("test");
+    EXPECT_FALSE(PyByteArray_Check(bytes_obj)) << "Failed for Bytes type";
+    Py_DECREF(bytes_obj);
+
+    // Test with Tuple
+    PyObject* tuple_obj = PyTuple_New(0);
+    EXPECT_FALSE(PyByteArray_Check(tuple_obj)) << "Failed for Tuple type";
+    Py_DECREF(tuple_obj);
+
+    // Test with List
+    PyObject* list_obj = PyList_New(0);
+    EXPECT_FALSE(PyByteArray_Check(list_obj)) << "Failed for List type";
+    Py_DECREF(list_obj);
+
+    // Test with Dict
+    PyObject* dict_obj = PyDict_New();
+    EXPECT_FALSE(PyByteArray_Check(dict_obj)) << "Failed for Dict type";
+    Py_DECREF(dict_obj);
+
+    // Test with Set
+    PyObject* set_obj = PySet_New(NULL);
+    EXPECT_FALSE(PyByteArray_Check(set_obj)) << "Failed for Set type";
+    Py_DECREF(set_obj);
+
+    PyObject* globals = PyDict_New();
+    PyObject* locals = PyDict_New();
+    if (PY_VERSION_HEX >= 0x03080000) {
+        // PyByteArray_Check is bugged for user-defined classes in Python 3.7
+        // Test with a user-defined class
+        PyRun_String(
+            "class TestClass:\n    pass\n\ntest_instance = TestClass()",
+            Py_file_input,
+            globals,
+            locals
+        );
+        PyObject* user_obj = PyDict_GetItemString(locals, "test_instance");
+        EXPECT_FALSE(PyByteArray_Check(user_obj)) << "Failed for user-defined class";
+    }
+
+    // Test with actual ByteArray (this should return true)
+    PyObject* bytearray_obj = PyByteArray_FromStringAndSize("test", 4);
+    EXPECT_TRUE(PyByteArray_Check(bytearray_obj)) << "Failed for actual ByteArray type";
+    Py_DECREF(bytearray_obj);
+
+    // Test with Complex number
+    PyObject* complex_obj = PyComplex_FromDoubles(1.0, 2.0);
+    EXPECT_FALSE(PyByteArray_Check(complex_obj)) << "Failed for Complex type";
+    Py_DECREF(complex_obj);
+
+    // Test with Function
+    PyRun_String(
+        "def test_function():\n    pass",
+        Py_file_input,
+        globals,
+        locals
+    );
+    PyObject* func_obj = PyDict_GetItemString(locals, "test_function");
+    EXPECT_FALSE(PyByteArray_Check(func_obj)) << "Failed for Function type";
+    Py_DECREF(globals);
+    Py_DECREF(locals);
+}

ddtrace/appsec/_iast/_taint_tracking/tests/test_stringutils.cpp

@@ -502,6 +502,29 @@ TEST_F(NewPyObjectIdCheck, NullObjectReturnsNull)
     EXPECT_EQ(new_id_obj, nullptr);
 }
 
+TEST_F(NewPyObjectIdCheck, NonTextObjectReturnsSameObject)
+{
+    PyObject* non_text_obj = PyLong_FromLong(42);
+    PyObject* new_id_obj = new_pyobject_id(non_text_obj);
+    EXPECT_EQ(new_id_obj, non_text_obj);
+    Py_DECREF(non_text_obj);
+}
+
+TEST_F(NewPyObjectIdCheck, WrongPointer)
+{
+    PyObject* wrong_object = reinterpret_cast<PyObject*>(0x12345);
+    PyObject* new_id_obj = new_pyobject_id(wrong_object);
+    EXPECT_EQ(new_id_obj, wrong_object);
+}
+
+TEST_F(NewPyObjectIdCheck, PyObjectNoType)
+{
+    PyObject* wrong_object = PyBytes_FromString("test");
+    wrong_object->ob_type = nullptr;
+    PyObject* new_id_obj = new_pyobject_id(wrong_object);
+    EXPECT_EQ(new_id_obj, wrong_object);
+}
+
 using GetPyObjectSizeCheck = PyEnvCheck;
 
 TEST_F(GetPyObjectSizeCheck, UnicodeReturnsCorrectSize)