Handle IOV_MAX limit in RandomAccess, add missing test (dotnet#109340)

* add test for Int32 overflow for WriteGather in RandomAccess

* add failing test fore more than IOV_MAX buffers

* fix both the native and managed parts

---------

Co-authored-by: Adeel Mujahid <[email protected]>
Co-authored-by: Stephen Toub <[email protected]>

Author: 3 people

src/libraries/System.Private.CoreLib/src/System/IO/RandomAccess.Unix.cs

@@ -196,9 +196,10 @@ internal static unsafe void WriteGatherAtOffset(SafeFileHandle handle, IReadOnly
                     }
 
                     long bytesWritten;
-                    fixed (Interop.Sys.IOVector* pinnedVectors = &MemoryMarshal.GetReference(vectors))
+                    Span<Interop.Sys.IOVector> left = vectors.Slice(buffersOffset);
+                    fixed (Interop.Sys.IOVector* pinnedVectors = &MemoryMarshal.GetReference(left))
                     {
-                        bytesWritten = Interop.Sys.PWriteV(handle, pinnedVectors, buffersCount, fileOffset);
+                        bytesWritten = Interop.Sys.PWriteV(handle, pinnedVectors, buffersCount - buffersOffset, fileOffset);
                     }
 
                     FileStreamHelpers.CheckFileCall(bytesWritten, handle.Path);
@@ -208,6 +209,8 @@ internal static unsafe void WriteGatherAtOffset(SafeFileHandle handle, IReadOnly
                     }
 
                     // The write completed successfully but for fewer bytes than requested.
+                    // We need to perform next write where the previous one has finished.
+                    fileOffset += bytesWritten;
                     // We need to try again for the remainder.
                     for (int i = 0; i < buffersCount; i++)
                     {

src/libraries/System.Runtime/tests/System.IO.FileSystem.Tests/RandomAccess/WriteGatherAsync.cs

@@ -6,6 +6,7 @@
 using System.Security.Cryptography;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.DotNet.XUnitExtensions;
 using Microsoft.Win32.SafeHandles;
 using Xunit;
 
@@ -133,5 +134,130 @@ public async Task DuplicatedBufferDuplicatesContentAsync(FileOptions options)
             Assert.Equal(repeatCount, actualContent.Length);
             Assert.All(actualContent, actual => Assert.Equal(value, actual));
         }
+
+        [OuterLoop("It consumes a lot of resources (disk space and memory).")]
+        [ConditionalTheory(typeof(PlatformDetection), nameof(PlatformDetection.Is64BitProcess), nameof(PlatformDetection.IsReleaseRuntime))]
+        [InlineData(false, false)]
+        [InlineData(false, true)]
+        [InlineData(true, true)]
+        [InlineData(true, false)]
+        public async Task NoInt32OverflowForLargeInputs(bool asyncFile, bool asyncMethod)
+        {
+            // We need to write more than Int32.MaxValue bytes to the disk to reproduce the problem.
+            // To reduce the number of used memory, we allocate only one write buffer and simply repeat it multiple times.
+            // For reading, we need unique buffers to ensure that all of them are getting populated with the right data.
+
+            const int BufferCount = 1002;
+            const int BufferSize = int.MaxValue / 1000;
+            const long FileSize = (long)BufferCount * BufferSize;
+            string filePath = GetTestFilePath();
+            ReadOnlyMemory<byte> writeBuffer = RandomNumberGenerator.GetBytes(BufferSize);
+            List<ReadOnlyMemory<byte>> writeBuffers = Enumerable.Repeat(writeBuffer, BufferCount).ToList();
+            List<Memory<byte>> readBuffers = new List<Memory<byte>>(BufferCount);
+
+            try
+            {
+                for (int i = 0; i < BufferCount; i++)
+                {
+                    readBuffers.Add(new byte[BufferSize]);
+                }
+            }
+            catch (OutOfMemoryException)
+            {
+                throw new SkipTestException("Not enough memory.");
+            }
+
+            FileOptions options = asyncFile ? FileOptions.Asynchronous : FileOptions.None; // we need to test both code paths
+            options |= FileOptions.DeleteOnClose;
+
+            SafeFileHandle? sfh;
+            try
+            {
+                sfh = File.OpenHandle(filePath, FileMode.CreateNew, FileAccess.ReadWrite, FileShare.None, options, preallocationSize: FileSize);
+            }
+            catch (IOException)
+            {
+                throw new SkipTestException("Not enough disk space.");
+            }
+
+            long fileOffset = 0, bytesRead = 0;
+            try
+            {
+                if (asyncMethod)
+                {
+                    await RandomAccess.WriteAsync(sfh, writeBuffers, fileOffset);
+                    bytesRead = await RandomAccess.ReadAsync(sfh, readBuffers, fileOffset);
+                }
+                else
+                {
+                    RandomAccess.Write(sfh, writeBuffers, fileOffset);
+                    bytesRead = RandomAccess.Read(sfh, readBuffers, fileOffset);
+                }
+            }
+            finally
+            {
+                sfh.Dispose(); // delete the file ASAP to avoid running out of resources in CI
+            }
+
+            Assert.Equal(FileSize, bytesRead);
+            for (int i = 0; i < BufferCount; i++)
+            {
+                Assert.Equal(writeBuffer, readBuffers[i]);
+            }
+        }
+
+        [Theory]
+        [InlineData(false, false)]
+        [InlineData(false, true)]
+        [InlineData(true, true)]
+        [InlineData(true, false)]
+        public async Task IovLimitsAreRespected(bool asyncFile, bool asyncMethod)
+        {
+            // We need to write and read more than IOV_MAX buffers at a time.
+            // IOV_MAX typical value is 1024.
+            const int BufferCount = 1026;
+            const int BufferSize = 1; // the less resources we use, the better
+            const int FileSize = BufferCount * BufferSize;
+            
+            ReadOnlyMemory<byte> writeBuffer = RandomNumberGenerator.GetBytes(BufferSize);
+            ReadOnlyMemory<byte>[] writeBuffers = Enumerable.Repeat(writeBuffer, BufferCount).ToArray();
+            Memory<byte>[] readBuffers = Enumerable.Range(0, BufferCount).Select(_ => new byte[BufferSize].AsMemory()).ToArray();
+
+            FileOptions options = asyncFile ? FileOptions.Asynchronous : FileOptions.None; // we need to test both code paths
+            options |= FileOptions.DeleteOnClose;
+
+            using SafeFileHandle sfh = File.OpenHandle(GetTestFilePath(), FileMode.CreateNew, FileAccess.ReadWrite, FileShare.None, options);
+
+            if (asyncMethod)
+            {
+                await RandomAccess.WriteAsync(sfh, writeBuffers, 0);
+            }
+            else
+            {
+                RandomAccess.Write(sfh, writeBuffers, 0);
+            }
+
+            Assert.Equal(FileSize, RandomAccess.GetLength(sfh));
+
+            long fileOffset = 0;
+            int bufferOffset = 0;
+            while (fileOffset < FileSize)
+            {
+                ArraySegment<Memory<byte>> left = new ArraySegment<Memory<byte>>(readBuffers, bufferOffset, readBuffers.Length - bufferOffset);
+
+                long bytesRead = asyncMethod
+                    ? await RandomAccess.ReadAsync(sfh, left, fileOffset)
+                    : RandomAccess.Read(sfh, left, fileOffset);
+
+                fileOffset += bytesRead;
+                // The following operation is correct only because the BufferSize is 1.
+                bufferOffset += (int)bytesRead;
+            }
+
+            for (int i = 0; i < BufferCount; ++i)
+            {
+                Assert.Equal(writeBuffers[i], readBuffers[i]);
+            }
+        }
     }
 }

src/native/libs/System.Native/pal_io.c

@@ -1883,6 +1883,24 @@ int32_t SystemNative_PWrite(intptr_t fd, void* buffer, int32_t bufferSize, int64
     return (int32_t)count;
 }
 
+#if (HAVE_PREADV || HAVE_PWRITEV) && !defined(TARGET_WASM)
+static int GetAllowedVectorCount(int32_t vectorCount)
+{
+    int allowedCount = (int)vectorCount;
+
+#if defined(IOV_MAX)
+    if (IOV_MAX < allowedCount)
+    {
+        // We need to respect the limit of items that can be passed in iov.
+        // In case of writes, the managed code is reponsible for handling incomplete writes.
+        // In case of reads, we simply returns the number of bytes read and it's up to the users.
+        allowedCount = IOV_MAX;
+    }
+#endif
+    return allowedCount;
+}
+#endif // (HAVE_PREADV || HAVE_PWRITEV) && !defined(TARGET_WASM)
+
 int64_t SystemNative_PReadV(intptr_t fd, IOVector* vectors, int32_t vectorCount, int64_t fileOffset)
 {
     assert(vectors != NULL);
@@ -1891,7 +1909,8 @@ int64_t SystemNative_PReadV(intptr_t fd, IOVector* vectors, int32_t vectorCount,
     int64_t count = 0;
     int fileDescriptor = ToFileDescriptor(fd);
 #if HAVE_PREADV && !defined(TARGET_WASM) // preadv is buggy on WASM
-    while ((count = preadv(fileDescriptor, (struct iovec*)vectors, (int)vectorCount, (off_t)fileOffset)) < 0 && errno == EINTR);
+    int allowedVectorCount = GetAllowedVectorCount(vectorCount);
+    while ((count = preadv(fileDescriptor, (struct iovec*)vectors, allowedVectorCount, (off_t)fileOffset)) < 0 && errno == EINTR);
 #else
     int64_t current;
     for (int i = 0; i < vectorCount; i++)
@@ -1931,7 +1950,8 @@ int64_t SystemNative_PWriteV(intptr_t fd, IOVector* vectors, int32_t vectorCount
     int64_t count = 0;
     int fileDescriptor = ToFileDescriptor(fd);
 #if HAVE_PWRITEV && !defined(TARGET_WASM) // pwritev is buggy on WASM
-    while ((count = pwritev(fileDescriptor, (struct iovec*)vectors, (int)vectorCount, (off_t)fileOffset)) < 0 && errno == EINTR);
+    int allowedVectorCount = GetAllowedVectorCount(vectorCount);
+    while ((count = pwritev(fileDescriptor, (struct iovec*)vectors, allowedVectorCount, (off_t)fileOffset)) < 0 && errno == EINTR);
 #else
     int64_t current;
     for (int i = 0; i < vectorCount; i++)