mirror of
https://github.com/apple/swift-nio-extras.git
synced 2025-05-14 00:42:41 +08:00
Motivation: Capturing all packets is expensive. Recording to a ring buffer and then outputting on a triggering event allows this cost to be reduced. Modifications: Add a new handler - NIOPCAPRingCaptureHandler. This derives from the existing NIOWritePCAPHandler and generates PCAP recordings. A ring buffer contained in this handler stores the captured packets until RecordPreviousPackets is received as a user message at which point they are flushed to the sink. Result: There is a new handler capable of outputting packet captured data only in the build up to a known event. Co-authored-by: Cory Benfield <lukasa@apple.com> Co-authored-by: George Barnett <gbrntt@gmail.com>
99 lines
4.0 KiB
Swift
99 lines
4.0 KiB
Swift
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This source file is part of the SwiftNIO open source project
|
|
//
|
|
// Copyright (c) 2020 Apple Inc. and the SwiftNIO project authors
|
|
// Licensed under Apache License v2.0
|
|
//
|
|
// See LICENSE.txt for license information
|
|
// See CONTRIBUTORS.txt for the list of SwiftNIO project authors
|
|
//
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
import NIO
|
|
|
|
// MARK: NIOPCAPRingBuffer
|
|
/// Storage for the most recent set of packets captured subject to constraints.
|
|
/// Use `addFragment` as the sink to a `NIOWritePCAPHandler` and call `emitPCAP`
|
|
/// when you wish to get the recorded data.
|
|
/// - Warning: This class is not thread safe so should only be called from one thread.
|
|
public class NIOPCAPRingBuffer {
|
|
private var pcapFragments: CircularBuffer<ByteBuffer>
|
|
private var pcapCurrentBytes: Int
|
|
private let maximumFragments: Int
|
|
private let maximumBytes: Int
|
|
|
|
/// Initialise the buffer, setting constraints.
|
|
/// - Parameters:
|
|
/// - maximumFragments: The maximum number of pcap fragments to store.
|
|
/// - maximumBytes: The maximum number of bytes to store - note, data written may exceed this by the header size.
|
|
public init(maximumFragments: Int, maximumBytes: Int) {
|
|
precondition(maximumFragments > 0)
|
|
precondition(maximumBytes > 0)
|
|
self.maximumFragments = maximumFragments
|
|
self.maximumBytes = maximumBytes
|
|
self.pcapCurrentBytes = 0
|
|
self.pcapFragments = CircularBuffer(initialCapacity: maximumFragments)
|
|
}
|
|
|
|
/// Initialise the buffer, setting constraints
|
|
/// - Parameter maximumBytes: The maximum number of bytes to store - note, data written may exceed this by the header size.
|
|
public convenience init(maximumBytes: Int) {
|
|
self.init(maximumFragments: .max, maximumBytes: maximumBytes)
|
|
}
|
|
|
|
/// Initialise the buffer, setting constraints
|
|
/// - Parameter maximumFragments: The maximum number of pcap fragments to store.
|
|
public convenience init(maximumFragments: Int) {
|
|
self.init(maximumFragments: maximumFragments, maximumBytes: .max)
|
|
}
|
|
|
|
@discardableResult
|
|
private func popFirst() -> ByteBuffer? {
|
|
let popped = self.pcapFragments.popFirst()
|
|
if let popped = popped {
|
|
self.pcapCurrentBytes -= popped.readableBytes
|
|
}
|
|
return popped
|
|
}
|
|
|
|
private func append(_ buffer: ByteBuffer) {
|
|
self.pcapFragments.append(buffer)
|
|
self.pcapCurrentBytes += buffer.readableBytes
|
|
assert(self.pcapFragments.count <= self.maximumFragments)
|
|
// It's expected that the caller will have made room if required
|
|
// for the fragment but we may well go over on bytes - they're
|
|
// expected to fix that afterwards.
|
|
}
|
|
|
|
/// Record a fragment into the buffer, making space if required.
|
|
/// - Parameter buffer: ByteBuffer containing a pcap fragment to store
|
|
public func addFragment(_ buffer: ByteBuffer) {
|
|
// Make sure we don't go over on the number of fragments.
|
|
if self.pcapFragments.count >= self.maximumFragments {
|
|
self.popFirst()
|
|
}
|
|
precondition(self.pcapFragments.count < self.maximumFragments)
|
|
|
|
// Add the new fragment
|
|
self.append(buffer)
|
|
|
|
// Trim if we've exceeded byte limit - this could remove multiple, and indeed all fragments.
|
|
while self.pcapCurrentBytes > self.maximumBytes {
|
|
self.popFirst()
|
|
}
|
|
precondition(self.pcapCurrentBytes <= self.maximumBytes)
|
|
}
|
|
|
|
/// Emit the captured data to a consuming function; then clear the captured data.
|
|
/// - Returns: A ciruclar buffer of captured fragments.
|
|
public func emitPCAP() -> CircularBuffer<ByteBuffer> {
|
|
let toReturn = self.pcapFragments // Copy before clearing.
|
|
self.pcapFragments.removeAll(keepingCapacity: true)
|
|
self.pcapCurrentBytes = 0
|
|
return toReturn
|
|
}
|
|
}
|