Fix QUIC ConnectionState NRE in HandleEventConnectionClose (#57655)

Moved ConnectionState assignment before msquic callback registration to avoid NRE in callback, in case Connection gets closed during Stream's ctor. Fixes #55815
dotnet · Aug 19, 2021 · d367827 · d367827
1 parent b75e55b
commit d367827
Showing 1 changed file with 16 additions and 17 deletions.
diff --git a/src/libraries/System.Net.Quic/src/System/Net/Quic/Implementations/MsQuic/MsQuicStream.cs b/src/libraries/System.Net.Quic/src/System/Net/Quic/Implementations/MsQuic/MsQuicStream.cs
@@ -95,6 +95,14 @@ public void Cleanup()
         // inbound.
         internal MsQuicStream(MsQuicConnection.State connectionState, SafeMsQuicStreamHandle streamHandle, QUIC_STREAM_OPEN_FLAGS flags)
         {
+            if (!connectionState.TryAddStream(this))
+            {
+                throw new ObjectDisposedException(nameof(QuicConnection));
+            }
+            // this assignment should be done before SetCallbackHandlerDelegate to prevent NRE in HandleEventConnectionClose
+            // but after TryAddStream to prevent unnecessary RemoveStream in finalizer
+            _state.ConnectionState = connectionState;
+
             _state.Handle = streamHandle;
             _canRead = true;
             _canWrite = !flags.HasFlag(QUIC_STREAM_OPEN_FLAGS.UNIDIRECTIONAL);
@@ -117,14 +125,6 @@ internal MsQuicStream(MsQuicConnection.State connectionState, SafeMsQuicStreamHa
                 throw;
             }
 
-            if (!connectionState.TryAddStream(this))
-            {
-                _state.StateGCHandle.Free();
-                throw new ObjectDisposedException(nameof(QuicConnection));
-            }
-
-            _state.ConnectionState = connectionState;
-
             _state.TraceId = MsQuicTraceHelper.GetTraceId(_state.Handle);
             if (NetEventSource.Log.IsEnabled())
             {
@@ -140,6 +140,14 @@ internal MsQuicStream(MsQuicConnection.State connectionState, QUIC_STREAM_OPEN_F
         {
             Debug.Assert(connectionState.Handle != null);
 
+            if (!connectionState.TryAddStream(this))
+            {
+                throw new ObjectDisposedException(nameof(QuicConnection));
+            }
+            // this assignment should be done before StreamOpenDelegate to prevent NRE in HandleEventConnectionClose
+            // but after TryAddStream to prevent unnecessary RemoveStream in finalizer
+            _state.ConnectionState = connectionState;
+
             _canRead = !flags.HasFlag(QUIC_STREAM_OPEN_FLAGS.UNIDIRECTIONAL);
             _canWrite = true;
 
@@ -170,15 +178,6 @@ internal MsQuicStream(MsQuicConnection.State connectionState, QUIC_STREAM_OPEN_F
                 throw;
             }
 
-            if (!connectionState.TryAddStream(this))
-            {
-                _state.Handle?.Dispose();
-                _state.StateGCHandle.Free();
-                throw new ObjectDisposedException(nameof(QuicConnection));
-            }
-
-            _state.ConnectionState = connectionState;
-
             _state.TraceId = MsQuicTraceHelper.GetTraceId(_state.Handle);
             if (NetEventSource.Log.IsEnabled())
             {