Add WireGuard RX/TX data statistics (#341)

Co-authored-by: Yevgeny <y.yezub@gmail.com>
Co-authored-by: Davide De Rosa <keeshux@gmail.com>

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.
-    }
-
 }

Package.swift

@@ -117,6 +117,7 @@ let package = Package(
             name: "TunnelKitWireGuardCore",
             dependencies: [
                 "__TunnelKitUtils",
+                "TunnelKitCore",
                 "WireGuardKit",
                 "SwiftyBeaver"
             ]),

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)
     }
 

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
+        }
+    }
 }

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
+        adapter.stop { [weak self] error in
+
             // BEGIN: TunnelKit
-            self.cfg._appexSetLastError(nil)
-            // END: TunnelKit
 
-            if let error = error {
-                wg_log(.error, message: "Failed to stop WireGuard adapter: \(error.localizedDescription)")
+            guard let self else {
+                completionHandler()
+                return
+            }
+            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()
             }
-            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<DataCount, Error>) -> 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 {

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 <http://www.gnu.org/licenses/>.
+//
+
+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))
+    }
+}

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)
+    }
 }