tcp_forwarder: fix error handling in data causing a crash

pymobiledevice3/tcp_forwarder.py

@@ -55,22 +55,29 @@ def start(self, address='0.0.0.0'):
             # as synchronous blocking
             readable, writable, exceptional = select.select(self.inputs, [], self.inputs, self.TIMEOUT)
             if self.stopped.is_set():
+                self.logger.debug("Closing since stopped is set")
                 break
 
             closed_sockets = set()
             for current_sock in readable:
+                self.logger.debug("Processing %r", current_sock)
                 if current_sock is self.server_socket:
                     self._handle_server_connection()
                 else:
                     if current_sock not in closed_sockets:
                         try:
                             self._handle_data(current_sock, closed_sockets)
                         except ConnectionResetError:
+                            self.logger.exception("Error when handling data")
                             self._handle_close_or_error(current_sock)
+                    else:
+                        self.logger.debug("Is closed")
 
             for current_sock in exceptional:
+                self.logger.error("Sock failed: %r", current_sock)
                 self._handle_close_or_error(current_sock)
 
+        self.logger.info("Closing everything")
         # on stop, close all currently opened sockets
         for current_sock in self.inputs:
             current_sock.close()
@@ -87,22 +94,39 @@ def _handle_close_or_error(self, from_sock):
         self.logger.info(f'connection {other_sock} was closed')
 
     def _handle_data(self, from_sock, closed_sockets):
-        data = from_sock.recv(1024)
-
-        if len(data) == 0:
-            # no data means socket was closed
+        self.logger.debug("Handling data from %s", from_sock)
+        data = None
+        try:
+            data = from_sock.recv(1024)
+        except OSError:
+            # Socket closing is handled in another if block
+            pass
+
+        if data is None or len(data) == 0:
+            if data is None:
+                # data is none means we had an error reading from socket
+                self.logger.debug("oserror when reading from_sock")
+            else:
+                # Empty data means socket was closed
+                self.logger.info("No data was read from the socket")
             self._handle_close_or_error(from_sock)
             closed_sockets.add(from_sock)
             closed_sockets.add(self.connections[from_sock])
             return
 
         # when data is received from one end, just forward it to the other
         other_sock = self.connections[from_sock]
-
-        # send the data in blocking manner
-        other_sock.setblocking(True)
-        other_sock.sendall(data)
-        other_sock.setblocking(False)
+        try:
+            # send the data in blocking manner
+            other_sock.setblocking(True)
+            other_sock.sendall(data)
+            other_sock.setblocking(False)
+        except OSError:
+            # Tried writing to closed socket
+            self.logger.exception("Exception when sending data to socket")
+            self._handle_close_or_error(other_sock)
+            closed_sockets.add(from_sock)
+            closed_sockets.add(self.connections[from_sock])
 
     @abstractmethod
     def _establish_remote_connection(self) -> socket.socket:
@@ -164,6 +188,7 @@ def _establish_remote_connection(self) -> socket.socket:
         # connect directly using usbmuxd
         mux_device = usbmux.select_device(self.serial, connection_type=self.usbmux_connection_type,
                                           usbmux_address=self.usbmux_address)
+        self.logger.debug("Selected device: %r", mux_device)
         if mux_device is None:
             raise ConnectionFailedError()
         return mux_device.connect(self.dst_port, usbmux_address=self.usbmux_address)