From 6aa697ea3deb8f14cce268367398eac436c4407b Mon Sep 17 00:00:00 2001
From: ShahanaFarooqui <shahana.farooqui@gmail.com>
Date: Fri, 8 Sep 2023 19:46:56 -0700
Subject: [PATCH] plugins/clnrest: Rune authentication for websocket server

---
 plugins/clnrest/clnrest.py          | 34 ++++++++++++++++++++---------
 plugins/clnrest/utilities/shared.py |  7 +++++-
 2 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/plugins/clnrest/clnrest.py b/plugins/clnrest/clnrest.py
index fbc6e5b6c..c24215a08 100755
--- a/plugins/clnrest/clnrest.py
+++ b/plugins/clnrest/clnrest.py
@@ -9,13 +9,13 @@ try:
     from gunicorn.workers import sync  # noqa: F401
 
     from pathlib import Path
-    from flask import Flask
+    from flask import Flask, request
     from flask_restx import Api
     from gunicorn.app.base import BaseApplication
     from multiprocessing import Process, Queue
-    from flask_socketio import SocketIO
+    from flask_socketio import SocketIO, disconnect
     from utilities.generate_certs import generate_certs
-    from utilities.shared import set_config
+    from utilities.shared import set_config, verify_rune
     from utilities.rpc_routes import rpcns
     from utilities.rpc_plugin import plugin
 except ModuleNotFoundError as err:
@@ -52,16 +52,30 @@ def broadcast_from_message_queue():
 socketio.start_background_task(broadcast_from_message_queue)
 
 
-@socketio.on("connect", namespace="/ws")
+@socketio.on("message")
+def handle_message(message):
+    plugin.log(f"Received message from client: {message}", "debug")
+    socketio.emit('message', {"client_message": message, "session": request.sid})
+
+
+@socketio.on("connect")
 def ws_connect():
-    plugin.log("Client Connected", "debug")
-    msgq.put("Client Connected")
+    try:
+        plugin.log("Client Connecting...", "debug")
+        is_valid_rune = verify_rune(plugin, request)
 
+        if "error" in is_valid_rune:
+            # Logging as error/warn emits the event for all clients
+            plugin.log(f"Error: {is_valid_rune}", "info")
+            raise Exception(is_valid_rune)
 
-@socketio.on("disconnect", namespace="/ws")
-def ws_disconnect():
-    plugin.log("Client Disconnected", "debug")
-    msgq.put("Client Disconnected")
+        plugin.log("Client Connected", "debug")
+        return True
+
+    except Exception as err:
+        # Logging as error/warn emits the event for all clients
+        plugin.log(f"{err}", "info")
+        disconnect()
 
 
 def create_app():
diff --git a/plugins/clnrest/utilities/shared.py b/plugins/clnrest/utilities/shared.py
index df36ca94a..9f139a72c 100644
--- a/plugins/clnrest/utilities/shared.py
+++ b/plugins/clnrest/utilities/shared.py
@@ -56,9 +56,14 @@ def verify_rune(plugin, request):
     else:
         rpc_params = request.form.to_dict()
 
+    try:
+        rpc_method = request.view_args["rpc_method"]
+    except Exception:
+        rpc_method = ""
+
     return call_rpc_method(plugin, "checkrune",
                            {"rune": rune,
-                            "method": request.view_args["rpc_method"],
+                            "method": rpc_method,
                             "params": rpc_params})