FIX: Access Token Handling, Shutdown Safety, and Connection Attribute Tests

mssql_python/pybind/connection/connection.cpp

@@ -173,16 +173,25 @@ SQLRETURN Connection::setAttribute(SQLINTEGER attribute, py::object value) {
     LOG("Setting SQL attribute");
     SQLPOINTER ptr = nullptr;
     SQLINTEGER length = 0;
+    std::string buffer; // to hold sensitive data temporarily
 
     if (py::isinstance<py::int_>(value)) {
         int intValue = value.cast<int>();
         ptr = reinterpret_cast<SQLPOINTER>(static_cast<uintptr_t>(intValue));
         length = SQL_IS_INTEGER;
     } else if (py::isinstance<py::bytes>(value) || py::isinstance<py::bytearray>(value)) {
-        static std::vector<std::string> buffers;
-        buffers.emplace_back(value.cast<std::string>());
-        ptr = const_cast<char*>(buffers.back().c_str());
-        length = static_cast<SQLINTEGER>(buffers.back().size());
+        buffer = value.cast<std::string>();  // stack buffer
+
+        // DEFENSIVE FIX: Protect against ODBC driver bug with short access tokens
+        // Microsoft ODBC Driver 18 crashes when given access tokens shorter than 32 bytes
+        // Real access tokens are typically 100+ bytes, so reject anything under 32 bytes
+        if (attribute == SQL_COPT_SS_ACCESS_TOKEN && buffer.size() < 32) {
+            LOG("Access token too short (< 32 bytes) - protecting against ODBC driver crash");
+            return SQL_ERROR;  // Return error instead of letting ODBC crash
+        }
+
+        ptr = buffer.data();
+        length = static_cast<SQLINTEGER>(buffer.size());
     } else {
         LOG("Unsupported attribute value type");
         return SQL_ERROR;
@@ -195,6 +204,11 @@ SQLRETURN Connection::setAttribute(SQLINTEGER attribute, py::object value) {
     else {
         LOG("Set attribute successfully");
     }
+
+    // Zero out sensitive data if used
+    if (!buffer.empty()) {
+        std::fill(buffer.begin(), buffer.end(), static_cast<char>(0));
+    }
     return ret;
 }
 

mssql_python/pybind/ddbc_bindings.cpp

@@ -665,24 +665,67 @@ void HandleZeroColumnSizeAtFetch(SQLULEN& columnSize) {
 
 }  // namespace
 
+// Helper function to check if Python is shutting down or finalizing
+// This centralizes the shutdown detection logic to avoid code duplication
+static bool is_python_finalizing() {
+    try {
+        if (Py_IsInitialized() == 0) {
+            return true; // Python is already shut down
+        }
+
+        py::gil_scoped_acquire gil;
+        py::object sys_module = py::module_::import("sys");
+        if (!sys_module.is_none()) {
+            // Check if the attribute exists before accessing it (for Python version compatibility)
+            if (py::hasattr(sys_module, "_is_finalizing")) {
+                py::object finalizing_func = sys_module.attr("_is_finalizing");
+                if (!finalizing_func.is_none() && finalizing_func().cast<bool>()) {
+                    return true; // Python is finalizing
+                }
+            }
+        }
+        return false;
+    } catch (...) {
+        std::cerr << "Error occurred while checking Python finalization state." << std::endl;
+        // Be conservative - don't assume shutdown on any exception
+        // Only return true if we're absolutely certain Python is shutting down
+        return false;
+    }
+}
+
 // TODO: Revisit GIL considerations if we're using python's logger
 template <typename... Args>
 void LOG(const std::string& formatString, Args&&... args) {
-    py::gil_scoped_acquire gil;  // <---- this ensures safe Python API usage
+    // Check if Python is shutting down to avoid crash during cleanup
+    if (is_python_finalizing()) {
+        return; // Python is shutting down or finalizing, don't log
+    }
+
+    try {
+        py::gil_scoped_acquire gil;  // <---- this ensures safe Python API usage
 
-    py::object logger = py::module_::import("mssql_python.logging_config").attr("get_logger")();
-    if (py::isinstance<py::none>(logger)) return;
+        py::object logger = py::module_::import("mssql_python.logging_config").attr("get_logger")();
+        if (py::isinstance<py::none>(logger)) return;
 
-    try {
-        std::string ddbcFormatString = "[DDBC Bindings log] " + formatString;
-        if constexpr (sizeof...(args) == 0) {
-            logger.attr("debug")(py::str(ddbcFormatString));
-        } else {
-            py::str message = py::str(ddbcFormatString).format(std::forward<Args>(args)...);
-            logger.attr("debug")(message);
+        try {
+            std::string ddbcFormatString = "[DDBC Bindings log] " + formatString;
+            if constexpr (sizeof...(args) == 0) {
+                logger.attr("debug")(py::str(ddbcFormatString));
+            } else {
+                py::str message = py::str(ddbcFormatString).format(std::forward<Args>(args)...);
+                logger.attr("debug")(message);
+            }
+        } catch (const std::exception& e) {
+            std::cerr << "Logging error: " << e.what() << std::endl;
         }
+    } catch (const py::error_already_set& e) {
+        // Python is shutting down or in an inconsistent state, silently ignore
+        (void)e; // Suppress unused variable warning
+        return;
     } catch (const std::exception& e) {
-        std::cerr << "Logging error: " << e.what() << std::endl;
+        // Any other error, ignore to prevent crash during cleanup
+        (void)e; // Suppress unused variable warning
+        return;
     }
 }
 
@@ -993,17 +1036,26 @@ SQLSMALLINT SqlHandle::type() const {
  */
 void SqlHandle::free() {
     if (_handle && SQLFreeHandle_ptr) {
-        const char* type_str = nullptr;
-        switch (_type) {
-            case SQL_HANDLE_ENV:  type_str = "ENV"; break;
-            case SQL_HANDLE_DBC:  type_str = "DBC"; break;
-            case SQL_HANDLE_STMT: type_str = "STMT"; break;
-            case SQL_HANDLE_DESC: type_str = "DESC"; break;
-            default:              type_str = "UNKNOWN"; break;
+        // Check if Python is shutting down using centralized helper function
+        bool pythonShuttingDown = is_python_finalizing();
+
+        // CRITICAL FIX: During Python shutdown, don't free STMT handles as their parent DBC may already be freed
+        // This prevents segfault when handles are freed in wrong order during interpreter shutdown
+        // Type 3 = SQL_HANDLE_STMT, Type 2 = SQL_HANDLE_DBC, Type 1 = SQL_HANDLE_ENV
+        if (pythonShuttingDown && _type == 3) {
+            _handle = nullptr; // Mark as freed to prevent double-free attempts
+            return;
         }
+
+        // Always clean up ODBC resources, regardless of Python state
         SQLFreeHandle_ptr(_type, _handle);
         _handle = nullptr;
-        // Don't log during destruction - it can cause segfaults during Python shutdown
+
+        // Only log if Python is not shutting down (to avoid segfault)
+        if (!pythonShuttingDown) {
+            // Don't log during destruction - even in normal cases it can be problematic
+            // If logging is needed, use explicit close() methods instead
+        }
     }
 }