segfault fix

Author: subrata-ms

mssql_python/connection.py

@@ -287,6 +287,10 @@ def __init__(
         # TODO: Think and implement scenarios for multi-threaded access
         # to cursors
         self._cursors = weakref.WeakSet()
+        
+        # Track if connection has been closed to prevent cursors from
+        # trying to free handles after connection is closed (prevents segfault)
+        self._connection_closed = False
 
         # Initialize output converters dictionary and its lock for thread safety
         self._output_converters = {}
@@ -1499,12 +1503,33 @@ def close(self) -> None:
         if self._closed:
             return
 
+        # CRITICAL: Set connection closed flag BEFORE closing anything
+        # This prevents cursors from trying to free handles during/after connection close
+        self._connection_closed = True
+
         # Close all cursors first, but don't let one failure stop the others
         if hasattr(self, "_cursors"):
             # Convert to list to avoid modification during iteration
             cursors_to_close = list(self._cursors)
             close_errors = []
 
+            # First pass: Invalidate cursor handles to prevent use-after-free
+            for cursor in cursors_to_close:
+                try:
+                    # CRITICAL: Mark cursor as invalidated BEFORE clearing handle
+                    # This tells cursor.close() to skip SQLFreeHandle entirely
+                    if hasattr(cursor, '_invalidated'):
+                        cursor._invalidated = True
+                    # Mark handles as freed before closing connection
+                    # This prevents cursors from trying to free already-freed handles
+                    if hasattr(cursor, '_handle_freed'):
+                        cursor._handle_freed = True
+                    if hasattr(cursor, 'hstmt'):
+                        cursor.hstmt = None
+                except Exception as e:  # pylint: disable=broad-exception-caught
+                    logger.warning("Error invalidating cursor handle: %s", e)
+
+            # Second pass: Close cursors
             for cursor in cursors_to_close:
                 try:
                     if not cursor.closed:

mssql_python/cursor.py

@@ -15,6 +15,8 @@
 import uuid
 import datetime
 import warnings
+import weakref
+import sys
 from typing import List, Union, Any, Optional, Tuple, Sequence, TYPE_CHECKING, Iterable
 from mssql_python.constants import ConstantsDDBC as ddbc_sql_const, SQLTypes
 from mssql_python.helpers import check_error
@@ -98,6 +100,13 @@ def __init__(self, connection: "Connection", timeout: int = 0) -> None:
         self._inputsizes: Optional[List[Union[int, Tuple[Any, ...]]]] = None
         # self.connection.autocommit = False
         self.hstmt: Optional[Any] = None
+        self._handle_freed: bool = False  # Track if statement handle has been freed
+        self._invalidated: bool = False  # Track if cursor was invalidated by connection close
+        
+        # Store weak reference to connection to check if it's closed
+        # This prevents segfault when trying to free handles after connection close
+        self._connection_ref: Any = weakref.ref(connection)
+        
         self._initialize_cursor()
         self.description: Optional[
             List[
@@ -728,14 +737,17 @@ def _reset_cursor(self) -> None:
         # Reinitialize the statement handle
         self._initialize_cursor()
 
-    def close(self) -> None:
+    def close(self, from_del: bool = False) -> None:
         """
         Close the connection now (rather than whenever .__del__() is called).
         Idempotent: subsequent calls have no effect and will be no-ops.
 
         The cursor will be unusable from this point forward; an InterfaceError
         will be raised if any operation (other than close) is attempted with the cursor.
         This is a deviation from pyodbc, which raises an exception if the cursor is already closed.
+        
+        Args:
+            from_del: Internal flag - when True, handle freeing is skipped to prevent segfaults during GC.
         """
         if self.closed:
             # Do nothing - not calling _check_closed() here since we want this to be idempotent
@@ -751,10 +763,81 @@ def close(self) -> None:
             except Exception as e:  # pylint: disable=broad-exception-caught
                 logger.warning("Error removing cursor from connection tracking: %s", e)
 
-        if self.hstmt:
-            self.hstmt.free()
-            self.hstmt = None
-            logger.debug("SQLFreeHandle succeeded")
+        # Free statement handle with protection against double-free
+        # CRITICAL SAFETY #0: If called from __del__, do NOTHING - just return
+        # During GC, ANY object access can trigger segfaults due to partially destroyed state
+        # Better to leak resources than crash - OS cleans up at process exit
+        if from_del:
+            return
+        
+        if self.hstmt and not self._handle_freed:
+            
+            # CRITICAL SAFETY #1: Check if cursor was explicitly invalidated by connection
+            # This happens when connection.close() invalidates all cursors BEFORE freeing connection handle
+            if hasattr(self, '_invalidated') and self._invalidated:
+                self.hstmt = None
+                self._handle_freed = True
+                return
+            
+            # CRITICAL SAFETY #2: Check if parent connection is closed via _connection attribute
+            # This is an additional safety check in case weak reference approach fails
+            try:
+                if hasattr(self, '_connection') and self._connection:
+                    if hasattr(self._connection, '_closed') and self._connection._closed:
+                        # Connection is closed, skip freeing handle
+                        self.hstmt = None
+                        self._handle_freed = True
+                        return
+                    if hasattr(self._connection, '_connection_closed') and self._connection._connection_closed:
+                        # Connection closed flag set, skip freeing handle
+                        self.hstmt = None
+                        self._handle_freed = True
+                        return
+            except (AttributeError, ReferenceError):
+                # Connection might be in invalid state, skip freeing to be safe
+                self.hstmt = None
+                self._handle_freed = True
+                return
+            
+            # CRITICAL SAFETY #3: Skip handle freeing during interpreter shutdown
+            # During shutdown, ODBC resources may already be freed, causing segfault
+            if sys.is_finalizing():
+                self.hstmt = None
+                self._handle_freed = True
+                return
+            
+            # Check if parent connection was closed or garbage collected
+            # First check if we even have a connection reference attribute
+            if not hasattr(self, '_connection_ref'):
+                # Old cursor without weak ref - conservatively skip free
+                self.hstmt = None
+                self._handle_freed = True
+                return
+            
+            # Try to get connection from weak reference
+            try:
+                conn = self._connection_ref()
+            except Exception:
+                # Weak ref might be invalid during cleanup - skip free to be safe
+                self.hstmt = None
+                self._handle_freed = True
+                return
+            
+            # Skip free if connection is closed or garbage collected
+            if conn is None or (hasattr(conn, '_connection_closed') and conn._connection_closed):
+                # Connection closed/GC'd - ODBC driver already freed handles
+                self.hstmt = None
+                self._handle_freed = True
+            else:
+                # Connection is still open - safe to free statement handle
+                try:
+                    self.hstmt.free()
+                    self._handle_freed = True
+                except Exception:
+                    # Silently handle errors during cleanup
+                    pass
+                finally:
+                    self.hstmt = None
         self._clear_rownumber()
         self.closed = True
 
@@ -2760,7 +2843,8 @@ def __del__(self):
         """
         if "closed" not in self.__dict__ or not self.closed:
             try:
-                self.close()
+                # Pass from_del=True to skip handle freeing (prevents segfaults during GC)
+                self.close(from_del=True)
             except Exception as e:  # pylint: disable=broad-exception-caught
                 # Don't raise an exception in __del__, just log it
                 # If interpreter is shutting down, we might not have logging set up

mssql_python/pybind/ddbc_bindings.cpp

@@ -1128,7 +1128,7 @@ void DriverLoader::loadDriver() {
 }
 
 // SqlHandle definition
-SqlHandle::SqlHandle(SQLSMALLINT type, SQLHANDLE rawHandle) : _type(type), _handle(rawHandle) {}
+SqlHandle::SqlHandle(SQLSMALLINT type, SQLHANDLE rawHandle) : _type(type), _handle(rawHandle), _freed(false) {}
 
 SqlHandle::~SqlHandle() {
     if (_handle) {
@@ -1152,32 +1152,65 @@ SQLSMALLINT SqlHandle::type() const {
  * If you need destruction logs, use explicit close() methods instead.
  */
 void SqlHandle::free() {
-    if (_handle && SQLFreeHandle_ptr) {
-        // Check if Python is shutting down using centralized helper function
-        bool pythonShuttingDown = is_python_finalizing();
-
-        // RESOURCE LEAK MITIGATION:
-        // When handles are skipped during shutdown, they are not freed, which could
-        // cause resource leaks. However, this is mitigated by:
-        // 1. Python-side atexit cleanup (in __init__.py) that explicitly closes all
-        //    connections before shutdown, ensuring handles are freed in correct order
-        // 2. OS-level cleanup at process termination recovers any remaining resources
-        // 3. This tradeoff prioritizes crash prevention over resource cleanup, which
-        //    is appropriate since we're already in shutdown sequence
-        if (pythonShuttingDown && (_type == SQL_HANDLE_STMT || _type == SQL_HANDLE_DBC)) {
-            _handle = nullptr;  // Mark as freed to prevent double-free attempts
-            return;
-        }
+    // Early return if handle was already explicitly freed
+    if (_freed) {
+        return;  // Already freed via explicit free() call
+    }
+    
+    // Early return if handle is nullptr
+    if (!_handle) {
+        _freed = true;
+        return;  // Already freed, nothing to do
+    }
+    
+    // Early return if SQLFreeHandle function is not loaded
+    if (!SQLFreeHandle_ptr) {
+        _handle = nullptr;  // Mark as freed to prevent future attempts
+        _freed = true;
+        return;
+    }
 
-        // Always clean up ODBC resources, regardless of Python state
-        SQLFreeHandle_ptr(_type, _handle);
-        _handle = nullptr;
+    // Check if Python is shutting down using centralized helper function
+    bool pythonShuttingDown = is_python_finalizing();
+
+    // RESOURCE LEAK MITIGATION:
+    // When handles are skipped during shutdown, they are not freed, which could
+    // cause resource leaks. However, this is mitigated by:
+    // 1. Python-side atexit cleanup (in __init__.py) that explicitly closes all
+    //    connections before shutdown, ensuring handles are freed in correct order
+    // 2. OS-level cleanup at process termination recovers any remaining resources
+    // 3. This tradeoff prioritizes crash prevention over resource cleanup, which
+    //    is appropriate since we're already in shutdown sequence
+    if (pythonShuttingDown && (_type == SQL_HANDLE_STMT || _type == SQL_HANDLE_DBC)) {
+        _handle = nullptr;  // Mark as freed to prevent double-free attempts
+        _freed = true;
+        return;
+    }
 
-        // Only log if Python is not shutting down (to avoid segfault)
+    // Additional safety: Check if handle is SQL_NULL_HANDLE before freeing
+    // This can happen if connection was closed and ODBC driver already freed statement handles
+    if (_handle == SQL_NULL_HANDLE) {
+        _freed = true;
+        return;  // Handle already null, nothing to free
+    }
+    
+    // USE-AFTER-FREE FIX: Always clean up ODBC resources with error handling
+    // to prevent segfaults when handle is already freed or invalid
+    SQLRETURN ret = SQLFreeHandle_ptr(_type, _handle);
+    
+    // Always clear the handle reference and mark as freed regardless of return code
+    // This prevents double-free attempts even if SQLFreeHandle fails
+    _handle = nullptr;
+    _freed = true;
+
+    // Handle errors gracefully - don't throw on invalid handle
+    // SQL_INVALID_HANDLE (-2) indicates handle was already freed or invalid
+    // This is expected during connection invalidation scenarios
+    if (ret != SQL_SUCCESS && ret != SQL_SUCCESS_WITH_INFO && ret != SQL_INVALID_HANDLE) {
+        // Only log non-critical errors if Python is not shutting down
         if (!pythonShuttingDown) {
-            // Don't log during destruction - even in normal cases it can be
-            // problematic If logging is needed, use explicit close() methods
-            // instead
+            // Log error but don't throw exception - prevents crashes during cleanup
+            // If logging is needed, use explicit close() methods instead of destructor
         }
     }
 }

mssql_python/pybind/ddbc_bindings.h

@@ -382,6 +382,7 @@ class SqlHandle {
   private:
     SQLSMALLINT _type;
     SQLHANDLE _handle;
+    bool _freed;  // Track if handle has been explicitly freed
 };
 using SqlHandlePtr = std::shared_ptr<SqlHandle>;