From 6b71ed5f0668df94d4bdb938c20d699a4cf051e9 Mon Sep 17 00:00:00 2001
From: Jackson Coxson <jkcoxson@gmail.com>
Date: Sat, 29 Mar 2025 13:35:02 -0600
Subject: [PATCH] Add VPN check before heartbeat

---
 StikJIT/StikJITApp.swift | 103 +++++++++++++++++++++++++++++++++------
 1 file changed, 88 insertions(+), 15 deletions(-)

diff --git a/StikJIT/StikJITApp.swift b/StikJIT/StikJITApp.swift
index ce72c82..5183294 100644
--- a/StikJIT/StikJITApp.swift
+++ b/StikJIT/StikJITApp.swift
@@ -6,6 +6,7 @@
 //
 
 import SwiftUI
+import Network
 import em_proxy
 import UniformTypeIdentifiers
 
@@ -44,26 +45,41 @@ struct HeartbeatApp: App {
                     .onAppear {
                         startProxy() { result, error in
                             if result {
-                                if FileManager.default.fileExists(atPath: URL.documentsDirectory.appendingPathComponent("pairingFile.plist").path) {
-                                    Timer.scheduledTimer(withTimeInterval: 0.5, repeats: true) { timer in
-                                        if pubHeartBeat {
-                                            isLoading = false
-                                            timer.invalidate()
-                                        } else {
-                                            if let error {
-                                                if error == InvalidHostID.rawValue {
-                                                    isPairing = true
+                                checkVPNConnection() { result, vpn_error in
+                                    if result {
+                                        if FileManager.default.fileExists(atPath: URL.documentsDirectory.appendingPathComponent("pairingFile.plist").path) {
+                                            Timer.scheduledTimer(withTimeInterval: 0.5, repeats: true) { timer in
+                                                if pubHeartBeat {
+                                                    isLoading = false
+                                                    timer.invalidate()
                                                 } else {
-                                                    startHeartbeatInBackground()
+                                                    if let error {
+                                                        if error == InvalidHostID.rawValue {
+                                                            isPairing = true
+                                                        } else {
+                                                            startHeartbeatInBackground()
+                                                        }
+                                                        self.error = nil
+                                                    }
                                                 }
-                                                self.error = nil
                                             }
+                                            
+                                            startHeartbeatInBackground()
+                                        } else {
+                                            isLoading = false
+                                        }
+                                    } else if let vpn_error {
+                                        showAlert(title: "Error", message: "EM Proxy failed to connect: \(vpn_error)", showOk: true) { _ in
+                                            // Add code here to exit the application
+                                            exit(0)  // This will terminate the app immediately
+                                            
+                                            // Or for a more graceful exit in iOS:
+                                            // UIApplication.shared.perform(#selector(NSXPCConnection.suspend))
+                                            // DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) {
+                                            //     exit(0)
+                                            // }
                                         }
                                     }
-                                    
-                                    startHeartbeatInBackground()
-                                } else {
-                                    isLoading = false
                                 }
                             } else if let error {
                                 showAlert(title: "Error", message: "EM Proxy Failed to start \(error)", showOk: true) { cool in
@@ -137,6 +153,63 @@ struct HeartbeatApp: App {
             }
         }
     }
+    
+    private func checkVPNConnection(callback: @escaping (Bool, String?) -> Void) {
+        let host = NWEndpoint.Host("10.7.0.1")
+        let port = NWEndpoint.Port(rawValue: 62078)!
+        
+        let connection = NWConnection(host: host, port: port, using: .tcp)
+        
+        // Create a variable to hold the timeout work item
+        var timeoutWorkItem: DispatchWorkItem?
+        
+        timeoutWorkItem = DispatchWorkItem { [weak connection] in
+            if connection?.state != .ready {
+                connection?.cancel()
+                DispatchQueue.main.async {
+                    // Only call back if we haven't already
+                    if timeoutWorkItem?.isCancelled == false {
+                        callback(false, "[TIMEOUT] Wireguard is not connected. Try closing this app, turn Wireguard off and back on.")
+                    }
+                }
+            }
+        }
+        
+        connection.stateUpdateHandler = { [weak connection] state in
+            switch state {
+            case .ready:
+                // Connection succeeded - cancel the timeout
+                timeoutWorkItem?.cancel()
+                connection?.cancel()
+                DispatchQueue.main.async {
+                    callback(true, nil)
+                }
+            case .failed(let error):
+                // Connection failed - cancel the timeout
+                timeoutWorkItem?.cancel()
+                connection?.cancel()
+                DispatchQueue.main.async {
+                    if error == NWError.posix(.ETIMEDOUT) {
+                        callback(false, "Wireguard is not connected. Try closing the app, turn it off and back on.")
+                    } else if error == NWError.posix(.ECONNREFUSED) {
+                        callback(false, "Wifi is not connected. StikJIT won't work on cellular data.")
+                    } else {
+                        callback(false, "em proxy check error: \(error.localizedDescription)")
+                    }
+                }
+            default:
+                break
+            }
+        }
+        
+        // Start the connection
+        connection.start(queue: .global())
+        
+        // Schedule the timeout
+        if let workItem = timeoutWorkItem {
+            DispatchQueue.global().asyncAfter(deadline: .now() + 10, execute: workItem)
+        }
+    }
 }
 
 var pubHeartBeat = false