fix(chat): 修复多个聊天错误 消息格式错误导致断连、新消息会通知到发件人、加强消息发送的成员验证

core/routers/chat/handlers.py

@@ -14,6 +14,7 @@ from core.db.models import ReadStatus, User, Chat, Message
 
 from .managers import SocketManager
 from . import schema
+from .services import get_chat
 
 logger = logging.getLogger(__name__)
 socket_manager = SocketManager()
@@ -28,7 +29,8 @@ async def send_new_message(
         incoming_message: schema.WSMessage,
 ):
     chat_guid: str = str(incoming_message.chat_guid)
-    chat: Chat | None = await CRUD(db_session, Chat).get(filters={'guid': chat_guid})
+    # chat: Chat | None = await CRUD(db_session, Chat).get(filters={'guid': chat_guid, 'is_deleted': False}
+    chat: Chat | None = await get_chat(chat_guid, current_user, db_session)
 
     # 判断是否为新聊天
     if not chats or chat_guid not in chats:
@@ -56,8 +58,6 @@ async def send_new_message(
 
         await db_session.flush()
         await db_session.commit()
-        # await db_session.refresh(message, attribute_names=["user", "chat"])
-        # await db_session.refresh(chat, attribute_names=["users"])
 
     except Exception as exc_info:
         await db_session.rollback()

core/routers/chat/managers.py

@@ -11,6 +11,7 @@ import redis.asyncio as aioredis
 import ujson
 from fastapi import WebSocket
 from fastapi.websockets import WebSocketState
+from simple_spider_tool import jsonpath
 
 from core.db.engine import get_cache
 
@@ -104,7 +105,9 @@ class SocketManager:
                             if socket.client_state != WebSocketState.CONNECTED:
                                 continue
                             data = message['data']
-                            await socket.send_text(data)
+                            if f',"user_id":{socket.user_id}' in data:
+                                continue
+                            await socket.send_json(data)
                 else:
                     if self._task_cancel is False:
                         break

core/routers/chat/routes.py

@@ -239,6 +239,7 @@ async def get_chat_websocket(
     await socket_manager.connect_socket(websocket=websocket)
     logger.info(f'Websocket connection is established: {login_user.user_id}')
     user_id = str(login_user.user_id)
+    websocket.user_id = login_user.user_id
     await socket_manager.add_user_socket_connection(user_id, websocket)
 
     ratelimit = WebSocketRateLimiter(times=50, seconds=10, callback=websocket_callback)
@@ -258,13 +259,13 @@ async def get_chat_websocket(
                 try:
                     message = schema.WSMessage(**incoming_message)
                 except ValueError:
-                    await socket_manager.send_error(websocket, 'The message type or format you sent is invalid')
+                    await socket_manager.send_error('The message type or format you sent is invalid', websocket)
                     continue
 
                 handler = socket_manager.handlers.get(message.type)
                 if not handler:
                     logger.error(f'No handler [{message.type}] exists')
-                    await socket_manager.send_error(websocket, f'Type: {message.type} was not found')
+                    await socket_manager.send_error(f'Type: {message.type} was not found', websocket)
                     continue
 
                 await handler(
@@ -277,14 +278,14 @@ async def get_chat_websocket(
 
             except (JSONDecodeError, AttributeError) as exc:
                 logger.exception(f'Websocket error, detail: {exc}')
-                await websocket.send_json({'status': 'error', 'message': 'Wrong message format'})
+                await socket_manager.send_error('Wrong message format', websocket)
 
             except ValueError as exc:
                 logger.exception(f'Websocket error, detail: {exc}')
-                await websocket.send_json({'status': 'error', 'message': 'Could not validate incoming message'})
+                await socket_manager.send_error('Could not validate incoming message', websocket)
 
             except WebsocketTooManyRequests:
                 logger.exception(f'User: {login_user} sent too many ws requests')
-                await websocket.send_json({'status': 'error', 'message': 'You have sent too many requests'})
+                await socket_manager.send_error('You have sent too many requests', websocket)
     except WebSocketDisconnect:
         logging.info(f'Websocket is disconnected: {login_user.user_id}')

core/routers/chat/services.py

@@ -11,12 +11,10 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import aliased, selectinload
 
 from core.base.exceptions import WebsocketTooManyRequests
-from core.db.models import User, Chat, ReadStatus, Message
+from core.db.models import User, Chat, ReadStatus, Message, chat_participant
 from core.routers.auth.schema import LoginUser
 from core.utils.helper import format_date
 
-from .schema import ChatUnreadOut
-
 
 async def websocket_callback(ws, pexpire):
     raise WebsocketTooManyRequests("Too many requests")
@@ -91,3 +89,22 @@ async def get_unread_messages_per_chat(
         ))
 
     return results
+
+
+async def get_chat(chat_guid, login_user: User, db_session: AsyncSession) -> Chat | None:
+    subquery = (
+        select(chat_participant.c.chat_id)
+        .where(chat_participant.c.user_id == login_user.user_id)
+    ).scalar_subquery()
+
+    query = (
+        select(Chat)
+        .where(and_(
+            Chat.id.in_(subquery),
+            Chat.guid == chat_guid,
+            Chat.is_deleted == False
+        ))
+    )
+    result = await db_session.execute(query)
+    await db_session.commit()
+    return result.scalar_one_or_none()