Merge f6c2bc355f073d437fe833a54b5874d7f04fd775 into f1f6f772198bee35d99dd145f1513d8581a54f2c

Sources/NIOExtras/WritePCAPHandler.swift

@@ -14,6 +14,7 @@
 
 import CNIOLinux
 import Dispatch
+import NIOConcurrencyHelpers
 import NIOCore
 
 #if canImport(Darwin)
@@ -810,6 +811,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

@@ -16,6 +16,7 @@ import CNIOLinux
 import Foundation
 import NIOCore
 import NIOEmbedded
+import NIOPosix
 import XCTest
 
 @testable import NIOExtras
@@ -816,6 +817,37 @@ 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
+            )
+
+        // Write test data directly using the file sink.
+        var buffer = ByteBufferAllocator().buffer(capacity: 64)
+        buffer.writeString("Test PCAP data")
+        try await fileSink.write(buffer: buffer)
+
+        // Sync and then close the file sink.
+        try await fileSink.asyncSync()
+        try await fileSink.close()
+
+        // Verify that the file exists and contains data.
+        let fileData = try Data(contentsOf: URL(fileURLWithPath: filePath))
+        XCTAssertGreaterThan(fileData.count, 0, "PCAP file should contain data")
+
+        // Clean up the temporary file.
+        try FileManager.default.removeItem(atPath: filePath)
+    }
 }
 
 struct PCAPRecord {