diff --git a/Sources/NIOExtras/WritePCAPHandler.swift b/Sources/NIOExtras/WritePCAPHandler.swift index 914309b3..47e6417a 100644 --- a/Sources/NIOExtras/WritePCAPHandler.swift +++ b/Sources/NIOExtras/WritePCAPHandler.swift @@ -14,6 +14,7 @@ import CNIOLinux import Dispatch +import NIOConcurrencyHelpers import NIOCore #if canImport(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 = 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 {} diff --git a/Tests/NIOExtrasTests/WritePCAPHandlerTest.swift b/Tests/NIOExtrasTests/WritePCAPHandlerTest.swift index 72f11a3a..e070765c 100644 --- a/Tests/NIOExtrasTests/WritePCAPHandlerTest.swift +++ b/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 {