Merge pull request #54 from jhjaggars/websockets

Adding Websockets

Author: jhjaggars

.gitignore

@@ -111,3 +111,4 @@ venv.bak/
 
 # mypy
 .mypy_cache/
+graph_*.dot

Pipfile

@@ -13,6 +13,7 @@ flake8 = "*"
 [packages]
 python-dateutil = "*"
 prometheus-client = "*"
+aiohttp = "*"
 
 [requires]
 python_version = "3.6"

Pipfile.lock

@@ -42,8 +42,8 @@ def main(args=None):
 
     try:
         config.go()
-    except Exception as e:
-        logger.error("An error occured while running receptor:\n%s" % (str(e),))
+    except Exception:
+        logger.exception("main: an error occured while running receptor")
         sys.exit(1)
 
 

receptor/__main__.py

@@ -5,6 +5,7 @@
 import sys
 
 from . import protocol
+from .ws import WSServer
 
 logger = logging.getLogger(__name__)
 
@@ -38,6 +39,13 @@ def mainloop(receptor, socket_path, loop=asyncio.get_event_loop()):
         config.controller_listen_address, config.controller_listen_port, ssl=config.get_server_ssl_context())
     logger.info("Serving on %s:%s", config.controller_listen_address, config.controller_listen_port)
     loop.create_task(listener)
+
+    ws_server = WSServer(receptor, loop)
+    ws_listener = loop.create_server(ws_server.app().make_handler(),
+        config.node_listen_address, config.node_listen_port + 1, ssl=config.get_server_ssl_context())
+    loop.create_task(ws_listener)
+    logger.info("Serving ws on %s:%s", config.node_listen_address, config.node_listen_port + 1)
+
     control_listener = loop.create_unix_server(
         lambda: protocol.BasicControllerProtocol(receptor, loop),
         path=socket_path

receptor/controller.py

@@ -27,6 +27,9 @@ def __init__(self, msg_id=None, header=None, payload=None):
         self.header = header
         self.payload = payload
 
+    def __repr__(self):
+        return f"FramedMessage(msg_id={self.msg_id}, header={self.header}, payload={self.payload})"
+
     def serialize(self):
         h = json.dumps(self.header).encode("utf-8")
         return b''.join([
@@ -89,7 +92,6 @@ async def handle_frame(self, data):
         await self.consume(rest)
 
     async def consume(self, data):
-        logger.debug("consuming %d bytes; to_read = %d bytes", len(data), self.to_read)
         data, rest = data[:self.to_read], data[self.to_read:]
         self.to_read -= len(data)
         self.bb += data
@@ -117,6 +119,9 @@ async def finish(self):
     async def get(self):
         return await self.q.get()
 
+    def get_nowait(self):
+        return self.q.get_nowait()
+
 
 class Frame:
     """

receptor/messages/envelope.py

@@ -1,10 +1,16 @@
 import asyncio
 import logging
+from urllib.parse import urlparse
 
 from .protocol import BasicProtocol, create_peer
+from .ws import WSClient, WSServer
 
 logger = logging.getLogger(__name__)
 
+def parse_peer(peer):
+    if "://" not in peer:
+        peer = f"receptor://{peer}"
+    return urlparse(peer)
 
 # FIXME: ping_interval is in the config, it shouldn't need to be passed as an arg here
 def mainloop(receptor, ping_interval=None, loop=asyncio.get_event_loop(), skip_run=False):
@@ -15,8 +21,23 @@ def mainloop(receptor, ping_interval=None, loop=asyncio.get_event_loop(), skip_r
             config.node_listen_address, config.node_listen_port, ssl=config.get_server_ssl_context())
         loop.create_task(listener)
         logger.info("Serving on %s:%s", config.node_listen_address, config.node_listen_port)
+
+        ws_server = WSServer(receptor, loop)
+        ws_listener = loop.create_server(ws_server.app().make_handler(),
+            config.node_listen_address, config.node_listen_port + 1, ssl=config.get_server_ssl_context())
+        loop.create_task(ws_listener)
+        logger.info("Serving ws on %s:%s", config.node_listen_address, config.node_listen_port + 1)
+
     for peer in config.node_peers:
-        loop.create_task(create_peer(receptor, loop, *peer.strip().split(":", 1)))
+        parsed = parse_peer(peer)
+        if parsed.scheme == "receptor":
+            loop.create_task(create_peer(receptor, loop, parsed.hostname, parsed.port))
+        elif parsed.scheme in ("ws", "wss"):
+            c = WSClient(receptor, loop)
+            loop.create_task(c.connect(peer))
+        else:
+            logger.warn(f"invalid peer: %s -> %s", peer, parsed)
+
     if ping_interval > 0:
         ping_time = (((int(loop.time()) + 1) // ping_interval) + 1) * ping_interval
         loop.call_at(ping_time, loop.create_task, send_pings_and_reschedule(receptor, loop, ping_time, ping_interval))
@@ -34,6 +55,6 @@ async def send_pings_and_reschedule(receptor, loop, ping_time, ping_interval):
     logger.debug(f'Scheduling mesh ping.')
     for node_id in receptor.router.get_nodes():
         await receptor.router.ping_node(node_id)
-    loop.call_at(ping_time + ping_interval, 
+    loop.call_at(ping_time + ping_interval,
                  loop.create_task, send_pings_and_reschedule(
                      receptor, loop, ping_time + ping_interval, ping_interval))

receptor/node.py

@@ -84,15 +84,7 @@ def handle_handshake(self, data):
         self.loop.create_task(self.receptor.message_handler(self.incoming_buffer))
 
     def send_handshake(self):
-        msg = envelope.CommandMessage(header={
-            "cmd": "HI",
-            "id": self.receptor.node_id,
-            "expire_time": time.time() + 10,
-            "meta": dict(capabilities=self.receptor.work_manager.get_capabilities(),
-                         groups=self.receptor.config.node_groups,
-                         work=self.receptor.work_manager.get_work())
-        })
-        self.transport.write(msg.serialize())
+        self.transport.write(self.receptor._say_hi().serialize())
 
 
 class BasicProtocol(BaseProtocol):

receptor/protocol.py

@@ -84,14 +84,6 @@ def update_connection_manifest(self, connection):
                             last=time.time()))
         self.write_connection_manifest(manifest)
 
-    def update_connections(self, protocol_obj):
-        self.router.register_edge(protocol_obj.id, self.node_id, 1)
-        if protocol_obj.id in self.connections:
-            self.connections[protocol_obj.id].append(protocol_obj)
-        else:
-            self.connections[protocol_obj.id] = [protocol_obj]
-        self.update_connection_manifest(protocol_obj.id)
-
     async def message_handler(self, buf):
         logger.debug("spawning message_handler")
         while True:
@@ -106,10 +98,27 @@ async def message_handler(self, buf):
                 else:
                     await self.handle_message(data)
 
+    def update_connections(self, protocol_obj, id_=None):
+        if id_ is None:
+            id_ = protocol_obj.id
+
+        self.router.register_edge(id_, self.node_id, 1)
+        if id_ in self.connections:
+            self.connections[id_].append(protocol_obj)
+        else:
+            self.connections[id_] = [protocol_obj]
+        self.update_connection_manifest(id_)
+
     def add_connection(self, protocol_obj):
         self.update_connections(protocol_obj)
 
-    def remove_connection(self, protocol_obj):
+    def remove_connection(self, protocol_obj, id_=None, loop=None):
+        if id_ is None:
+            id_ = protocol_obj.id
+
+        if loop is None:
+            loop = protocol_obj.loop
+
         notify_connections = []
         for connection_node in self.connections:
             if protocol_obj in self.connections[connection_node]:
@@ -120,14 +129,24 @@ def remove_connection(self, protocol_obj):
                 self.router.debug_router()
                 self.update_connection_manifest(connection_node)
             notify_connections += self.connections[connection_node]
-        protocol_obj.loop.create_task(self.send_route_advertisement(self.router.get_edges()))
+        loop.create_task(self.send_route_advertisement(self.router.get_edges()))
 
     async def shutdown_handler(self):
         while True:
             if self.stop:
                 return
             await asyncio.sleep(1)
 
+    def _say_hi(self):
+        return envelope.CommandMessage(header={
+            "cmd": "HI",
+            "id": self.node_id,
+            "expire_time": time.time() + 10,
+            "meta": dict(capabilities=self.work_manager.get_capabilities(),
+                         groups=self.config.node_groups,
+                         work=self.work_manager.get_work())
+        })
+
     async def handle_route_advertisement(self, data):
         self.router.add_edges(data["edges"])
         await self.send_route_advertisement(data["edges"], data["seen"])

receptor/receptor.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import uuid
 
@@ -119,10 +120,6 @@ async def test_malformed_frame(framed_buffer, msg_id):
         )
 
 
-@pytest.mark.skip(
-    reason="""
-    This test illustrates that sending an incomplete stream corrupts the transport"""
-)
 @pytest.mark.asyncio
 async def test_too_short(framed_buffer, msg_id):
     f1 = Frame(Frame.Types.HEADER, 1, 100, 1, 1)
@@ -133,4 +130,5 @@ async def test_too_short(framed_buffer, msg_id):
     await framed_buffer.put(f1.serialize() + too_short_header)
     await framed_buffer.put(f2.serialize() + too_short_payload)
 
-    await framed_buffer.get()
+    with pytest.raises(asyncio.QueueEmpty):
+        framed_buffer.get_nowait()