diff --git a/app/main.py b/app/main.py index 71d694d..af7cde5 100644 --- a/app/main.py +++ b/app/main.py @@ -3,7 +3,7 @@ import sys import traceback from contextlib import asynccontextmanager -from fastapi import FastAPI, Request +from fastapi import FastAPI, Request, WebSocket from fastapi.encoders import jsonable_encoder from fastapi.exceptions import HTTPException, RequestValidationError from loguru import logger @@ -12,19 +12,20 @@ from redis.asyncio import Redis from starlette import status from starlette.middleware.cors import CORSMiddleware from starlette.responses import JSONResponse, RedirectResponse +from starlette.websockets import WebSocketDisconnect from app.api.config import config as dconfig from app.api.error_report.report import Report from app.api.models import ApiErrors, ApiStartupStatus, ErrorMessage, StartupState -from app.api.utils import SSE, broadcast_sse_msg, build_sse_event, sse_mgr +from app.api.utils import Event, broadcast_msg from app.api.warmup import ( get_bitcoin_client_warmup_data, get_full_client_warmup_data, get_full_client_warmup_data_bitcoinonly, ) +from app.api.ws_manager import ws_mgr from app.apps.router import register_app_status_update_handlers from app.apps.router import router as app_router -from app.auth.auth_bearer import JWTBearer from app.auth.auth_handler import ( handle_local_cookie, register_cookie_updater, @@ -97,7 +98,7 @@ async def lifespan(app: FastAPI): raise RuntimeError("Redis not initialized correctly, got a Sentinel") register_cookie_updater() - await broadcast_sse_msg(SSE.SYSTEM_STARTUP_INFO, api_startup_status.model_dump()) + await broadcast_msg(Event.SYSTEM_STARTUP_INFO, api_startup_status.model_dump()) btc_task = asyncio.create_task(_initialize_bitcoin()) ln_task = asyncio.create_task(_initialize_lightning()) await register_all_handlers() @@ -191,7 +192,7 @@ async def _set_startup_status( api_startup_status.lightning_msg = lightning_msg asyncio.create_task(warmup_new_connections()) - await broadcast_sse_msg(SSE.SYSTEM_STARTUP_INFO, api_startup_status.model_dump()) + await broadcast_msg(Event.SYSTEM_STARTUP_INFO, api_startup_status.model_dump()) @logger.catch @@ -278,45 +279,32 @@ def index(req: Request): new_connections = [] -async def _send_sse_event(id, event, data): - return await sse_mgr.send_to_single(id, build_sse_event(event, data)) +async def _send_ws_event(id, event, data): + return await ws_mgr.send_to_single(id, event, data) -@app.get( - "/sse/subscribe", - status_code=status.HTTP_200_OK, -) -async def stream(request: Request): - token = request.cookies.get("access_token") - if not token: - # No token in cookies found, try to get it from the Authorization header - token = request.headers.get("authorization") +@app.websocket("/ws") +async def stream(websocket: WebSocket): + conn_id, authed = await ws_mgr.connect(websocket) + if not authed: + return # ws_mgr already closed the socket - if not token: - # Raise an exception as there is not token found - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail="No authorization code.", - ) - - token = token.replace("Bearer ", "") - if not JWTBearer().verify_jwt(jwtoken=token): - # token is invalid - raise HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail="Invalid authorization code.", - ) - - event_source, id = sse_mgr.add_connection(request) - new_connections.append(id) - - await _send_sse_event( - id, SSE.SYSTEM_STARTUP_INFO, jsonable_encoder(api_startup_status.model_dump()) + new_connections.append(conn_id) + await _send_ws_event( + conn_id, + Event.SYSTEM_STARTUP_INFO, + jsonable_encoder(api_startup_status.model_dump()), ) - asyncio.create_task(warmup_new_connections()) - return event_source + try: + while True: + # server-push only; receive loop just detects disconnect + await websocket.receive_text() + except WebSocketDisconnect: + ws_mgr.disconnect(conn_id) + if conn_id in new_connections: + new_connections.remove(conn_id) warmup_running = False @@ -331,11 +319,11 @@ async def warmup_new_connections(): async def _handle(id, event, res): match res: case BaseModel(): - return await _send_sse_event(id, event, res.model_dump()) + return await _send_ws_event(id, event, res.model_dump()) case dict() | list(): - return await _send_sse_event(id, event, res) + return await _send_ws_event(id, event, res) case Ok(data) if data and isinstance(data, BaseModel): - return await _send_sse_event(id, event, data.model_dump()) + return await _send_ws_event(id, event, data.model_dump()) case Ok(data) if not data: logger.debug(f"No data to send for warmup event {event}") return @@ -346,7 +334,7 @@ async def warmup_new_connections(): ) logger.error(f"Error while fetching warmup_data for {event}: {res}") - return await _send_sse_event(id, event, {"error": f"{res}"}) + return await _send_ws_event(id, event, {"error": f"{res}"}) global new_connections if len(new_connections) == 0: @@ -368,13 +356,13 @@ async def warmup_new_connections(): for id in new_connections: await asyncio.gather( *[ - _handle(id, SSE.SYSTEM_INFO, res[0]), - _handle(id, SSE.BTC_INFO, res[1]), - _handle(id, SSE.LN_INFO, res[2]), - _handle(id, SSE.LN_FEE_REVENUE, res[3]), - _handle(id, SSE.WALLET_BALANCE, res[4]), - _handle(id, SSE.APP_STATE_MESSAGE, res[5]), - _handle(id, SSE.HARDWARE_INFO, res[6]), + _handle(id, Event.SYSTEM_INFO, res[0]), + _handle(id, Event.BTC_INFO, res[1]), + _handle(id, Event.LN_INFO, res[2]), + _handle(id, Event.LN_FEE_REVENUE, res[3]), + _handle(id, Event.WALLET_BALANCE, res[4]), + _handle(id, Event.APP_STATE_MESSAGE, res[5]), + _handle(id, Event.HARDWARE_INFO, res[6]), ] ) @@ -384,10 +372,10 @@ async def warmup_new_connections(): for id in new_connections: await asyncio.gather( *[ - _handle(id, SSE.SYSTEM_INFO, res[0]), - _handle(id, SSE.BTC_INFO, res[1]), - _handle(id, SSE.APP_STATE_MESSAGE, res[2]), - _handle(id, SSE.HARDWARE_INFO, res[3]), + _handle(id, Event.SYSTEM_INFO, res[0]), + _handle(id, Event.BTC_INFO, res[1]), + _handle(id, Event.APP_STATE_MESSAGE, res[2]), + _handle(id, Event.HARDWARE_INFO, res[3]), ] ) @@ -402,8 +390,8 @@ async def warmup_new_connections(): for id in new_connections: await asyncio.gather( *[ - _handle(id, SSE.BTC_INFO, res[0]), - _handle(id, SSE.HARDWARE_INFO, res[1]), + _handle(id, Event.BTC_INFO, res[0]), + _handle(id, Event.HARDWARE_INFO, res[1]), ] ) @@ -418,7 +406,7 @@ async def warmup_new_connections(): # Bitcoin Core and Lightning running res = await get_hardware_info() for id in new_connections: - await _send_sse_event(id, SSE.HARDWARE_INFO, res) + await _send_ws_event(id, Event.HARDWARE_INFO, res) # don't clear new_connections, # we'll try again later when api is initialized diff --git a/tests/test_hw_gatherer_resilience.py b/tests/test_hw_gatherer_resilience.py index 02c150a..4f9a28c 100644 --- a/tests/test_hw_gatherer_resilience.py +++ b/tests/test_hw_gatherer_resilience.py @@ -33,7 +33,7 @@ async def test_gatherer_survives_get_hardware_info_error(monkeypatch): raise asyncio.CancelledError() monkeypatch.setattr(service, "get_hardware_info", flaky_hardware_info) - monkeypatch.setattr(service, "broadcast_sse_msg", fake_broadcast) + monkeypatch.setattr(service, "broadcast_msg", fake_broadcast) monkeypatch.setattr(service.asyncio, "sleep", fake_sleep) with pytest.raises(asyncio.CancelledError): diff --git a/tests/test_warmup_events.py b/tests/test_warmup_events.py index 7cb0994..49600a5 100644 --- a/tests/test_warmup_events.py +++ b/tests/test_warmup_events.py @@ -7,7 +7,7 @@ name the WebUI does not listen to. """ from app.api.models import ApiStartupStatus, StartupState -from app.api.utils import SSE +from app.api.utils import Event from app.apps.constants import AppManagementProcessState from app.apps.models import AppStatusQueryResult, AppStatusUpdateTaskMessage from app.external.result_type.src.result.result import Ok @@ -37,7 +37,7 @@ def _patch_common(monkeypatch, sent, *, node_type: str, lightning: StartupState) "api_startup_status", ApiStartupStatus(bitcoin=StartupState.DONE, lightning=lightning), ) - monkeypatch.setattr(main, "_send_sse_event", fake_send) + monkeypatch.setattr(main, "_send_ws_event", fake_send) monkeypatch.setattr( main, "get_bitcoin_client_warmup_data", fake_bitcoin_warmup_data ) @@ -46,7 +46,7 @@ def _patch_common(monkeypatch, sent, *, node_type: str, lightning: StartupState) async def test_bitcoinonly_warmup_sends_app_status_as_app_state_message(monkeypatch): - """The WebUI only listens for SSE.APP_STATE_MESSAGE to populate the Apps + """The WebUI only listens for Event.APP_STATE_MESSAGE to populate the Apps tab. In bitcoin-only mode the warmup app status must be sent under that event, exactly like in lightning mode.""" sent = [] @@ -69,8 +69,8 @@ async def test_bitcoinonly_warmup_sends_app_status_as_app_state_message(monkeypa await main.warmup_new_connections() events = [event for (_, event, _) in sent] - assert SSE.APP_STATE_MESSAGE in events, ( - f"expected app status under event '{SSE.APP_STATE_MESSAGE}', " + assert Event.APP_STATE_MESSAGE in events, ( + f"expected app status under event '{Event.APP_STATE_MESSAGE}', " f"got events: {events}" ) diff --git a/tests/test_ws_endpoint.py b/tests/test_ws_endpoint.py new file mode 100644 index 0000000..188ca4a --- /dev/null +++ b/tests/test_ws_endpoint.py @@ -0,0 +1,29 @@ +import json + +from starlette.testclient import TestClient +from starlette.websockets import WebSocketDisconnect + +from app.auth.auth_handler import sign_jwt +from app.main import app + +client = TestClient(app) + + +def test_ws_valid_auth_receives_startup_frame(): + # conftest.py sets BAPI_JWT_SECRET/ALGORITHM, so sign_jwt() is a valid token + token = sign_jwt() + with client.websocket_connect("/ws") as ws: + ws.send_text(json.dumps({"type": "auth", "token": token})) + frame = json.loads(ws.receive_text()) + assert frame["event"] == "system_startup_info" + assert "data" in frame + + +def test_ws_bad_token_is_closed_4401(): + with client.websocket_connect("/ws") as ws: + ws.send_text(json.dumps({"type": "auth", "token": "not-a-jwt"})) + try: + ws.receive_text() + assert False, "expected the server to close the connection" + except WebSocketDisconnect as e: + assert e.code == 4401