FIX: Validate caller-controlled native buffer sizes

mssql_python/pybind/param_detect.hpp

@@ -140,6 +140,26 @@ inline constexpr int MAX_INLINE_CHAR = 4000;
 // Binary data longer than this uses DAE streaming (SQL Server max for non-MAX types)
 inline constexpr int MAX_INLINE_BINARY = 8000;
 
+inline SQLULEN DAEColumnSize(SQLSMALLINT sqlType, SQLULEN actualSize) {
+    switch (sqlType) {
+        case SQL_CHAR:
+        case SQL_VARCHAR:
+            return actualSize > MAX_INLINE_CHAR ? 0 : actualSize;
+        case SQL_WCHAR:
+        case SQL_WVARCHAR:
+            return actualSize > MAX_INLINE_CHAR ? 0 : actualSize;
+        case SQL_BINARY:
+        case SQL_VARBINARY:
+            return actualSize > MAX_INLINE_BINARY ? 0 : actualSize;
+        case SQL_LONGVARCHAR:
+        case SQL_WLONGVARCHAR:
+        case SQL_LONGVARBINARY:
+            return actualSize;
+        default:
+            return actualSize;
+    }
+}
+
 // SQL Server maximum numeric precision
 inline constexpr int MAX_NUMERIC_PRECISION = 38;
 
@@ -190,6 +210,45 @@ inline bool PyLongGreaterThan(PyObject* value, long long threshold) {
     return overflow > 0 || (overflow == 0 && result > threshold);
 }
 
+inline Py_ssize_t UnicodeUtf16Length(PyObject* value) {
+    const Py_ssize_t length = PyUnicode_GET_LENGTH(value);
+    if (PyUnicode_KIND(value) <= PyUnicode_2BYTE_KIND) {
+        return length;
+    }
+
+    Py_ssize_t utf16Length = 0;
+    const Py_UCS4* data = PyUnicode_4BYTE_DATA(value);
+    for (Py_ssize_t index = 0; index < length; ++index) {
+        utf16Length += data[index] > 0xFFFF ? 2 : 1;
+    }
+    return utf16Length;
+}
+
+inline Py_ssize_t EncodedUnicodeLength(PyObject* value, const std::string& encoding) {
+    py::object encoderFactory =
+        py::module_::import("codecs").attr("getincrementalencoder")(encoding);
+    py::object encoder = encoderFactory("strict");
+    const Py_ssize_t length = PyUnicode_GET_LENGTH(value);
+    constexpr Py_ssize_t chunkSize = 4096;
+    Py_ssize_t total = 0;
+    for (Py_ssize_t offset = 0; offset < length; offset += chunkSize) {
+        const Py_ssize_t end = std::min(offset + chunkSize, length);
+        py::object chunk = steal(PyUnicode_Substring(value, offset, end));
+        if (!chunk) throw py::error_already_set();
+        py::object encoded = encoder.attr("encode")(chunk, end == length);
+        const Py_ssize_t encodedSize = PyBytes_GET_SIZE(encoded.ptr());
+        if (encodedSize > MAX_INLINE_BINARY - total) {
+            return MAX_INLINE_BINARY + 1;
+        }
+        total += encodedSize;
+    }
+    if (length == 0) {
+        py::object encoded = encoder.attr("encode")(py::str(), true);
+        total = PyBytes_GET_SIZE(encoded.ptr());
+    }
+    return total;
+}
+
 inline PyObject* FormatDecimalParam(PyObject* params, Py_ssize_t index, PyObject* value) {
     py::object formatted = steal(PyObject_CallMethod(value, "__format__", "s", "f"));
     if (!formatted) throw py::error_already_set();
@@ -215,8 +274,91 @@ inline void NormalizeTimeParam(PyObject* params, Py_ssize_t index, SQLULEN& colu
     }
 }
 
+inline long long ValidatedInputSizeInteger(PyObject* value, const char* fieldName) {
+    if (!PyLong_Check(value) || PyBool_Check(value)) {
+        throw py::type_error(std::string(fieldName) + " must be an integer");
+    }
+    int overflow = 0;
+    const long long result = PyLong_AsLongLongAndOverflow(value, &overflow);
+    if ((result == -1 && PyErr_Occurred()) || overflow != 0) {
+        PyErr_Clear();
+        throw py::value_error(std::string(fieldName) + " is out of range");
+    }
+    return result;
+}
+
+inline void ValidateInputSizes(PyObject* inputSizes) {
+    if (inputSizes == Py_None) {
+        return;
+    }
+    if (!PyList_Check(inputSizes)) {
+        throw py::type_error("inputSizes must be None or a list");
+    }
+
+    const Py_ssize_t count = PyList_GET_SIZE(inputSizes);
+    for (Py_ssize_t index = 0; index < count; ++index) {
+        PyObject* entry = PyList_GET_ITEM(inputSizes, index);
+        if (!PyTuple_Check(entry) || PyTuple_GET_SIZE(entry) != 4) {
+            throw py::type_error("each inputSizes entry must be a four-item tuple");
+        }
+
+        const long long sqlType =
+            ValidatedInputSizeInteger(PyTuple_GET_ITEM(entry, 0), "SQL type");
+        const long long cType =
+            ValidatedInputSizeInteger(PyTuple_GET_ITEM(entry, 1), "C type");
+        PyObject* columnSize = PyTuple_GET_ITEM(entry, 2);
+        PyObject* decimalDigits = PyTuple_GET_ITEM(entry, 3);
+        if (!PyLong_Check(columnSize) || PyBool_Check(columnSize) ||
+            !PyLong_Check(decimalDigits) || PyBool_Check(decimalDigits)) {
+            throw py::type_error("column size and decimal digits must be integers");
+        }
+
+        if (sqlType < std::numeric_limits<SQLSMALLINT>::min() ||
+            sqlType > std::numeric_limits<SQLSMALLINT>::max() ||
+            cType < std::numeric_limits<SQLSMALLINT>::min() ||
+            cType > std::numeric_limits<SQLSMALLINT>::max()) {
+            throw py::value_error("SQL and C types must fit in SQLSMALLINT");
+        }
+        const bool isNumeric = sqlType == SQL_DECIMAL || sqlType == SQL_NUMERIC;
+        py::int_ zero(0);
+        const int negativeColumnSize =
+            PyObject_RichCompareBool(columnSize, zero.ptr(), Py_LT);
+        const int negativeDecimalDigits =
+            PyObject_RichCompareBool(decimalDigits, zero.ptr(), Py_LT);
+        if (negativeColumnSize == -1 || negativeDecimalDigits == -1) {
+            throw py::error_already_set();
+        }
+        if (negativeColumnSize == 1) {
+            throw py::value_error("column size must be non-negative");
+        }
+        if (negativeDecimalDigits == 1) {
+            throw py::value_error("decimal digits must be non-negative");
+        }
+        if (!isNumeric) {
+            const unsigned long long requestedSize = PyLong_AsUnsignedLongLong(columnSize);
+            if (requestedSize == static_cast<unsigned long long>(-1) && PyErr_Occurred()) {
+                PyErr_Clear();
+                throw py::value_error("column size is out of range");
+            }
+            if (requestedSize > std::numeric_limits<SQLULEN>::max()) {
+                throw py::value_error("column size is out of range");
+            }
+            const unsigned long long requestedDigits =
+                PyLong_AsUnsignedLongLong(decimalDigits);
+            if (requestedDigits == static_cast<unsigned long long>(-1) && PyErr_Occurred()) {
+                PyErr_Clear();
+                throw py::value_error("decimal digits are out of range");
+            }
+            if (requestedDigits >
+                static_cast<unsigned long long>(std::numeric_limits<SQLSMALLINT>::max())) {
+                throw py::value_error("decimal digits are out of range");
+            }
+        }
+    }
+}
+
 inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssize_t index,
-                                   ParamInfo& info) {
+                                   ParamInfo& info, const std::string& charEncoding) {
     py::tuple values = borrow<py::tuple>(inputSize);
     info.paramSQLType = values[0].cast<SQLSMALLINT>();
     info.paramCType = values[1].cast<SQLSMALLINT>();
@@ -257,10 +399,42 @@ inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssi
         }
     }
 
-    info.isDAE =
-        (PyUnicode_Check(obj) && PyLongGreaterThan(columnSize, MAX_INLINE_CHAR)) ||
-        ((PyBytes_Check(obj) || PyByteArray_Check(obj)) &&
-         PyLongGreaterThan(columnSize, MAX_INLINE_BINARY));
+    if ((PyBytes_Check(obj) || PyByteArray_Check(obj)) &&
+        (info.paramSQLType == SQL_CHAR || info.paramSQLType == SQL_VARCHAR ||
+         info.paramSQLType == SQL_LONGVARCHAR)) {
+        info.paramCType = SQL_C_CHAR;
+    }
+
+    if (info.paramCType == SQL_C_WCHAR &&
+        (PyBytes_Check(obj) || PyByteArray_Check(obj))) {
+        throw py::type_error("bytes values cannot be bound as SQL_C_WCHAR");
+    }
+
+    Py_ssize_t actualTextLength = 0;
+    if (PyUnicode_Check(obj)) {
+        actualTextLength = info.paramCType == SQL_C_CHAR
+                               ? EncodedUnicodeLength(obj, charEncoding)
+                               : UnicodeUtf16Length(obj);
+    }
+    const bool textNeedsDAE =
+        !isNumeric && PyUnicode_Check(obj) && actualTextLength > MAX_INLINE_CHAR;
+    const bool binaryNeedsDAE =
+        (PyBytes_Check(obj) || PyByteArray_Check(obj)) &&
+        (PyBytes_Check(obj) ? PyBytes_GET_SIZE(obj) : PyByteArray_GET_SIZE(obj)) >
+            MAX_INLINE_BINARY;
+    info.isDAE = textNeedsDAE || binaryNeedsDAE;
+    if (!isNumeric && PyUnicode_Check(obj)) {
+        const SQLULEN actualSize = static_cast<SQLULEN>(actualTextLength);
+        info.columnSize =
+            info.isDAE ? DAEColumnSize(info.paramSQLType, actualSize)
+                       : std::max(info.columnSize, actualSize);
+    } else if (PyBytes_Check(obj) || PyByteArray_Check(obj)) {
+        const SQLULEN actualSize = static_cast<SQLULEN>(
+            PyBytes_Check(obj) ? PyBytes_GET_SIZE(obj) : PyByteArray_GET_SIZE(obj));
+        info.columnSize =
+            info.isDAE ? DAEColumnSize(info.paramSQLType, actualSize)
+                       : std::max(info.columnSize, actualSize);
+    }
 
     if (PyTime_Check(obj) && info.paramCType == PARAM_C_TYPE_TEXT) {
         NormalizeTimeParam(params, index, info.columnSize);
@@ -296,9 +470,15 @@ inline void ApplyInputSizeOverride(PyObject* params, PyObject* inputSize, Py_ssi
 //
 // Takes raw PyObject* lists. Caller guarantees params is a fresh copy (cursor.py
 // does list(actual_params)), so in-place mutation via PyList_SetItem is safe.
-inline std::vector<ParamInfo> DetectParamTypes(PyObject* params, PyObject* inputSizes) {
+inline std::vector<ParamInfo> DetectParamTypes(PyObject* params, PyObject* inputSizes,
+                                               const std::string& charEncoding = "utf-8") {
     PyTypeCache::initialize();
 
+    if (!PyList_Check(params)) {
+        throw py::type_error("params must be a list");
+    }
+    ValidateInputSizes(inputSizes);
+
     const Py_ssize_t n = PyList_GET_SIZE(params);
     const Py_ssize_t inputSizeCount = inputSizes == Py_None ? 0 : PyList_GET_SIZE(inputSizes);
     std::vector<ParamInfo> infos(n);
@@ -312,7 +492,7 @@ inline std::vector<ParamInfo> DetectParamTypes(PyObject* params, PyObject* input
         info.isDAE = false;
 
         if (i < inputSizeCount) {
-            ApplyInputSizeOverride(params, PyList_GET_ITEM(inputSizes, i), i, info);
+            ApplyInputSizeOverride(params, PyList_GET_ITEM(inputSizes, i), i, info, charEncoding);
             continue;
         }
 
@@ -398,16 +578,7 @@ inline std::vector<ParamInfo> DetectParamTypes(PyObject* params, PyObject* input
             unsigned int kind = PyUnicode_KIND(obj);
             const void* udata = PyUnicode_DATA(obj);
 
-            Py_ssize_t utf16_len;
-            if (kind <= PyUnicode_2BYTE_KIND) {
-                utf16_len = length;
-            } else {
-                utf16_len = 0;
-                const Py_UCS4* data = PyUnicode_4BYTE_DATA(obj);
-                for (Py_ssize_t j = 0; j < length; ++j) {
-                    utf16_len += (data[j] > 0xFFFF) ? 2 : 1;
-                }
-            }
+            const Py_ssize_t utf16_len = UnicodeUtf16Length(obj);
 
             // Detect whether the string needs wide-char (NVARCHAR) or narrow (VARCHAR) binding.
             // PyUnicode_IS_COMPACT_ASCII is a struct field check (O(1)), not a content scan.
@@ -439,16 +610,14 @@ inline std::vector<ParamInfo> DetectParamTypes(PyObject* params, PyObject* input
                 // Strings > 4000 UTF-16 code units exceed SQL Server's inline NVARCHAR(MAX)
                 // threshold. Switch to data-at-execution (DAE) streaming: ODBC driver pulls
                 // data in chunks via SQLPutData, avoiding a single massive buffer allocation.
-                // DAE path: match slow-path types exactly.
-                // Non-unicode (ASCII) → SQL_VARCHAR + PARAM_C_TYPE_TEXT, which is
-                //   SQL_C_WCHAR and matches the slow path's SQL_C_CHAR (numerically
-                //   -8 == SQL_C_WCHAR — a long-standing alias in the Python layer).
-                // Unicode → SQL_WVARCHAR + SQL_C_WCHAR (wide-char streaming)
+                // Use the validated payload size when it is legal fixed-width metadata;
+                // larger values use the driver's MAX-length sentinel.
                 info.isDAE = true;
-                info.columnSize = 0;
+                const SQLSMALLINT sqlType = is_unicode ? SQL_WVARCHAR : SQL_VARCHAR;
+                info.columnSize = DAEColumnSize(sqlType, utf16_len);
                 info.utf16Len = utf16_len;
                 info.dataPtr = borrow(obj);
-                info.paramSQLType = is_unicode ? SQL_WVARCHAR : SQL_VARCHAR;
+                info.paramSQLType = sqlType;
                 info.paramCType = is_unicode ? SQL_C_WCHAR : PARAM_C_TYPE_TEXT;
             } else {
                 info.columnSize = is_unicode ? utf16_len : length;
@@ -467,7 +636,8 @@ inline std::vector<ParamInfo> DetectParamTypes(PyObject* params, PyObject* input
             info.decimalDigits = 0;
             if (length > MAX_INLINE_BINARY) {
                 info.isDAE = true;
-                info.columnSize = 0;
+                info.columnSize =
+                    DAEColumnSize(SQL_VARBINARY, static_cast<SQLULEN>(length));
                 info.dataPtr = borrow(obj);
             } else {
                 info.columnSize = std::max<SQLULEN>(length, 1);

tests/test_004_cursor.py

@@ -1693,6 +1693,64 @@ def test_arraysize(cursor):
     assert cursor.arraysize == 5, "Arraysize mismatch after change"
 
 
+@pytest.mark.parametrize("value", [0, -1, 1_000_001])
+def test_arraysize_rejects_out_of_range_values(value):
+    cursor = mssql_python.Cursor.__new__(mssql_python.Cursor)
+    with pytest.raises(ValueError, match="arraysize"):
+        cursor.arraysize = value
+
+
+@pytest.mark.parametrize("value", [True, False, 1.5, "10"])
+def test_arraysize_rejects_non_integer_values(value):
+    cursor = mssql_python.Cursor.__new__(mssql_python.Cursor)
+    with pytest.raises(TypeError, match="arraysize"):
+        cursor.arraysize = value
+
+
+@pytest.mark.parametrize(
+    "size_info",
+    [
+        (True,),
+        (mssql_python.SQL_WVARCHAR, True, 0),
+        (mssql_python.SQL_DECIMAL, 18, True),
+    ],
+)
+def test_setinputsizes_rejects_boolean_sizes(size_info):
+    cursor = mssql_python.Cursor.__new__(mssql_python.Cursor)
+    with pytest.raises(ValueError):
+        cursor.setinputsizes([size_info])
+
+
+def test_setinputsizes_rejects_excessive_column_size():
+    cursor = mssql_python.Cursor.__new__(mssql_python.Cursor)
+    with pytest.raises(ValueError, match="column size"):
+        cursor.setinputsizes([(mssql_python.SQL_VARCHAR, 1 << 40, 0)])
+
+
+def test_executemany_rejects_excessive_cumulative_parameter_buffers(cursor):
+    cursor.setinputsizes([(mssql_python.SQL_WVARCHAR, 50_000_000, 0)] * 3)
+    with pytest.raises(RuntimeError, match="Parameter buffers exceed the 256 MiB allocation limit"):
+        cursor.executemany("SELECT ?, ?, ?", [(None, None, None)])
+
+
+def test_executemany_rejects_text_for_binary_parameter(cursor):
+    cursor.setinputsizes([(mssql_python.SQL_VARBINARY, 10, 0)])
+    with pytest.raises(RuntimeError, match="object type does not match"):
+        cursor.executemany("SELECT CAST(? AS VARBINARY(10))", [("text",)])
+
+
+def test_fetchmany_rejects_excessive_native_buffer(cursor):
+    cursor.execute("SELECT CAST('x' AS VARCHAR(8000))")
+    with pytest.raises(RuntimeError, match="256 MiB allocation limit"):
+        cursor.fetchmany(100_000)
+    assert cursor.fetchone()[0] == "x"
+
+
+def test_fetchall_clamps_wide_result_batch_to_native_buffer_budget(cursor):
+    cursor.execute("SELECT " + ", ".join("CAST(N'x' AS NVARCHAR(4000))" for _ in range(34)))
+    assert [tuple(row) for row in cursor.fetchall()] == [("x",) * 34]
+
+
 def test_description(cursor):
     """Test description"""
     cursor.execute("SELECT * FROM #pytest_all_data_types WHERE id = 1")

tests/test_004_cursor_arrow.py

@@ -225,11 +225,25 @@ def test_arrow_empty_fetch(cursor: mssql_python.Cursor):
         _test_arrow_test_data(cursor, [col_data], fetch_length=0)
 
 
+def test_arrow_wide_schema_only_and_single_row_fit_native_buffer_budget(
+    cursor: mssql_python.Cursor,
+):
+    columns = ", ".join(f"CAST(N'x' AS NVARCHAR(4000)) AS col_{index}" for index in range(600))
+    cursor.execute(f"SELECT {columns}")
+
+    schema_batch = cursor.arrow_batch(0)
+    assert schema_batch.num_rows == 0
+    assert schema_batch.num_columns == 600
+
+    data_batch = cursor.arrow_batch(1)
+    assert data_batch.num_rows == 1
+    assert data_batch.num_columns == 600
+
+
 def test_arrow_table_batchsize_negative(cursor: mssql_python.Cursor):
-    tbl = cursor.execute("select 1 a").arrow(batch_size=-42)
-    assert type(tbl) is pa.Table
-    assert tbl.num_rows == 0
-    assert tbl.num_columns == 1
+    cursor.execute("select 1 a")
+    with pytest.raises(ValueError, match="batch_size"):
+        cursor.arrow(batch_size=-42)
     assert cursor.fetchone()[0] == 1