Merge pull request #63 from ImDevinC/fix-socket-read

fix(sockets): improve socket connection reliability and callback cleanup

Author: ImDevinC

actions/DiscordCore.py

@@ -47,11 +47,13 @@ def register_backend_callback(self, key: str, callback: callable):
         """Register a callback and track it for cleanup."""
         self.backend.register_callback(key, callback)
         self._registered_callbacks.append((key, callback))
+        self.plugin_base.add_callback(key, callback)
 
     def cleanup_callbacks(self):
         """Unregister all tracked callbacks to prevent memory leaks."""
         for key, callback in self._registered_callbacks:
             self.backend.unregister_callback(key, callback)
+            self.plugin_base.remove_callback(key, callback)
         self._registered_callbacks.clear()
 
     def __del__(self):

discordrpc/asyncdiscord.py

@@ -5,7 +5,7 @@
 import requests
 from loguru import logger as log
 
-from .sockets import UnixPipe
+from .sockets import UnixPipe, SOCKET_BAD_BUFFER_SIZE, SOCKET_DISCONNECTED
 from .commands import *
 from .exceptions import *
 from .constants import MAX_SOCKET_RETRY_ATTEMPTS
@@ -39,15 +39,27 @@ def is_connected(self):
     def connect(self, callback: callable):
         tries = 0
         while tries < MAX_SOCKET_RETRY_ATTEMPTS:
-            log.debug(
-                f"Attempting to connect to socket, attempt {tries + 1}/{MAX_SOCKET_RETRY_ATTEMPTS}"
-            )
-            self.rpc.connect()
-            self.rpc.send({"v": 1, "client_id": self.client_id}, OP_HANDSHAKE)
-            _, resp = self.rpc.receive()
-            if resp:
+            try:
+                log.debug(
+                    f"Attempting to connect to socket, attempt {tries + 1}/{MAX_SOCKET_RETRY_ATTEMPTS}"
+                )
+                self.rpc.connect()
                 break
-            tries += 1
+            except Exception as ex:
+                log.error(f"failed to connect to socket. {ex}")
+                tries += 1
+
+        self.rpc.send({"v": "1", "client_id": self.client_id}, OP_HANDSHAKE)
+        code, resp = self.rpc.receive()
+
+        if not resp:
+            log.error("no response from discord client")
+            raise RPCException
+
+        if code == SOCKET_BAD_BUFFER_SIZE:
+            log.error("bad buffer size when receiving data from socket")
+            raise RPCException
+
         try:
             data = json.loads(resp)
         except Exception as ex:
@@ -74,7 +86,9 @@ def poll_callback(self, callback: callable):
                 log.error(f"error receiving data from socket. {ex}")
                 self.disconnect()
                 return
-            if val[0] == -1:
+            if val[0] == SOCKET_BAD_BUFFER_SIZE:
+                log.debug("bad buffer size when receiving data from socket")
+            if val[0] == SOCKET_DISCONNECTED:
                 self.disconnect()
             callback(val[0], val[1])
 

discordrpc/sockets.py

@@ -11,16 +11,19 @@
 from .constants import MAX_IPC_SOCKET_RANGE, SOCKET_SELECT_TIMEOUT, SOCKET_BUFFER_SIZE
 
 SOCKET_DISCONNECTED: int = -1
-
+SOCKET_BAD_BUFFER_SIZE: int = -2
+SOCKET_SEND_TIMEOUT: int = 5
+SOCKET_CONNECT_TIMEOUT: int = 2
+SOCKET_RECEIVE_TIMEOUT: int = 5
 
 class UnixPipe:
     def __init__(self):
         self.socket: socket.socket = None
 
     def connect(self):
         if self.socket is None:
-            self.socket = socket.socket(socket.AF_UNIX)
-        self.socket.setblocking(False)
+            self.socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        self.socket.settimeout(SOCKET_CONNECT_TIMEOUT)
         base_path = path = (
             os.environ.get("XDG_RUNTIME_DIR")
             or os.environ.get("TMPDIR")
@@ -44,6 +47,7 @@ def connect(self):
                 pass
         else:
             raise DiscordNotOpened
+        self.socket.setblocking(False)
 
     def disconnect(self):
         if self.socket is None:
@@ -60,17 +64,15 @@ def disconnect(self):
         self.socket = None  # Reset so connect() creates a fresh socket
 
     def send(self, payload, op):
-        payload = json.dumps(payload).encode("UTF-8")
-        payload = struct.pack("<ii", op, len(payload)) + payload
-        size = 0
-        while size == 0 or size < len(payload):
-            res = self.socket.send(payload[size:])
-            size += res
+        log.debug(f"Sending payload: {payload} with op: {op}")
+        payload_bytes = json.dumps(payload).encode("UTF-8")
+        header = struct.pack("<ii", op, len(payload_bytes))
+        message = header + payload_bytes
+        self.socket.settimeout(SOCKET_SEND_TIMEOUT)
+        self.socket.sendall(message)
 
     def receive(self) -> (int, str):
-        ready = select.select([self.socket], [], [], SOCKET_SELECT_TIMEOUT)
-        if not ready[0]:
-            return 0, {}
+        self.socket.settimeout(SOCKET_RECEIVE_TIMEOUT)
         data = self.socket.recv(SOCKET_BUFFER_SIZE)
         if len(data) == 0:
             return SOCKET_DISCONNECTED, {}
@@ -80,7 +82,7 @@ def receive(self) -> (int, str):
         all_data = data[8:]
         buffer_size = length - len(all_data)
         if buffer_size < 0:
-            return 0, {}
+            return SOCKET_BAD_BUFFER_SIZE, {}
         data = self.socket.recv(length - len(all_data))
         all_data += data
         return code, all_data.decode("UTF-8")

main.py

@@ -190,3 +190,12 @@ def on_auth_callback(self, success: bool, message: str = None):
 
     def get_settings_area(self):
         return self._settings_manager.get_settings_area()
+
+    def clear_callbacks(self, key: str, callback: callable):
+        callbacks = self.callbacks.get(key, [])
+        if callback in callbacks:
+            callbacks.remove(callback)
+            if callbacks:
+                self.callbacks[key] = callbacks
+            else:
+                del self.callbacks[key]

manifest.json

@@ -1,5 +1,5 @@
 {
-  "version": "1.9.0",
+  "version": "1.9.1",
   "thumbnail": "store/thumbnail.png",
   "id": "com_imdevinc_StreamControllerDiscordPlugin",
   "name": "Discord - Debug",