Send buffers and handles crash fixes

src/libraries/System.Net.Quic/src/System/Net/Quic/Internal/MsQuicBuffers.cs

@@ -1,9 +1,11 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Diagnostics;
 using System.Collections.Generic;
 using System.Runtime.InteropServices;
 using Microsoft.Quic;
+using System.Threading;
 
 namespace System.Net.Quic;
 
@@ -18,14 +20,27 @@ internal unsafe struct MsQuicBuffers : IDisposable
     private QUIC_BUFFER* _buffers;
     // Number of QUIC_BUFFER instance currently allocated in _buffers, so that we can reuse the memory instead of reallocating.
     private int _count;
+    private bool _initialized;
+    private bool _disposed;
 
     public MsQuicBuffers()
     {
         _buffers = null;
         _count = 0;
+        _initialized=false;
+        _disposed = false;
+    }
+
+    public QUIC_BUFFER* Buffers
+    {
+        get
+        {
+            ObjectDisposedException.ThrowIf(_disposed, typeof(MsQuicBuffers));
+            Debug.Assert(_initialized);
+            return _buffers;
+        }
     }
 
-    public QUIC_BUFFER* Buffers => _buffers;
     public int Count => _count;
 
     private void FreeNativeMemory()
@@ -48,6 +63,7 @@ private void Reserve(int count)
 
     private void SetBuffer(int index, ReadOnlyMemory<byte> buffer)
     {
+        Debug.Assert(_buffers[index].Buffer == null);
         _buffers[index].Buffer = (byte*)NativeMemory.Alloc((nuint)buffer.Length, (nuint)sizeof(byte));
         _buffers[index].Length = (uint)buffer.Length;
         buffer.Span.CopyTo(_buffers[index].Span);
@@ -62,11 +78,14 @@ private void SetBuffer(int index, ReadOnlyMemory<byte> buffer)
     /// <typeparam name="T">The type of the inputs.</typeparam>
     public void Initialize<T>(IList<T> inputs, Func<T, ReadOnlyMemory<byte>> toBuffer)
     {
+        ObjectDisposedException.ThrowIf(_disposed, typeof(MsQuicBuffers));
+        Debug.Assert(!_initialized);
         Reserve(inputs.Count);
         for (int i = 0; i < inputs.Count; ++i)
         {
             SetBuffer(i, toBuffer(inputs[i]));
         }
+        _initialized = true;
     }
 
     /// <summary>
@@ -76,15 +95,21 @@ public void Initialize<T>(IList<T> inputs, Func<T, ReadOnlyMemory<byte>> toBuffe
     /// <param name="buffer">Buffer to be passed to MsQuic as QUIC_BUFFER*.</param>
     public void Initialize(ReadOnlyMemory<byte> buffer)
     {
+        ObjectDisposedException.ThrowIf(_disposed, typeof(MsQuicBuffers));
+        Debug.Assert(!_initialized);
         Reserve(1);
         SetBuffer(0, buffer);
+        _initialized = true;
     }
 
     /// <summary>
     /// Release the native memory of individual buffers and allows reuse of this struct.
     /// </summary>
-    public void Reset()
+    public void Reset() => Reset(disposing: false);
+
+    private void Reset(bool disposing)
     {
+        ObjectDisposedException.ThrowIf(_disposed && !disposing, typeof(MsQuicBuffers));
         for (int i = 0; i < _count; ++i)
         {
             if (_buffers[i].Buffer is null)
@@ -96,14 +121,16 @@ public void Reset()
             NativeMemory.Free(buffer);
             _buffers[i].Length = 0;
         }
+        _initialized = false;
     }
 
     /// <summary>
     /// Releases all the native memory.
     /// </summary>
     public void Dispose()
     {
-        Reset();
+        _disposed = true;
+        Reset(disposing: true);
         FreeNativeMemory();
     }
 }

src/libraries/System.Net.Quic/src/System/Net/Quic/Internal/MsQuicSafeHandle.cs

@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Diagnostics;
 using System.Runtime.InteropServices;
 using Microsoft.Quic;
 
@@ -25,7 +26,14 @@ internal unsafe class MsQuicSafeHandle : SafeHandle
 
     public override bool IsInvalid => handle == IntPtr.Zero;
 
-    public QUIC_HANDLE* QuicHandle => (QUIC_HANDLE*)DangerousGetHandle();
+    public QUIC_HANDLE* QuicHandle
+    {
+        get
+        {
+            ObjectDisposedException.ThrowIf(IsInvalid, this);
+            return (QUIC_HANDLE*)DangerousGetHandle();
+        }
+    }
 
     public MsQuicSafeHandle(QUIC_HANDLE* handle, delegate* unmanaged[Cdecl]<QUIC_HANDLE*, void> releaseAction, SafeHandleType safeHandleType)
         : base((IntPtr)handle, ownsHandle: true)
@@ -41,8 +49,9 @@ public MsQuicSafeHandle(QUIC_HANDLE* handle, delegate* unmanaged[Cdecl]<QUIC_HAN
 
     protected override bool ReleaseHandle()
     {
-        _releaseAction(QuicHandle);
+        QUIC_HANDLE* quicHandle = QuicHandle;
         SetHandle(IntPtr.Zero);
+        _releaseAction(quicHandle);
 
         if (NetEventSource.Log.IsEnabled())
         {
@@ -52,6 +61,44 @@ protected override bool ReleaseHandle()
         return true;
     }
 
+    public TResult SafeCall<TResult>(Func<MsQuicSafeHandle, TResult> call)
+    {
+        ObjectDisposedException.ThrowIf(IsInvalid, this);
+        bool success = false;
+        try
+        {
+            DangerousAddRef(ref success);
+            Debug.Assert(success);
+            return call(this);
+        }
+        finally
+        {
+            if (success)
+            {
+                DangerousRelease();
+            }
+        }
+    }
+
+    public void SafeCall(Action<MsQuicSafeHandle> call)
+    {
+        ObjectDisposedException.ThrowIf(IsInvalid, this);
+        bool success = false;
+        try
+        {
+            DangerousAddRef(ref success);
+            Debug.Assert(success);
+            call(this);
+        }
+        finally
+        {
+            if (success)
+            {
+                DangerousRelease();
+            }
+        }
+    }
+
     public override string ToString() => _traceId ??= $"[{s_typeName[(int)_type]}][0x{DangerousGetHandle():X11}]";
 }
 

src/libraries/System.Net.Quic/src/System/Net/Quic/QuicConnection.cs

@@ -242,7 +242,7 @@ private async ValueTask FinishConnectAsync(QuicClientConnectionOptions options,
             if (address is not null)
             {
                 QuicAddr quicAddress = new IPEndPoint(address, port).ToQuicAddr();
-                MsQuicHelpers.SetMsQuicParameter(_handle, QUIC_PARAM_CONN_REMOTE_ADDRESS, quicAddress);
+                _handle.SafeCall(handle => MsQuicHelpers.SetMsQuicParameter(handle, QUIC_PARAM_CONN_REMOTE_ADDRESS, quicAddress));
             }
             // RemoteEndPoint is DnsEndPoint containing hostname that is different from requested SNI.
             // --> Resolve the hostname and set the IP directly, use requested SNI in ConnectionStart.
@@ -257,7 +257,7 @@ private async ValueTask FinishConnectAsync(QuicClientConnectionOptions options,
                 }
 
                 QuicAddr quicAddress = new IPEndPoint(addresses[0], port).ToQuicAddr();
-                MsQuicHelpers.SetMsQuicParameter(_handle, QUIC_PARAM_CONN_REMOTE_ADDRESS, quicAddress);
+                _handle.SafeCall(handle => MsQuicHelpers.SetMsQuicParameter(handle, QUIC_PARAM_CONN_REMOTE_ADDRESS, quicAddress));
             }
             // RemoteEndPoint is DnsEndPoint containing hostname that is the same as the requested SNI.
             // --> Let MsQuic resolve the hostname/SNI, give address family hint is specified in DnsEndPoint.
@@ -276,7 +276,7 @@ private async ValueTask FinishConnectAsync(QuicClientConnectionOptions options,
             if (options.LocalEndPoint is not null)
             {
                 QuicAddr quicAddress = options.LocalEndPoint.ToQuicAddr();
-                MsQuicHelpers.SetMsQuicParameter(_handle, QUIC_PARAM_CONN_LOCAL_ADDRESS, quicAddress);
+                _handle.SafeCall(handle => MsQuicHelpers.SetMsQuicParameter(handle, QUIC_PARAM_CONN_LOCAL_ADDRESS, quicAddress));
             }
 
             _sslConnectionOptions = new SslConnectionOptions(
@@ -294,13 +294,13 @@ private async ValueTask FinishConnectAsync(QuicClientConnectionOptions options,
             {
                 unsafe
                 {
-                    ThrowHelper.ThrowIfMsQuicError(MsQuicApi.Api.ApiTable->ConnectionStart(
-                        _handle.QuicHandle,
+                    _handle.SafeCall(handle => ThrowHelper.ThrowIfMsQuicError(MsQuicApi.Api.ApiTable->ConnectionStart(
+                        handle.QuicHandle,
                         _configuration.QuicHandle,
                         (ushort)addressFamily,
                         (sbyte*)targetHostPtr,
                         (ushort)port),
-                        "ConnectionStart failed");
+                        "ConnectionStart failed"));
                 }
             }
             finally
@@ -334,10 +334,10 @@ internal ValueTask FinishHandshakeAsync(QuicServerConnectionOptions options, str
 
             unsafe
             {
-                ThrowHelper.ThrowIfMsQuicError(MsQuicApi.Api.ApiTable->ConnectionSetConfiguration(
-                    _handle.QuicHandle,
+                _handle.SafeCall(handle => ThrowHelper.ThrowIfMsQuicError(MsQuicApi.Api.ApiTable->ConnectionSetConfiguration(
+                    handle.QuicHandle,
                     _configuration.QuicHandle),
-                    "ConnectionSetConfiguration failed");
+                    "ConnectionSetConfiguration failed"));
             }
         }
 
@@ -359,7 +359,7 @@ public async ValueTask<QuicStream> OpenOutboundStreamAsync(QuicStreamType type,
         QuicStream? stream = null;
         try
         {
-            stream = new QuicStream(_handle, type, _defaultStreamErrorCode);
+            stream = _handle.SafeCall(handle => new QuicStream((MsQuicContextSafeHandle)handle, type, _defaultStreamErrorCode));
             await stream.StartAsync(cancellationToken).ConfigureAwait(false);
         }
         catch
@@ -430,10 +430,10 @@ public ValueTask CloseAsync(long errorCode, CancellationToken cancellationToken
         {
             unsafe
             {
-                MsQuicApi.Api.ApiTable->ConnectionShutdown(
-                    _handle.QuicHandle,
+                _handle.SafeCall(handle => MsQuicApi.Api.ApiTable->ConnectionShutdown(
+                    handle.QuicHandle,
                     QUIC_CONNECTION_SHUTDOWN_FLAGS.NONE,
-                    (ulong)errorCode);
+                    (ulong)errorCode));
             }
         }
 
@@ -492,8 +492,8 @@ private unsafe int HandleEventShutdownComplete(ref SHUTDOWN_COMPLETE_DATA data)
             NetEventSource.Info(this, $"{this} Received event SHUTDOWN_INITIATED_BY_PEER_DATA");
         }
 
-        _shutdownTcs.TrySetResult();
         _acceptQueue.Writer.TryComplete(ExceptionDispatchInfo.SetCurrentStackTrace(ThrowHelper.GetOperationAbortedException()));
+        _shutdownTcs.TrySetResult();
         return QUIC_STATUS_SUCCESS;
     }
 
@@ -631,10 +631,10 @@ public async ValueTask DisposeAsync()
         {
             unsafe
             {
-                MsQuicApi.Api.ApiTable->ConnectionShutdown(
-                    _handle.QuicHandle,
+                _handle.SafeCall(handle => MsQuicApi.Api.ApiTable->ConnectionShutdown(
+                    handle.QuicHandle,
                     QUIC_CONNECTION_SHUTDOWN_FLAGS.NONE,
-                    (ulong)_defaultCloseErrorCode);
+                    (ulong)_defaultCloseErrorCode));
             }
         }