From cae52a71bca51256f91d22a58034e91f98905428 Mon Sep 17 00:00:00 2001
From: Benjamin Braatz <bb@bbraatz.eu>
Date: Tue, 25 Jul 2023 11:31:05 +0200
Subject: [PATCH] Improve connection handling, especially in case of
 conflicting names.

---
 controlpi_plugins/wsserver.py | 41 +++++++++++++++++++++--------------
 setup.py                      |  2 +-
 2 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/controlpi_plugins/wsserver.py b/controlpi_plugins/wsserver.py
index aefbdbf..2eaf215 100644
--- a/controlpi_plugins/wsserver.py
+++ b/controlpi_plugins/wsserver.py
@@ -39,6 +39,7 @@ class Connection:
                            [MessageTemplate()],
                            [([MessageTemplate()],
                              self._receive)])
+        self._registered = True
 
     async def _receive(self, message: Message) -> None:
         """Receive messages from bus and relay to websocket."""
@@ -61,41 +62,49 @@ class Connection:
                         message['command'] == 'configure websocket' and
                         'target' in message and
                         message['target'] == ''):
-                    await self._bus.send(Message(self._name,
-                                                 {'event':
-                                                  'connection closed'}))
+                    self._registered = False
                     self._bus.unregister(self._name)
+                    new_name = self._name
                     if 'name' in message:
-                        self._name = message['name']
+                        new_name = message['name']
                     sends = []
                     sends.append(MessageTemplate({'event':
                                                   {'const':
                                                    'connection opened'}}))
+                    sends.append(MessageTemplate({'event':
+                                                  {'const':
+                                                   'connection configured'}}))
                     sends.append(MessageTemplate({'event':
                                                   {'const':
                                                    'connection closed'}}))
                     for template in message['up filter']:
                         sends.append(template)
                     try:
-                        self._bus.register(self._name, 'WSServer', sends,
+                        self._bus.register(new_name, 'WSServer', sends,
                                            [(message['down filter'],
                                              self._receive)])
+                        self._registered = True
+                        self._name = new_name
+                        configure_message = Message(self._name)
+                        configure_message['event'] = 'connection configured'
+                        configure_message['address'] = self._address
+                        configure_message['port'] = self._port
+                        if 'mac' in message:
+                            configure_message['mac'] = message['mac']
+                        await self._bus.send(configure_message)
                     except BusException as e:
-                        print(f"Unable to register client '{self._name}'"
+                        print(f"Unable to register client '{new_name}'"
                               f"on bus: {e}")
-                    open_message = Message(self._name)
-                    open_message['event'] = 'connection opened'
-                    open_message['address'] = self._address
-                    open_message['port'] = self._port
-                    if 'mac' in message:
-                        open_message['mac'] = message['mac']
-                    await self._bus.send(open_message)
+                        await self._websocket.close()
                 else:
-                    await self._bus.send(Message(self._name, message))
+                    if self._registered:
+                        await self._bus.send(Message(self._name, message))
         except ConnectionClosed:
             pass
-        await self._bus.send(Message(self._name,
-                                     {'event': 'connection closed'}))
+        if self._registered:
+            await self._bus.send(Message(self._name,
+                                         {'event': 'connection closed'}))
+        self._registered = False
         self._bus.unregister(self._name)
 
 
diff --git a/setup.py b/setup.py
index ae38e84..7c09372 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@ with open("README.md", "r") as readme_file:
 
 setuptools.setup(
     name="controlpi-wsserver",
-    version="0.3.0",
+    version="0.3.1",
     author="Graph-IT GmbH",
     author_email="info@graph-it.com",
     description="ControlPi Plugin for Websocket Servers",
-- 
2.34.1