socket auth

Author: mariusandra

backend/app/fastapi.py

@@ -12,9 +12,9 @@
 from alembic.config import Config as AlembicConfig
 from alembic import command as alembic_command
 
-from app.services.ws_broadcast import register_ws_routes, redis_listener
+from app.websockets import register_ws_routes, redis_listener
 from app.config import get_config
-from backend.app.utils.sentry import initialize_sentry
+from app.utils.sentry import initialize_sentry
 
 def lifespan(app: FastAPI):
     initialize_sentry()

backend/app/models/frame.py

@@ -11,7 +11,7 @@
 from app.models.apps import get_app_configs
 from app.models.settings import get_settings_dict
 from app.utils.token import secure_token
-from app.services.ws_broadcast import publish_message
+from app.websockets import publish_message
 
 
 # NB! Update frontend/src/types.tsx if you change this

backend/app/models/log.py

@@ -8,7 +8,7 @@
 from app.database import Base
 from sqlalchemy import Integer, String, DateTime, ForeignKey, Text, func
 from sqlalchemy.orm import relationship, backref, Session, mapped_column
-from app.services.ws_broadcast import publish_message
+from app.websockets import publish_message
 
 class Log(Base):
     __tablename__ = 'log'

backend/app/models/metrics.py

@@ -3,7 +3,7 @@
 from sqlalchemy import Integer, String, ForeignKey, DateTime, func
 from app.database import Base
 from sqlalchemy.orm import relationship, backref, Session, mapped_column
-from app.services.ws_broadcast import publish_message
+from app.websockets import publish_message
 
 class Metrics(Base):
     __tablename__ = 'metrics'

backend/app/services/ws_broadcast.py renamed to backend/app/websockets.py

@@ -1,11 +1,15 @@
 import asyncio
 import json
 import uuid
+from jose import jwt, JWTError
 from typing import List
 from redis.asyncio import from_url as create_redis
-from fastapi import WebSocket, WebSocketDisconnect
+from fastapi import WebSocket, WebSocketDisconnect, Depends
+from sqlalchemy.orm import Session
+from app.database import get_db
 
 from app.config import get_config
+from app.models.user import User
 
 redis_pub = None
 redis_sub = None
@@ -82,17 +86,35 @@ async def publish_message(event: str, data: dict):
 
 def register_ws_routes(app):
     @app.websocket("/ws")
-    async def websocket_endpoint(websocket: WebSocket):
+    async def websocket_endpoint(websocket: WebSocket, db: Session = Depends(get_db)):
+        token = websocket.query_params.get('token')
+        if not token:
+            await websocket.close(code=1008, reason="Missing token")
+            return
+
+        try:
+            from app.api.auth import ALGORITHM, SECRET_KEY
+            payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+            user_email = payload.get("sub")
+            if not user_email:
+                raise ValueError("Invalid token")
+        except JWTError:
+            await websocket.close(code=1008, reason="Invalid token")
+            return
+
+        user = db.query(User).filter(User.email == user_email).first()
+        if user is None:
+            await websocket.close(code=1008, reason="User not found")
+            return
+
         await manager.connect(websocket)
         try:
             while True:
                 data = await websocket.receive_text()
-                # Handle incoming messages
-                await manager.send_personal_message(f"You said: {data}", websocket)
+                # Optionally handle incoming messages
+                await manager.send_personal_message("You said: " + data, websocket)
         except WebSocketDisconnect:
             await manager.disconnect(websocket)
-        except Exception as e:
-            print(f"Error: {e}")
-            await manager.disconnect(websocket)
+
 
     return manager

frontend/src/scenes/socketLogic.tsx

@@ -14,7 +14,7 @@ export const socketLogic = kea<socketLogicType>([
     newMetrics: (metrics: Record<string, any>) => ({ metrics }),
   }),
   afterMount(({ actions, cache }) => {
-    cache.ws = new WebSocket('/ws')
+    cache.ws = new WebSocket('/ws?token=' + localStorage.getItem('token'))
     cache.ws.onopen = function (event: any) {
       console.log('🔵 Connected to the WebSocket server.')
     }

backend/gpt.sh renamed to gpt.sh

@@ -7,11 +7,13 @@ OUTPUT="gpt.txt"
 
 # Define the patterns you want to collect files from:
 patterns=(
-    "app/*.py"
-    "app/schemas/*.py"
-    "app/models/*.py"
-    "app/api/*.py"
-    "app/services/*.py"
+    "backend/app/*.py"
+    "backend/app/schemas/*.py"
+    "backend/app/models/*.py"
+    "backend/app/api/*.py"
+    "backend/app/services/*.py"
+    "backend/app/services/ws_broadcast.py"
+    "frontend/src/scenes/socketLogic.tsx"
 )
 
 for pattern in "${patterns[@]}"; do