Track data count in shared UserDefaults

Default disabled (dataCountInterval = 0).

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.

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

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 {
+            if let proxy = proxy, let dataCount = proxy.dataCount() {
-                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()
     }

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