From 2329f712cfd6a19177a204fe26854c9b9f2aa206 Mon Sep 17 00:00:00 2001
From: Andrej Mihajlov <and@mullvad.net>
Date: Wed, 2 Dec 2020 13:48:24 +0100
Subject: [PATCH] WireGuardKit: Pass logHandler via constructor

Signed-off-by: Andrej Mihajlov <and@mullvad.net>
---
 .../PacketTunnelProvider.swift                | 11 +++---
 .../WireGuardKit/WireGuardAdapter.swift       | 35 ++++++-------------
 2 files changed, 15 insertions(+), 31 deletions(-)

diff --git a/WireGuard/WireGuardNetworkExtension/PacketTunnelProvider.swift b/WireGuard/WireGuardNetworkExtension/PacketTunnelProvider.swift
index 28eb986..f88d117 100644
--- a/WireGuard/WireGuardNetworkExtension/PacketTunnelProvider.swift
+++ b/WireGuard/WireGuardNetworkExtension/PacketTunnelProvider.swift
@@ -8,7 +8,11 @@ import os
 
 class PacketTunnelProvider: NEPacketTunnelProvider {
 
-    private lazy var adapter = WireGuardAdapter(with: self)
+    private lazy var adapter: WireGuardAdapter = {
+        return WireGuardAdapter(with: self) { logLevel, message in
+            wg_log(logLevel.osLogLevel, message: message)
+        }
+    }()
 
     override func startTunnel(options: [String: NSObject]?, completionHandler: @escaping (Error?) -> Void) {
         let activationAttemptId = options?["activationAttemptId"] as? String
@@ -25,11 +29,6 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
             return
         }
 
-        // Setup WireGuard logger
-        adapter.setLogHandler { logLevel, message in
-            wg_log(logLevel.osLogLevel, message: message)
-        }
-
         // Start the tunnel
         adapter.start(tunnelConfiguration: tunnelConfiguration) { adapterError in
             guard let adapterError = adapterError else {
diff --git a/WireGuardKit/Sources/WireGuardKit/WireGuardAdapter.swift b/WireGuardKit/Sources/WireGuardKit/WireGuardAdapter.swift
index 749def3..a19939f 100644
--- a/WireGuardKit/Sources/WireGuardKit/WireGuardAdapter.swift
+++ b/WireGuardKit/Sources/WireGuardKit/WireGuardAdapter.swift
@@ -35,7 +35,7 @@ public class WireGuardAdapter {
     private weak var packetTunnelProvider: NEPacketTunnelProvider?
 
     /// Log handler closure.
-    private var logHandler: LogHandler?
+    private let logHandler: LogHandler
 
     /// WireGuard internal handle returned by `wgTurnOn` that's used to associate the calls
     /// with the specific WireGuard tunnel.
@@ -91,14 +91,18 @@ public class WireGuardAdapter {
     /// Designated initializer.
     /// - Parameter packetTunnelProvider: an instance of `NEPacketTunnelProvider`. Internally stored
     ///   as a weak reference.
-    public init(with packetTunnelProvider: NEPacketTunnelProvider) {
+    /// - Parameter logHandler: a log handler closure.
+    public init(with packetTunnelProvider: NEPacketTunnelProvider, logHandler: @escaping LogHandler) {
         self.packetTunnelProvider = packetTunnelProvider
+        self.logHandler = logHandler
+
+        setupLogHandler()
     }
 
     deinit {
-        // Force deactivate logger to make sure that no further calls to the instance of this class
+        // Force remove logger to make sure that no further calls to the instance of this class
         // can happen after deallocation.
-        deactivateLogHandler()
+        wgSetLogger(nil, nil)
 
         // Cancel network monitor
         networkMonitor?.cancel()
@@ -129,20 +133,6 @@ public class WireGuardAdapter {
         }
     }
 
-    /// Set log handler.
-    /// - Parameter logHandler: log handler closure
-    public func setLogHandler(_ logHandler: LogHandler?) {
-        workQueue.async {
-            self.logHandler = logHandler
-        }
-
-        if logHandler == nil {
-            deactivateLogHandler()
-        } else {
-            activateLogHandler()
-        }
-    }
-
     /// Start the tunnel tunnel.
     /// - Parameters:
     ///   - tunnelConfiguration: tunnel configuration.
@@ -247,8 +237,8 @@ public class WireGuardAdapter {
 
     // MARK: - Private methods
 
-    /// Install WireGuard log handler.
-    private func activateLogHandler() {
+    /// Setup WireGuard log handler.
+    private func setupLogHandler() {
         let context = Unmanaged.passUnretained(self).toOpaque()
         wgSetLogger(context) { (context, logLevel, message) in
             guard let context = context, let message = message else { return }
@@ -263,11 +253,6 @@ public class WireGuardAdapter {
         }
     }
 
-    /// Uninstall WireGuard log handler.
-    private func deactivateLogHandler() {
-        wgSetLogger(nil, nil)
-    }
-
     /// Resolve endpoints and update network configuration.
     /// - Parameters:
     ///   - tunnelConfiguration: tunnel configuration