From 0c07df4b731aec3068a79bcc7d567cc315f25cf7 Mon Sep 17 00:00:00 2001
From: Antoine Millet <antoine.millet@smartjog.com>
Date: Wed, 22 Dec 2010 10:16:53 +0100
Subject: [PATCH] Now handle properly disconnections of a client.

---
 ccserver/ccserver.py | 14 +++++++++++++-
 ccserver/handlers.py |  9 +++++++--
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/ccserver/ccserver.py b/ccserver/ccserver.py
index 399b341..a24184f 100644
--- a/ccserver/ccserver.py
+++ b/ccserver/ccserver.py
@@ -47,7 +47,8 @@ class CCServer(object):
         # Create the connection manager:
         self.manager = SimpleSslRpcServer(sock, certfile=certfile,
                                           keyfile=keyfile,
-                                          default_handler=WelcomeHandler(self))
+                                          default_handler=WelcomeHandler(self),
+                                          on_disconnect='on_disconnect')
 
     def iterrole(self, role):
         '''
@@ -71,6 +72,17 @@ class CCServer(object):
         
         self._connected[login] = CCClient(login, role, self, connection)
 
+    def unregister(self, connection):
+        '''
+        Unregister an already connected account.
+
+        :param connection: the connection of the client to unregister
+        '''
+        
+        client = self.search_client_by_connection(connection)
+        if client.login in self._connected:
+            del self._connected[client.login]
+
     def search_client_by_connection(self, connection):
         '''
         Search a connected client by it connection. If no client is found,
diff --git a/ccserver/handlers.py b/ccserver/handlers.py
index c97f291..48c10a2 100644
--- a/ccserver/handlers.py
+++ b/ccserver/handlers.py
@@ -42,14 +42,19 @@ class CCHandler(RpcHandler):
 
         return cmd_list
 
-class HypervisorHandler(CCHandler):
+class OnlineCCHandler(CCHandler):
+
+    def on_disconnect(self, connection):
+        self._server.unregister(connection)
+
+class HypervisorHandler(OnlineCCHandler):
     '''
     Handler binded to 'node' role.
     '''
     role_name = 'hypervisor'
 
 
-class ClientHandler(CCHandler):
+class ClientHandler(OnlineCCHandler):
     '''
     Handler binded to 'cli' role.
     '''
-- 
GitLab