diff --git a/Cargo.lock b/Cargo.lock index 1f732b0..b649c17 100644 --- a/Cargo.lock +++ b/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" diff --git a/Cargo.toml b/Cargo.toml index 0a5a099..3072733 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/protocol/initialize.rs b/src/protocol/initialize.rs index f03b58f..2bea02a 100644 --- a/src/protocol/initialize.rs +++ b/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 { + pub async fn initialize(server: &Arc, mut socket: WebSocket) -> anyhow::Result { 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, }) } } diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index e7ead39..4c517c9 100644 --- a/src/protocol/mod.rs +++ b/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, }, diff --git a/src/server.rs b/src/server.rs index c57c2b3..0eb6d92 100644 --- a/src/server.rs +++ b/src/server.rs @@ -24,8 +24,10 @@ impl Server { impl Server { pub async fn ws_handler(self: &Arc, 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}"); diff --git a/src/signature.rs b/src/signature.rs index 2c95e85..851b9eb 100644 --- a/src/signature.rs +++ b/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 { @@ -19,3 +19,25 @@ pub async fn get() -> anyhow::Result { )) } } + +pub fn to_string(key: &VerifyingKey) -> String { + bs58::encode(key.to_bytes()).into_string() +} + +pub fn from_string(key: &str) -> anyhow::Result { + 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 { + Ok(Signature::from_bytes( + &bs58::decode(signature) + .into_vec()? + .try_into() + .map_err(|_| anyhow::anyhow!("Invalid public key"))?, + )) +}