From 39ce5f4a4c987425795ea57b7d5879c9e073c883 Mon Sep 17 00:00:00 2001
From: doronz88 <doron88@gmail.com>
Date: Wed, 5 Mar 2025 13:55:20 +0200
Subject: [PATCH] tunneld: Fix possible FD leaks during connection establish

---
 .../remote/remote_service_discovery.py        | 43 +++++++++++--------
 pymobiledevice3/remote/remotexpc.py           |  6 ++-
 pymobiledevice3/remote/tunnel_service.py      |  2 +
 3 files changed, 32 insertions(+), 19 deletions(-)

diff --git a/pymobiledevice3/remote/remote_service_discovery.py b/pymobiledevice3/remote/remote_service_discovery.py
index d1ec598ae..4d38f5bfc 100644
--- a/pymobiledevice3/remote/remote_service_discovery.py
+++ b/pymobiledevice3/remote/remote_service_discovery.py
@@ -57,15 +57,18 @@ def set_language(self, language: str) -> None:
 
     async def connect(self) -> None:
         await self.service.connect()
-        self.peer_info = await self.service.receive_response()
-        self.udid = self.peer_info['Properties']['UniqueDeviceID']
-        self.product_type = self.peer_info['Properties']['ProductType']
         try:
-            self.lockdown = create_using_remote(self.start_lockdown_service('com.apple.mobile.lockdown.remote.trusted'))
-        except InvalidServiceError:
-            self.lockdown = create_using_remote(
-                self.start_lockdown_service('com.apple.mobile.lockdown.remote.untrusted'))
-        self.all_values = self.lockdown.all_values
+            self.peer_info = await self.service.receive_response()
+            self.udid = self.peer_info['Properties']['UniqueDeviceID']
+            self.product_type = self.peer_info['Properties']['ProductType']
+            try:
+                self.lockdown = create_using_remote(self.start_lockdown_service('com.apple.mobile.lockdown.remote.trusted'))
+            except InvalidServiceError:
+                self.lockdown = create_using_remote(
+                    self.start_lockdown_service('com.apple.mobile.lockdown.remote.untrusted'))
+            self.all_values = self.lockdown.all_values
+        except Exception:  # noqa: E722
+            await self.close()
 
     def get_value(self, domain: Optional[str] = None, key: Optional[str] = None) -> Any:
         return self.lockdown.get_value(domain, key)
@@ -75,16 +78,20 @@ def start_lockdown_service_without_checkin(self, name: str) -> ServiceConnection
 
     def start_lockdown_service(self, name: str, include_escrow_bag: bool = False) -> ServiceConnection:
         service = self.start_lockdown_service_without_checkin(name)
-        checkin = {'Label': 'pymobiledevice3', 'ProtocolVersion': '2', 'Request': 'RSDCheckin'}
-        if include_escrow_bag:
-            pairing_record = get_local_pairing_record(get_remote_pairing_record_filename(self.udid), get_home_folder())
-            checkin['EscrowBag'] = base64.b64decode(pairing_record['remote_unlock_host_key'])
-        response = service.send_recv_plist(checkin)
-        if response['Request'] != 'RSDCheckin':
-            raise PyMobileDevice3Exception(f'Invalid response for RSDCheckIn: {response}. Expected "RSDCheckIn"')
-        response = service.recv_plist()
-        if response['Request'] != 'StartService':
-            raise PyMobileDevice3Exception(f'Invalid response for RSDCheckIn: {response}. Expected "ServiceService"')
+        try:
+            checkin = {'Label': 'pymobiledevice3', 'ProtocolVersion': '2', 'Request': 'RSDCheckin'}
+            if include_escrow_bag:
+                pairing_record = get_local_pairing_record(get_remote_pairing_record_filename(self.udid), get_home_folder())
+                checkin['EscrowBag'] = base64.b64decode(pairing_record['remote_unlock_host_key'])
+            response = service.send_recv_plist(checkin)
+            if response['Request'] != 'RSDCheckin':
+                raise PyMobileDevice3Exception(f'Invalid response for RSDCheckIn: {response}. Expected "RSDCheckIn"')
+            response = service.recv_plist()
+            if response['Request'] != 'StartService':
+                raise PyMobileDevice3Exception(f'Invalid response for RSDCheckIn: {response}. Expected "ServiceService"')
+        except Exception:  # noqa: E722
+            service.close()
+            raise
         return service
 
     async def aio_start_lockdown_service(
diff --git a/pymobiledevice3/remote/remotexpc.py b/pymobiledevice3/remote/remotexpc.py
index caef34ddb..6fae27a4c 100644
--- a/pymobiledevice3/remote/remotexpc.py
+++ b/pymobiledevice3/remote/remotexpc.py
@@ -53,7 +53,11 @@ async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
 
     async def connect(self) -> None:
         self._reader, self._writer = await asyncio.open_connection(self.address[0], self.address[1])
-        await self._do_handshake()
+        try:
+            await self._do_handshake()
+        except Exception:  # noqa: E722
+            await self.close()
+            raise
 
     async def close(self) -> None:
         if self._writer is None:
diff --git a/pymobiledevice3/remote/tunnel_service.py b/pymobiledevice3/remote/tunnel_service.py
index 89bc533ff..6d4c0be52 100644
--- a/pymobiledevice3/remote/tunnel_service.py
+++ b/pymobiledevice3/remote/tunnel_service.py
@@ -963,6 +963,8 @@ async def create_core_device_tunnel_service_using_rsd(
     service = CoreDeviceTunnelService(rsd)
     try:
         await service.connect(autopair=autopair)
+    except Exception:  # noqa: E722
+        await service.close()
     except RemotePairingCompletedError:
         # The connection must be reestablished upon pairing is completed
         await service.close()