refactored CMClient disconnect handling

steam/core/cm.py

@@ -56,6 +56,8 @@ class CMClient(EventEmitter):
         else:
             raise ValueError("Only TCP is supported")
 
+        self.connection.event_connected.rawlink(self._handle_disconnect)
+
         self.on(EMsg.ChannelEncryptRequest, self._handle_encrypt_request),
         self.on(EMsg.Multi, self._handle_multi),
         self.on(EMsg.ClientLogOnResponse, self._handle_logon),
@@ -69,23 +71,25 @@ class CMClient(EventEmitter):
         logger.debug("Connect initiated.")
 
         while True:
-            with gevent.Timeout(15, False):
-                server_addr = random.choice(server_list)
-                self.connection.connect(server_addr)
+            server_addr = random.choice(server_list)
 
-            if not self.connection.event_connected.is_set():
-                logger.debug("Failed to connect. Retrying...")
-                continue
+            if self.connection.connect(server_addr):
+                break
 
-            break
+            logger.debug("Failed to connect. Retrying...")
 
         self.connected = True
         self.emit("connected")
         self._recv_loop = gevent.spawn(self._recv_messages)
 
+    def _handle_disconnect(self, event):
+        if not event.is_set():
+            gevent.spawn(self.disconnect)
+
     def disconnect(self, reconnect=False):
-        if reconnect:
-            gevent.spawn(self.connect)
+        if not self.connected:
+            return
+        self.connected = False
 
         self.connection.disconnect()
 
@@ -96,6 +100,9 @@ class CMClient(EventEmitter):
         self._init_attributes()
         self.emit('disconnected')
 
+        if reconnect:
+            gevent.spawn(self.connect)
+
     def _init_attributes(self):
         self.connected = False
 
@@ -133,14 +140,9 @@ class CMClient(EventEmitter):
         self.connection.put_message(data)
 
     def _recv_messages(self):
-        while True:
-            try:
-                message = self.connection.get_message(timeout=1)
-            except queue.Empty:
-                if not self.connection.event_connected.is_set():
-                    gevent.spawn(self.disconnect)
-                    return
-                continue
+        for message in self.connection:
+            if not self.connected:
+                break
 
             if self.key:
                 if self.hmac_secret:

steam/core/connection.py

@@ -32,7 +32,11 @@ class Connection:
         self._new_socket()
 
         logger.debug("Attempting connection to %s", str(server_addr))
-        self._connect(server_addr)
+
+        try:
+            self._connect(server_addr)
+        except socket.error:
+            return False
 
         self.server_addr = server_addr
 
@@ -41,6 +45,7 @@ class Connection:
 
         logger.debug("Connected.")
         self.event_connected.set()
+        return True
 
     def disconnect(self):
         if not self.event_connected.is_set():
@@ -62,8 +67,15 @@ class Connection:
         self.socket.close()
 
         logger.debug("Disconnected.")
+
+        # hack
+        # rawlink will only call the callback on set
+        self.event_connected.set()
         self.event_connected.clear()
 
+    def __iter__(self):
+        return self.recv_queue
+
     def get_message(self, block=True, timeout=None):
         return self.recv_queue.get(block, timeout)