chore: handle waking from device sleep

Author: ethanndickson

Coder Desktop/Coder Desktop/Views/VPNMenu.swift

@@ -5,13 +5,6 @@ struct VPNMenu<VPN: VPNService, S: Session>: View {
     @EnvironmentObject var session: S
     @Environment(\.openSettings) private var openSettings
 
-    // There appears to be a race between the VPN service reporting itself as disconnected,
-    // and the system extension process exiting. When the VPN is toggled off and on quickly,
-    // an error is shown: "The VPN session failed because an internal error occurred".
-    // This forces the user to wait a few seconds before they can toggle the VPN back on.
-    @State private var waitCleanup = false
-    private var waitCleanupDuration: Duration = .seconds(6)
-
     let inspection = Inspection<Self>()
 
     var body: some View {
@@ -23,7 +16,7 @@ struct VPNMenu<VPN: VPNService, S: Session>: View {
                     Toggle(isOn: Binding(
                         get: { vpn.state == .connected || vpn.state == .connecting },
                         set: { isOn in Task {
-                            if isOn { await vpn.start() } else { await stop() }
+                            if isOn { await vpn.start() } else { await vpn.stop() }
                         }
                         }
                     )) {
@@ -93,21 +86,11 @@ struct VPNMenu<VPN: VPNService, S: Session>: View {
     }
 
     private var vpnDisabled: Bool {
-        waitCleanup ||
-            !session.hasSession ||
+        !session.hasSession ||
             vpn.state == .connecting ||
             vpn.state == .disconnecting ||
             vpn.state == .failed(.systemExtensionError(.needsUserApproval))
     }
-
-    private func stop() async {
-        await vpn.stop()
-        waitCleanup = true
-        Task {
-            try? await Task.sleep(for: waitCleanupDuration)
-            waitCleanup = false
-        }
-    }
 }
 
 func openSystemExtensionSettings() {

Coder Desktop/VPN/Manager.swift

@@ -27,7 +27,10 @@ actor Manager {
             fatalError("unknown architecture")
         #endif
         do {
-            try await download(src: dylibPath, dest: dest)
+            let sessionConfig = URLSessionConfiguration.default
+            // The tunnel might be asked to start before the network interfaces have woken up from sleep
+            sessionConfig.waitsForConnectivity = true
+            try await download(src: dylibPath, dest: dest, urlSession: URLSession(configuration: sessionConfig))
         } catch {
             throw .download(error)
         }

Coder Desktop/VPN/PacketTunnelProvider.swift

@@ -47,6 +47,11 @@ class PacketTunnelProvider: NEPacketTunnelProvider, @unchecked Sendable {
         options _: [String: NSObject]?, completionHandler: @escaping (Error?) -> Void
     ) {
         logger.info("startTunnel called")
+        start(completionHandler)
+    }
+
+    // called by `startTunnel` and on `wake`
+    func start(_ completionHandler: @escaping (Error?) -> Void) {
         guard manager == nil else {
             logger.error("startTunnel called with non-nil Manager")
             completionHandler(makeNSError(suffix: "PTP", desc: "Already running"))
@@ -95,8 +100,13 @@ class PacketTunnelProvider: NEPacketTunnelProvider, @unchecked Sendable {
         with _: NEProviderStopReason, completionHandler: @escaping () -> Void
     ) {
         logger.debug("stopTunnel called")
+        teardown(completionHandler)
+    }
+
+    // called by `stopTunnel` and `sleep`
+    func teardown(_ completionHandler: @escaping () -> Void) {
         guard let manager else {
-            logger.error("stopTunnel called with nil Manager")
+            logger.error("teardown called with nil Manager")
             completionHandler()
             return
         }
@@ -121,15 +131,25 @@ class PacketTunnelProvider: NEPacketTunnelProvider, @unchecked Sendable {
         }
     }
 
+    // sleep and wake reference: https://developer.apple.com/forums/thread/95988
     override func sleep(completionHandler: @escaping () -> Void) {
-        // Add code here to get ready to sleep.
         logger.debug("sleep called")
-        completionHandler()
+        teardown(completionHandler)
     }
 
     override func wake() {
-        // Add code here to wake up.
         logger.debug("wake called")
+        reasserting = true
+        currentSettings = .init(tunnelRemoteAddress: "127.0.0.1")
+        setTunnelNetworkSettings(nil)
+        start { error in
+            if let error {
+                self.logger.error("error starting tunnel after wake: \(error.localizedDescription)")
+                self.cancelTunnelWithError(error)
+            } else {
+                self.reasserting = false
+            }
+        }
     }
 
     // Wrapper around `setTunnelNetworkSettings` that supports merging updates

Coder Desktop/VPNLib/Download.swift

@@ -101,7 +101,7 @@ public class SignatureValidator {
     }
 }
 
-public func download(src: URL, dest: URL) async throws(DownloadError) {
+public func download(src: URL, dest: URL, urlSession: URLSession) async throws(DownloadError) {
     var req = URLRequest(url: src)
     if FileManager.default.fileExists(atPath: dest.path) {
         if let existingFileData = try? Data(contentsOf: dest, options: .mappedIfSafe) {
@@ -112,7 +112,7 @@ public func download(src: URL, dest: URL) async throws(DownloadError) {
     let tempURL: URL
     let response: URLResponse
     do {
-        (tempURL, response) = try await URLSession.shared.download(for: req)
+        (tempURL, response) = try await urlSession.download(for: req)
     } catch {
         throw .networkError(error)
     }

Coder Desktop/VPNLibTests/DownloadTests.swift

@@ -13,7 +13,7 @@ struct DownloadTests {
         let fileURL = URL(string: "http://example.com/test1.txt")!
         Mock(url: fileURL, contentType: .html, statusCode: 200, data: [.get: testData]).register()
 
-        try await download(src: fileURL, dest: destinationURL)
+        try await download(src: fileURL, dest: destinationURL, urlSession: URLSession.shared)
 
         try #require(FileManager.default.fileExists(atPath: destinationURL.path))
         defer { try? FileManager.default.removeItem(at: destinationURL) }
@@ -32,7 +32,7 @@ struct DownloadTests {
 
         Mock(url: fileURL, contentType: .html, statusCode: 200, data: [.get: testData]).register()
 
-        try await download(src: fileURL, dest: destinationURL)
+        try await download(src: fileURL, dest: destinationURL, urlSession: URLSession.shared)
         try #require(FileManager.default.fileExists(atPath: destinationURL.path))
         let downloadedData = try Data(contentsOf: destinationURL)
         #expect(downloadedData == testData)
@@ -44,7 +44,7 @@ struct DownloadTests {
         }
         mock.register()
 
-        try await download(src: fileURL, dest: destinationURL)
+        try await download(src: fileURL, dest: destinationURL, urlSession: URLSession.shared)
         let unchangedData = try Data(contentsOf: destinationURL)
         #expect(unchangedData == testData)
         #expect(etagIncluded)
@@ -61,7 +61,7 @@ struct DownloadTests {
 
         Mock(url: fileURL, contentType: .html, statusCode: 200, data: [.get: ogData]).register()
 
-        try await download(src: fileURL, dest: destinationURL)
+        try await download(src: fileURL, dest: destinationURL, urlSession: URLSession.shared)
         try #require(FileManager.default.fileExists(atPath: destinationURL.path))
         var downloadedData = try Data(contentsOf: destinationURL)
         #expect(downloadedData == ogData)
@@ -73,7 +73,7 @@ struct DownloadTests {
         }
         mock.register()
 
-        try await download(src: fileURL, dest: destinationURL)
+        try await download(src: fileURL, dest: destinationURL, urlSession: URLSession.shared)
         downloadedData = try Data(contentsOf: destinationURL)
         #expect(downloadedData == newData)
         #expect(etagIncluded)