@@ -542,7 +542,7 @@ async def chat(request: ChatRequest): |
542 | 542 | # For now, just return dummy text |
543 | 543 | return {"response": "This is a dummy response from the Feast feature server."} |
544 | 544 |
|
545 | | - @app.post("/read-document") |
| 545 | + @app.post("/read-document", dependencies=[Depends(inject_user_details)]) |
546 | 546 | async def read_document_endpoint(request: ReadDocumentRequest): |
547 | 547 | try: |
548 | 548 | import os |
@@ -557,7 +557,7 @@ async def read_document_endpoint(request: ReadDocumentRequest): |
557 | 557 | except Exception as e: |
558 | 558 | return {"error": str(e)} |
559 | 559 |
|
560 | | - @app.post("/save-document") |
| 560 | + @app.post("/save-document", dependencies=[Depends(inject_user_details)]) |
561 | 561 | async def save_document_endpoint(request: SaveDocumentRequest): |
562 | 562 | try: |
563 | 563 | import json |
@@ -668,12 +668,41 @@ async def send_message(self, message: str, websocket: WebSocket): |
668 | 668 |
|
669 | 669 | manager = ConnectionManager() |
670 | 670 |
|
| 671 | + MAX_WS_CONNECTIONS = 5 |
| 672 | + MAX_MESSAGE_SIZE = 4096 |
| 673 | + MAX_MESSAGES_PER_MINUTE = 60 |
| 674 | + WS_READ_TIMEOUT_SEC = 60 |
| 675 | + |
671 | 676 | @app.websocket("/ws/chat") |
672 | 677 | async def websocket_endpoint(websocket: WebSocket): |
| 678 | + if len(manager.active_connections) >= MAX_WS_CONNECTIONS: |
| 679 | + await websocket.close(code=status.WS_1008_POLICY_VIOLATION) |
| 680 | + return |
| 681 | + |
673 | 682 | await manager.connect(websocket) |
| 683 | + message_timestamps: List[float] = [] |
674 | 684 | try: |
675 | 685 | while True: |
676 | | - message = await websocket.receive_text() |
| 686 | + try: |
| 687 | + message = await asyncio.wait_for( |
| 688 | + websocket.receive_text(), timeout=WS_READ_TIMEOUT_SEC |
| 689 | + ) |
| 690 | + except asyncio.TimeoutError: |
| 691 | + await websocket.close(code=status.WS_1001_GOING_AWAY) |
| 692 | + return |
| 693 | + |
| 694 | + if len(message) > MAX_MESSAGE_SIZE: |
| 695 | + await websocket.close(code=status.WS_1009_MESSAGE_TOO_BIG) |
| 696 | + return |
| 697 | + |
| 698 | + now = time.time() |
| 699 | + cutoff = now - 60 |
| 700 | + message_timestamps = [ts for ts in message_timestamps if ts >= cutoff] |
| 701 | + if len(message_timestamps) >= MAX_MESSAGES_PER_MINUTE: |
| 702 | + await websocket.close(code=status.WS_1008_POLICY_VIOLATION) |
| 703 | + return |
| 704 | + message_timestamps.append(now) |
| 705 | + |
677 | 706 | # Process the received message (currently unused but kept for future implementation) |
678 | 707 | # For now, just return dummy text |
679 | 708 | response = f"You sent: '{message}'. This is a dummy response from the Feast feature server." |
|