From 696df071ddb8073922a909f05b7c1391fd448c6d Mon Sep 17 00:00:00 2001
From: Klesti Selimaj <klesti.selimay@icloud.com>
Date: Fri, 14 Aug 2026 10:42:00 -0400
Subject: [PATCH] Improved messages

---
 src/protocol/initialize.rs | 22 ++++++++++++++++++++--
 src/protocol/mod.rs        | 12 ++++++------
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/src/protocol/initialize.rs b/src/protocol/initialize.rs
index 75dfe9a..2c6dff4 100644
--- a/src/protocol/initialize.rs
+++ b/src/protocol/initialize.rs
@@ -37,15 +37,33 @@ impl UserConnections {
             ));
         };
 
-        let public_key = crate::signature::from_string(&public_key_string)?;
+        let Ok(public_key) = crate::signature::from_string(&public_key_string) else {
+            send_socket(
+                &mut socket,
+                &ClientMethod::Error {
+                    error: Cow::Borrowed("Invalid public key"),
+                },
+            )
+            .await?;
+
+            return Err(anyhow::anyhow!("Invalid public key"));
+        };
 
         if public_key
             .verify_strict(
                 format!("{timestamp}@{hostname}").as_bytes(),
                 &crate::signature::from_string_sig(&signature)?,
             )
-            .is_ok()
+            .is_err()
         {
+            send_socket(
+                &mut socket,
+                &ClientMethod::Error {
+                    error: Cow::Borrowed("Invalid signature"),
+                },
+            )
+            .await?;
+
             return Err(anyhow::anyhow!("Invalid signature"));
         }
 
diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs
index 94bd18a..0503573 100644
--- a/src/protocol/mod.rs
+++ b/src/protocol/mod.rs
@@ -76,21 +76,21 @@ pub async fn read_loop(socket: &Arc<Mutex<WebSocket>>) -> anyhow::Result<()> {
 
 pub async fn read_socket(socket: &mut WebSocket) -> anyhow::Result<Option<ServerMethod>> {
     match socket.recv().await.transpose()? {
-        Some(Message::Text(text)) => {
-            if let Ok(msg) = serde_json::from_str(&text.to_string()) {
-                Ok(Some(msg))
-            } else {
+        Some(Message::Text(text)) => match serde_json::from_str(&text.to_string()) {
+            Ok(msg) => Ok(Some(msg)),
+
+            Err(e) => {
                 send_socket(
                     socket,
                     &ClientMethod::Error {
-                        error: Cow::Borrowed("Unable to parse message: {text}"),
+                        error: Cow::Owned(format!("Unable to parse message: {e}")),
                     },
                 )
                 .await?;
 
                 Ok(None)
             }
-        }
+        },
 
         Some(Message::Ping(v)) => {
             socket.send(Message::Pong(v)).await?;