Repository: apple/swift-nio-extras
Pull request #75 (2020-01-17 15:15:57)

NIOWritePCAPHandler: make pcap issuing configurable
### Motivation:

In the past, we would issue the written data immediately on `write` event. As of #74 , we switched to issuing the data when the write completed which reflects what happens on the wire much more closely. We did however notice that in certain cases, the old behaviour is preferable.

### Modifications:

- make the issuing behaviour configurable, defaulting to when the write completes
- make `NIOWritePCAPHandler` removable

### Result:

- more configurability


Patch:
diff --git a/Sources/NIOExtras/WritePCAPHandler.swift b/Sources/NIOExtras/WritePCAPHandler.swift
index db352f4..2d7d388 100644
--- a/Sources/NIOExtras/WritePCAPHandler.swift
+++ b/Sources/NIOExtras/WritePCAPHandler.swift
@@ -121,12 +121,48 @@ struct PCAPRecordHeader {
 ///
 /// `NIOWritePCAPHandler` will also work with Unix Domain Sockets in which case it will still synthesize a TCP packet
 /// capture with local address `111.111.111.111` (port `1111`) and remote address `222.222.222.222` (port `2222`).
-public class NIOWritePCAPHandler {
+public class NIOWritePCAPHandler: RemovableChannelHandler {
     public enum Mode {
         case client
         case server
     }
 
+    /// Settings for `NIOWritePCAPHandler`.
+    public struct Settings {
+        /// When to issue data into the `.pcap` file.
+        public enum EmitPCAP {
+            /// Write the data immediately when `NIOWritePCAPHandler` saw the event on the `ChannelPipeline`.
+            ///
+            /// For writes this means when the `write` event is triggered. Please note that this will write potentially
+            /// unflushed data into the `.pcap` file.
+            ///
+            /// If in doubt, prefer `.whenCompleted`.
+            case whenIssued
+
+            /// Write the data when the event completed.
+            ///
+            /// For writes this means when the `write` promise is succeeded. The `whenCompleted` mode mirrors most
+            /// closely what's actually sent over the wire.
+            case whenCompleted
+        }
+
+        /// When to emit the data from the `write` event into the `.pcap` file.
+        public var emitPCAPWrites: EmitPCAP
+
+        /// Default settings for the `NIOWritePCAPHandler`.
+        public init() {
+            self = .init(emitPCAPWrites: .whenCompleted)
+        }
+
+        /// Settings with customization.
+        ///
+        /// - parameters:
+        ///    - emitPCAPWrites: When to issue the writes into the `.pcap` file, see `EmitPCAP`.
+        public init(emitPCAPWrites: EmitPCAP) {
+            self.emitPCAPWrites = emitPCAPWrites
+        }
+    }
+
     private enum CloseState {
         case notClosing
         case closedInitiatorLocal
@@ -136,11 +172,12 @@ public class NIOWritePCAPHandler {
     private let fileSink: (ByteBuffer) -> Void
     private let mode: Mode
     private let maxPayloadSize = Int(UInt16.max - 40 /* needs to fit into the IPv4 header which adds 40 */)
+    private let settings: Settings
     private var buffer: ByteBuffer!
     private var readInboundBytes = 0
     private var writtenOutboundBytes = 0
     private var closeState = CloseState.notClosing
-    
+
     private static let fakeLocalAddress = try! SocketAddress(ipAddress: "111.111.111.111", port: 1111)
     private static let fakeRemoteAddress = try! SocketAddress(ipAddress: "222.222.222.222", port: 2222)
     
@@ -158,13 +195,16 @@ public class NIOWritePCAPHandler {
     /// - parameters:
     ///     - fakeLocalAddress: Allows you to optionally override the local address to be different from the real one.
     ///     - fakeRemoteAddress: Allows you to optionally override the remote address to be different from the real one.
+    ///     - settings: The settings for the `NIOWritePCAPHandler`.
     ///     - fileSink: The `fileSink` closure is called every time a new chunk of the `.pcap` file is ready to be
     ///                 written to disk or elsewhere. See `NIOSynchronizedFileSink` for a convenient way to write to
     ///                 disk.
     public init(mode: Mode,
                 fakeLocalAddress: SocketAddress? = nil,
                 fakeRemoteAddress: SocketAddress? = nil,
+                settings: Settings,
                 fileSink: @escaping (ByteBuffer) -> Void) {
+        self.settings = settings
         self.fileSink = fileSink
         self.mode = mode
         if let fakeLocalAddress = fakeLocalAddress {
@@ -174,6 +214,25 @@ public class NIOWritePCAPHandler {
             self.remoteAddress = fakeRemoteAddress
         }
     }
+
+    /// Initialize a `NIOWritePCAPHandler` with default settings.
+    ///
+    /// - parameters:
+    ///     - fakeLocalAddress: Allows you to optionally override the local address to be different from the real one.
+    ///     - fakeRemoteAddress: Allows you to optionally override the remote address to be different from the real one.
+    ///     - fileSink: The `fileSink` closure is called every time a new chunk of the `.pcap` file is ready to be
+    ///                 written to disk or elsewhere. See `NIOSynchronizedFileSink` for a convenient way to write to
+    ///                 disk.
+    public convenience init(mode: Mode,
+                            fakeLocalAddress: SocketAddress? = nil,
+                            fakeRemoteAddress: SocketAddress? = nil,
+                            fileSink: @escaping (ByteBuffer) -> Void) {
+        self.init(mode: mode,
+                  fakeLocalAddress: fakeLocalAddress,
+                  fakeRemoteAddress: fakeRemoteAddress,
+                  settings: Settings(),
+                  fileSink: fileSink)
+    }
     
     private func writeBuffer(_ buffer: ByteBuffer) {
         self.fileSink(buffer)
@@ -361,12 +420,10 @@ extension NIOWritePCAPHandler: ChannelDuplexHandler {
     }
     
     public func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
-        let promise = promise ?? context.eventLoop.makePromise()
-        let buffer = self.unwrapInboundIn(data)
+        var buffer = self.unwrapInboundIn(data)
 
-        promise.futureResult.whenSuccess {
+        func emitWrites() {
             do {
-                var buffer = buffer
                 self.buffer.clear()
                 while var payloadToSend = self.takeSensiblySizedPayload(buffer: &buffer) {
                     try self.buffer.writePCAPRecord(.init(payloadLength: payloadToSend.readableBytes,
@@ -385,7 +442,18 @@ extension NIOWritePCAPHandler: ChannelDuplexHandler {
                 context.fireErrorCaught(error)
             }
         }
-        context.write(data, promise: promise)
+
+        switch self.settings.emitPCAPWrites {
+        case .whenCompleted:
+            let promise = promise ?? context.eventLoop.makePromise()
+            promise.futureResult.whenSuccess {
+                emitWrites()
+            }
+            context.write(data, promise: promise)
+        case .whenIssued:
+            emitWrites()
+            context.write(data, promise: promise)
+        }
     }
     
     public func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) {


Test patch:
diff --git a/Tests/NIOExtrasTests/WritePCAPHandlerTest+XCTest.swift b/Tests/NIOExtrasTests/WritePCAPHandlerTest+XCTest.swift
index fc2ea65..b5dd03b 100644
--- a/Tests/NIOExtrasTests/WritePCAPHandlerTest+XCTest.swift
+++ b/Tests/NIOExtrasTests/WritePCAPHandlerTest+XCTest.swift
@@ -39,6 +39,7 @@ extension WritePCAPHandlerTest {
                 ("testOversizedOutboundDataComesAsTwoPacketsIPv6", testOversizedOutboundDataComesAsTwoPacketsIPv6),
                 ("testUnflushedOutboundDataIsNotWritten", testUnflushedOutboundDataIsNotWritten),
                 ("testDataWrittenAfterCloseIsDiscarded", testDataWrittenAfterCloseIsDiscarded),
+                ("testUnflushedOutboundDataIsWrittenWhenEmittingWritesOnIssue", testUnflushedOutboundDataIsWrittenWhenEmittingWritesOnIssue),
            ]
    }
 }
diff --git a/Tests/NIOExtrasTests/WritePCAPHandlerTest.swift b/Tests/NIOExtrasTests/WritePCAPHandlerTest.swift
index 8ed3194..720a2cd 100644
--- a/Tests/NIOExtrasTests/WritePCAPHandlerTest.swift
+++ b/Tests/NIOExtrasTests/WritePCAPHandlerTest.swift
@@ -42,12 +42,13 @@ class WritePCAPHandlerTest: XCTestCase {
     
     override func setUp() {
         self.accumulatedPackets = []
-        self.channel = EmbeddedChannel(handler: NIOWritePCAPHandler(mode: .client,
-                                                                    fakeLocalAddress: nil,
-                                                                    fakeRemoteAddress: nil,
-                                                                    fileSink: {
-            self.accumulatedPackets.append($0)
-        }))
+        self.channel = EmbeddedChannel()
+        XCTAssertNoThrow(try self.channel.pipeline.addHandler(NIOWritePCAPHandler(mode: .client,
+                                                                                  fakeLocalAddress: nil,
+                                                                                  fakeRemoteAddress: nil,
+                                                                                  fileSink: {
+                                                                                    self.accumulatedPackets.append($0)
+            }), name: "NIOWritePCAPHandler").wait())
         self.scratchBuffer = self.channel.allocator.buffer(capacity: 128)
     }
     
@@ -615,6 +616,30 @@ class WritePCAPHandlerTest: XCTestCase {
         XCTAssertEqual(20 /* just the TCP header */, ipPackets[2].payloadLength)
     }
 
+    func testUnflushedOutboundDataIsWrittenWhenEmittingWritesOnIssue() throws {
+        XCTAssertNoThrow(try self.channel.pipeline.removeHandler(name: "NIOWritePCAPHandler").wait())
+        let settings = NIOWritePCAPHandler.Settings(emitPCAPWrites: .whenIssued)
+        XCTAssertNoThrow(try self.channel.pipeline.addHandler(NIOWritePCAPHandler(mode: .client,
+                                                                                  fakeLocalAddress: nil,
+                                                                                  fakeRemoteAddress: nil,
+                                                                                  settings: settings,
+                                                                                  fileSink: {
+                                                                                    self.accumulatedPackets.append($0)
+        })).wait())
+        self.channel.localAddress = try! SocketAddress(ipAddress: "1.2.3.4", port: 1111)
+        self.channel.remoteAddress = try! SocketAddress(ipAddress: "9.8.7.6", port: 2222)
+        self.scratchBuffer.writeStaticString("hello")
+        XCTAssertNoThrow(try self.channel.writeOutbound(self.scratchBuffer))
+
+        // this is unflushed, yet we check it'll still be written because we set `settings.emitPCAPWrites = .whenIssued`
+        self.channel.write(self.scratchBuffer, promise: nil)
+        XCTAssertEqual(2, self.accumulatedPackets.count)
+        var packet1Bytes = self.accumulatedPackets.first
+        XCTAssertNotNil(packet1Bytes?.readPCAPRecord())
+        var packet2Bytes = self.accumulatedPackets.dropFirst().first
+        XCTAssertNotNil(packet2Bytes?.readPCAPRecord())
+    }
+
 }
 
 struct PCAPRecord {