| FazBrowse GitHub Viewer | Trending | | Home |
| Tools: [Download Repo ZIP] [Original HTTPS Page] |
1 parent 8b9bb50 commit 2081b55
1 file changed
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -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." | |
| Back | FazBrowse Home | New Git URL |
0 commit comments