From 44fb5a5b4859d6c8131b9bd14dc7bfcb97e5db30 Mon Sep 17 00:00:00 2001 From: Davide De Rosa Date: Thu, 28 Mar 2019 00:04:35 +0100 Subject: [PATCH] Track data count in shared UserDefaults Default disabled (dataCountInterval = 0). --- CHANGELOG.md | 4 +++ .../TunnelKitProvider+Configuration.swift | 36 ++++++++++++++++++- .../AppExtension/TunnelKitProvider.swift | 30 ++++++++++++++-- TunnelKit/Sources/Core/SessionProxy.swift | 5 ++- 4 files changed, 71 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index eb18d5e..4f223ca 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## Unreleased +### Added + +- Optional data count report via `TunnelKitProvider.Configuration.dataCount(in:)`. + ### Fixed - `checksEKU` not propagated to TunnelKitProvider. diff --git a/TunnelKit/Sources/AppExtension/TunnelKitProvider+Configuration.swift b/TunnelKit/Sources/AppExtension/TunnelKitProvider+Configuration.swift index 74029fc..6680032 100644 --- a/TunnelKit/Sources/AppExtension/TunnelKitProvider+Configuration.swift +++ b/TunnelKit/Sources/AppExtension/TunnelKitProvider+Configuration.swift @@ -305,7 +305,9 @@ extension TunnelKitProvider { static let debugLogFilename = "debug.log" - static let lastErrorKey = "LastTunnelKitError" + static let lastErrorKey = "TunnelKitLastError" + + fileprivate static let dataCountKey = "TunnelKitDataCount" /** Returns the URL of the latest debug log. @@ -358,6 +360,22 @@ extension TunnelKitProvider { UserDefaults(suiteName: appGroup)?.removeObject(forKey: Configuration.lastErrorKey) } + /** + Returns the most recent (received, sent) count in bytes. + + - Parameter in: The app group where to locate the count pair. + - Returns: The bytes count pair, if any. + */ + public func dataCount(in appGroup: String) -> (Int, Int)? { + guard let rawValue = UserDefaults(suiteName: appGroup)?.dataCountArray else { + return nil + } + guard rawValue.count == 2 else { + return nil + } + return (rawValue[0], rawValue[1]) + } + // MARK: API /** @@ -574,3 +592,19 @@ extension EndpointProtocol: Codable { try container.encode(rawValue) } } + +/// :nodoc: +public extension UserDefaults { + @objc public var dataCountArray: [Int]? { + get { + return array(forKey: TunnelKitProvider.Configuration.dataCountKey) as? [Int] + } + set { + set(newValue, forKey: TunnelKitProvider.Configuration.dataCountKey) + } + } + + public func removeDataCountArray() { + removeObject(forKey: TunnelKitProvider.Configuration.dataCountKey) + } +} diff --git a/TunnelKit/Sources/AppExtension/TunnelKitProvider.swift b/TunnelKit/Sources/AppExtension/TunnelKitProvider.swift index 6e8592f..48e6ca9 100644 --- a/TunnelKit/Sources/AppExtension/TunnelKitProvider.swift +++ b/TunnelKit/Sources/AppExtension/TunnelKitProvider.swift @@ -72,6 +72,9 @@ open class TunnelKitProvider: NEPacketTunnelProvider { /// The number of link failures after which the tunnel is expected to die. public var maxLinkFailures = 3 + /// The number of milliseconds between data count updates. Set to 0 to disable updates (default). + public var dataCountInterval = 0 + // MARK: Constants private let memoryLog = MemoryDestination() @@ -112,6 +115,8 @@ open class TunnelKitProvider: NEPacketTunnelProvider { private var pendingStopHandler: (() -> Void)? + private var isCountingData = false + // MARK: NEPacketTunnelProvider (XPC queue) /// :nodoc: @@ -194,6 +199,7 @@ open class TunnelKitProvider: NEPacketTunnelProvider { let proxy: SessionProxy do { proxy = try SessionProxy(queue: tunnelQueue, configuration: cfg.sessionConfiguration, cachesURL: cachesURL) + refreshDataCount() } catch let e { completionHandler(e) return @@ -244,8 +250,7 @@ open class TunnelKitProvider: NEPacketTunnelProvider { response = memoryLog.description.data(using: .utf8) case .dataCount: - if let proxy = proxy { - let dataCount = proxy.dataCount() + if let proxy = proxy, let dataCount = proxy.dataCount() { response = Data() response?.append(UInt64(dataCount.0)) // inbound response?.append(UInt64(dataCount.1)) // outbound @@ -340,6 +345,22 @@ open class TunnelKitProvider: NEPacketTunnelProvider { cancelTunnelWithError(error) } } + + // MARK: Data counter (tunnel queue) + + private func refreshDataCount() { + guard dataCountInterval > 0 else { + return + } + tunnelQueue.schedule(after: .milliseconds(dataCountInterval)) { [weak self] in + self?.refreshDataCount() + } + guard isCountingData, let proxy = proxy, let dataCount = proxy.dataCount() else { + defaults?.removeDataCountArray() + return + } + defaults?.dataCountArray = [dataCount.0, dataCount.1] + } } extension TunnelKitProvider: GenericSocketDelegate { @@ -461,12 +482,17 @@ extension TunnelKitProvider: SessionProxyDelegate { self.pendingStartHandler?(nil) self.pendingStartHandler = nil } + + isCountingData = true } /// :nodoc: public func sessionDidStop(_: SessionProxy, shouldReconnect: Bool) { log.info("Session did stop") + isCountingData = false + refreshDataCount() + reasserting = shouldReconnect socket?.shutdown() } diff --git a/TunnelKit/Sources/Core/SessionProxy.swift b/TunnelKit/Sources/Core/SessionProxy.swift index cbb04de..70968da 100644 --- a/TunnelKit/Sources/Core/SessionProxy.swift +++ b/TunnelKit/Sources/Core/SessionProxy.swift @@ -309,7 +309,10 @@ public class SessionProxy { - Returns: The current data bytes count as a pair, inbound first. */ - public func dataCount() -> (Int, Int) { + public func dataCount() -> (Int, Int)? { + guard let _ = link else { + return nil + } return controlChannel.currentDataCount() }