nex/college-bot-v1
1
0
Fork 0
mirror of https://github.com/nexy7574/LCC-bot.git synced 2024-09-19 18:16:34 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix websocket lock

Browse source
This commit is contained in:
Nexus 2023-11-07 13:43:39 +00:00
parent 98c916bb9b
commit f79247928c
Signed by: nex
GPG key ID: 0FA334385D0B689F
1 changed files with 20 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
web/server.py
View file

@ -11,7 +11,8 @@ from pathlib import Path
import discord
import httpx
from config import guilds
from fastapi import FastAPI, Header, HTTPException, Request
from asyncio import Lock
from fastapi import FastAPI, Header, HTTPException, Request, WebSocketException as _WSException
from websockets.exceptions import WebSocketException
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse
from starlette.websockets import WebSocket, WebSocketDisconnect
@ -55,7 +56,7 @@ except ImportError:
bot = None
app.state.last_sender = None
app.state.last_sender_ts = datetime.utcnow()
app.state.ws_connected = False
app.state.ws_connected = Lock()
@app.middleware("http")
@ -290,14 +291,14 @@ async def bridge(req: Request):
@app.websocket("/bridge/recv")
async def bridge_recv(ws: WebSocket, secret: str = Header(None)):
await ws.accept()
if secret != app.state.bot.http.token:
raise HTTPException(status_code=401, detail="Invalid secret.")
if app.state.ws_connected:
raise HTTPException(status_code=409, detail="Already connected.")
raise _WSException(code=1008, reason="Invalid Secret")
if app.state.ws_connected.locked():
raise _WSException(code=1008, reason="Already connected.")
queue: asyncio.Queue = app.state.bot.bridge_queue
await ws.accept()
app.state.ws_connected = True
async with app.state.ws_connected:
while True:
try:
data = queue.get_nowait()
@ -311,4 +312,3 @@ async def bridge_recv(ws: WebSocket, secret: str = Header(None)):
break
finally:
queue.task_done()
app.state.ws_connected = False