Tighten up some ArrayPool.Return usage (#56229)
authorStephen Toub <stoub@microsoft.com>
Tue, 27 Jul 2021 11:40:36 +0000 (07:40 -0400)
committerGitHub <noreply@github.com>
Tue, 27 Jul 2021 11:40:36 +0000 (07:40 -0400)
Avoid potential problems if Return were to throw an exception after having already put the array back into the pool.

20 files changed:
src/libraries/Common/src/System/Text/Json/PooledByteBufferWriter.cs
src/libraries/System.Drawing.Common/src/System/Drawing/Icon.Windows.cs
src/libraries/System.Net.Http.Json/src/System/Net/Http/Json/TranscodingReadStream.cs
src/libraries/System.Net.Http.Json/src/System/Net/Http/Json/TranscodingWriteStream.cs
src/libraries/System.Net.Security/src/System/Net/Security/SslStream.Implementation.cs
src/libraries/System.Net.Security/tests/StressTests/SslStress/StressOperations.cs
src/libraries/System.Net.WebSockets/src/System/Net/WebSockets/Compression/WebSocketDeflater.cs
src/libraries/System.Net.WebSockets/src/System/Net/WebSockets/Compression/WebSocketInflater.cs
src/libraries/System.Net.WebSockets/src/System/Net/WebSockets/ManagedWebSocket.cs
src/libraries/System.Private.CoreLib/src/System/Collections/Generic/ValueListBuilder.cs
src/libraries/System.Private.CoreLib/src/System/IO/Enumeration/FileSystemEnumerator.Unix.cs
src/libraries/System.Private.CoreLib/src/System/IO/File.cs
src/libraries/System.Private.CoreLib/src/System/IO/FileSystem.Windows.cs
src/libraries/System.Private.CoreLib/src/System/TimeZoneInfo.Unix.NonAndroid.cs
src/libraries/System.Security.Cryptography.Primitives/src/System/Security/Cryptography/CryptoStream.cs
src/libraries/System.Security.Cryptography.X509Certificates/src/Internal/Cryptography/Pal.Unix/UnixPkcs12Reader.cs
src/libraries/System.Security.Cryptography.X509Certificates/src/Internal/Cryptography/Pal.iOS/AppleCertificatePal.Pem.cs
src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Object/ObjectWithParameterizedConstructorConverter.Large.cs
src/libraries/System.Text.Json/src/System/Text/Json/Serialization/Converters/Object/ObjectWithParameterizedConstructorConverter.cs
src/libraries/System.Text.Json/src/System/Text/Json/Serialization/ReadBufferState.cs

index 80e7bc2..8694a78 100644 (file)
@@ -86,8 +86,9 @@ namespace System.Text.Json
             }
 
             ClearHelper();
-            ArrayPool<byte>.Shared.Return(_rentedBuffer);
+            byte[] toReturn = _rentedBuffer;
             _rentedBuffer = null!;
+            ArrayPool<byte>.Shared.Return(toReturn);
         }
 
         public void Advance(int count)
index 4acd637..73e65b3 100644 (file)
@@ -188,14 +188,17 @@ namespace System.Drawing
             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;
index 894ea8c..cc8ef87 100644 (file)
@@ -208,17 +208,20 @@ namespace System.Net.Http.Json
             {
                 _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();
             }
index 51c73d0..2217bae 100644 (file)
@@ -133,8 +133,9 @@ namespace System.Net.Http.Json
             if (!_disposed)
             {
                 _disposed = true;
-                ArrayPool<char>.Shared.Return(_charBuffer);
+                char[] toReturn = _charBuffer;
                 _charBuffer = null!;
+                ArrayPool<char>.Shared.Return(toReturn);
             }
         }
 
index 6b1b1e0..0cd8884 100644 (file)
@@ -742,8 +742,10 @@ namespace System.Net.Security
                     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 @@ namespace System.Net.Security
         // 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)
             {
index 8c2ec6f..eea046d 100644 (file)
@@ -24,7 +24,7 @@ namespace SslStress
 {
     public struct DataSegment
     {
-        private readonly byte[] _buffer;
+        private byte[] _buffer;
 
         public DataSegment(int length)
         {
@@ -37,7 +37,12 @@ namespace SslStress
         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)
index f3ecf27..7d8eb83 100644 (file)
@@ -35,10 +35,10 @@ namespace System.Net.WebSockets.Compression
 
         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 @@ namespace System.Net.WebSockets.Compression
                 // 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);
index d47e646..5f580e9 100644 (file)
@@ -111,9 +111,11 @@ namespace System.Net.WebSockets.Compression
                     {
                         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 @@ namespace System.Net.WebSockets.Compression
 
         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);
             }
         }
 
index 864b742..297b4c5 100644 (file)
@@ -1452,11 +1452,10 @@ namespace System.Net.WebSockets
         {
             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);
             }
         }
 
index ec23d53..a82372b 100644 (file)
@@ -59,10 +59,11 @@ namespace System.Collections.Generic
         [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);
             }
         }
 
index 04f8282..a2c83ab 100644 (file)
@@ -236,12 +236,17 @@ namespace System.IO.Enumeration
 
                     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);
+                    }
                 }
             }
 
index 668e78c..ff2a7bd 100644 (file)
@@ -796,8 +796,11 @@ namespace System.IO
 
                         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);
index 698e3bc..191106a 100644 (file)
@@ -568,8 +568,9 @@ namespace System.IO
                 // 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);
                 }
index c5ad4b0..d7e54e1 100644 (file)
@@ -181,10 +181,9 @@ namespace System
             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 @@ namespace System
             }
             finally
             {
-                if (dirBuffer != null)
-                    ArrayPool<byte>.Shared.Return(dirBuffer);
+                ArrayPool<byte>.Shared.Return(dirBuffer);
             }
         }
 
index 9a67b46..c4c7c05 100644 (file)
@@ -637,7 +637,7 @@ namespace System.Security.Cryptography
                 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 @@ namespace System.Security.Cryptography
                     }
 
                     ArrayPool<byte>.Shared.Return(rentedBuffer);
-                    rentedBuffer = null!;
+                    rentedBuffer = null;
                     return result;
                 }
             }
@@ -667,7 +667,7 @@ namespace System.Security.Cryptography
             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 @@ namespace System.Security.Cryptography
                 }
             }
             ArrayPool<byte>.Shared.Return(rentedBuffer);
-            rentedBuffer = null!;
+            rentedBuffer = null;
         }
 
         /// <inheritdoc/>
@@ -699,7 +699,7 @@ namespace System.Security.Cryptography
         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 @@ namespace System.Security.Cryptography
                 pinHandle.Free();
             }
             ArrayPool<byte>.Shared.Return(rentedBuffer);
-            rentedBuffer = null!;
+            rentedBuffer = null;
         }
 
         private void CheckCopyToArguments(Stream destination, int bufferSize)
index b227a05..8cd7abe 100644 (file)
@@ -167,6 +167,7 @@ namespace Internal.Cryptography.Pal
                 if (contentType == DecryptedSentinel)
                 {
                     ReadOnlyMemory<byte> content = rentedContents[i].Content;
+                    rentedContents[i].Content = default;
 
                     if (!MemoryMarshal.TryGetArray(content, out ArraySegment<byte> segment))
                     {
@@ -174,7 +175,6 @@ namespace Internal.Cryptography.Pal
                     }
 
                     CryptoPool.Return(segment);
-                    rentedContents[0].Content = default;
                 }
             }
 
index 4fae469..73d0e7b 100644 (file)
@@ -63,8 +63,9 @@ namespace Internal.Cryptography.Pal
                             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)
                         {
index 0514750..02c7149 100644 (file)
@@ -32,6 +32,7 @@ namespace System.Text.Json.Serialization.Converters
         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;
 
index 867b74f..528ac02 100644 (file)
@@ -76,8 +76,9 @@ namespace System.Text.Json.Serialization.Converters
                         ReadPropertyValue(obj, ref state, ref tempReader, jsonPropertyInfo, useExtensionProperty);
                     }
 
-                    ArrayPool<FoundProperty>.Shared.Return(argumentState.FoundProperties!, clearArray: true);
+                    FoundProperty[] toReturn = argumentState.FoundProperties!;
                     argumentState.FoundProperties = null;
+                    ArrayPool<FoundProperty>.Shared.Return(toReturn, clearArray: true);
                 }
             }
             else
@@ -133,8 +134,9 @@ namespace System.Text.Json.Serialization.Converters
                         }
                     }
 
-                    ArrayPool<FoundPropertyAsync>.Shared.Return(argumentState.FoundPropertiesAsync!, clearArray: true);
+                    FoundPropertyAsync[] toReturn = argumentState.FoundPropertiesAsync!;
                     argumentState.FoundPropertiesAsync = null;
+                    ArrayPool<FoundPropertyAsync>.Shared.Return(toReturn, clearArray: true);
                 }
             }
 
@@ -238,9 +240,10 @@ namespace System.Text.Json.Serialization.Converters
 
                             argumentState.FoundProperties.CopyTo(newCache, 0);
 
-                            ArrayPool<FoundProperty>.Shared.Return(argumentState.FoundProperties, clearArray: true);
-
+                            FoundProperty[] toReturn = argumentState.FoundProperties;
                             argumentState.FoundProperties = newCache!;
+
+                            ArrayPool<FoundProperty>.Shared.Return(toReturn, clearArray: true);
                         }
 
                         argumentState.FoundProperties[argumentState.FoundPropertyCount++] = (
@@ -436,9 +439,10 @@ namespace System.Text.Json.Serialization.Converters
 
                 argumentState.FoundPropertiesAsync!.CopyTo(newCache, 0);
 
-                ArrayPool<FoundPropertyAsync>.Shared.Return(argumentState.FoundPropertiesAsync!, clearArray: true);
-
+                FoundPropertyAsync[] toReturn = argumentState.FoundPropertiesAsync!;
                 argumentState.FoundPropertiesAsync = newCache!;
+
+                ArrayPool<FoundPropertyAsync>.Shared.Return(toReturn, clearArray: true);
             }
 
             // Cache the property name and value.
index 395c5f9..d909716 100644 (file)
@@ -25,8 +25,11 @@ namespace System.Text.Json.Serialization
         {
             // Clear only what we used and return the buffer to the pool
             new Span<byte>(Buffer, 0, ClearMax).Clear();
-            ArrayPool<byte>.Shared.Return(Buffer);
+
+            byte[] toReturn = Buffer;
             Buffer = null!;
+
+            ArrayPool<byte>.Shared.Return(toReturn);
         }
     }
 }