Harden SqlConnection state transitions with CAS guards (#4267)

paulmedynski · web-flow · commit c7ee148cb9fa · 2026-06-26T19:57:21.000Z
diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/Common/AdapterUtil.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/Common/AdapterUtil.cs
@@ -1106,6 +1106,10 @@ internal enum ConnectionError
             GetConnectionReturnsNull,
             ConnectionOptionsMissing,
             CouldNotSwitchToClosedPreviouslyOpenedState,
+            InvalidSourceStateForSetInnerConnectionTo,
+            FailedToCommitSetInnerConnectionTo,
+            InvalidSourceStateForSetInnerConnectionEvent,
+            FailedToCommitSetInnerConnectionEvent,
         }
 
         internal static Exception InternalConnectionError(ConnectionError internalError)
diff --git a/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlConnection.cs b/src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlConnection.cs
@@ -2236,7 +2236,16 @@ private bool TryOpen(TaskCompletionSource<DbConnectionInternal> retry, SqlConnec
             return result;
         }
 
-        private bool TryOpenInner(TaskCompletionSource<DbConnectionInternal> retry)
+        /// <summary>
+        /// Completes the inner open/replace operation and initializes parser state for the active inner connection.
+        /// </summary>
+        /// <param name="retry">Retry continuation used by async open paths.</param>
+        /// <returns><see langword="true"/> when open initialization completed synchronously; otherwise <see langword="false"/>.</returns>
+        /// <remarks>
+        /// The inner connection is snapshotted after the open call so downstream parser access uses a single observed
+        /// instance and does not rely on a second racy read of <see cref="InnerConnection"/>.
+        /// </remarks>
+        internal bool TryOpenInner(TaskCompletionSource<DbConnectionInternal> retry)
         {
             // Create a single TimeoutTimer that represents the overall connection-open budget for this
             // attempt. The timer is threaded through every layer (inner connection state transitions,
@@ -2408,11 +2417,32 @@ internal void RemoveWeakReference(object value)
 
         // ClosedBusy->Closed (never opened)
         // Connecting->Closed (exception during open, return to previous closed state)
+        /// <summary>
+        /// Finalizes a non-evented transition from a transitional inner state to a stable closed state.
+        /// </summary>
+        /// <param name="to">The target inner connection state object.</param>
+        /// <remarks>
+        /// This path is restricted to transition sources that intentionally serialize external callers while state is
+        /// being updated. A compare-and-swap prevents stale writers from overwriting a newer inner connection state.
+        /// </remarks>
         internal void SetInnerConnectionTo(DbConnectionInternal to)
         {
             Debug.Assert(_innerConnection != null, "null InnerConnection");
             Debug.Assert(to != null, "to null InnerConnection");
-            _innerConnection = to;
+
+            DbConnectionInternal originalInnerConnection = _innerConnection;
+            if (originalInnerConnection != DbConnectionClosedBusy.SingletonInstance &&
+                originalInnerConnection != DbConnectionClosedConnecting.SingletonInstance)
+            {
+                // SetInnerConnectionTo is only valid while leaving a known transitional state.
+                throw ADP.InternalConnectionError(ADP.ConnectionError.InvalidSourceStateForSetInnerConnectionTo);
+            }
+
+            if (originalInnerConnection != Interlocked.CompareExchange(ref _innerConnection, to, originalInnerConnection))
+            {
+                // Reject stale writers when another thread already advanced the state machine.
+                throw ADP.InternalConnectionError(ADP.ConnectionError.FailedToCommitSetInnerConnectionTo);
+            }
         }
 
         // Closed->Connecting: prevent set_ConnectionString during Open
@@ -2430,14 +2460,37 @@ internal bool SetInnerConnectionFrom(DbConnectionInternal to, DbConnectionIntern
 
         // OpenBusy->Closed (previously opened)
         // Connecting->Open
+        /// <summary>
+        /// Finalizes an evented state transition from a transitional inner state to a stable target state.
+        /// </summary>
+        /// <param name="to">The target inner connection state object.</param>
+        /// <remarks>
+        /// This path is intentionally restricted to the expected transition sources (<see cref="DbConnectionOpenBusy"/>
+        /// and <see cref="DbConnectionClosedConnecting"/>). A compare-and-swap enforces that the transition commits only
+        /// if no concurrent writer has already changed <see cref="_innerConnection"/>.
+        /// </remarks>
         internal void SetInnerConnectionEvent(DbConnectionInternal to)
         {
             Debug.Assert(_innerConnection != null, "null InnerConnection");
             Debug.Assert(to != null, "to null InnerConnection");
 
-            ConnectionState originalState = _innerConnection.State & ConnectionState.Open;
+            DbConnectionInternal originalInnerConnection = _innerConnection;
+            if (originalInnerConnection != DbConnectionOpenBusy.SingletonInstance &&
+                originalInnerConnection != DbConnectionClosedConnecting.SingletonInstance)
+            {
+                // Evented transitions are valid only from the two expected transitional sources.
+                throw ADP.InternalConnectionError(ADP.ConnectionError.InvalidSourceStateForSetInnerConnectionEvent);
+            }
+
+            ConnectionState originalState = originalInnerConnection.State & ConnectionState.Open;
             ConnectionState currentState = to.State & ConnectionState.Open;
 
+            if (originalInnerConnection != Interlocked.CompareExchange(ref _innerConnection, to, originalInnerConnection))
+            {
+                // Ensure event and close-count side effects are tied to a committed transition only.
+                throw ADP.InternalConnectionError(ADP.ConnectionError.FailedToCommitSetInnerConnectionEvent);
+            }
+
             if ((originalState != currentState) && (ConnectionState.Closed == currentState))
             {
                 unchecked
@@ -2446,8 +2499,6 @@ internal void SetInnerConnectionEvent(DbConnectionInternal to)
                 }
             }
 
-            _innerConnection = to;
-
             if (ConnectionState.Closed == originalState && ConnectionState.Open == currentState)
             {
                 OnStateChange(DbConnectionInternal.StateChangeOpen);
@@ -2574,9 +2625,10 @@ private void ConnectionString_Set(ConnectionPoolKey key)
                 flag = SetInnerConnectionFrom(DbConnectionClosedBusy.SingletonInstance, connectionInternal);
                 if (flag)
                 {
+                    // Publish option/pool metadata before releasing the busy state back to a stable closed state.
                     _userConnectionOptions = connectionOptions;
                     _poolGroup = poolGroup;
-                    _innerConnection = DbConnectionClosedNeverOpened.SingletonInstance;
+                    SetInnerConnectionTo(DbConnectionClosedNeverOpened.SingletonInstance);
                 }
             }
             if (!flag)
diff --git a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/SqlConnectionConcurrentOpenTests.cs b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/SqlConnectionConcurrentOpenTests.cs
@@ -36,7 +36,11 @@ private static DbConnectionInternal GetConnectingSingleton()
 
         private static void ForceInnerConnection(SqlConnection connection, DbConnectionInternal innerConnectionValue)
         {
-            connection.SetInnerConnectionTo(innerConnectionValue);
+            // Use SetInnerConnectionFrom to transition from the current state via CAS.
+            // SetInnerConnectionTo requires the connection to already be in a transitional state.
+            DbConnectionInternal current = connection.InnerConnection;
+            bool swapped = connection.SetInnerConnectionFrom(innerConnectionValue, current);
+            Assert.True(swapped, $"CAS failed: expected current state {current.GetType().Name} but it changed concurrently.");
         }
 
         private static bool InvokeTryOpenInner(SqlConnection connection, TaskCompletionSource<DbConnectionInternal> retry)
diff --git a/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/SqlConnectionStateTransitionTests.cs b/src/Microsoft.Data.SqlClient/tests/UnitTests/Microsoft/Data/SqlClient/SqlConnectionStateTransitionTests.cs

Original file line number	Diff line number	Diff line change
`@@ -36,7 +36,11 @@ private static DbConnectionInternal GetConnectingSingleton()`
`36`	`36`
`37`	`37`	`private static void ForceInnerConnection(SqlConnection connection, DbConnectionInternal innerConnectionValue)`
`38`	`38`	`{`
`39`		`- connection.SetInnerConnectionTo(innerConnectionValue);`
	`39`	`+ // Use SetInnerConnectionFrom to transition from the current state via CAS.`
	`40`	`+ // SetInnerConnectionTo requires the connection to already be in a transitional state.`
	`41`	`+ DbConnectionInternal current = connection.InnerConnection;`
	`42`	`+ bool swapped = connection.SetInnerConnectionFrom(innerConnectionValue, current);`
	`43`	`+ Assert.True(swapped, $"CAS failed: expected current state {current.GetType().Name} but it changed concurrently.");`
`40`	`44`	`}`
`41`	`45`
`42`	`46`	`private static bool InvokeTryOpenInner(SqlConnection connection, TaskCompletionSource<DbConnectionInternal> retry)`