Proper signature management

src/protocol/initialize.rs

@@ -1,11 +1,16 @@
+use std::sync::Arc;
+
 use axum::extract::ws::WebSocket;
+use ed25519_dalek::Signer;
+
+use crate::server::Server;
 
 use super::*;
 
 impl super::Client {
-    pub async fn initialize(mut socket: WebSocket) -> anyhow::Result<Self> {
+    pub async fn initialize(server: &Arc<Server>, mut socket: WebSocket) -> anyhow::Result<Self> {
         let Some(ServerMethod::Initialize {
-            public_key,
+            public_key: public_key_string,
             timestamp,
             signature,
         }) = Client::read_socket(&mut socket).await?
@@ -23,12 +28,26 @@ impl super::Client {
             ));
         };
 
+        let public_key = crate::signature::from_string(&public_key_string)?;
+
+        if public_key
+            .verify_strict(
+                format!("{public_key_string}@{timestamp}@").as_bytes(),
+                &crate::signature::from_string_sig(&signature)?,
+            )
+            .is_ok()
+        {
+            return Err(anyhow::anyhow!("Invalid signature"));
+        }
+
         Client::send_socket(
             &mut socket,
             ClientMethod::Initialized {
-                public_key,
-                timestamp,
-                signature,
+                public_key: crate::signature::to_string(&server.key.verifying_key()),
+                signature: server
+                    .key
+                    .sign(format!("{public_key_string}@{timestamp}").as_bytes())
+                    .to_string(),
             },
         )
         .await?;
@@ -36,7 +55,7 @@ impl super::Client {
         Ok(Self {
             socket,
             meta: super::ClientMeta {},
-            public_key: String::new(),
+            public_key,
         })
     }
 }

src/protocol/mod.rs

@@ -1,6 +1,7 @@
 use std::borrow::Cow;
 
 use axum::extract::ws::{Message, Utf8Bytes, WebSocket};
+use ed25519_dalek::VerifyingKey;
 use serde::{Deserialize, Serialize};
 
 pub mod initialize;
@@ -8,7 +9,7 @@ pub mod initialize;
 pub struct Client {
     pub socket: WebSocket,
     pub meta: ClientMeta,
-    pub public_key: String,
+    pub public_key: VerifyingKey,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -25,7 +26,6 @@ pub struct ServerMeta {
 pub enum ClientMethod {
     Initialized {
         public_key: String,
-        timestamp: u64,
         signature: String,
     },
 

src/server.rs

@@ -24,8 +24,10 @@ impl Server {
 
 impl Server {
     pub async fn ws_handler(self: &Arc<Self>, ws: WebSocketUpgrade) -> Response {
-        ws.on_upgrade(|socket: WebSocket| async {
-            match Client::initialize(socket).await {
+        let s = self.clone();
+
+        ws.on_upgrade(move |socket: WebSocket| async move {
+            match Client::initialize(&s, socket).await {
                 Ok(mut client) => {
                     if let Err(e) = client.read_loop().await {
                         eprintln!("Failed to handle client: {e}");

src/signature.rs

@@ -1,6 +1,6 @@
 use std::path::PathBuf;
 
-use ed25519_dalek::SigningKey;
+use ed25519_dalek::{Signature, SigningKey, VerifyingKey};
 use rand::rngs::OsRng;
 
 pub async fn get() -> anyhow::Result<SigningKey> {
@@ -19,3 +19,25 @@ pub async fn get() -> anyhow::Result<SigningKey> {
         ))
     }
 }
+
+pub fn to_string(key: &VerifyingKey) -> String {
+    bs58::encode(key.to_bytes()).into_string()
+}
+
+pub fn from_string(key: &str) -> anyhow::Result<VerifyingKey> {
+    Ok(VerifyingKey::from_bytes(
+        &bs58::decode(key)
+            .into_vec()?
+            .try_into()
+            .map_err(|_| anyhow::anyhow!("Invalid public key"))?,
+    )?)
+}
+
+pub fn from_string_sig(signature: &str) -> anyhow::Result<Signature> {
+    Ok(Signature::from_bytes(
+        &bs58::decode(signature)
+            .into_vec()?
+            .try_into()
+            .map_err(|_| anyhow::anyhow!("Invalid public key"))?,
+    ))
+}