Tighten up some ArrayPool.Return usage (#56229)

Avoid potential problems if Return were to throw an exception after having already put the array back into the pool.
dotnet · Jul 27, 2021 · 59a5610 · 59a5610
1 parent f32b55f
commit 59a5610
Show file tree

Hide file tree

Showing 20 changed files with 94 additions and 59 deletions.
diff --git a/src/libraries/Common/src/System/Text/Json/PooledByteBufferWriter.cs b/src/libraries/Common/src/System/Text/Json/PooledByteBufferWriter.cs
@@ -86,8 +86,9 @@ public void Dispose()
             }
 
             ClearHelper();
-            ArrayPool<byte>.Shared.Return(_rentedBuffer);
+            byte[] toReturn = _rentedBuffer;
             _rentedBuffer = null!;
+            ArrayPool<byte>.Shared.Return(toReturn);
         }
 
         public void Advance(int count)

diff --git a/src/libraries/System.Drawing.Common/src/System/Drawing/Icon.Windows.cs b/src/libraries/System.Drawing.Common/src/System/Drawing/Icon.Windows.cs
@@ -188,14 +188,17 @@ void ISerializable.GetObjectData(SerializationInfo si, StreamingContext context)
             filePath.CopyTo(0, buffer, 0, filePath.Length);
             buffer[filePath.Length] = '\0';
 
+            IntPtr hIcon;
             fixed (char* b = buffer)
             {
-                IntPtr hIcon = Interop.Shell32.ExtractAssociatedIcon(NativeMethods.NullHandleRef, b, ref index);
-                ArrayPool<char>.Shared.Return(buffer);
-                if (hIcon != IntPtr.Zero)
-                {
-                    return new Icon(hIcon, true);
-                }
+                hIcon = Interop.Shell32.ExtractAssociatedIcon(NativeMethods.NullHandleRef, b, ref index);
+            }
+
+            ArrayPool<char>.Shared.Return(buffer);
+
+            if (hIcon != IntPtr.Zero)
+            {
+                return new Icon(hIcon, true);
             }
 
             return null;

diff --git a/src/libraries/System.Net.Http.Json/src/System/Net/Http/Json/TranscodingReadStream.cs b/src/libraries/System.Net.Http.Json/src/System/Net/Http/Json/TranscodingReadStream.cs
@@ -208,17 +208,20 @@ protected override void Dispose(bool disposing)
             {
                 _disposed = true;
 
-                Debug.Assert(_charBuffer.Array != null);
-                ArrayPool<char>.Shared.Return(_charBuffer.Array);
+                char[]? charBuffer = _charBuffer.Array;
+                Debug.Assert(charBuffer != null);
                 _charBuffer = default;
+                ArrayPool<char>.Shared.Return(charBuffer);
 
-                Debug.Assert(_byteBuffer.Array != null);
-                ArrayPool<byte>.Shared.Return(_byteBuffer.Array);
+                byte[]? byteBuffer = _byteBuffer.Array;
+                Debug.Assert(byteBuffer != null);
                 _byteBuffer = default;
+                ArrayPool<byte>.Shared.Return(byteBuffer);
 
-                Debug.Assert(_overflowBuffer.Array != null);
-                ArrayPool<byte>.Shared.Return(_overflowBuffer.Array);
+                byte[]? overflowBuffer = _overflowBuffer.Array;
+                Debug.Assert(overflowBuffer != null);
                 _overflowBuffer = default;
+                ArrayPool<byte>.Shared.Return(overflowBuffer);
 
                 _stream.Dispose();
             }

diff --git a/src/libraries/System.Net.Http.Json/src/System/Net/Http/Json/TranscodingWriteStream.cs b/src/libraries/System.Net.Http.Json/src/System/Net/Http/Json/TranscodingWriteStream.cs
@@ -133,8 +133,9 @@ protected override void Dispose(bool disposing)
             if (!_disposed)
             {
                 _disposed = true;
-                ArrayPool<char>.Shared.Return(_charBuffer);
+                char[] toReturn = _charBuffer;
                 _charBuffer = null!;
+                ArrayPool<char>.Shared.Return(toReturn);
             }
         }
 

diff --git a/src/libraries/System.Net.Security/src/System/Net/Security/SslStream.Implementation.cs b/src/libraries/System.Net.Security/src/System/Net/Security/SslStream.Implementation.cs
@@ -742,8 +742,10 @@ async ValueTask WaitAndWriteAsync(TIOAdapter writeAdapter, ReadOnlyMemory<byte>
                     if (status.ErrorCode == SecurityStatusPalErrorCode.TryAgain)
                     {
                         // No need to hold on the buffer any more.
-                        ArrayPool<byte>.Shared.Return(bufferToReturn);
+                        byte[] tmp = bufferToReturn;
                         bufferToReturn = null;
+                        ArrayPool<byte>.Shared.Return(tmp);
+
                         // Call WriteSingleChunk() recursively to avoid code duplication.
                         // This should be extremely rare in cases when second renegotiation happens concurrently with Write.
                         await WriteSingleChunk(writeAdapter, buffer).ConfigureAwait(false);
@@ -788,14 +790,12 @@ async ValueTask CompleteWriteAsync(ValueTask writeTask, byte[] bufferToReturn)
         // actually contains no decrypted or encrypted bytes
         private void ReturnReadBufferIfEmpty()
         {
-            if (_internalBuffer != null && _decryptedBytesCount == 0 && _internalBufferCount == 0)
+            if (_internalBuffer is byte[] internalBuffer && _decryptedBytesCount == 0 && _internalBufferCount == 0)
             {
-                ArrayPool<byte>.Shared.Return(_internalBuffer);
                 _internalBuffer = null;
-                _internalBufferCount = 0;
                 _internalOffset = 0;
-                _decryptedBytesCount = 0;
                 _decryptedBytesOffset = 0;
+                ArrayPool<byte>.Shared.Return(internalBuffer);
             }
             else if (_decryptedBytesCount == 0)
             {

diff --git a/src/libraries/System.Net.Security/tests/StressTests/SslStress/StressOperations.cs b/src/libraries/System.Net.Security/tests/StressTests/SslStress/StressOperations.cs
@@ -24,7 +24,7 @@ namespace SslStress
 {
     public struct DataSegment
     {
-        private readonly byte[] _buffer;
+        private byte[] _buffer;
 
         public DataSegment(int length)
         {
@@ -37,7 +37,12 @@ public DataSegment(int length)
         public Span<byte> AsSpan() => new Span<byte>(_buffer, 0, Length);
 
         public ulong Checksum => CRC.CalculateCRC(AsSpan());
-        public void Return() => ArrayPool<byte>.Shared.Return(_buffer);
+        public void Return()
+        {
+            byte[] toReturn = _buffer;
+            _buffer = null;
+            ArrayPool<byte>.Shared.Return(toReturn);
+        }
 
         /// Create and populate a segment with random data
         public static DataSegment CreateRandom(Random random, int maxLength)

diff --git a/...ibraries/System.Net.WebSockets/src/System/Net/WebSockets/Compression/WebSocketDeflater.cs b/...ibraries/System.Net.WebSockets/src/System/Net/WebSockets/Compression/WebSocketDeflater.cs
@@ -35,10 +35,10 @@ public void Dispose()
 
         public void ReleaseBuffer()
         {
-            if (_buffer is not null)
+            if (_buffer is byte[] toReturn)
             {
-                ArrayPool<byte>.Shared.Return(_buffer);
                 _buffer = null;
+                ArrayPool<byte>.Shared.Return(toReturn);
             }
         }
 
@@ -70,8 +70,11 @@ public ReadOnlySpan<byte> Deflate(ReadOnlySpan<byte> payload, bool endOfMessage)
                 // Rent a 30% bigger buffer
                 byte[] newBuffer = ArrayPool<byte>.Shared.Rent((int)(_buffer.Length * 1.3));
                 _buffer.AsSpan(0, position).CopyTo(newBuffer);
-                ArrayPool<byte>.Shared.Return(_buffer);
+
+                byte[] toReturn = _buffer;
                 _buffer = newBuffer;
+
+                ArrayPool<byte>.Shared.Return(toReturn);
             }
 
             return new ReadOnlySpan<byte>(_buffer, 0, position);

diff --git a/...ibraries/System.Net.WebSockets/src/System/Net/WebSockets/Compression/WebSocketInflater.cs b/...ibraries/System.Net.WebSockets/src/System/Net/WebSockets/Compression/WebSocketInflater.cs
@@ -111,9 +111,11 @@ public void AddBytes(int totalBytesReceived, bool endOfMessage)
                     {
                         byte[] newBuffer = ArrayPool<byte>.Shared.Rent(_available + FlushMarkerLength);
                         _buffer.AsSpan(0, _available).CopyTo(newBuffer);
-                        ArrayPool<byte>.Shared.Return(_buffer);
 
+                        byte[] toReturn = _buffer;
                         _buffer = newBuffer;
+
+                        ArrayPool<byte>.Shared.Return(toReturn);
                     }
 
                     FlushMarker.CopyTo(_buffer.AsSpan(_available));
@@ -202,12 +204,13 @@ private unsafe bool Finish(Span<byte> output, ref int written)
 
         private void ReleaseBuffer()
         {
-            if (_buffer is not null)
+            if (_buffer is byte[] toReturn)
             {
-                ArrayPool<byte>.Shared.Return(_buffer);
                 _buffer = null;
                 _available = 0;
                 _position = 0;
+
+                ArrayPool<byte>.Shared.Return(toReturn);
             }
         }
 

diff --git a/src/libraries/System.Net.WebSockets/src/System/Net/WebSockets/ManagedWebSocket.cs b/src/libraries/System.Net.WebSockets/src/System/Net/WebSockets/ManagedWebSocket.cs
@@ -1452,11 +1452,10 @@ private void ReleaseSendBuffer()
         {
             Debug.Assert(_sendFrameAsyncLock.CurrentCount == 0, "Caller should hold the _sendFrameAsyncLock");
 
-            byte[]? old = _sendBuffer;
-            if (old != null)
+            if (_sendBuffer is byte[] toReturn)
             {
                 _sendBuffer = null;
-                ArrayPool<byte>.Shared.Return(old);
+                ArrayPool<byte>.Shared.Return(toReturn);
             }
         }
 

diff --git a/src/libraries/System.Private.CoreLib/src/System/Collections/Generic/ValueListBuilder.cs b/src/libraries/System.Private.CoreLib/src/System/Collections/Generic/ValueListBuilder.cs
@@ -59,10 +59,11 @@ public ReadOnlySpan<T> AsSpan()
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public void Dispose()
         {
-            if (_arrayFromPool != null)
+            T[]? toReturn = _arrayFromPool;
+            if (toReturn != null)
             {
-                ArrayPool<T>.Shared.Return(_arrayFromPool);
                 _arrayFromPool = null;
+                ArrayPool<T>.Shared.Return(toReturn);
             }
         }
 

diff --git a/src/libraries/System.Private.CoreLib/src/System/IO/Enumeration/FileSystemEnumerator.Unix.cs b/src/libraries/System.Private.CoreLib/src/System/IO/Enumeration/FileSystemEnumerator.Unix.cs
@@ -236,12 +236,17 @@ private void InternalDispose(bool disposing)
 
                     CloseDirectoryHandle();
 
-                    if (_pathBuffer != null)
-                        ArrayPool<char>.Shared.Return(_pathBuffer);
-                    _pathBuffer = null;
-                    if (_entryBuffer != null)
-                        ArrayPool<byte>.Shared.Return(_entryBuffer);
-                    _entryBuffer = null;
+                    if (_pathBuffer is char[] pathBuffer)
+                    {
+                        _pathBuffer = null;
+                        ArrayPool<char>.Shared.Return(pathBuffer);
+                    }
+
+                    if (_entryBuffer is byte[] entryBuffer)
+                    {
+                        _entryBuffer = null;
+                        ArrayPool<byte>.Shared.Return(entryBuffer);
+                    }
                 }
             }
 

diff --git a/src/libraries/System.Private.CoreLib/src/System/IO/File.cs b/src/libraries/System.Private.CoreLib/src/System/IO/File.cs
@@ -796,8 +796,11 @@ private static async Task<byte[]> InternalReadAllBytesUnknownLengthAsync(FileStr
 
                         byte[] tmp = ArrayPool<byte>.Shared.Rent((int)newLength);
                         Buffer.BlockCopy(rentedArray, 0, tmp, 0, bytesRead);
-                        ArrayPool<byte>.Shared.Return(rentedArray);
+
+                        byte[] toReturn = rentedArray;
                         rentedArray = tmp;
+
+                        ArrayPool<byte>.Shared.Return(toReturn);
                     }
 
                     Debug.Assert(bytesRead < rentedArray.Length);

diff --git a/src/libraries/System.Private.CoreLib/src/System/IO/FileSystem.Windows.cs b/src/libraries/System.Private.CoreLib/src/System/IO/FileSystem.Windows.cs
@@ -568,8 +568,9 @@ internal static void CreateSymbolicLink(string path, string pathToTarget, bool i
                 // the return value is the required buffer size, in TCHARs. This value includes the size of the terminating null character.
                 if (result > buffer.Length)
                 {
-                    ArrayPool<char>.Shared.Return(buffer);
+                    char[] toReturn = buffer;
                     buffer = ArrayPool<char>.Shared.Rent((int)result);
+                    ArrayPool<char>.Shared.Return(toReturn);
 
                     result = GetFinalPathNameByHandle(handle, buffer);
                 }

diff --git a/src/libraries/System.Private.CoreLib/src/System/TimeZoneInfo.Unix.NonAndroid.cs b/src/libraries/System.Private.CoreLib/src/System/TimeZoneInfo.Unix.NonAndroid.cs
@@ -181,10 +181,9 @@ private static unsafe void EnumerateFilesRecursively(string path, Predicate<stri
             List<string>? toExplore = null; // List used as a stack
 
             int bufferSize = Interop.Sys.GetReadDirRBufferSize();
-            byte[]? dirBuffer = null;
+            byte[] dirBuffer = ArrayPool<byte>.Shared.Rent(bufferSize);
             try
             {
-                dirBuffer = ArrayPool<byte>.Shared.Rent(bufferSize);
                 string currentPath = path;
 
                 fixed (byte* dirBufferPtr = dirBuffer)
@@ -266,8 +265,7 @@ private static unsafe void EnumerateFilesRecursively(string path, Predicate<stri
             }
             finally
             {
-                if (dirBuffer != null)
-                    ArrayPool<byte>.Shared.Return(dirBuffer);
+                ArrayPool<byte>.Shared.Return(dirBuffer);
             }
         }
 

diff --git a/.../System.Security.Cryptography.Primitives/src/System/Security/Cryptography/CryptoStream.cs b/.../System.Security.Cryptography.Primitives/src/System/Security/Cryptography/CryptoStream.cs
@@ -637,7 +637,7 @@ unsafe static int TransformBlock(ICryptoTransform transform, ReadOnlyMemory<byte
                 else
                 {
                     // Use ArrayPool.Shared instead of CryptoPool because the array is passed out.
-                    byte[] rentedBuffer = ArrayPool<byte>.Shared.Rent(inputBuffer.Length);
+                    byte[]? rentedBuffer = ArrayPool<byte>.Shared.Rent(inputBuffer.Length);
                     int result = default;
 
                     // Pin the rented buffer for security.
@@ -655,7 +655,7 @@ unsafe static int TransformBlock(ICryptoTransform transform, ReadOnlyMemory<byte
                     }
 
                     ArrayPool<byte>.Shared.Return(rentedBuffer);
-                    rentedBuffer = null!;
+                    rentedBuffer = null;
                     return result;
                 }
             }
@@ -667,7 +667,7 @@ public unsafe override void CopyTo(Stream destination, int bufferSize)
             CheckCopyToArguments(destination, bufferSize);
 
             // Use ArrayPool<byte>.Shared instead of CryptoPool because the array is passed out.
-            byte[] rentedBuffer = ArrayPool<byte>.Shared.Rent(bufferSize);
+            byte[]? rentedBuffer = ArrayPool<byte>.Shared.Rent(bufferSize);
             // Pin the array for security.
             fixed (byte* _ = &rentedBuffer[0])
             {
@@ -686,7 +686,7 @@ public unsafe override void CopyTo(Stream destination, int bufferSize)
                 }
             }
             ArrayPool<byte>.Shared.Return(rentedBuffer);
-            rentedBuffer = null!;
+            rentedBuffer = null;
         }
 
         /// <inheritdoc/>
@@ -699,7 +699,7 @@ public override Task CopyToAsync(Stream destination, int bufferSize, Cancellatio
         private async Task CopyToAsyncInternal(Stream destination, int bufferSize, CancellationToken cancellationToken)
         {
             // Use ArrayPool<byte>.Shared instead of CryptoPool because the array is passed out.
-            byte[] rentedBuffer = ArrayPool<byte>.Shared.Rent(bufferSize);
+            byte[]? rentedBuffer = ArrayPool<byte>.Shared.Rent(bufferSize);
             // Pin the array for security.
             GCHandle pinHandle = GCHandle.Alloc(rentedBuffer, GCHandleType.Pinned);
             try
@@ -717,7 +717,7 @@ private async Task CopyToAsyncInternal(Stream destination, int bufferSize, Cance
                 pinHandle.Free();
             }
             ArrayPool<byte>.Shared.Return(rentedBuffer);
-            rentedBuffer = null!;
+            rentedBuffer = null;
         }
 
         private void CheckCopyToArguments(Stream destination, int bufferSize)

diff --git a/...rity.Cryptography.X509Certificates/src/Internal/Cryptography/Pal.Unix/UnixPkcs12Reader.cs b/...rity.Cryptography.X509Certificates/src/Internal/Cryptography/Pal.Unix/UnixPkcs12Reader.cs
@@ -167,14 +167,14 @@ private static void ReturnRentedContentInfos(ContentInfoAsn[] rentedContents)
                 if (contentType == DecryptedSentinel)
                 {
                     ReadOnlyMemory<byte> content = rentedContents[i].Content;
+                    rentedContents[i].Content = default;
 
                     if (!MemoryMarshal.TryGetArray(content, out ArraySegment<byte> segment))
                     {
                         Debug.Fail("Couldn't unpack decrypted buffer.");
                     }
 
                     CryptoPool.Return(segment);
-                    rentedContents[0].Content = default;
                 }
             }
 

diff --git a/...ryptography.X509Certificates/src/Internal/Cryptography/Pal.iOS/AppleCertificatePal.Pem.cs b/...ryptography.X509Certificates/src/Internal/Cryptography/Pal.iOS/AppleCertificatePal.Pem.cs
@@ -63,8 +63,9 @@ internal static bool TryDecodePem(ReadOnlySpan<byte> rawData, DerCallback derCal
                             X509ContentType.Pkcs7;
                         bool cont = derCallback(certBytes.AsSpan(0, bytesWritten), contentType);
 
-                        CryptoPool.Return(certBytes, clearSize: 0);
+                        byte[] toReturn = certBytes;
                         certBytes = null;
+                        CryptoPool.Return(toReturn, clearSize: 0);
 
                         if (!cont)
                         {

diff --git a/...Json/Serialization/Converters/Object/ObjectWithParameterizedConstructorConverter.Large.cs b/...Json/Serialization/Converters/Object/ObjectWithParameterizedConstructorConverter.Large.cs
@@ -32,6 +32,7 @@ protected override bool ReadAndCacheConstructorArgument(ref ReadStack state, ref
         protected override object CreateObject(ref ReadStackFrame frame)
         {
             object[] arguments = (object[])frame.CtorArgumentState!.Arguments;
+            frame.CtorArgumentState.Arguments = null!;
 
             var createObject = (JsonTypeInfo.ParameterizedConstructorDelegate<T>?)frame.JsonTypeInfo.CreateObjectWithArgs;