From d003730bd8640971318357f6fc77be3befeeea23 Mon Sep 17 00:00:00 2001
From: cnkwocha <ci_nkwocha@apple.com>
Date: Wed, 9 Oct 2024 15:35:06 +0100
Subject: [PATCH] [NIOFileSystem] Avoid crash when reading more than
 `ByteBuffer` capacity

Motivation:

As described in issue (#2878)[https://github.com/apple/swift-nio/issues/2878], NIOFileSystem crashes when reading more than `ByteBuffer` capacity.

Modifications:

- Add a new static property, `byteBufferCapacity`, to `ByteCount` representing the maximum amount of bytes that can be written to `ByteBuffer`.
- Throw a `FileSystemError` in `ReadableFileHandleProtocol.readToEnd(fromAbsoluteOffset:maximumSizeAllowed:)` when `maximumSizeAllowed` is more than `ByteCount.byteBufferCapacity`.

Result:

NIOFileSystem will `throw` instead of crashing.
---
 Sources/NIOFileSystem/ByteCount.swift           | 15 +++++++++++++++
 Sources/NIOFileSystem/FileHandleProtocol.swift  | 17 ++++++++++++++++-
 .../FileSystemTests.swift                       | 15 +++++++++++++++
 3 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/Sources/NIOFileSystem/ByteCount.swift b/Sources/NIOFileSystem/ByteCount.swift
index 64fe76a309..983b7c0380 100644
--- a/Sources/NIOFileSystem/ByteCount.swift
+++ b/Sources/NIOFileSystem/ByteCount.swift
@@ -80,6 +80,21 @@ public struct ByteCount: Hashable, Sendable {
     }
 }
 
+extension ByteCount {
+    #if arch(arm) || arch(i386) || arch(arm64_32)
+    // on 32-bit platforms we can't make use of a whole UInt32.max (as it doesn't fit in an Int)
+    private static let byteBufferMaxIndex = UInt32(Int.max)
+    #else
+    // on 64-bit platforms we're good
+    private static let byteBufferMaxIndex = UInt32.max
+    #endif
+
+    /// A ``ByteCount`` for the maximum amount of bytes that can be written to `ByteBuffer`.
+    internal static var byteBufferCapacity: ByteCount {
+        ByteCount(bytes: Int64(byteBufferMaxIndex))
+    }
+}
+
 extension ByteCount: AdditiveArithmetic {
     public static var zero: ByteCount { ByteCount(bytes: 0) }
 
diff --git a/Sources/NIOFileSystem/FileHandleProtocol.swift b/Sources/NIOFileSystem/FileHandleProtocol.swift
index d58e33065a..94dc4a349b 100644
--- a/Sources/NIOFileSystem/FileHandleProtocol.swift
+++ b/Sources/NIOFileSystem/FileHandleProtocol.swift
@@ -322,7 +322,9 @@ extension ReadableFileHandleProtocol {
     ///
     /// - Important: This method checks whether the file is seekable or not (i.e., whether it's a socket,
     /// pipe or FIFO), and will throw ``FileSystemError/Code-swift.struct/unsupported`` if
-    /// an offset other than zero is passed.
+    /// an offset other than zero is passed. Also, it will throw
+    /// ``FileSystemError/Code-swift.struct/resourceExhausted`` if `maximumSizeAllowed` is more than can be
+    /// written to `ByteBuffer`.
     ///
     /// - Parameters:
     ///   - offset: The absolute offset into the file to read from. Defaults to zero.
@@ -340,6 +342,19 @@ extension ReadableFileHandleProtocol {
         let fileSize = Int64(info.size)
         let readSize = max(Int(fileSize - offset), 0)
 
+        if maximumSizeAllowed > .byteBufferCapacity {
+            throw FileSystemError(
+                code: .resourceExhausted,
+                message: """
+                    The maximum size allowed (\(maximumSizeAllowed)) is more than the maximum \
+                    amount of bytes that can be written to ByteBuffer \
+                    (\(ByteCount.byteBufferCapacity)).
+                    """,
+                cause: nil,
+                location: .here()
+            )
+        }
+
         if readSize > maximumSizeAllowed.bytes {
             throw FileSystemError(
                 code: .resourceExhausted,
diff --git a/Tests/NIOFileSystemIntegrationTests/FileSystemTests.swift b/Tests/NIOFileSystemIntegrationTests/FileSystemTests.swift
index 397e9b76ca..f0e6aee74a 100644
--- a/Tests/NIOFileSystemIntegrationTests/FileSystemTests.swift
+++ b/Tests/NIOFileSystemIntegrationTests/FileSystemTests.swift
@@ -1803,6 +1803,21 @@ extension FileSystemTests {
             XCTAssertEqual(byteCount, Int(size))
         }
     }
+
+    func testReadMoreThanByteBufferCapacity() async throws {
+        let path = try await self.fs.temporaryFilePath()
+
+        try await self.fs.withFileHandle(forReadingAndWritingAt: path) { fileHandle in
+            await XCTAssertThrowsFileSystemErrorAsync {
+                // Set `maximumSizeAllowed` to 1 byte more than can be written to `ByteBuffer`.
+                try await fileHandle.readToEnd(
+                    maximumSizeAllowed: .byteBufferCapacity + .bytes(1)
+                )
+            } onError: { error in
+                XCTAssertEqual(error.code, .resourceExhausted)
+            }
+        }
+    }
 }
 
 #if !canImport(Darwin) && swift(<5.9.2)