From ae8326f591ab70d1efd083a9dc442d6d8356dacb Mon Sep 17 00:00:00 2001
From: Rossen Georgiev <rossen@rgp.io>
Date: Wed, 30 Mar 2016 07:10:39 +0100
Subject: [PATCH] reworked CMClient conn drop handling; fix #21

---
 steam/client/__init__.py | 8 +++++++-
 steam/core/cm.py         | 8 ++------
 steam/core/connection.py | 6 +-----
 3 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/steam/client/__init__.py b/steam/client/__init__.py
index 6d42887..475423d 100644
--- a/steam/client/__init__.py
+++ b/steam/client/__init__.py
@@ -35,6 +35,7 @@ class SteamClient(EventEmitter, FeatureBase):
         # register listners
         self.cm.on(None, self._handle_cm_events)
         self.cm.on("disconnected", self._handle_disconnect)
+        self.cm.on("reconnect", self._handle_disconnect)
         self.on(EMsg.ClientLogOnResponse, self._handle_logon)
         self.on(EMsg.ClientUpdateMachineAuth, self._handle_update_machine_auth)
 
@@ -86,6 +87,7 @@ class SteamClient(EventEmitter, FeatureBase):
         """
         Close connection
         """
+        self.logged_on = False
         self.cm.disconnect()
 
     def _handle_cm_events(self, event, *args):
@@ -102,7 +104,11 @@ class SteamClient(EventEmitter, FeatureBase):
             if jobid not in (-1, 18446744073709551615):
                 self.emit("job_%d" % jobid, *args)
 
-    def _handle_disconnect(self):
+    def _handle_disconnect(self, *args):
+        # only disconnect, reconnect has 1 arg
+        if self.logged_on and len(args) == 0:
+            gevent.spawn(self.connect)
+
         self.username = None
         self.logged_on = False
         self.current_jobid = 0
diff --git a/steam/core/cm.py b/steam/core/cm.py
index 0541c48..603c488 100644
--- a/steam/core/cm.py
+++ b/steam/core/cm.py
@@ -60,8 +60,6 @@ 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),
@@ -95,10 +93,6 @@ class CMClient(EventEmitter):
         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, nodelay=False):
         """Close connection
 
@@ -208,6 +202,8 @@ class CMClient(EventEmitter):
             gevent.spawn(self._parse_message, message)
             gevent.idle()
 
+        gevent.spawn(self.disconnect)
+
     def _parse_message(self, message):
         emsg_id, = struct.unpack_from("<I", message)
         emsg = EMsg(clear_proto_bit(emsg_id))
diff --git a/steam/core/connection.py b/steam/core/connection.py
index 8dfcfef..ba4ae17 100644
--- a/steam/core/connection.py
+++ b/steam/core/connection.py
@@ -51,6 +51,7 @@ class Connection(object):
     def disconnect(self):
         if not self.event_connected.is_set():
             return
+        self.event_connected.clear()
 
         self.server_addr = None
 
@@ -70,11 +71,6 @@ class Connection(object):
 
         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