SqlClient handle socket network error gracefully

src/System.Data.SqlClient/src/System/Data/SqlClient/SNI/SNIMarsConnection.cs

@@ -44,7 +44,7 @@ public SNIMarsConnection(SNIHandle lowerHandle)
             _lowerHandle.SetAsyncCallbacks(HandleReceiveComplete, HandleSendComplete);
         }
 
-        public SNIMarsHandle CreateMarsSession(object callbackObject, bool async)
+        public SNIMarsHandle CreateMarsSession(TdsParserStateObject callbackObject, bool async)
         {
             lock (this)
             {
@@ -126,12 +126,12 @@ public uint CheckConnection()
         /// <summary>
         /// Process a receive error
         /// </summary>
-        public void HandleReceiveError(SNIPacket packet)
+        public void HandleReceiveError(SNIPacket packet, uint sniErrorCode)
         {
             Debug.Assert(Monitor.IsEntered(this), "HandleReceiveError was called without being locked.");
             foreach (SNIMarsHandle handle in _sessions.Values)
             {
-                handle.HandleReceiveError(packet);
+                handle.HandleReceiveError(packet, sniErrorCode);
             }
             packet?.Dispose();
         }
@@ -161,7 +161,7 @@ public void HandleReceiveComplete(SNIPacket packet, uint sniErrorCode)
             {
                 lock (this)
                 {
-                    HandleReceiveError(packet);
+                    HandleReceiveError(packet, sniErrorCode);
                     return;
                 }
             }
@@ -192,7 +192,7 @@ public void HandleReceiveComplete(SNIPacket packet, uint sniErrorCode)
                                     return;
                                 }
 
-                                HandleReceiveError(packet);
+                                HandleReceiveError(packet, sniErrorCode);
                                 return;
                             }
                         }
@@ -223,7 +223,7 @@ public void HandleReceiveComplete(SNIPacket packet, uint sniErrorCode)
                                     return;
                                 }
 
-                                HandleReceiveError(packet);
+                                HandleReceiveError(packet, sniErrorCode);
                                 return;
                             }
                         }
@@ -234,7 +234,7 @@ public void HandleReceiveComplete(SNIPacket packet, uint sniErrorCode)
                     if (!_sessions.ContainsKey(_currentHeader.sessionId))
                     {
                         SNILoadHandle.SingletonInstance.LastError = new SNIError(SNIProviders.SMUX_PROV, 0, SNICommon.InvalidParameterError, string.Empty);
-                        HandleReceiveError(packet);
+                        HandleReceiveError(packet, sniErrorCode);
                         _lowerHandle.Dispose();
                         _lowerHandle = null;
                         return;
@@ -280,7 +280,7 @@ public void HandleReceiveComplete(SNIPacket packet, uint sniErrorCode)
                             return;
                         }
 
-                        HandleReceiveError(packet);
+                        HandleReceiveError(packet, sniErrorCode);
                         return;
                     }
                 }

src/System.Data.SqlClient/src/System/Data/SqlClient/SNI/SNIMarsHandle.cs

@@ -19,7 +19,7 @@ internal class SNIMarsHandle : SNIHandle
         private readonly uint _status = TdsEnums.SNI_UNINITIALIZED;
         private readonly Queue<SNIPacket> _receivedPacketQueue = new Queue<SNIPacket>();
         private readonly Queue<SNIMarsQueuedPacket> _sendPacketQueue = new Queue<SNIMarsQueuedPacket>();
-        private readonly object _callbackObject;
+        private readonly TdsParserStateObject _callbackObject;
         private readonly Guid _connectionId = Guid.NewGuid();
         private readonly ushort _sessionId;
         private readonly ManualResetEventSlim _packetEvent = new ManualResetEventSlim(false);
@@ -78,7 +78,7 @@ public override void Dispose()
         /// <param name="sessionId">MARS session ID</param>
         /// <param name="callbackObject">Callback object</param>
         /// <param name="async">true if connection is asynchronous</param>
-        public SNIMarsHandle(SNIMarsConnection connection, ushort sessionId, object callbackObject, bool async)
+        public SNIMarsHandle(SNIMarsConnection connection, ushort sessionId, TdsParserStateObject callbackObject, bool async)
         {
             _sessionId = sessionId;
             _connection = connection;
@@ -295,15 +295,23 @@ public override uint ReceiveAsync(ref SNIPacket packet)
         /// <summary>
         /// Handle receive error
         /// </summary>
-        public void HandleReceiveError(SNIPacket packet)
+        public void HandleReceiveError(SNIPacket packet, uint sniErrorCode)
         {
             lock (_receivedPacketQueue)
             {
                 _connectionError = SNILoadHandle.SingletonInstance.LastError;
                 _packetEvent.Set();
             }
-
-            ((TdsParserStateObject)_callbackObject).ReadAsyncCallback(PacketHandle.FromManagedPacket(packet), 1);
+            if (sniErrorCode == TdsEnums.SNI_WSAECONNRESET)
+            {
+                TdsParser parser = _callbackObject.Parser;
+                parser.State = TdsParserState.Broken;
+                parser.Connection.BreakConnection();
+            }
+            else
+            {
+                _callbackObject.ReadAsyncCallback(PacketHandle.FromManagedPacket(packet), sniErrorCode);
+            }
         }
 
         /// <summary>
@@ -317,7 +325,7 @@ public void HandleSendComplete(SNIPacket packet, uint sniErrorCode)
             {
                 Debug.Assert(_callbackObject != null);
 
-                ((TdsParserStateObject)_callbackObject).WriteAsyncCallback(PacketHandle.FromManagedPacket(packet), sniErrorCode);
+                _callbackObject.WriteAsyncCallback(PacketHandle.FromManagedPacket(packet), sniErrorCode);
             }
         }
 

src/System.Data.SqlClient/src/System/Data/SqlClient/SNI/SNIPacket.NetCoreApp.cs

@@ -4,6 +4,7 @@
 
 using System.Buffers;
 using System.IO;
+using System.Net.Sockets;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -18,31 +19,39 @@ internal partial class SNIPacket
         /// <param name="callback">Completion callback</param>
         public void ReadFromStreamAsync(Stream stream, SNIAsyncCallback callback)
         {
-            // Treat local function as a static and pass all params otherwise as async will allocate
-            async Task ReadFromStreamAsync(SNIPacket packet, SNIAsyncCallback cb, ValueTask<int> valueTask)
+            static async Task ReadFromStreamAsync(SNIPacket packet, SNIAsyncCallback cb, ValueTask<int> valueTask)
             {
-                bool error = false;
+                uint errorCode = TdsEnums.SNI_SUCCESS;
                 try
                 {
                     packet._length = await valueTask.ConfigureAwait(false);
                     if (packet._length == 0)
                     {
                         SNILoadHandle.SingletonInstance.LastError = new SNIError(SNIProviders.TCP_PROV, 0, SNICommon.ConnTerminatedError, string.Empty);
-                        error = true;
+                        errorCode = TdsEnums.SNI_WSAECONNRESET;
                     }
                 }
+                catch (IOException ioException) when ( 
+                    ioException?.InnerException is SocketException socketException && 
+                    socketException!=null && 
+                    socketException.SocketErrorCode==SocketError.OperationAborted
+                )
+                {
+                    SNILoadHandle.SingletonInstance.LastError = new SNIError(SNIProviders.TCP_PROV, SNICommon.InternalExceptionError, socketException);
+                    errorCode = TdsEnums.SNI_WSAECONNRESET;
+                }
                 catch (Exception ex)
                 {
                     SNILoadHandle.SingletonInstance.LastError = new SNIError(SNIProviders.TCP_PROV, SNICommon.InternalExceptionError, ex);
-                    error = true;
+                    errorCode = TdsEnums.SNI_ERROR;
                 }
 
-                if (error)
+                if (errorCode != TdsEnums.SNI_SUCCESS)
                 {
                     packet.Release();
                 }
 
-                cb(packet, error ? TdsEnums.SNI_ERROR : TdsEnums.SNI_SUCCESS);
+                cb(packet, errorCode);
             }
 
             ValueTask<int> vt = stream.ReadAsync(new Memory<byte>(_data, 0, _capacity), CancellationToken.None);
@@ -70,8 +79,7 @@ async Task ReadFromStreamAsync(SNIPacket packet, SNIAsyncCallback cb, ValueTask<
         /// <param name="stream">Stream to write to</param>
         public void WriteToStreamAsync(Stream stream, SNIAsyncCallback callback, SNIProviders provider, bool disposeAfterWriteAsync = false)
         {
-            // Treat local function as a static and pass all params otherwise as async will allocate
-            async Task WriteToStreamAsync(SNIPacket packet, SNIAsyncCallback cb, SNIProviders providers, bool disposeAfter, ValueTask valueTask)
+            static async Task WriteToStreamAsync(SNIPacket packet, SNIAsyncCallback cb, SNIProviders providers, bool disposeAfter, ValueTask valueTask)
             {
                 uint status = TdsEnums.SNI_SUCCESS;
                 try

src/System.Data.SqlClient/src/System/Data/SqlClient/SNI/SNIPacket.NetStandard.cs

@@ -4,6 +4,7 @@
 
 using System.Buffers;
 using System.IO;
+using System.Net.Sockets;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -18,31 +19,39 @@ internal partial class SNIPacket
         /// <param name="callback">Completion callback</param>
         public void ReadFromStreamAsync(Stream stream, SNIAsyncCallback callback)
         {
-            // Treat local function as a static and pass all params otherwise as async will allocate
-            async Task ReadFromStreamAsync(SNIPacket packet, SNIAsyncCallback cb, Task<int> task)
+            static async Task ReadFromStreamAsync(SNIPacket packet, SNIAsyncCallback cb, Task<int> task)
             {
-                bool error = false;
+                uint errorCode = TdsEnums.SNI_SUCCESS;
                 try
                 {
                     packet._length = await task.ConfigureAwait(false);
                     if (packet._length == 0)
                     {
                         SNILoadHandle.SingletonInstance.LastError = new SNIError(SNIProviders.TCP_PROV, 0, SNICommon.ConnTerminatedError, string.Empty);
-                        error = true;
+                        errorCode = TdsEnums.SNI_WSAECONNRESET;
                     }
                 }
+                catch (IOException ioException) when (
+                    ioException?.InnerException is SocketException socketException &&
+                    socketException != null &&
+                    socketException.SocketErrorCode == SocketError.OperationAborted
+                )
+                {
+                    SNILoadHandle.SingletonInstance.LastError = new SNIError(SNIProviders.TCP_PROV, SNICommon.InternalExceptionError, socketException);
+                    errorCode = TdsEnums.SNI_WSAECONNRESET;
+                }
                 catch (Exception ex)
                 {
                     SNILoadHandle.SingletonInstance.LastError = new SNIError(SNIProviders.TCP_PROV, SNICommon.InternalExceptionError, ex);
-                    error = true;
+                    errorCode = TdsEnums.SNI_ERROR;
                 }
 
-                if (error)
+                if (errorCode != TdsEnums.SNI_SUCCESS)
                 {
                     packet.Release();
                 }
 
-                cb(packet, error ? TdsEnums.SNI_ERROR : TdsEnums.SNI_SUCCESS);
+                cb(packet, errorCode);
             }
 
             Task<int> t = stream.ReadAsync(_data, 0, _capacity, CancellationToken.None);

src/System.Data.SqlClient/src/System/Data/SqlClient/TdsEnums.cs

@@ -538,7 +538,7 @@ public enum FedAuthLibrary : byte
         public const uint SNI_SUCCESS_IO_PENDING = 997;      // Overlapped I/O operation is in progress.
 
         // Windows Sockets Error Codes
-        public const short SNI_WSAECONNRESET = 10054;    // An existing connection was forcibly closed by the remote host.
+        public const uint SNI_WSAECONNRESET = 10054;    // An existing connection was forcibly closed by the remote host.
 
         // SNI internal errors (shouldn't overlap with Win32 / socket errors)
         public const uint SNI_QUEUE_FULL = 1048576;		 // Packet queue is full

src/System.Data.SqlClient/src/System/Data/SqlClient/TdsParserStateObject.cs

@@ -2691,7 +2691,7 @@ public void ReadAsyncCallback(IntPtr key, PacketHandle packet, uint error)
 
             // The mars physical connection can get a callback
             // with a packet but no result after the connection is closed.
-            if (source == null && _parser._pMarsPhysicalConObj == this)
+            if (source == null)
             {
                 return;
             }

src/System.Data.SqlClient/src/System/Data/SqlClient/TdsParserStateObjectManaged.cs

@@ -47,7 +47,7 @@ protected override void CreateSessionHandle(TdsParserStateObject physicalConnect
             _sessionHandle = managedSNIObject.CreateMarsSession(this, async);
         }
 
-        internal SNIMarsHandle CreateMarsSession(object callbackObject, bool async)
+        internal SNIMarsHandle CreateMarsSession(TdsParserStateObject callbackObject, bool async)
         {
             return _marsConnection.CreateMarsSession(callbackObject, async);
         }

src/System.Data.SqlClient/tests/ManualTests/SQL/ConnectionPoolTest/ConnectionPoolTest.cs

@@ -15,13 +15,13 @@ public static class ConnectionPoolTest
         private static readonly string _tcpMarsConnStr = (new SqlConnectionStringBuilder(DataTestUtility.TcpConnStr) { MultipleActiveResultSets = true, Pooling = true }).ConnectionString;
 
 
-        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup), /* [ActiveIssue(33930)]: */ nameof(DataTestUtility.IsUsingNativeSNI))]
+        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup))]
         public static void ConnectionPool_NonMars()
         {
             RunDataTestForSingleConnString(_tcpConnStr);
         }
 
-        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup), /* [ActiveIssue(33930)] */ nameof(DataTestUtility.IsUsingNativeSNI))]
+        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup))]
         public static void ConnectionPool_Mars()
         {
             RunDataTestForSingleConnString(_tcpMarsConnStr);

src/System.Data.SqlClient/tests/ManualTests/SQL/ConnectionPoolTest/PoolBlockPeriodTest.netcoreapp.cs

@@ -21,7 +21,7 @@ public class PoolBlockPeriodTest
         private const int ConnectionTimeout = 15;
         private const int CompareMargin = 2;
 
-        [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), /* [ActiveIssue(33930)] */ nameof(DataTestUtility.IsUsingNativeSNI))]
+        [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
         [InlineData("Azure with Default Policy must Disable blocking (*.database.windows.net)", new object[] { AzureEndpointSample })]
         [InlineData("Azure with Default Policy must Disable blocking (*.database.chinacloudapi.cn)", new object[] { AzureChinaEnpointSample })]
         [InlineData("Azure with Default Policy must Disable blocking (*.database.usgovcloudapi.net)", new object[] { AzureUSGovernmentEndpointSample })]
@@ -45,7 +45,7 @@ public void TestAzureBlockingPeriod(string description, object[] Params)
             PoolBlockingPeriodAzureTest(connString, policy);
         }
 
-        [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), /* [ActiveIssue(33930)] */ nameof(DataTestUtility.IsUsingNativeSNI))]
+        [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
         [InlineData("NonAzure with Default Policy must Enable blocking", new object[] { NonExistentServer })]
         [InlineData("NonAzure with Auto Policy must Enable Blocking", new object[] { NonExistentServer, PoolBlockingPeriod.Auto })]
         [InlineData("NonAzure with Always Policy must Enable Blocking", new object[] { NonExistentServer, PoolBlockingPeriod.AlwaysBlock })]
@@ -66,7 +66,7 @@ public void TestNonAzureBlockingPeriod(string description, object[] Params)
             PoolBlockingPeriodNonAzureTest(connString, policy);
         }
 
-        [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), /* [ActiveIssue(33930)] */ nameof(DataTestUtility.IsUsingNativeSNI))]
+        [ConditionalTheory(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
         [InlineData("Test policy with Auto (lowercase)", "auto")]
         [InlineData("Test policy with Auto (PascalCase)", "Auto")]
         [InlineData("Test policy with Always (lowercase)", "alwaysblock")]

src/System.Data.SqlClient/tests/ManualTests/SQL/DataStreamTest/DataStreamTest.cs

@@ -64,10 +64,10 @@ private static void RunAllTestsForSingleServer(string connectionString, bool usi
             // These tests fail with named pipes, since they try to do DNS lookups on named pipe paths.
             if (!usingNamePipes)
             {
-                //if (DataTestUtility.IsUsingNativeSNI()) /* [ActiveIssue(33930)] */
-                //{
-                //    TimeoutDuringReadAsyncWithClosedReaderTest(connectionString);
-                //}
+                if (DataTestUtility.IsUsingNativeSNI())
+                {
+                    TimeoutDuringReadAsyncWithClosedReaderTest(connectionString);
+                }
                 NonFatalTimeoutDuringRead(connectionString);
             }
         }

src/System.Data.SqlClient/tests/ManualTests/SQL/MARSTest/MARSTest.cs

@@ -33,7 +33,7 @@ public static void NamedPipesMARSTest()
         }
 
 #if DEBUG
-        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup), /* [ActiveIssue(33930)] */ nameof(DataTestUtility.IsUsingNativeSNI))]
+        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup))]
         public static void MARSAsyncTimeoutTest()
         {
             using (SqlConnection connection = new SqlConnection(_connStr))
@@ -73,7 +73,7 @@ public static void MARSAsyncTimeoutTest()
             }
         }
 
-        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup), /* [ActiveIssue(33930)] */ nameof(DataTestUtility.IsUsingNativeSNI))]
+        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup))]
         public static void MARSSyncTimeoutTest()
         {
             using (SqlConnection connection = new SqlConnection(_connStr))

src/System.Data.SqlClient/tests/ManualTests/SQL/SqlCredentialTest/SqlCredentialTest.cs

@@ -17,7 +17,7 @@ namespace System.Data.SqlClient.ManualTesting.Tests
     public static class SqlCredentialTest
     {
 
-        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup), /* [ActiveIssue(33930)] */ nameof(DataTestUtility.IsUsingNativeSNI))]
+        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup))]
         public static void CreateSqlConnectionWithCredential()
         {
             var user = "u" + Guid.NewGuid().ToString().Replace("-", "");
@@ -49,7 +49,7 @@ public static void CreateSqlConnectionWithCredential()
             }
         }
 
-        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup), /* [ActiveIssue(33930)] */ nameof(DataTestUtility.IsUsingNativeSNI))]
+        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup))]
         public static void SqlConnectionChangePasswordPlaintext()
         {
             var user = "u" + Guid.NewGuid().ToString().Replace("-", "");
@@ -82,7 +82,7 @@ public static void SqlConnectionChangePasswordPlaintext()
             }
         }
 
-        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup), /* [ActiveIssue(33930)] */ nameof(DataTestUtility.IsUsingNativeSNI))]
+        [ConditionalFact(typeof(DataTestUtility),nameof(DataTestUtility.AreConnStringsSetup))]
         public static void SqlConnectionChangePasswordSecureString()
         {
             var user = "u" + Guid.NewGuid().ToString().Replace("-", "");
@@ -122,7 +122,7 @@ public static void SqlConnectionChangePasswordSecureString()
             }
         }
 
-        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup), /* [ActiveIssue(33930)] */ nameof(DataTestUtility.IsUsingNativeSNI))]
+        [ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
         public static void OldCredentialsShouldFail()
         {
             String user = "u" + Guid.NewGuid().ToString().Replace("-", "");