Refactor WebSocket console code to work with FastAPI.

Fix endpoint routes.

gns3server/endpoints/controller/appliances.py

@@ -25,7 +25,7 @@ from typing import Optional
 router = APIRouter()
 
 
-@router.get("/")
+@router.get("")
 async def get_appliances(update: Optional[bool] = None, symbol_theme:  Optional[str] = "Classic"):
     """
     Return all appliances known by the controller.

gns3server/endpoints/controller/computes.py

@@ -35,7 +35,7 @@ responses = {
 }
 
 
-@router.post("/",
+@router.post("",
              status_code=status.HTTP_201_CREATED,
              response_model=schemas.Compute,
              responses={404: {"model": ErrorMessage, "description": "Could not connect to compute"},
@@ -64,7 +64,7 @@ def get_compute(compute_id: Union[str, UUID]):
     return compute.__json__()
 
 
-@router.get("/",
+@router.get("",
             response_model=List[schemas.Compute],
             response_model_exclude_unset=True)
 async def get_computes():

gns3server/endpoints/controller/drawings.py

@@ -35,7 +35,7 @@ responses = {
 }
 
 
-@router.get("/",
+@router.get("",
             response_model=List[Drawing],
             response_model_exclude_unset=True)
 async def get_drawings(project_id: UUID):
@@ -47,7 +47,7 @@ async def get_drawings(project_id: UUID):
     return [v.__json__() for v in project.drawings.values()]
 
 
-@router.post("/",
+@router.post("",
              status_code=status.HTTP_201_CREATED,
              response_model=Drawing,
              responses=responses)

gns3server/endpoints/controller/gns3vm.py

@@ -48,8 +48,7 @@ async def get_vms(engine: str):
     return vms
 
 
-@router.get("/",
-            response_model=GNS3VM)
+@router.get("", response_model=GNS3VM)
 async def get_gns3vm_settings():
     """
     Return the GNS3 VM settings.
@@ -58,9 +57,7 @@ async def get_gns3vm_settings():
     return Controller.instance().gns3vm.__json__()
 
 
-@router.put("/",
-            response_model=GNS3VM,
-            response_model_exclude_unset=True)
+@router.put("", response_model=GNS3VM, response_model_exclude_unset=True)
 async def update_gns3vm_settings(gns3vm_data: GNS3VM):
     """
     Update the GNS3 VM settings.

gns3server/endpoints/controller/links.py

@@ -48,7 +48,7 @@ async def dep_link(project_id: UUID, link_id: UUID):
     return link
 
 
-@router.get("/",
+@router.get("",
             response_model=List[schemas.Link],
             response_model_exclude_unset=True)
 async def get_links(project_id: UUID):
@@ -60,7 +60,7 @@ async def get_links(project_id: UUID):
     return [v.__json__() for v in project.links.values()]
 
 
-@router.post("/",
+@router.post("",
              status_code=status.HTTP_201_CREATED,
              response_model=schemas.Link,
              responses={404: {"model": ErrorMessage, "description": "Could not find project"},
@@ -214,4 +214,4 @@ async def reset_link(link: Link = Depends(dep_link)):
 #                     break
 #                 await proxied_response.write(data)
 #
-#     #return StreamingResponse(file_like, media_type="video/mp4"))
+#     #return StreamingResponse(file_like, media_type="video/mp4"))

gns3server/endpoints/controller/nodes.py

@@ -19,9 +19,10 @@
 API endpoints for nodes.
 """
 
+import aiohttp
 import asyncio
 
-from fastapi import APIRouter, Depends, Request, Response, status
+from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect, Request, Response, status
 from fastapi.encoders import jsonable_encoder
 from fastapi.routing import APIRoute
 from typing import List, Callable
@@ -35,6 +36,9 @@ from gns3server.controller.controller_error import ControllerForbiddenError
 from gns3server.endpoints.schemas.common import ErrorMessage
 from gns3server.endpoints import schemas
 
+import logging
+log = logging.getLogger(__name__)
+
 node_locks = {}
 
 
@@ -97,7 +101,7 @@ async def dep_node(node_id: UUID, project: Project = Depends(dep_project)):
     return node
 
 
-@router.post("/",
+@router.post("",
              status_code=status.HTTP_201_CREATED,
              response_model=schemas.Node,
              responses={404: {"model": ErrorMessage, "description": "Could not find project"},
@@ -117,7 +121,7 @@ async def create_node(node_data: schemas.Node, project: Project = Depends(dep_pr
     return node.__json__()
 
 
-@router.get("/",
+@router.get("",
             response_model=List[schemas.Node],
             response_model_exclude_unset=True)
 async def get_nodes(project: Project = Depends(dep_project)):
@@ -367,59 +371,47 @@ async def post_file(file_path: str, request: Request, node: Node = Depends(dep_n
                                   raw=True)
 
 
-# @Route.get(
-#     r"/projects/{project_id}/nodes/{node_id}/console/ws",
-#     parameters={
-#         "project_id": "Project UUID",
-#         "node_id": "Node UUID"
-#     },
-#     description="Connect to WebSocket console",
-#     status_codes={
-#         200: "File returned",
-#         403: "Permission denied",
-#         404: "The file doesn't exist"
-#     })
-# async def ws_console(request, response):
-#
-#     project = await Controller.instance().get_loaded_project(request.match_info["project_id"])
-#     node = project.get_node(request.match_info["node_id"])
-#     compute = node.compute
-#     ws = aiohttp.web.WebSocketResponse()
-#     await ws.prepare(request)
-#     request.app['websockets'].add(ws)
-#
-#     ws_console_compute_url = "ws://{compute_host}:{compute_port}/v2/compute/projects/{project_id}/{node_type}/nodes/{node_id}/console/ws".format(compute_host=compute.host,
-#                                                                                                                                                  compute_port=compute.port,
-#                                                                                                                                                  project_id=project.id,
-#                                                                                                                                                  node_type=node.node_type,
-#                                                                                                                                                  node_id=node.id)
-#
-#     async def ws_forward(ws_client):
-#         async for msg in ws:
-#             if msg.type == aiohttp.WSMsgType.TEXT:
-#                 await ws_client.send_str(msg.data)
-#             elif msg.type == aiohttp.WSMsgType.BINARY:
-#                 await ws_client.send_bytes(msg.data)
-#             elif msg.type == aiohttp.WSMsgType.ERROR:
-#                 break
-#
-#     try:
-#         async with aiohttp.ClientSession(connector=aiohttp.TCPConnector(limit=None, force_close=True)) as session:
-#             async with session.ws_connect(ws_console_compute_url) as ws_client:
-#                 asyncio.ensure_future(ws_forward(ws_client))
-#                 async for msg in ws_client:
-#                     if msg.type == aiohttp.WSMsgType.TEXT:
-#                         await ws.send_str(msg.data)
-#                     elif msg.type == aiohttp.WSMsgType.BINARY:
-#                         await ws.send_bytes(msg.data)
-#                     elif msg.type == aiohttp.WSMsgType.ERROR:
-#                         break
-#     finally:
-#         if not ws.closed:
-#             await ws.close()
-#         request.app['websockets'].discard(ws)
-#
-#     return ws
+@router.websocket("/{node_id}/console/ws")
+async def ws_console(websocket: WebSocket, node: Node = Depends(dep_node)):
+    """
+    WebSocket console.
+    """
+
+    compute = node.compute
+    await websocket.accept()
+    log.info(f"New client {websocket.client.host}:{websocket.client.port} has connected to controller console WebSocket")
+    ws_console_compute_url = f"ws://{compute.host}:{compute.port}/v2/compute/projects/" \
+                             f"{node.project.id}/{node.node_type}/nodes/{node.id}/console/ws"
+
+    async def ws_receive(ws_console_compute):
+        """
+        Receive WebSocket data from client and forward to compute console WebSocket.
+        """
+
+        try:
+            while True:
+                data = await websocket.receive_text()
+                if data:
+                    await ws_console_compute.send_str(data)
+        except WebSocketDisconnect:
+            await ws_console_compute.close()
+            log.info(f"Client {websocket.client.host}:{websocket.client.port} has disconnected from controller"
+                     f" console WebSocket")
+
+    try:
+        # receive WebSocket data from compute console WebSocket and forward to client.
+        async with aiohttp.ClientSession(connector=aiohttp.TCPConnector(limit=None, force_close=True)) as session:
+            async with session.ws_connect(ws_console_compute_url) as ws_console_compute:
+                asyncio.ensure_future(ws_receive(ws_console_compute))
+                async for msg in ws_console_compute:
+                    if msg.type == aiohttp.WSMsgType.TEXT:
+                        await websocket.send_text(msg.data)
+                    elif msg.type == aiohttp.WSMsgType.BINARY:
+                        await websocket.send_bytes(msg.data)
+                    elif msg.type == aiohttp.WSMsgType.ERROR:
+                        break
+    except aiohttp.client_exceptions.ClientResponseError as e:
+        log.error(f"Client response error received when forwarding to compute console WebSocket: {e}")
 
 
 @router.post("/console/reset",

gns3server/endpoints/controller/notifications.py

@@ -19,51 +19,58 @@
 API endpoints for controller notifications.
 """
 
+import asyncio
+
+from fastapi import APIRouter, WebSocket
+from fastapi.responses import StreamingResponse
+from starlette.endpoints import WebSocketEndpoint
 
-from fastapi import APIRouter, Request, Response, WebSocket, WebSocketDisconnect
-from websockets.exceptions import WebSocketException
 from gns3server.controller import Controller
 
-router = APIRouter()
-
 import logging
 log = logging.getLogger(__name__)
 
-
-# @router.get("/")
-# async def notification(request: Request):
-#     """
-#     Receive notifications about the controller from HTTP
-#     """
-#
-#     controller = Controller.instance()
-#
-#     await response.prepare(request)
-#     response = Response(content, media_type="application/json")
-#
-#     with controller.notification.controller_queue() as queue:
-#         while True:
-#             msg = await queue.get_json(5)
-#             await response.write(("{}\n".format(msg)).encode("utf-8"))
-#
-#
-#             await response(scope, receive, send)
+router = APIRouter()
 
 
-@router.websocket("/ws")
-async def notification_ws(websocket: WebSocket):
+@router.get("")
+async def http_notification():
     """
-    Receive notifications about the controller from a Websocket
+    Receive controller notifications about the controller from HTTP stream.
     """
 
-    controller = Controller.instance()
-    await websocket.accept()
-    log.info("New client has connected to controller WebSocket")
-    try:
-        with controller.notification.controller_queue() as queue:
+    async def event_stream():
+
+        with Controller.instance().notification.controller_queue() as queue:
+            while True:
+                msg = await queue.get_json(5)
+                yield ("{}\n".format(msg)).encode("utf-8")
+
+    return StreamingResponse(event_stream(), media_type="application/json")
+
+
+@router.websocket_route("/ws")
+class ControllerWebSocketNotifications(WebSocketEndpoint):
+    """
+    Receive controller notifications about the controller from WebSocket stream.
+    """
+
+    async def on_connect(self, websocket: WebSocket) -> None:
+
+        await websocket.accept()
+        log.info(f"New client {websocket.client.host}:{websocket.client.port} has connected to controller WebSocket")
+
+        self._notification_task = asyncio.ensure_future(self._stream_notifications(websocket=websocket))
+
+    async def on_disconnect(self, websocket: WebSocket, close_code: int) -> None:
+
+        self._notification_task.cancel()
+        log.info(f"Client {websocket.client.host}:{websocket.client.port} has disconnected from controller WebSocket"
+                 f" with close code {close_code}")
+
+    async def _stream_notifications(self, websocket: WebSocket) -> None:
+
+        with Controller.instance().notifications.queue() as queue:
             while True:
                 notification = await queue.get_json(5)
                 await websocket.send_text(notification)
-    except (WebSocketException, WebSocketDisconnect):
-        log.info("Client has disconnected from controller WebSocket")
-        await websocket.close()

gns3server/endpoints/controller/projects.py

@@ -32,7 +32,7 @@ log = logging.getLogger()
 from fastapi import APIRouter, Depends, Request, Body, HTTPException, status, WebSocket, WebSocketDisconnect
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import StreamingResponse, FileResponse
-from websockets.exceptions import WebSocketException
+from websockets.exceptions import ConnectionClosed, WebSocketException
 from typing import List
 from uuid import UUID
 
@@ -66,7 +66,19 @@ def dep_project(project_id: UUID):
 CHUNK_SIZE = 1024 * 8  # 8KB
 
 
-@router.post("/",
+@router.get("",
+            response_model=List[schemas.Project],
+            response_model_exclude_unset=True)
+def get_projects():
+    """
+    Return all projects.
+    """
+
+    controller = Controller.instance()
+    return [p.__json__() for p in controller.projects.values()]
+
+
+@router.post("",
              status_code=status.HTTP_201_CREATED,
              response_model=schemas.Project,
              response_model_exclude_unset=True,
@@ -78,22 +90,9 @@ async def create_project(project_data: schemas.ProjectCreate):
 
     controller = Controller.instance()
     project = await controller.add_project(**jsonable_encoder(project_data, exclude_unset=True))
-    print(project.__json__()["variables"])
     return project.__json__()
 
 
-@router.get("/",
-            response_model=List[schemas.Project],
-            response_model_exclude_unset=True)
-def get_projects():
-    """
-    Return all projects.
-    """
-
-    controller = Controller.instance()
-    return [p.__json__() for p in controller.projects.values()]
-
-
 @router.get("/{project_id}",
             response_model=schemas.Project,
             responses=responses)
@@ -193,52 +192,57 @@ async def load_project(path: str = Body(..., embed=True)):
     return project.__json__()
 
 
-# @router.get("/projects/{project_id}/notifications",
-#             summary="Receive notifications about projects",
-#             responses={404: {"model": ErrorMessage, "description": "Could not find project"}})
-# async def notification(project_id: UUID):
-#
-#     controller = Controller.instance()
-#     project = controller.get_project(str(project_id))
-#     #response.content_type = "application/json"
-#     #response.set_status(200)
-#     #response.enable_chunked_encoding()
-#     #await response.prepare(request)
-#     log.info("New client has connected to the notification stream for project ID '{}' (HTTP long-polling method)".format(project.id))
-#
-#     try:
-#         with controller.notification.project_queue(project.id) as queue:
-#             while True:
-#                 msg = await queue.get_json(5)
-#                 await response.write(("{}\n".format(msg)).encode("utf-8"))
-#     finally:
-#         log.info("Client has disconnected from notification for project ID '{}' (HTTP long-polling method)".format(project.id))
-#         if project.auto_close:
-#             # To avoid trouble with client connecting disconnecting we sleep few seconds before checking
-#             # if someone else is not connected
-#             await asyncio.sleep(5)
-#             if not controller.notification.project_has_listeners(project.id):
-#                 log.info("Project '{}' is automatically closing due to no client listening".format(project.id))
-#                 await project.close()
+@router.get("/{project_id}/notifications")
+async def notification(project_id: UUID):
+    """
+    Receive project notifications about the controller from HTTP stream.
+    """
+
+    controller = Controller.instance()
+    project = controller.get_project(str(project_id))
+
+    log.info("New client has connected to the notification stream for project ID '{}' (HTTP steam method)".format(project.id))
+
+    async def event_stream():
+
+        try:
+            with controller.notification.project_queue(project.id) as queue:
+                while True:
+                    msg = await queue.get_json(5)
+                    yield ("{}\n".format(msg)).encode("utf-8")
+        finally:
+            log.info("Client has disconnected from notification for project ID '{}' (HTTP stream method)".format(project.id))
+            if project.auto_close:
+                # To avoid trouble with client connecting disconnecting we sleep few seconds before checking
+                # if someone else is not connected
+                await asyncio.sleep(5)
+                if not controller.notification.project_has_listeners(project.id):
+                    log.info("Project '{}' is automatically closing due to no client listening".format(project.id))
+                    await project.close()
+
+    return StreamingResponse(event_stream(), media_type="application/json")
 
 
 @router.websocket("/{project_id}/notifications/ws")
 async def notification_ws(project_id: UUID, websocket: WebSocket):
+    """
+    Receive project notifications about the controller from WebSocket.
+    """
 
     controller = Controller.instance()
     project = controller.get_project(str(project_id))
     await websocket.accept()
 
-    #request.app['websockets'].add(ws)
-    #asyncio.ensure_future(process_websocket(ws))
     log.info("New client has connected to the notification stream for project ID '{}' (WebSocket method)".format(project.id))
     try:
         with controller.notification.project_queue(project.id) as queue:
             while True:
                 notification = await queue.get_json(5)
                 await websocket.send_text(notification)
-    except (WebSocketException, WebSocketDisconnect):
+    except (ConnectionClosed, WebSocketDisconnect):
         log.info("Client has disconnected from notification stream for project ID '{}' (WebSocket method)".format(project.id))
+    except WebSocketException as e:
+        log.warning("Error while sending to project event to WebSocket client: '{}'".format(e))
     finally:
         await websocket.close()
         if project.auto_close:

gns3server/endpoints/controller/snapshots.py

@@ -19,7 +19,6 @@
 API endpoints for snapshots.
 """
 
-
 import logging
 log = logging.getLogger()
 
@@ -48,7 +47,7 @@ def dep_project(project_id: UUID):
     return project
 
 
-@router.post("/",
+@router.post("",
              status_code=status.HTTP_201_CREATED,
              response_model=schemas.Snapshot,
              responses=responses)
@@ -61,7 +60,7 @@ async def create_snapshot(snapshot_data: schemas.SnapshotCreate, project: Projec
     return snapshot.__json__()
 
 
-@router.get("/",
+@router.get("",
             response_model=List[schemas.Snapshot],
             response_model_exclude_unset=True,
             responses=responses)

gns3server/endpoints/controller/symbols.py

@@ -35,7 +35,7 @@ log = logging.getLogger(__name__)
 router = APIRouter()
 
 
-@router.get("/")
+@router.get("")
 def get_symbols():
 
     controller = Controller.instance()