From 5b0f078f7a1f7f93fb09a1a824394db1abf3a219 Mon Sep 17 00:00:00 2001
From: Joscha <brummer@migejolise.de>
Date: Fri, 8 Sep 2017 08:00:50 +0000
Subject: [PATCH] Add timeout to Controller's connect()

---
 yaboli/connection.py |  8 ++++----
 yaboli/controller.py | 21 ++++++++++++++-------
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/yaboli/connection.py b/yaboli/connection.py
index 315af0a..6fc2cb7 100644
--- a/yaboli/connection.py
+++ b/yaboli/connection.py
@@ -77,6 +77,7 @@ class Connection:
 			except:
 				pass # errors are not useful here
 			
+			self._pingtask.cancel()
 			await self._pingtask # should stop now that the ws is closed
 			self._ws = None
 	
@@ -91,14 +92,13 @@ class Connection:
 				wait_for_reply = await self._ws.ping()
 				await asyncio.wait_for(wait_for_reply, self.ping_timeout)
 				logger.debug("Pinged!")
+				await asyncio.sleep(self.ping_delay)
 			except asyncio.TimeoutError:
 				logger.warning("Ping timed out.")
 				await self._ws.close()
 				break
-			except (websockets.ConnectionClosed, ConnectionResetError):
-				break
-			else:
-				await asyncio.sleep(self.ping_delay)
+			except (websockets.ConnectionClosed, ConnectionResetError, asyncio.CancelledError):
+				return
 	
 	async def stop(self):
 		"""
diff --git a/yaboli/controller.py b/yaboli/controller.py
index 212e371..5c84ff8 100644
--- a/yaboli/controller.py
+++ b/yaboli/controller.py
@@ -27,11 +27,12 @@ class Controller:
 	
 	"""
 	
-	def __init__(self, nick, human=False, cookie=None):
+	def __init__(self, nick, human=False, cookie=None, connect_timeout=10):
 		"""
-		roomname  - name of room to connect to
-		human     - whether the human flag should be set on connections
-		cookie    - cookie to use in HTTP request, if any
+		roomname        - name of room to connect to
+		human           - whether the human flag should be set on connections
+		cookie          - cookie to use in HTTP request, if any
+		connect_timeout - time for authentication to complete
 		"""
 		self.nick = nick
 		self.human = human
@@ -41,6 +42,7 @@ class Controller:
 		self.password = None
 		
 		self.room = None
+		self.connect_timeout = connect_timeout # in seconds
 		self._connect_result = None
 	
 	def _create_room(self, roomname):
@@ -70,6 +72,7 @@ class Controller:
 		  "no password"    = password needed to connect to room
 		  "wrong password" = password given does not work
 		  "disconnected"   = connection closed before client could access the room
+		  "timeout"        = timed out while waiting for server
 		  "success"        = no failure
 		"""
 		
@@ -97,9 +100,13 @@ class Controller:
 		
 		# connection succeeded, now we need to know whether we can log in
 		# wait for success/authentication/disconnect
-		# TODO: add a timeout
-		await self._connect_result
-		result = self._connect_result.result()
+		try:
+			await asyncio.wait_for(self._connect_result, self.connect_timeout)
+		except asyncio.TimeoutError:
+			result = "timeout"
+		else:
+			result = self._connect_result.result()
+		
 		logger.debug(f"&{roomname}._connect_result: {result!r}")
 		
 		# deal with result