From 49360b08fe2f8e916237b5ad18c92cccf13dee2f Mon Sep 17 00:00:00 2001
From: Diego Rey Mendez <dreymendez@duckduckgo.com>
Date: Mon, 9 Dec 2024 20:05:13 -0300
Subject: [PATCH] Improves handling of VPN uninstall errors

---
 .../Sources/UDSHelper/UDSClient.swift         | 32 +++++++++++++++++--
 1 file changed, 30 insertions(+), 2 deletions(-)

diff --git a/LocalPackages/UDSHelper/Sources/UDSHelper/UDSClient.swift b/LocalPackages/UDSHelper/Sources/UDSHelper/UDSClient.swift
index 58dfde0afe..ffdf5f2dda 100644
--- a/LocalPackages/UDSHelper/Sources/UDSHelper/UDSClient.swift
+++ b/LocalPackages/UDSHelper/Sources/UDSHelper/UDSClient.swift
@@ -24,9 +24,33 @@ public actor UDSClient {
 
     public typealias PayloadHandler = (Data) async throws -> Void
 
-    enum ConnectionError: Error {
+    enum ConnectionError: CustomNSError {
         case cancelled
         case failure(_ error: Error)
+
+        /// This means there's either no peer to connect to, or the connection timed out.
+        ///
+        /// For our local UDS connection either is a clear problem.
+        ///
+        case timeout(_ error: Error)
+
+        var errorCode: Int {
+            switch self {
+            case .cancelled: return 1
+            case .failure: return 2
+            case .timeout: return 3
+            }
+        }
+
+        var errorUserInfo: [String: Any] {
+            switch self {
+            case .cancelled:
+                return [:]
+            case .failure(let error),
+                    .timeout(let error):
+                return [NSUnderlyingErrorKey: error as NSError]
+            }
+        }
     }
 
     private var internalConnection: NWConnection?
@@ -71,7 +95,9 @@ public actor UDSClient {
     ///
     private func connect() async throws -> NWConnection {
         let endpoint = NWEndpoint.unix(path: socketFileURL.path)
-        let parameters = NWParameters.tcp
+        let options = NWProtocolTCP.Options()
+        options.connectionTimeout = 5
+        let parameters = NWParameters(tls: nil, tcp: options)
         let connection = NWConnection(to: endpoint, using: parameters)
 
         connection.stateUpdateHandler = { state in
@@ -91,6 +117,8 @@ public actor UDSClient {
                 throw ConnectionError.cancelled
             case .failed(let error):
                 throw ConnectionError.failure(error)
+            case .waiting(let error):
+                throw ConnectionError.timeout(error)
             default:
                 try await Task.sleep(nanoseconds: 200 * MSEC_PER_SEC)
             }