From 1022b07e0e2cd6f34d62626d8ad361a07acd89db Mon Sep 17 00:00:00 2001
From: TheSpookyCat <24483290+TheSpookyCat@users.noreply.github.com>
Date: Thu, 25 May 2023 17:39:17 +0100
Subject: [PATCH] - Introduce `connection_timeout` and `response_timeout`
 kwargs - - Loosely based on #195 - Calls `get_ipc_path()` each time a
 handshake is begun - - (Windows) DiscordNotFound will no longer be raised on
 object initialization, instead raises on handshake - `InvalidID` is no longer
 raised when the pipe is closed for any reason. `PipeClosed` is now raised
 instead

---
 pypresence/baseclient.py | 46 +++++++++++++++++++++++-----------------
 pypresence/exceptions.py | 29 +++++++++++++++++++------
 2 files changed, 49 insertions(+), 26 deletions(-)

diff --git a/pypresence/baseclient.py b/pypresence/baseclient.py
index ac3fae7..78e7d11 100644
--- a/pypresence/baseclient.py
+++ b/pypresence/baseclient.py
@@ -16,16 +16,14 @@
 class BaseClient:
 
     def __init__(self, client_id: str, **kwargs):
-        pipe = kwargs.get('pipe', None)
         loop = kwargs.get('loop', None)
         handler = kwargs.get('handler', None)
+        self.pipe = kwargs.get('pipe', None)
         self.isasync = kwargs.get('isasync', False)
+        self.connection_timeout = kwargs.get('connection_timeout', 30)
+        self.response_timeout = kwargs.get('response_timeout', 10)
 
         client_id = str(client_id)
-        self.ipc_path = get_ipc_path(pipe)
-
-        if not self.ipc_path:
-            raise DiscordNotFound
 
         if loop is not None:
             self.update_event_loop(loop)
@@ -78,11 +76,13 @@ async def _async_err_handle(self, loop, context: dict):
 
     async def read_output(self):
         try:
-            preamble = await self.sock_reader.read(8)
+            preamble = await asyncio.wait_for(self.sock_reader.read(8), self.response_timeout)
             status_code, length = struct.unpack('<II', preamble[:8])
-            data = await self.sock_reader.read(length)
-        except BrokenPipeError:
-            raise InvalidID
+            data = await asyncio.wait_for(self.sock_reader.read(length), self.response_timeout)
+        except (BrokenPipeError, struct.error):
+            raise PipeClosed
+        except asyncio.TimeoutError:
+            raise ResponseTimeout
         payload = json.loads(data.decode('utf-8'))
         if payload["evt"] == "ERROR":
             raise ServerError(payload["data"]["message"])
@@ -103,21 +103,29 @@ def send_data(self, op: int, payload: Union[dict, Payload]):
             payload.encode('utf-8'))
 
     async def handshake(self):
-        if sys.platform == 'linux' or sys.platform == 'darwin':
-            self.sock_reader, self.sock_writer = await asyncio.open_unix_connection(self.ipc_path)
-        elif sys.platform == 'win32' or sys.platform == 'win64':
-            self.sock_reader = asyncio.StreamReader(loop=self.loop)
-            reader_protocol = asyncio.StreamReaderProtocol(
-                self.sock_reader, loop=self.loop)
-            try:
-                self.sock_writer, _ = await self.loop.create_pipe_connection(lambda: reader_protocol, self.ipc_path)
-            except FileNotFoundError:
-                raise InvalidPipe
+        ipc_path = get_ipc_path(self.pipe)
+        if not ipc_path:
+            raise DiscordNotFound
+
+        try:
+            if sys.platform == 'linux' or sys.platform == 'darwin':
+                self.sock_reader, self.sock_writer = await asyncio.wait_for(asyncio.open_unix_connection(ipc_path), self.connection_timeout)
+            elif sys.platform == 'win32' or sys.platform == 'win64':
+                self.sock_reader = asyncio.StreamReader(loop=self.loop)
+                reader_protocol = asyncio.StreamReaderProtocol(self.sock_reader, loop=self.loop)
+                self.sock_writer, _ = await asyncio.wait_for(self.loop.create_pipe_connection(lambda: reader_protocol, ipc_path), self.connection_timeout)
+        except FileNotFoundError:
+            raise InvalidPipe
+        except asyncio.TimeoutError:
+            raise ConnectionTimeout
+
         self.send_data(0, {'v': 1, 'client_id': self.client_id})
         preamble = await self.sock_reader.read(8)
         code, length = struct.unpack('<ii', preamble)
         data = json.loads(await self.sock_reader.read(length))
         if 'code' in data:
+            if data['message'] == 'Invalid Client ID':
+                raise InvalidID
             raise DiscordError(data['code'], data['message'])
         if self._events_on:
             self.sock_reader.feed_data = self.on_event
diff --git a/pypresence/exceptions.py b/pypresence/exceptions.py
index 281b061..5fb8d97 100644
--- a/pypresence/exceptions.py
+++ b/pypresence/exceptions.py
@@ -10,11 +10,6 @@ def __init__(self):
         super().__init__('Could not find Discord installed and running on this machine.')
 
 
-class InvalidID(PyPresenceException):
-    def __init__(self):
-        super().__init__('Client ID is Invalid')
-
-
 class InvalidPipe(PyPresenceException):
     def __init__(self):
         super().__init__('Pipe Not Found - Is Discord Running?')
@@ -34,10 +29,15 @@ def __init__(self, message: str):
 
 
 class DiscordError(PyPresenceException):
-    def __init__(self, code: int, message: str):
+    def __init__(self, code: int, message: str, override=False):
         self.code = code
         self.message = message
-        super().__init__('Error Code: {0} Message: {1}'.format(code, message))
+        super().__init__('Error Code: {0} Message: {1}'.format(code, message) if not override else message)
+
+
+class InvalidID(DiscordError):
+    def __init__(self):
+        super().__init__(4000, 'Client ID is Invalid')
 
 
 class ArgumentError(PyPresenceException):
@@ -48,3 +48,18 @@ def __init__(self):
 class EventNotFound(PyPresenceException):
     def __init__(self, event):
         super().__init__('No event with name {0} exists.'.format(event))
+
+
+class PipeClosed(PyPresenceException):
+    def __init__(self):
+        super().__init__('The pipe was closed. Catch this exception and re-connect your instance.')
+
+
+class ResponseTimeout(PyPresenceException):
+    def __init__(self):
+        super().__init__('No response was received from the pipe in time')
+
+
+class ConnectionTimeout(PyPresenceException):
+    def __init__(self):
+        super().__init__('Unable to create a connection to the pipe in time')