improve websocket connection

app.py

@@ -1,17 +1,37 @@
-from fastapi import FastAPI, Response, WebSocket, WebSocketDisconnect, Depends, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
+import asyncio
+import logging
+from fastapi.concurrency import asynccontextmanager
 import uvicorn
-from dotenv import load_dotenv
 import os
+from dotenv import load_dotenv
+from fastapi import FastAPI, Response, WebSocket, WebSocketDisconnect, status, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+
+from models.connection_manager import ConnectionManager
+from models.request_payload import RequestPayload
+from utils.package_manager import PackageManager
 
 # Load environment variables from .env file
 load_dotenv()
 IS_DEV = os.environ.get('ENV', 'DEV') != 'PROD'
-SECURE_TOKEN = os.getenv("SECURE_TOKEN")
+WEBSOCKET_SECURE_TOKEN = os.getenv("SECURE_TOKEN")
+WHITELIST_CHANNEL_IDS = os.getenv('WHITELIST_CHANNEL_IDS')
 X_REQUEST_USER = os.environ.get('X_REQUEST_USER')
 X_API_KEY = os.environ.get('X_API_KEY')
 
+WHITELIST_CHANNEL_IDS = WHITELIST_CHANNEL_IDS.split(',') if WHITELIST_CHANNEL_IDS is not None else []
+
 app = FastAPI()
+# Initialize the connection manager
+manager = ConnectionManager()
+package = PackageManager()
+
+logging.basicConfig(
+    level=logging.WARNING,
+    format='%(asctime)s %(name)s %(levelname)-8s  %(message)s',
+    datefmt='(%H:%M:%S)'
+)
 
 # CORS Middleware: restrict access to only trusted origins
 app.add_middleware(
@@ -23,31 +43,68 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-
-# Simple token-based authentication dependency
-async def authenticate_token(token: str):
-    if token != f"Bearer {SECURE_TOKEN}":
-        raise HTTPException(status_code=401, detail="Invalid token")
-
 @app.get("/")
 def root():
-    return Response(status=200, data='ok')
+    return Response(status_code=status.HTTP_200_OK, data='ok')
         
 @app.get("/health")
 def healthcheck():
-    return Response(status=200, data='ok')
+    return Response(status_code=status.HTTP_200_OK, data='ok')
+
+@app.post("/hi_mlhub")
+async def hi_mlhub(payload: RequestPayload):
+    if manager.available is not None:
+        request_id, compressed_data = package.gzip(payload)
+        
+        # Send binary data to all connected WebSocket clients
+        await manager.send_bytes(manager.available, compressed_data)
         
+        try:
+            # Wait for the response with a timeout (e.g., 10 seconds)
+            data = await asyncio.wait_for(manager.listen(manager.available, request_id), timeout=10.0)
+            return JSONResponse(status_code=status.HTTP_200_OK, content=data)
+        except Exception:
+            return JSONResponse(status_code=status.HTTP_504_GATEWAY_TIMEOUT, content={ "error": "Timeout" })
+    else:
+        return JSONResponse(status_code=status.HTTP_502_BAD_GATEWAY, content={ "error": "MLaaS is not available." })
+
+# Simple token-based authentication dependency
+def is_valid_token(token: str):
+    return token == WEBSOCKET_SECURE_TOKEN
+
+def is_valid_apikey(channel_id: str):
+    return channel_id is not None and channel_id in WHITELIST_CHANNEL_IDS
+
 # WebSocket endpoint
 @app.websocket("/ws")
-async def websocket_endpoint(websocket: WebSocket, token: str = Depends(authenticate_token)):
-    await websocket.accept()
+async def websocket_endpoint(websocket: WebSocket):
+    headers = websocket.headers
+    token = headers.get("x-token")
+    channel_id = headers.get("x-channel-id")
+
+    if not is_valid_token(token):
+        return HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token")
+    if not is_valid_apikey(channel_id):
+        return HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="No permission")
+
+    await manager.connect(channel_id, websocket)
+    
     try:
         while True:
-            data = await websocket.receive_text()
-            print(f"Message from client: {data}")
-            await websocket.send_text(f"Server received: {data}")
+            # Common receiver
+            data = await manager.receive_text(channel_id)
+            print(f"Message from MLaaS: {data}")
+
+            # Notify the manager that a message was received
+            await manager.notify(channel_id, data)
+
+            # Broadcast the message to all clients
+            #await manager.broadcast(f"Client {channel_id} says: {data}")
     except WebSocketDisconnect:
-        print("Client disconnected")
+        manager.disconnect(channel_id)
+        await manager.broadcast(f"A client has disconnected with ID: {channel_id}")
+
+    return None
 
 def is_valid(u, p):
     return u == X_REQUEST_USER and p == X_API_KEY

models/connection_manager.py

@@ -0,0 +1,96 @@
+import asyncio
+import json
+import logging
+from fastapi import HTTPException, WebSocket, status
+from typing import Dict
+
+class InRequest:
+    def __init__(self):
+        self.responses: Dict[str, asyncio.Future] = {}
+
+class ConnectionManager:
+    
+    def __init__(self):
+        self.available = None
+        self.active_connections: Dict[str, WebSocket] = {}  # Maps socket ID to WebSocket connection
+        self.in_request: Dict[str, InRequest] = {}  # Store pending response futures
+
+    async def connect(self, socket_id: str, websocket: WebSocket):
+        await websocket.accept()
+        self.active_connections[socket_id] = websocket
+        if self.available is None:
+            self.available = socket_id
+        return socket_id
+
+    def disconnect(self, socket_id: str):
+        if socket_id in self.active_connections:
+            del self.active_connections[socket_id]
+            if self.available == socket_id:
+                self.available = None
+
+    async def broadcast(self, message: str):
+        for connection in self.active_connections.values():
+            await connection.send_text(message)
+
+    async def receive_text(self, socket_id: str):
+        websocket = self.active_connections.get(socket_id)
+        if websocket:
+            return await websocket.receive_text()
+        else:
+            raise HTTPException(
+                status_code=status.HTTP_502_BAD_GATEWAY, 
+                detail=f"Socket ID {socket_id} not connected")
+
+    async def send_text(self, socket_id: str, message: str):
+        websocket = self.active_connections.get(socket_id)
+        if websocket:
+            await websocket.send_text(message)
+        else:
+            raise HTTPException(
+                status_code=status.HTTP_502_BAD_GATEWAY, 
+                detail="WebSocket connection not found.")
+
+    async def send_bytes(self, socket_id: str, binary_data: bytes):
+        websocket = self.active_connections.get(socket_id)
+        if websocket:
+            await websocket.send_bytes(binary_data)  # Send binary data
+        else:
+            raise HTTPException(
+                status_code=status.HTTP_502_BAD_GATEWAY, 
+                detail=f"Socket ID {socket_id} not connected")
+
+    async def listen(self, socket_id:str, request_id:str) -> str:
+        req = InRequest()
+        # Create a Future for waiting for the response
+        future = asyncio.get_event_loop().create_future()
+        req.responses[request_id] = future
+        self.in_request[socket_id] = req
+        try:
+            return await future  # Await the future until it's set with a response
+        except asyncio.CancelledError:
+            raise HTTPException(
+                status_code=status.HTTP_502_BAD_GATEWAY, 
+                detail=f"Socket ID {socket_id} not connected or canceled")
+
+    async def notify(self, socket_id: str, message: str):
+        logging.debug(message)
+        # If there is a pending future for this socket, set the result
+        if socket_id in self.in_request:
+            request_id, payload = self.extract_message(message)
+            if request_id is not None:
+                self.in_request[socket_id].responses[request_id].set_result(payload)
+                self.in_request.pop(socket_id, None)
+
+    def extract_message(self, message:str):
+        request_id = None
+        payload = None
+        logging.debug(message)
+
+        try:
+            o = json.loads(message)
+            if o is not None:
+                request_id, payload = o.get('request_id'), o.get('payload')
+        except Exception as e:
+            logging.warning(f"extract_message error: {str(e)}")
+
+        return request_id, payload

models/gateway_payload.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel
+
+class GwPayload(BaseModel):
+    request_id: str
+    payload: str

models/request_payload.py

@@ -0,0 +1,4 @@
+from pydantic import BaseModel
+
+class RequestPayload(BaseModel):
+    request: object

requirements.txt

@@ -1,4 +1,5 @@
 python-dotenv==1.0.1
 fastapi==0.115.2
 uvicorn==0.31.1
-websockets==13.1
+websockets==13.1
+pydantic==2.9.2

utils/package_manager.py

@@ -0,0 +1,19 @@
+import gzip
+import json
+import uuid
+
+from models.request_payload import RequestPayload
+
+class PackageManager:
+    def gzip(self, payload: RequestPayload):
+        request_id = str(uuid.uuid4())
+    
+        json_object = {
+            "request_id": request_id,
+            "payload": payload.request  # Use the incoming data
+        }
+
+        # Convert JSON object to binary then zip it
+        binary_data = json.dumps(json_object).encode('utf-8')
+        compressed_data = gzip.compress(binary_data)
+        return request_id, compressed_data