Proper signature management

Cargo.lock

@@ -90,6 +90,15 @@ dependencies = [
  "generic-array",
 ]
 
+[[package]]
+name = "bs58"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bf88ba1141d185c399bee5288d850d63b8369520c1eafc32a0430b5b6c287bf4"
+dependencies = [
+ "tinyvec",
+]
+
 [[package]]
 name = "bytes"
 version = "1.12.1"
@@ -211,6 +220,7 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "axum",
+ "bs58",
  "ed25519-dalek",
  "rand 0.8.7",
  "serde",
@@ -757,6 +767,21 @@ dependencies = [
  "syn 3.0.3",
 ]
 
+[[package]]
+name = "tinyvec"
+version = "1.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bb4ebadaa0af04fab11ae01eb5f9fdb5f9c5b875506e210e71c07873528baa7f"
+dependencies = [
+ "tinyvec_macros",
+]
+
+[[package]]
+name = "tinyvec_macros"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
+
 [[package]]
 name = "tokio"
 version = "1.53.1"

Cargo.toml

@@ -11,3 +11,4 @@ serde_json = "1.0.151"
 tokio = { version = "1.53.1", features = ["rt", "rt-multi-thread", "macros", "sync", "fs"] }
 ed25519-dalek = { version = "2", features = ["rand_core"] }
 rand = "0.8"
+bs58 = "0.5.1"

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"))?,
+    ))
+}