webui: Save connected WebSocket clients and allow broadcast

2 files changed, 38 insertions, 12 deletions

diff --git a/fg21sim/webui/app.py b/fg21sim/webui/app.py
index 3793fda..199e2d6 100644
--- a/fg21sim/webui/app.py
+++ b/fg21sim/webui/app.py
@@ -22,14 +22,32 @@ from ..configs import ConfigManager
 
 
 class Application(tornado.web.Application):
-    configmanager = ConfigManager()
+    """
+    Application of the "fg21sim" Web UI.
+
+    Attributes
+    ----------
+    configmanager : `~fg21sim.configs.ConfigManager`
+        A ``ConfigManager`` instance, which saves the current configurations
+        status.  The configuration operations (e.g., "set", "get", "load")
+        are performed on this instance, which is also passed to the
+        foregrounds simulation programs.
+    ws_clients : set
+        Current connected clients through WebSocket.
+        When a new WebSocket connection established, it is added to this
+        list, which is also removed from this list when the connection lost.
+    """
 
     def __init__(self, **kwargs):
+        self.configmanager = ConfigManager()
+        self.ws_clients = set()
+        # URL handlers
         handlers = [
             url(r"/", IndexHandler, name="index"),
             url(r"/login", LoginHandler, name="login"),
             url(r"/ws", FG21simWSHandler),
         ]
+        # Application settings
         settings = {
             # The static files will be served from the default "/static/" URI.
             # Recommend to use `{{ static_url(filepath) }}` in the templates.
diff --git a/fg21sim/webui/handlers/websocket.py b/fg21sim/webui/handlers/websocket.py
index 6031109..941dd39 100644
--- a/fg21sim/webui/handlers/websocket.py
+++ b/fg21sim/webui/handlers/websocket.py
@@ -59,18 +59,15 @@ class FG21simWSHandler(tornado.websocket.WebSocketHandler):
     from_localhost = None
 
     def check_origin(self, origin):
-        """Check the origin of the WebSocket access.
+        """
+        Check the origin of the WebSocket connection to determine whether
+        the access is allowed.
 
         Attributes
         ----------
         from_localhost : bool
-            Set to ``True`` if the access is from the localhost,
-            otherwise ``False``.
-
-        NOTE
-        ----
-        Currently, only allow access from the ``localhost``
-        (i.e., 127.0.0.1) and local LAN.
+            Set to ``True`` if the access is from the "localhost" (i.e.,
+            127.0.0.1), otherwise ``False``.
         """
         self.from_localhost = False
         logger.info("WebSocket: {0}: origin: {1}".format(self.name, origin))
@@ -96,9 +93,11 @@ class FG21simWSHandler(tornado.websocket.WebSocketHandler):
 
     def open(self):
         """Invoked when a new WebSocket is opened by the client."""
+        # Add to the set of current connected clients
+        self.application.ws_clients.add(self)
+        logger.warning("Added new WebSocket client: {0}".format(self))
         # FIXME:
         # * better to move to the `Application` class ??
-        # * or create a ``ConfigsHandler`` similar to the ``ConsoleHandler``
         self.configs = self.application.configmanager
         self.console_handler = ConsoleHandler(websocket=self)
         self.configs_handler = ConfigsHandler(configs=self.configs)
@@ -108,12 +107,15 @@ class FG21simWSHandler(tornado.websocket.WebSocketHandler):
 
     def on_close(self):
         """Invoked when a new WebSocket is closed by the client."""
+        # Remove from the set of current connected clients
+        self.application.ws_clients.remove(self)
+        logger.warning("Removed WebSocket client: {0}".format(self))
         code, reason = None, None
         if hasattr(self, "close_code"):
             code = self.close_code
         if hasattr(self, "close_reason"):
             reason = self.close_reason
-        logger.info("WebSocket: {0}: closed by client: {1}, {2}".format(
+        logger.warning("WebSocket: {0}: closed by client: {1}, {2}".format(
             self.name, code, reason))
 
     # FIXME/XXX:
@@ -123,7 +125,8 @@ class FG21simWSHandler(tornado.websocket.WebSocketHandler):
     # [1] https://stackoverflow.com/a/35543856/4856091
     # [2] https://stackoverflow.com/a/33724486/4856091
     def on_message(self, message):
-        """Handle incoming messages and dispatch task according to the
+        """
+        Handle incoming messages and dispatch task according to the
         message type.
 
         NOTE
@@ -184,6 +187,11 @@ class FG21simWSHandler(tornado.websocket.WebSocketHandler):
         msg_response = json.dumps(response)
         self.write_message(msg_response)
 
+    def broadcast(self, message):
+        """Broadcast/push the given message to all connected clients."""
+        for ws in self.application.ws_clients:
+            ws.write_message(message)
+
     def _handle_results(self, msg):
         # Got a message of supported types
         msg_type = msg["type"]