Fix cancelability of DNS queries (#104420)

* Fix cancelability of DNS queries * Properly cascade cancelation info for connect by name exceptions
dotnet · Jul 12, 2024 · 45f3250 · 45f3250
1 parent 42b2fc8
commit 45f3250
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/src/libraries/System.Net.Sockets/src/System/Net/Sockets/Socket.Tasks.cs b/src/libraries/System.Net.Sockets/src/System/Net/Sockets/Socket.Tasks.cs
@@ -97,7 +97,7 @@ public ValueTask ConnectAsync(EndPoint remoteEP, CancellationToken cancellationT
 
             saea.RemoteEndPoint = remoteEP;
 
-            ValueTask connectTask = saea.ConnectAsync(this);
+            ValueTask connectTask = saea.ConnectAsync(this, saeaCancelable: cancellationToken.CanBeCanceled);
             if (connectTask.IsCompleted || !cancellationToken.CanBeCanceled)
             {
                 // Avoid async invocation overhead
@@ -1210,11 +1210,11 @@ public ValueTask<int> SendToAsync(Socket socket, CancellationToken cancellationT
                     ValueTask.FromException<int>(CreateException(error));
             }
 
-            public ValueTask ConnectAsync(Socket socket)
+            public ValueTask ConnectAsync(Socket socket, bool saeaCancelable)
             {
                 try
                 {
-                    if (socket.ConnectAsync(this, userSocket: true, saeaCancelable: false))
+                    if (socket.ConnectAsync(this, userSocket: true, saeaCancelable: saeaCancelable))
                     {
                         return new ValueTask(this, _mrvtsc.Version);
                     }

diff --git a/src/libraries/System.Net.Sockets/src/System/Net/Sockets/SocketAsyncEventArgs.cs b/src/libraries/System.Net.Sockets/src/System/Net/Sockets/SocketAsyncEventArgs.cs
@@ -428,6 +428,12 @@ internal void SetResults(Exception exception, int bytesTransferred, SocketFlags
                 {
                     _socketError = socketException.SocketErrorCode;
                 }
+                else if (exception is OperationCanceledException)
+                {
+                    // Preserve information about the cancellation when it is canceled at non Socket operation.
+                    // It is used to throw the right exception later in the stack.
+                    _socketError = SocketError.OperationAborted;
+                }
                 else
                 {
                     _socketError = SocketError.SocketError;