From bda84bf569792fbb702d0173de3c9c58768f9153 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 15 Dec 2023 05:01:26 +0800 Subject: [PATCH] Add WireGuard RX/TX data statistics (#341) Co-authored-by: Yevgeny Co-authored-by: Davide De Rosa --- Demo/Host/ViewController.swift | 11 -- Package.swift | 1 + .../NetworkExtensionVPN.swift | 2 +- .../TunnelKitManager/VPNNotification.swift | 15 +++ .../WireGuardTunnelProvider.swift | 101 +++++++++++++++--- .../DataCount+WireGuard.swift | 60 +++++++++++ .../WireGuard+ProviderConfiguration.swift | 45 ++++++++ 7 files changed, 211 insertions(+), 24 deletions(-) create mode 100644 Sources/TunnelKitWireGuardManager/DataCount+WireGuard.swift diff --git a/Demo/Host/ViewController.swift b/Demo/Host/ViewController.swift index 5ecce47..eb2eb13 100644 --- a/Demo/Host/ViewController.swift +++ b/Demo/Host/ViewController.swift @@ -37,15 +37,4 @@ import UIKit class ViewController: UIViewController { - - override func viewDidLoad() { - super.viewDidLoad() - // Do any additional setup after loading the view, typically from a nib. - } - - override func didReceiveMemoryWarning() { - super.didReceiveMemoryWarning() - // Dispose of any resources that can be recreated. - } - } diff --git a/Package.swift b/Package.swift index de80840..64dcc4a 100644 --- a/Package.swift +++ b/Package.swift @@ -117,6 +117,7 @@ let package = Package( name: "TunnelKitWireGuardCore", dependencies: [ "__TunnelKitUtils", + "TunnelKitCore", "WireGuardKit", "SwiftyBeaver" ]), diff --git a/Sources/TunnelKitManager/NetworkExtensionVPN.swift b/Sources/TunnelKitManager/NetworkExtensionVPN.swift index 3e09699..3e59466 100644 --- a/Sources/TunnelKitManager/NetworkExtensionVPN.swift +++ b/Sources/TunnelKitManager/NetworkExtensionVPN.swift @@ -240,11 +240,11 @@ public class NetworkExtensionVPN: VPN { } let bundleId = connection.manager.tunnelBundleIdentifier log.debug("VPN status did change (\(bundleId ?? "?")): isEnabled=\(connection.manager.isEnabled), status=\(connection.status.rawValue)") - var notification = Notification(name: VPNNotification.didChangeStatus) notification.vpnBundleIdentifier = bundleId notification.vpnIsEnabled = connection.manager.isEnabled notification.vpnStatus = connection.status.wrappedStatus + notification.connectionDate = connection.connectedDate NotificationCenter.default.post(notification) } diff --git a/Sources/TunnelKitManager/VPNNotification.swift b/Sources/TunnelKitManager/VPNNotification.swift index 40d5a9a..b4b63e1 100644 --- a/Sources/TunnelKitManager/VPNNotification.swift +++ b/Sources/TunnelKitManager/VPNNotification.swift @@ -99,4 +99,19 @@ extension Notification { userInfo = newInfo } } + + /// The current VPN connection date. + public var connectionDate: Date? { + get { + guard let date = userInfo?["ConnectionDate"] as? Date else { + fatalError("Notification has no connectionDate") + } + return date + } + set { + var newInfo = userInfo ?? [:] + newInfo["ConnectionDate"] = newValue + userInfo = newInfo + } + } } diff --git a/Sources/TunnelKitWireGuardAppExtension/WireGuardTunnelProvider.swift b/Sources/TunnelKitWireGuardAppExtension/WireGuardTunnelProvider.swift index 5bfc1f9..64fb5a3 100644 --- a/Sources/TunnelKitWireGuardAppExtension/WireGuardTunnelProvider.swift +++ b/Sources/TunnelKitWireGuardAppExtension/WireGuardTunnelProvider.swift @@ -1,3 +1,4 @@ +import TunnelKitCore import TunnelKitWireGuardCore import TunnelKitWireGuardManager import WireGuardKit @@ -14,6 +15,14 @@ import os open class WireGuardTunnelProvider: NEPacketTunnelProvider { private var cfg: WireGuard.ProviderConfiguration! + /// The number of milliseconds between data count updates. Set to 0 to disable updates (default). + public var dataCountInterval = 0 + + /// Once the tunnel starts, enable this property to update connection stats + private var tunnelIsStarted = false + + private let tunnelQueue = DispatchQueue(label: WireGuardTunnelProvider.description(), qos: .utility) + private lazy var adapter: WireGuardAdapter = { return WireGuardAdapter(with: self) { logLevel, message in wg_log(logLevel.osLogLevel, message: message) @@ -45,12 +54,20 @@ open class WireGuardTunnelProvider: NEPacketTunnelProvider { // END: TunnelKit // Start the tunnel - adapter.start(tunnelConfiguration: tunnelConfiguration) { adapterError in + adapter.start(tunnelConfiguration: tunnelConfiguration) { [weak self] adapterError in + guard let self else { + completionHandler(nil) + return + } + guard let adapterError = adapterError else { let interfaceName = self.adapter.interfaceName ?? "unknown" wg_log(.info, message: "Tunnel interface is \(interfaceName)") - + self.tunnelQueue.async { + self.tunnelIsStarted = true + self.refreshDataCount() + } completionHandler(nil) return } @@ -88,15 +105,24 @@ open class WireGuardTunnelProvider: NEPacketTunnelProvider { open override func stopTunnel(with reason: NEProviderStopReason, completionHandler: @escaping () -> Void) { wg_log(.info, staticMessage: "Stopping tunnel") - adapter.stop { error in - // BEGIN: TunnelKit - self.cfg._appexSetLastError(nil) - // END: TunnelKit + adapter.stop { [weak self] error in - if let error = error { - wg_log(.error, message: "Failed to stop WireGuard adapter: \(error.localizedDescription)") + // BEGIN: TunnelKit + + guard let self else { + completionHandler() + return } - completionHandler() + self.tunnelQueue.async { + self.cfg._appexSetLastError(nil) + self.tunnelIsStarted = false + if let error = error { + wg_log(.error, message: "Failed to stop WireGuard adapter: \(error.localizedDescription)") + } + completionHandler() + } + + // END: TunnelKit #if os(macOS) // HACK: This is a filthy hack to work around Apple bug 32073323 (dup'd by us as 47526107). @@ -108,7 +134,9 @@ open class WireGuardTunnelProvider: NEPacketTunnelProvider { } open override func handleAppMessage(_ messageData: Data, completionHandler: ((Data?) -> Void)? = nil) { - guard let completionHandler = completionHandler else { return } + guard let completionHandler = completionHandler else { + return + } if messageData.count == 1 && messageData[0] == 0 { adapter.getRuntimeConfiguration { settings in @@ -122,10 +150,48 @@ open class WireGuardTunnelProvider: NEPacketTunnelProvider { completionHandler(nil) } } + + // MARK: Data counter (tunnel queue) + + // XXX: thread-safety here is poor, but we know that: + // + // - dataCountInterval is virtually constant, set on tunnel creation + // - cfg only modifies UserDefaults, which is thread-safe + // - adapter, used in fetchDataCount, is thread-safe + // + private func refreshDataCount() { + guard dataCountInterval > 0 else { + return + } + + tunnelQueue.schedule(after: DispatchTimeInterval.milliseconds(dataCountInterval)) { [weak self] in + self?.refreshDataCount() + } + + guard tunnelIsStarted else { + cfg._appexSetDataCount(nil) + return + } + fetchDataCount { [weak self] result in + guard let self else { + return + } + switch result { + case .success(let dataCount): + self.cfg._appexSetDataCount(dataCount) + case .failure(let error): + wg_log(.error, message: "Failed to refresh data count \(error.localizedDescription)") + } + } + } } -extension WireGuardTunnelProvider { - private func configureLogging() { +private extension WireGuardTunnelProvider { + enum StatsError: Error { + case parseFailure + } + + func configureLogging() { let logLevel: SwiftyBeaver.Level = (cfg.shouldDebug ? .debug : .info) let logFormat = cfg.debugLogFormat ?? "$Dyyyy-MM-dd HH:mm:ss.SSS$d $L $N.$F:$l - $M" @@ -146,6 +212,17 @@ extension WireGuardTunnelProvider { // store path for clients cfg._appexSetDebugLogPath() } + + func fetchDataCount(completiondHandler: @escaping (Result) -> Void) { + adapter.getRuntimeConfiguration { configurationString in + if let configurationString = configurationString, + let wireGuardDataCount = DataCount.from(wireGuardString: configurationString) { + completiondHandler(.success(wireGuardDataCount)) + } else { + completiondHandler(.failure(StatsError.parseFailure)) + } + } + } } extension WireGuardLogLevel { diff --git a/Sources/TunnelKitWireGuardManager/DataCount+WireGuard.swift b/Sources/TunnelKitWireGuardManager/DataCount+WireGuard.swift new file mode 100644 index 0000000..748d53f --- /dev/null +++ b/Sources/TunnelKitWireGuardManager/DataCount+WireGuard.swift @@ -0,0 +1,60 @@ +// +// DataCount+WireGuard.swift +// Passepartout +// +// Created by Yevgeny Yezub on 11/17/23. +// Copyright (c) 2023 Yevgeny Yezub. All rights reserved. +// +// https://github.com/passepartoutvpn +// +// This file is part of Passepartout. +// +// Passepartout is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// Passepartout is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with Passepartout. If not, see . +// + +import Foundation +import TunnelKitCore + +extension DataCount { + public static func from(wireGuardString string: String) -> DataCount? { + var bytesReceived: UInt? + var bytesSent: UInt? + + string.enumerateLines { line, stop in + if bytesReceived == nil, let value = line.getPrefix("rx_bytes=") { + bytesReceived = value + } else if bytesSent == nil, let value = line.getPrefix("tx_bytes=") { + bytesSent = value + } + if bytesReceived != nil, bytesSent != nil { + stop = true + } + } + + guard let bytesReceived, let bytesSent else { + return nil + } + + return DataCount(bytesReceived, bytesSent) + } +} + +private extension String { + func getPrefix(_ prefixKey: String) -> UInt? { + guard hasPrefix(prefixKey) else { + return nil + } + return UInt(dropFirst(prefixKey.count)) + } +} diff --git a/Sources/TunnelKitWireGuardManager/WireGuard+ProviderConfiguration.swift b/Sources/TunnelKitWireGuardManager/WireGuard+ProviderConfiguration.swift index 9bf2ab3..e0d29b8 100644 --- a/Sources/TunnelKitWireGuardManager/WireGuard+ProviderConfiguration.swift +++ b/Sources/TunnelKitWireGuardManager/WireGuard+ProviderConfiguration.swift @@ -25,6 +25,7 @@ import Foundation import NetworkExtension +import TunnelKitCore import TunnelKitManager import TunnelKitWireGuardCore import WireGuardKit @@ -41,6 +42,8 @@ extension WireGuard { case logPath = "WireGuard.LogPath" case lastError = "WireGuard.LastError" + + case dataCount = "WireGuard.DataCount" } public let title: String @@ -91,6 +94,12 @@ extension WireGuard.ProviderConfiguration: NetworkExtensionConfiguration { // MARK: Shared data extension WireGuard.ProviderConfiguration { + + /// The most recent (received, sent) count in bytes. + public var dataCount: DataCount? { + return defaults?.wireGuardDataCount + } + public var lastError: TunnelKitWireGuardError? { return defaults?.wireGuardLastError } @@ -102,9 +111,14 @@ extension WireGuard.ProviderConfiguration { private var defaults: UserDefaults? { return UserDefaults(suiteName: appGroup) } + } extension WireGuard.ProviderConfiguration { + public func _appexSetDataCount(_ newValue: DataCount?) { + defaults?.wireGuardDataCount = newValue + } + public func _appexSetLastError(_ newValue: TunnelKitWireGuardError?) { defaults?.wireGuardLastError = newValue } @@ -146,4 +160,35 @@ extension UserDefaults { set(newValue.rawValue, forKey: WireGuard.ProviderConfiguration.Keys.lastError.rawValue) } } + + public fileprivate(set) var wireGuardDataCount: DataCount? { + get { + guard let rawValue = wireGuardDataCountArray else { + return nil + } + guard rawValue.count == 2 else { + return nil + } + return DataCount(rawValue[0], rawValue[1]) + } + set { + guard let newValue = newValue else { + wireGuardRemoveDataCountArray() + return + } + wireGuardDataCountArray = [newValue.received, newValue.sent] + } + } + + @objc private var wireGuardDataCountArray: [UInt]? { + get { + return array(forKey: WireGuard.ProviderConfiguration.Keys.dataCount.rawValue) as? [UInt] + } + set { + set(newValue, forKey: WireGuard.ProviderConfiguration.Keys.dataCount.rawValue) + } + } + private func wireGuardRemoveDataCountArray() { + removeObject(forKey: WireGuard.ProviderConfiguration.Keys.dataCount.rawValue) + } }