Add AsynchronizedFileSink for non-blocking PCAP recording

Sources/NIOExtras/WritePCAPHandler.swift

@@ -15,6 +15,7 @@
 import CNIOLinux
 import Dispatch
 import NIOCore
+import NIOConcurrencyHelpers
 
 #if canImport(Darwin)
 import Darwin
@@ -808,6 +809,103 @@ extension NIOWritePCAPHandler {
             }
         }
     }
+    public final class AsynchronizedFileSink {
+        private let fileHandle: NIOFileHandle
+        private let eventLoop: EventLoop
+        private let errorHandler: @Sendable (Swift.Error) -> Void
+        private let state: NIOLockedValueBox<State> = NIOLockedValueBox(.running)
+
+        public enum FileWritingMode {
+            case appendToExistingPCAPFile
+            case createNewPCAPFile
+        }
+
+        public struct Error: Swift.Error {
+            public var errorCode: Int
+
+            internal enum ErrorCode: Int {
+                case cannotOpenFileError = 1
+                case cannotWriteToFileError
+            }
+        }
+
+        private enum State {
+            case running
+            case error(Swift.Error)
+        }
+
+        /// Creates an AsynchronizedFileSink for writing to a .pcap file at the given path.
+        /// If fileWritingMode is `.createNewPCAPFile`, a file header is written.
+        public static func fileSinkWritingToFile(
+            path: String,
+            fileWritingMode: FileWritingMode = .createNewPCAPFile,
+            errorHandler: @escaping @Sendable (Swift.Error) -> Void,
+            on eventLoop: EventLoop
+        ) async throws -> AsynchronizedFileSink {
+            let oflag: CInt = fileWritingMode == .createNewPCAPFile ? (O_TRUNC | O_CREAT) : O_APPEND
+            let fd: CInt = path.withCString { pathPtr in
+                open(pathPtr, O_WRONLY | oflag, 0o600)
+            }
+            if fd < 0 {
+                throw Error(errorCode: Error.ErrorCode.cannotOpenFileError.rawValue)
+            }
+
+            /// Write PCAP file header
+            if fileWritingMode == .createNewPCAPFile {
+                let writeOk: Bool = NIOWritePCAPHandler.pcapFileHeader.withUnsafeReadableBytes { ptr in
+                    sysWrite(fd, ptr.baseAddress, ptr.count) == ptr.count
+                }
+                if !writeOk {
+                    throw Error(errorCode: Error.ErrorCode.cannotWriteToFileError.rawValue)
+                }
+            }
+
+            let fileHandle: NIOFileHandle = NIOFileHandle(_deprecatedTakingOwnershipOfDescriptor: fd)
+            return AsynchronizedFileSink(fileHandle: fileHandle, eventLoop: eventLoop, errorHandler: errorHandler)
+        }
+
+        private init(
+            fileHandle: NIOFileHandle,
+            eventLoop: EventLoop,
+            errorHandler: @escaping @Sendable (Swift.Error) -> Void
+        ) {
+            self.fileHandle = fileHandle
+            self.eventLoop = eventLoop
+            self.errorHandler = errorHandler
+        }
+
+        public func write(buffer: ByteBuffer) async throws {
+            try self.fileHandle.withUnsafeFileDescriptor { fd in
+                var buffer = buffer
+                while buffer.readableBytes > 0 {
+                    try buffer.readWithUnsafeReadableBytes { dataPtr in
+                        let written = sysWrite(fd, dataPtr.baseAddress, dataPtr.count)
+                        guard written > 0 else {
+                            throw Error(errorCode: Error.ErrorCode.cannotWriteToFileError.rawValue)
+                        }
+                        return written
+                    }
+                }
+            }
+        }
+
+        /// Syncs the file to disk using fsync.
+        public func asyncSync() async throws {
+            try self.fileHandle.withUnsafeFileDescriptor { fd in
+                let result: CInt = fsync(fd)
+                if result != 0 {
+                    throw Error(errorCode: Error.ErrorCode.cannotWriteToFileError.rawValue)
+                }
+            }
+        }
+
+        /// Closes the file sink.
+        public func close() async throws {
+            try self.fileHandle.close()
+            print("File successfully closed.")
+        }
+    }
 }
 
 extension NIOWritePCAPHandler.SynchronizedFileSink: @unchecked Sendable {}
+extension NIOWritePCAPHandler.AsynchronizedFileSink: Sendable {}

Tests/NIOExtrasTests/WritePCAPHandlerTest.swift

@@ -17,6 +17,7 @@ import Foundation
 import NIOCore
 import NIOEmbedded
 import XCTest
+import NIOPosix
 
 @testable import NIOExtras
 
@@ -816,6 +817,61 @@ class WritePCAPHandlerTest: XCTestCase {
         XCTAssertNoThrow(XCTAssertTrue(try channel.finish().isClean))
     }
 
+    func testAsynchronizedFileSinkWritesDataToFile() async throws {
+        let testHostname: String = "testhost"
+        let filePath: String = "/tmp/packets-\(testHostname)-\(UUID())-\(getpid())-\(Int(Date().timeIntervalSince1970)).pcap"
+
+        let eventLoop: EmbeddedEventLoop = EmbeddedEventLoop()
+
+        let fileSink: NIOWritePCAPHandler.AsynchronizedFileSink = try await NIOWritePCAPHandler.AsynchronizedFileSink.fileSinkWritingToFile(
+            path: filePath,
+            fileWritingMode: .createNewPCAPFile,
+            errorHandler: { error in XCTFail("PCAP logging error: \(error)") },
+            on: eventLoop
+        )
+
+        let channel: EmbeddedChannel = EmbeddedChannel(handler: NIOWritePCAPHandler(
+            mode: .client,
+            fakeLocalAddress: nil,
+            fakeRemoteAddress: nil,
+            fileSink: { buffer in
+                Task.detached {
+                    do {
+                        try await fileSink.write(buffer: buffer)
+                    } catch {
+                        XCTFail("Failed to write to file sink: \(error)")
+                    }
+                }
+            }
+        ))
+
+        var buffer: ByteBuffer = channel.allocator.buffer(capacity: 64)
+        buffer.writeString("Test PCAP data")
+        try await fileSink.write(buffer: buffer)
+        try await channel.closeFuture.get()
+
+        // Flush any buffered data to disk and close the file.
+        try await fileSink.asyncSync()
+        // try await fileSink.close()
+        let expectation: XCTestExpectation = XCTestExpectation(description: "File sink should close")
+        Task {
+            do {
+                try await fileSink.close()
+                expectation.fulfill()
+            } catch {
+                XCTFail("Failed to close file sink: \(error)")
+            }
+        }
+        await fulfillment(of: [expectation], timeout: 5)
+
+        // Verify that the file exists and contains data.
+        let fileManager: FileManager = FileManager.default
+        let fileData: Data = try Data(contentsOf: URL(fileURLWithPath: filePath))
+        XCTAssertGreaterThan(fileData.count, 0, "PCAP file should contain data")
+
+        // Clean up the temporary file.
+        try fileManager.removeItem(atPath: filePath)
+    }
 }
 
 struct PCAPRecord {