Fixed blocking socket connect freezing IDA

idarling/interface/widget.py

@@ -262,7 +262,7 @@ def create_servers_group(servers):
 
             for server in servers:
                 is_connected = (
-                    self._plugin.network.connected
+                    current_server is not None
                     and server["host"] == current_server["host"]
                     and server["port"] == current_server["port"]
                 )
@@ -278,14 +278,12 @@ def server_action_triggered(server_action):
                 Called when a action is clicked. Connects to the new server
                 or disconnects from the old server.
                 """
-                was_connected = (
-                    self._plugin.network.connected
-                    and self._plugin.network.server == server
-                )
+                server = server_action._server
+                was_connected = self._plugin.network.server == server
                 self._plugin.network.stop_server()
                 self._plugin.network.disconnect()
                 if not was_connected:
-                    self._plugin.network.connect(server_action._server)
+                    self._plugin.network.connect(server)
 
             servers_group.triggered.connect(server_action_triggered)
 

idarling/network/client.py

@@ -49,11 +49,6 @@ def __init__(self, plugin, parent=None):
             DownloadFile.Query: self._handle_download_file,
         }
 
-    def disconnect(self, err=None):
-        self._plugin.logger.info("Connection lost")
-        ClientSocket.disconnect(self, err)
-        self._plugin.network.disconnect()
-
     def recv_packet(self, packet):
         if isinstance(packet, Command):
             # Call the corresponding handler
@@ -85,6 +80,26 @@ def send_packet(self, packet):
             packet.tick = self._plugin.core.tick
         return ClientSocket.send_packet(self, packet)
 
+    def disconnect(self, err=None):
+        ret = ClientSocket.disconnect(self, err)
+        self._plugin.network._client = None
+        self._plugin.network._server = None
+
+        # Update the user interface
+        self._plugin.interface.update()
+        self._plugin.interface.clear_invites()
+        return ret
+
+    def _check_socket(self):
+        was_connected = self._connected
+        ret = ClientSocket._check_socket(self)
+        if not was_connected and self._connected:
+            # Update the user interface
+            self._plugin.interface.update()
+            # Subscribe to the events
+            self._plugin.core.join_session()
+        return ret
+
     def _handle_join_session(self, packet):
         # Add the user to the navbar
         self._plugin.interface.painter.paint(

idarling/network/network.py

@@ -10,6 +10,7 @@
 
 # You should have received a copy of the GNU General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
+import errno
 import socket
 import ssl
 
@@ -66,69 +67,50 @@ def _uninstall(self):
     def connect(self, server):
         """Connect to the specified server."""
         # Make sure we're not already connected
-        if self.connected:
-            return False
+        if self._client:
+            return
 
+        self._client = Client(self._plugin)
         self._server = server.copy()  # Make a copy
         host = self._server["host"]
         if host == "0.0.0.0":  # Windows can't connect to 0.0.0.0
             host = "127.0.0.1"
         port = self._server["port"]
         no_ssl = self._server["port"]
 
-        # Do the actual connection process
-        self._client = Client(self._plugin)
-        self._plugin.logger.info("Connecting to %s:%d..." % (host, port))
         # Update the user interface
         self._plugin.interface.update()
+        self._plugin.logger.info("Connecting to %s:%d..." % (host, port))
 
+        # Create a new socket
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
         # Wrap the socket in a SSL tunnel
         if not no_ssl:
             ctx = ssl.create_default_context()
             sock = ctx.wrap_socket(sock, server_hostname=host)
-
-        try:
-            sock.connect((host, port))
-        except socket.error as e:
-            self._plugin.logger.warning("Connection failed")
-            self._plugin.logger.exception(e)
-            self._client = None
-            self._server = None
-
-            # Update the user interface
-            self._plugin.interface.update()
-            return False
-        sock.settimeout(0)  # No timeout
-        sock.setblocking(0)  # No blocking
-        self._client.connect(sock)
+        self._client.wrap_socket(sock)
 
         # Set TCP keep-alive options
         cnt = self._plugin.config["keep"]["cnt"]
         intvl = self._plugin.config["keep"]["intvl"]
         idle = self._plugin.config["keep"]["idle"]
         self._client.set_keep_alive(cnt, intvl, idle)
 
-        self._plugin.logger.info("Connected")
-        # Update the user interface
-        self._plugin.interface.update()
-        # Subscribe to the events
-        self._plugin.core.join_session()
-        return True
+        # Connect the socket
+        sock.settimeout(0)  # No timeout
+        sock.setblocking(0)  # No blocking
+        ret = sock.connect_ex((host, port))
+        if ret != 0 and ret != errno.EINPROGRESS and ret != errno.EWOULDBLOCK:
+            self._client.disconnect()
 
     def disconnect(self):
         """Disconnect from the current server."""
-        # Do the actual disconnection process
-        self._plugin.logger.info("Disconnecting...")
-        if self.connected:
-            self._client.disconnect()
-        self._client = None
-        self._server = None
+        # Make sure we aren't already disconnected
+        if not self._client:
+            return
 
-        # Update the user interface
-        self._plugin.interface.update()
-        self._plugin.interface.clear_invites()
-        return True
+        self._plugin.logger.info("Disconnecting...")
+        self._client.disconnect()
 
     def send_packet(self, packet):
         """Send a packet to the server."""
@@ -139,12 +121,12 @@ def send_packet(self, packet):
     def start_server(self):
         """Start the integrated server."""
         if self._integrated:
-            return False
+            return
 
         self._plugin.logger.info("Starting the integrated server...")
         server = IntegratedServer(self._plugin)
         if not server.start("0.0.0.0"):
-            return False  # Couldn't start the server
+            return  # Couldn't start the server
         self._integrated = server
         integrated_arg = {
             "host": "0.0.0.0",
@@ -153,15 +135,14 @@ def start_server(self):
         }
         # Connect the client to the server
         self.disconnect()
-        return self.connect(integrated_arg)
+        self.connect(integrated_arg)
 
     def stop_server(self):
         """Stop the integrated server."""
         if not self._integrated:
-            return False
+            return
 
         self._plugin.logger.info("Stopping the integrated server...")
         self.disconnect()
         self._integrated.stop()
         self._integrated = None
-        return True

idarling/shared/server.py

@@ -70,8 +70,8 @@ def color(self):
     def ea(self):
         return self._ea
 
-    def connect(self, sock):
-        ClientSocket.connect(self, sock)
+    def wrap_socket(self, sock):
+        ClientSocket.wrap_socket(self, sock)
 
         # Setup command handlers
         self._handlers = {
@@ -351,7 +351,7 @@ def _accept(self, sock):
 
         sock.settimeout(0)  # No timeout
         sock.setblocking(0)  # No blocking
-        client.connect(sock)
+        client.wrap_socket(sock)
         self._clients.append(client)
 
     def reject(self, client):

idarling/shared/sockets.py

@@ -13,6 +13,7 @@
 import collections
 import errno
 import json
+import os
 import socket
 import ssl
 import sys
@@ -67,8 +68,8 @@ def connected(self):
         """Is the underlying socket connected?"""
         return self._connected
 
-    def connect(self, sock):
-        """Sets the underlying socket to utilize."""
+    def wrap_socket(self, sock):
+        """Sets the underlying socket to use."""
         self._read_notifier = QSocketNotifier(
             sock.fileno(), QSocketNotifier.Read, self
         )
@@ -79,17 +80,17 @@ def connect(self, sock):
             sock.fileno(), QSocketNotifier.Write, self
         )
         self._write_notifier.activated.connect(self._notify_write)
-        self._write_notifier.setEnabled(False)
+        self._write_notifier.setEnabled(True)
 
         self._socket = sock
-        self._connected = True
 
     def disconnect(self, err=None):
         """Terminates the current connection."""
         if not self._socket:
             return
+
+        self._logger.debug("Disconnected")
         if err:
-            self._logger.warning("Connection lost")
             self._logger.exception(err)
         self._read_notifier.setEnabled(False)
         self._write_notifier.setEnabled(False)
@@ -135,8 +136,27 @@ def set_keep_alive(self, cnt, intvl, idle):
                 sio_keeplive_vals, (1, idle * 1000, intvl * 1000)
             )
 
+    def _check_socket(self):
+        """Check if the connection has been established yet."""
+        # Ignore if you're already connected
+        if self._connected:
+            return True
+
+        # Check if the connection was successful
+        ret = self._socket.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
+        if ret != 0 and ret != errno.EINPROGRESS and ret != errno.EWOULDBLOCK:
+            self.disconnect(socket.error(ret, os.strerror(ret)))
+            return False
+        else:
+            self._connected = True
+            self._logger.debug("Connected")
+            return True
+
     def _notify_read(self):
         """Callback called when some data is ready to be read on the socket."""
+        if not self._check_socket():
+            return
+
         # Read as much data as is available
         while True:
             try:
@@ -202,6 +222,9 @@ def _notify_read(self):
 
     def _notify_write(self):
         """Callback called when some data is ready to written on the socket."""
+        if not self._check_socket():
+            return
+
         while True:
             if not self._write_buffer:
                 if not self._outgoing: