Protect controlle notification endpoints. Ref #1888

(WebSocket endpoint is not secured, it takes an optional token).
pull/1906/head
grossmj 3 years ago
parent a02aede1b0
commit 49be4146d6

@ -18,12 +18,15 @@
API routes for controller notifications.
"""
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from fastapi import APIRouter, Depends, Query, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.responses import StreamingResponse
from websockets.exceptions import ConnectionClosed, WebSocketException
from gns3server.services import auth_service
from gns3server.controller import Controller
from .dependencies.authentication import get_current_active_user
import logging
log = logging.getLogger(__name__)
@ -31,7 +34,7 @@ log = logging.getLogger(__name__)
router = APIRouter()
@router.get("")
@router.get("", dependencies=[Depends(get_current_active_user)])
async def http_notification() -> StreamingResponse:
"""
Receive controller notifications about the controller from HTTP stream.
@ -41,18 +44,26 @@ async def http_notification() -> StreamingResponse:
with Controller.instance().notification.controller_queue() as queue:
while True:
msg = await queue.get_json(5)
yield (f"{msg}\n").encode("utf-8")
yield f"{msg}\n".encode("utf-8")
return StreamingResponse(event_stream(), media_type="application/json")
@router.websocket("/ws")
async def notification_ws(websocket: WebSocket) -> None:
async def notification_ws(websocket: WebSocket, token: str = Query(None)) -> None:
"""
Receive project notifications about the controller from WebSocket.
"""
await websocket.accept()
if token:
try:
username = auth_service.get_username_from_token(token)
except HTTPException:
log.error("Invalid token received")
await websocket.close(code=1008)
return
log.info(f"New client {websocket.client.host}:{websocket.client.port} has connected to controller WebSocket")
try:
with Controller.instance().notification.controller_queue() as queue:

Loading…
Cancel
Save