From 49be4146d65a39cba6e355fb96cbea63a3af1d4a Mon Sep 17 00:00:00 2001
From: grossmj <grossmj@gns3.net>
Date: Mon, 26 Apr 2021 16:18:18 +0930
Subject: [PATCH] Protect controlle notification endpoints. Ref #1888
 (WebSocket endpoint is not secured, it takes an optional token).

---
 .../api/routes/controller/notifications.py    | 21 ++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/gns3server/api/routes/controller/notifications.py b/gns3server/api/routes/controller/notifications.py
index 3634a7d4..79e99328 100644
--- a/gns3server/api/routes/controller/notifications.py
+++ b/gns3server/api/routes/controller/notifications.py
@@ -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: