From b0e646b97db501150d15310af6e319374a46d7bf Mon Sep 17 00:00:00 2001
From: grossmj <grossmj@gns3.net>
Date: Tue, 13 Sep 2022 22:10:01 +0200
Subject: [PATCH] Fix some issues with HTTP notification streams

---
 .../api/routes/controller/notifications.py    | 21 ++++++++++++-------
 gns3server/api/routes/controller/projects.py  |  7 ++++---
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/gns3server/api/routes/controller/notifications.py b/gns3server/api/routes/controller/notifications.py
index 2618efae..52d29e74 100644
--- a/gns3server/api/routes/controller/notifications.py
+++ b/gns3server/api/routes/controller/notifications.py
@@ -18,7 +18,7 @@
 API routes for controller notifications.
 """
 
-from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect
+from fastapi import APIRouter, Request, Depends, WebSocket, WebSocketDisconnect
 from fastapi.responses import StreamingResponse
 from websockets.exceptions import ConnectionClosed, WebSocketException
 
@@ -35,17 +35,24 @@ router = APIRouter()
 
 
 @router.get("", dependencies=[Depends(get_current_active_user)])
-async def controller_http_notifications() -> StreamingResponse:
+async def controller_http_notifications(request: Request) -> StreamingResponse:
     """
     Receive controller notifications about the controller from HTTP stream.
     """
 
-    async def event_stream():
-        with Controller.instance().notification.controller_queue() as queue:
-            while True:
-                msg = await queue.get_json(5)
-                yield f"{msg}\n".encode("utf-8")
+    from gns3server.api.server import app
+    log.info(f"New client {request.client.host}:{request.client.port} has connected to controller HTTP "
+             f"notification stream")
 
+    async def event_stream():
+        try:
+            with Controller.instance().notification.controller_queue() as queue:
+                while not app.state.exiting:
+                    msg = await queue.get_json(5)
+                    yield f"{msg}\n".encode("utf-8")
+        finally:
+            log.info(f"Client {request.client.host}:{request.client.port} has disconnected from controller HTTP "
+                     f"notification stream")
     return StreamingResponse(event_stream(), media_type="application/json")
 
 
diff --git a/gns3server/api/routes/controller/projects.py b/gns3server/api/routes/controller/projects.py
index 91021f64..302173d0 100644
--- a/gns3server/api/routes/controller/projects.py
+++ b/gns3server/api/routes/controller/projects.py
@@ -210,18 +210,19 @@ async def project_http_notifications(project_id: UUID) -> StreamingResponse:
     Receive project notifications about the controller from HTTP stream.
     """
 
+    from gns3server.api.server import app
     controller = Controller.instance()
     project = controller.get_project(str(project_id))
 
-    log.info(f"New client has connected to the notification stream for project ID '{project.id}' (HTTP steam method)")
+    log.info(f"New client has connected to the notification stream for project ID '{project.id}' (HTTP stream method)")
 
     async def event_stream():
 
         try:
             with controller.notification.project_queue(project.id) as queue:
-                while True:
+                while not app.state.exiting:
                     msg = await queue.get_json(5)
-                    yield (f"{msg}\n").encode("utf-8")
+                    yield f"{msg}\n".encode("utf-8")
         finally:
             log.info(f"Client has disconnected from notification for project ID '{project.id}' (HTTP stream method)")
             if project.auto_close: