[QUIC] Throw ODE if connection/listener is disposed (dotnet#92215)

* AcceptConnection/StreamAsync now throw ODE in case the listener/connection was stopped by DisposeAsync. * Fix exception type and make behavior stable for disposal
Artromskiy · Sep 20, 2023 · 521e1e6 · 521e1e6
1 parent d411f50
commit 521e1e6
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 13 deletions.
diff --git a/src/libraries/System.Net.Quic/src/System/Net/Quic/QuicConnection.cs b/src/libraries/System.Net.Quic/src/System/Net/Quic/QuicConnection.cs
@@ -382,6 +382,9 @@ public async ValueTask<QuicStream> OpenOutboundStreamAsync(QuicStreamType type,
             {
                 await stream.DisposeAsync().ConfigureAwait(false);
             }
+
+            // Propagate ODE if disposed in the meantime.
+            ObjectDisposedException.ThrowIf(_disposed == 1, this);
             // Propagate connection error if present.
             if (_acceptQueue.Reader.Completion.IsFaulted)
             {
@@ -485,7 +488,7 @@ private unsafe int HandleEventShutdownInitiatedByPeer(ref SHUTDOWN_INITIATED_BY_
     }
     private unsafe int HandleEventShutdownComplete()
     {
-        Exception exception = ExceptionDispatchInfo.SetCurrentStackTrace(ThrowHelper.GetOperationAbortedException());
+        Exception exception = ExceptionDispatchInfo.SetCurrentStackTrace(_disposed == 1 ? new ObjectDisposedException(GetType().FullName) : ThrowHelper.GetOperationAbortedException());
         _acceptQueue.Writer.TryComplete(exception);
         _connectedTcs.TrySetException(exception);
         _shutdownTcs.TrySetResult();
@@ -622,7 +625,7 @@ public async ValueTask DisposeAsync()
         }
 
         // Flush the queue and dispose all remaining streams.
-        _acceptQueue.Writer.TryComplete(ExceptionDispatchInfo.SetCurrentStackTrace(ThrowHelper.GetOperationAbortedException()));
+        _acceptQueue.Writer.TryComplete(ExceptionDispatchInfo.SetCurrentStackTrace(new ObjectDisposedException(GetType().FullName)));
         while (_acceptQueue.Reader.TryRead(out QuicStream? stream))
         {
             await stream.DisposeAsync().ConfigureAwait(false);

diff --git a/src/libraries/System.Net.Quic/src/System/Net/Quic/QuicListener.cs b/src/libraries/System.Net.Quic/src/System/Net/Quic/QuicListener.cs
@@ -377,8 +377,8 @@ public async ValueTask DisposeAsync()
         _handle.Dispose();
 
         // Flush the queue and dispose all remaining connections.
-        _disposeCts.Cancel();
-        _acceptQueue.Writer.TryComplete(ExceptionDispatchInfo.SetCurrentStackTrace(ThrowHelper.GetOperationAbortedException()));
+        await _disposeCts.CancelAsync().ConfigureAwait(false);
+        _acceptQueue.Writer.TryComplete(ExceptionDispatchInfo.SetCurrentStackTrace(new ObjectDisposedException(GetType().FullName)));
         while (_acceptQueue.Reader.TryRead(out object? item))
         {
             if (item is QuicConnection connection)

diff --git a/src/libraries/System.Net.Quic/tests/FunctionalTests/QuicConnectionTests.cs b/src/libraries/System.Net.Quic/tests/FunctionalTests/QuicConnectionTests.cs
@@ -2,7 +2,6 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Collections.Generic;
-using System.Diagnostics.CodeAnalysis;
 using System.Net.Sockets;
 using System.Security.Cryptography.X509Certificates;
 using System.Threading;
@@ -88,7 +87,6 @@ await RunClientServer(
                     await AssertThrowsQuicExceptionAsync(QuicError.OperationAborted, () => connectTask);
 
                     // Subsequent attempts should fail
-                    // TODO: Which exception is correct?
                     await AssertThrowsQuicExceptionAsync(QuicError.OperationAborted, async () => await serverConnection.AcceptInboundStreamAsync());
                     await Assert.ThrowsAsync<QuicException>(() => OpenAndUseStreamAsync(serverConnection));
                 });
@@ -117,11 +115,10 @@ await RunClientServer(
                     sync.Release();
 
                     // Pending ops should fail
-                    await AssertThrowsQuicExceptionAsync(QuicError.OperationAborted, () => acceptTask);
-                    await AssertThrowsQuicExceptionAsync(QuicError.OperationAborted, () => connectTask);
+                    await Assert.ThrowsAsync<ObjectDisposedException>(async () => await acceptTask);
+                    await Assert.ThrowsAsync<ObjectDisposedException>(async () => await connectTask);
 
                     // Subsequent attempts should fail
-                    // TODO: Should these be QuicOperationAbortedException, to match above? Or vice-versa?
                     await Assert.ThrowsAsync<ObjectDisposedException>(async () => await serverConnection.AcceptInboundStreamAsync());
                     await Assert.ThrowsAsync<ObjectDisposedException>(async () => await OpenAndUseStreamAsync(serverConnection));
                 });
@@ -312,6 +309,21 @@ await RunClientServer(
                 _ => Task.CompletedTask);
         }
 
+        [Fact]
+        public async Task AcceptStreamAsync_ConnectionDisposed_Throws()
+        {
+            (QuicConnection clientConnection, QuicConnection serverConnection) = await CreateConnectedQuicConnection();
+
+            // One task issues before the disposal.
+            ValueTask<QuicStream> acceptTask1 = serverConnection.AcceptInboundStreamAsync();
+            await serverConnection.DisposeAsync();
+            // Another task issued after the disposal.
+            ValueTask<QuicStream> acceptTask2 = serverConnection.AcceptInboundStreamAsync();
+
+            var accept1Exception = await Assert.ThrowsAsync<ObjectDisposedException>(async () => await acceptTask1);
+            var accept2Exception = await Assert.ThrowsAsync<ObjectDisposedException>(async () => await acceptTask2);
+        }
+
         [Theory]
         [InlineData(true)]
         [InlineData(false)]

diff --git a/src/libraries/System.Net.Quic/tests/FunctionalTests/QuicListenerTests.cs b/src/libraries/System.Net.Quic/tests/FunctionalTests/QuicListenerTests.cs
@@ -196,10 +196,8 @@ public async Task AcceptConnectionAsync_ListenerDisposed_Throws()
             await listener.DisposeAsync();
             serverDisposed.SetResult();
 
-            var accept1Exception = await AssertThrowsQuicExceptionAsync(QuicError.OperationAborted, async () => await acceptTask1);
-            var accept2Exception = await AssertThrowsQuicExceptionAsync(QuicError.OperationAborted, async () => await acceptTask2);
-
-            Assert.Equal(accept1Exception, accept2Exception);
+            var accept1Exception = await Assert.ThrowsAsync<ObjectDisposedException>(async () => await acceptTask1);
+            var accept2Exception = await Assert.ThrowsAsync<ObjectDisposedException>(async () => await acceptTask2);
 
             // Connect attempt should be stopped with "UserCanceled".
             var connectException = await Assert.ThrowsAsync<AuthenticationException>(async () => await connectTask);