From 1debf7ea5b64b119e7eb7dcd8a496fa46fdc2c78 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Thu, 1 May 2025 20:57:40 +0100
Subject: [PATCH] Got a websocket + dispatcher class that works

---
 .../trustgraph/gateway/endpoint/flows.py      | 30 +++++-
 .../trustgraph/gateway/endpoint/stream.py     | 92 ++++++++++++++++---
 2 files changed, 109 insertions(+), 13 deletions(-)

diff --git a/trustgraph-flow/trustgraph/gateway/endpoint/flows.py b/trustgraph-flow/trustgraph/gateway/endpoint/flows.py
index 34396518..b0114219 100644
--- a/trustgraph-flow/trustgraph/gateway/endpoint/flows.py
+++ b/trustgraph-flow/trustgraph/gateway/endpoint/flows.py
@@ -1,4 +1,6 @@
 
+import asyncio
+
 from . endpoint import ServiceEndpoint
 
 from . flow_endpoint import FlowEndpoint
@@ -25,6 +27,32 @@ class FlowEndpointManager:
         self.services = {
         }
 
+        class DispInst:
+            def __init__(self, ws, running):
+                self.ws = ws
+                self.running = running
+                self.num = 1
+            async def destroy(self):
+                print("Destroy..")
+                await self.ws.close()
+                self.running.stop()
+            async def run(self):
+                while self.running.get():
+                    await asyncio.sleep(1)
+                    await self.ws.send_json({"number": self.num})
+                    self.num += 1
+                    print("Tick")
+            async def receive(self, msg):
+                print("Message...")
+                print(msg.data)
+
+        class Dispatcher:
+            def __init__(self):
+                pass
+            async def create(self, ws, running, request):
+                print("Create")
+                return DispInst(ws, running)
+        
         self.endpoints = [
             FlowEndpoint(
                 endpoint_path = "/api/v1/flow/{flow}/{kind}",
@@ -34,7 +62,7 @@ class FlowEndpointManager:
             StreamEndpoint(
                 endpoint_path = "/api/v1/flow/{flow}/stream/{kind}",
                 auth = auth,
-                dispatchers = self.services,
+                dispatcher = Dispatcher(),
             ),
         ]
 
diff --git a/trustgraph-flow/trustgraph/gateway/endpoint/stream.py b/trustgraph-flow/trustgraph/gateway/endpoint/stream.py
index 9cdd56fb..e3d1a97d 100644
--- a/trustgraph-flow/trustgraph/gateway/endpoint/stream.py
+++ b/trustgraph-flow/trustgraph/gateway/endpoint/stream.py
@@ -4,42 +4,110 @@ from aiohttp import web, WSMsgType
 import logging
 
 from .. running import Running
-from . socket import SocketEndpoint
 
 logger = logging.getLogger("socket")
 logger.setLevel(logging.INFO)
 
-class StreamEndpoint(SocketEndpoint):
+class StreamEndpoint:
 
     def __init__(
-            self, endpoint_path, auth, dispatchers,
+            self, endpoint_path, auth, dispatcher,
     ):
 
-        super(StreamEndpoint, self).__init__(endpoint_path, auth)
+        self.path = endpoint_path
+        self.auth = auth
+        self.operation = "socket"
 
-    async def listener(self, ws):
+        self.dispatcher = dispatcher
+
+    async def worker(self, ws, dispatcher, running):
+
+        await dispatcher.run()
+
+    async def listener(self, ws, dispatcher, running):
 
         async for msg in ws:
+
             # On error, finish
             if msg.type == WSMsgType.TEXT:
-                print("Received:", msg)
+                await dispatcher.receive(msg)
                 continue
             elif msg.type == WSMsgType.BINARY:
-                # Ignore incoming message
+                await dispatcher.receive(msg)
                 continue
             else:
                 break
 
-        self.running.stop()
+        running.stop()
+        await ws.close()
         
-    async def dispatcher(self, ws):
+    async def handle(self, request):
 
-        for i in range(0, 10):
+        try:
+            token = request.query['token']
+        except:
+            token = ""
 
-            await asyncio.sleep(1)
+        if not self.auth.permitted(token, self.operation):
+            return web.HTTPUnauthorized()
+        
+        # 50MB max message size
+        ws = web.WebSocketResponse(max_msg_size=52428800)
 
-            await ws.send_json({"number": i})
+        await ws.prepare(request)
+
+        try:
+
+            async with asyncio.TaskGroup() as tg:
+
+                running = Running()
+
+                print("Create...")
+                print(self.dispatcher)
+                dispinst = await self.dispatcher.create(ws, running, request)
+
+                print("Create worker...")
+                worker_task = tg.create_task(
+                    self.worker(ws, dispinst, running)
+                )
+
+                print("Create listener")
+                lsnr_task = tg.create_task(
+                    self.listener(ws, dispinst, running)
+                )
+
+                print("Created taskgroup, waiting...")
+
+                # Wait for threads to complete
+
+            print("Task group closed")
+
+            # Finally?
+            await dispinst.destroy()
+
+        except ExceptionGroup as e:
+
+            print("Exception group:", flush=True)
+
+            for se in e.exceptions:
+                print("  Type:", type(se), flush=True)
+                print(f"  Exception: {se}", flush=True)
+        except Exception as e:
+            print("Socket exception:", e, flush=True)
 
         await ws.close()
+
+        return ws
+
+    async def start(self):
+        pass
+
+    async def stop(self):
         self.running.stop()
 
+    def add_routes(self, app):
+
+        app.add_routes([
+            web.get(self.path, self.handle),
+        ])
+