Fix incorrect bounds check in ArrayMemoryPoolBuffer.Pin (#28032)

Author: GrabYourPitchforks

src/System.Memory/src/System/Buffers/ArrayMemoryPool.ArrayMemoryPoolBuffer.cs

@@ -72,11 +72,22 @@ public sealed override MemoryHandle Pin(int byteOffset = 0)
                 {
                     Retain(); // this checks IsDisposed
 
-                    if (byteOffset != 0 && (((uint)byteOffset) - 1) / Unsafe.SizeOf<T>() >= _array.Length)
-                        ThrowHelper.ThrowArgumentOutOfRangeException(ExceptionArgument.byteOffset);
+                    try
+                    {
+                        if ((IntPtr.Size == 4 && (uint)byteOffset > (uint)_array.Length * (uint)Unsafe.SizeOf<T>())
+                            || (IntPtr.Size != 4 && (ulong)byteOffset > (uint)_array.Length * (ulong)Unsafe.SizeOf<T>()))
+                        {
+                            ThrowHelper.ThrowArgumentOutOfRangeException(ExceptionArgument.byteOffset);
+                        }
 
-                    GCHandle handle = GCHandle.Alloc(_array, GCHandleType.Pinned);
-                    return new MemoryHandle(this, ((byte*)handle.AddrOfPinnedObject()) + byteOffset, handle);
+                        GCHandle handle = GCHandle.Alloc(_array, GCHandleType.Pinned);
+                        return new MemoryHandle(this, ((byte*)handle.AddrOfPinnedObject()) + byteOffset, handle);
+                    }
+                    catch
+                    {
+                        Release();
+                        throw;
+                    }
                 }
             }
 

src/System.Memory/tests/Memory/CustomMemoryForTest.cs

@@ -51,11 +51,24 @@ public override MemoryHandle Pin(int byteOffset = 0)
         {
             unsafe
             {
-                Retain();
-                if (byteOffset != 0 && (((uint)byteOffset) - 1) / Unsafe.SizeOf<T>() >= _array.Length)
-                    throw new ArgumentOutOfRangeException(nameof(byteOffset));
-                var handle = GCHandle.Alloc(_array, GCHandleType.Pinned);
-                return new MemoryHandle(this, Unsafe.Add<byte>((void*)handle.AddrOfPinnedObject(), _offset + byteOffset), handle);
+                Retain(); // this checks IsDisposed
+
+                try
+                {
+                    if ((IntPtr.Size == 4 && (uint)byteOffset > (uint)_array.Length * (uint)Unsafe.SizeOf<T>())
+                        || (IntPtr.Size != 4 && (ulong)byteOffset > (uint)_array.Length * (ulong)Unsafe.SizeOf<T>()))
+                    {
+                        throw new ArgumentOutOfRangeException(nameof(byteOffset));
+                    }
+
+                    var handle = GCHandle.Alloc(_array, GCHandleType.Pinned);
+                    return new MemoryHandle(this, Unsafe.Add<byte>((void*)handle.AddrOfPinnedObject(), _offset + byteOffset), handle);
+                }
+                catch
+                {
+                    Release();
+                    throw;
+                }
             }
         }
 

src/System.Memory/tests/Memory/OwnedMemory.cs

@@ -117,6 +117,15 @@ public static void OwnedMemoryPinArray()
             handle.Dispose();
         }
 
+        [Fact]
+        [OuterLoop]
+        public static void OwnedMemoryPinLargeArray()
+        {
+            int[] array = new int[0x2000_0000]; // will produce array with total byte length > 2 GB
+            OwnedMemory<int> owner = new CustomMemoryForTest<int>(array);
+            Assert.Throws<ArgumentOutOfRangeException>(() => owner.Pin(int.MinValue));
+        }
+
         [Fact]
         public static void MemoryFromOwnedMemoryAfterDispose()
         {