From c49a198bab4f875fb278e2fe4615fe974fceab7c Mon Sep 17 00:00:00 2001
From: tan01 <jonnytan@google.com>
Date: Thu, 27 Jun 2024 11:41:09 -0700
Subject: [PATCH] Remove sockJS and use websocket for the front-end.

This attempts to remove both tornado-sockjs from the HTF servers and sockJS from the front-end getting served.

SockJS is a websocket emulator used as a compatibility shim in environments that don't support the websocket API. However, the TornadoSockJS library is no longer being maintained.

PiperOrigin-RevId: 647392280
---
 openhtf/output/servers/dashboard_server.py |  6 +++---
 openhtf/output/servers/pub_sub.py          | 18 ++++++++++++------
 openhtf/output/servers/station_server.py   | 17 +++++++++--------
 3 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/openhtf/output/servers/dashboard_server.py b/openhtf/output/servers/dashboard_server.py
index a1276dacf..cd131bb9f 100644
--- a/openhtf/output/servers/dashboard_server.py
+++ b/openhtf/output/servers/dashboard_server.py
@@ -32,7 +32,6 @@
 from openhtf.output.web_gui import web_launcher
 from openhtf.util import data
 from openhtf.util import multicast
-import sockjs.tornado
 import tornado.web
 
 _LOG = logging.getLogger(__name__)
@@ -128,10 +127,11 @@ class DashboardServer(web_gui_server.WebGuiServer):
   """Serves a list of known stations and an Angular frontend."""
 
   def __init__(self, port):
-    dash_router = sockjs.tornado.SockJSRouter(DashboardPubSub, '/sub/dashboard')
-    routes = dash_router.urls + [
+    routes = [
         ('/station_list', StationListHandler),
+        ('/sub/dashboard', DashboardPubSub),
     ]
+    _LOG.info('Initializing dashboard server.')
     super(DashboardServer, self).__init__(routes, port)
 
   def _get_config(self):
diff --git a/openhtf/output/servers/pub_sub.py b/openhtf/output/servers/pub_sub.py
index dc2acfdcd..7e27c7e78 100644
--- a/openhtf/output/servers/pub_sub.py
+++ b/openhtf/output/servers/pub_sub.py
@@ -16,12 +16,12 @@
 import logging
 
 from openhtf import util as htf_util
-import sockjs.tornado
+import tornado.websocket
 
 _LOG = logging.getLogger(__name__)
 
 
-class PubSub(sockjs.tornado.SockJSConnection):
+class PubSub(tornado.websocket.WebSocketHandler):
   """Generic pub/sub based on SockJS connections."""
 
   @htf_util.classproperty
@@ -53,18 +53,20 @@ def publish(cls, message, client_filter=None):
         if (not client_filter) or client_filter(client):
           client.send(message)
 
-  def on_open(self, info):
-    _LOG.debug('New subscriber from %s.', info.ip)
+  def open(self, *args, **kwargs):
     with self._lock:  # pylint: disable=not-context-manager
       self.subscribers.add(self)
-    self.on_subscribe(info)
+    self.on_subscribe(None)
 
   def on_close(self):
-    _LOG.debug('A client unsubscribed.')
+    _LOG.info('A client unsubscribed.')
     with self._lock:  # pylint: disable=not-context-manager
       self.subscribers.remove(self)
     self.on_unsubscribe()
 
+  def send(self, message):
+    self.write_message(message)
+
   def on_subscribe(self, info):
     """Called when new clients subscribe. Subclasses can override."""
     pass
@@ -72,3 +74,7 @@ def on_subscribe(self, info):
   def on_unsubscribe(self):
     """Called when clients unsubscribe. Subclasses can override."""
     pass
+
+  def check_origin(self, origin: str) -> bool:
+    """Allows all cross-origin traffic. Dangerous."""
+    return True
diff --git a/openhtf/output/servers/station_server.py b/openhtf/output/servers/station_server.py
index 756f0bf39..cf0766a48 100644
--- a/openhtf/output/servers/station_server.py
+++ b/openhtf/output/servers/station_server.py
@@ -39,7 +39,7 @@
 from openhtf.util import functions
 from openhtf.util import multicast
 from openhtf.util import timeouts
-import sockjs.tornado
+import tornado
 
 CONF = configuration.CONF
 
@@ -225,7 +225,7 @@ def _to_dict_with_event(cls, test_state):
     return test_state_dict, event
 
 
-class DashboardPubSub(sockjs.tornado.SockJSConnection):
+class DashboardPubSub(tornado.websocket.WebSocketHandler):
   """WebSocket endpoint for the list of available stations.
 
   In this case, there is always exactly one available station: the station
@@ -244,9 +244,9 @@ def for_port(cls, port):
     """Returns a new subclass with the port set."""
     return type(cls.__name__, (cls,), {'port': port})
 
-  def on_open(self, unused_info):
+  def open(self, unused_info):
     """Called by the base class when a client connects."""
-    self.send(self._make_message())
+    self.write_message(self._make_message())
 
   @classmethod
   def _make_message(cls):
@@ -595,11 +595,12 @@ def __init__(
     station_watcher = StationWatcher(StationPubSub.publish_update)
     station_watcher.start()
 
-    # Set up the SockJS endpoints.
+    # Set up the endpoints.
     dashboard_class = DashboardPubSub.for_port(port)
-    dash_router = sockjs.tornado.SockJSRouter(dashboard_class, '/sub/dashboard')
-    station_router = sockjs.tornado.SockJSRouter(StationPubSub, '/sub/station')
-    routes = dash_router.urls + station_router.urls
+    routes = [
+        ('/sub/dashboard', dashboard_class),
+        ('/sub/station', StationPubSub),
+    ]
 
     # Set up the other endpoints.
     routes.extend((