Server structure

src/main.rs

@@ -1,38 +1,34 @@
 pub mod protocol;
+pub mod server;
+pub mod signature;
+
+use std::sync::Arc;
 
 use axum::{
     Router,
-    extract::{WebSocketUpgrade, ws::WebSocket},
+    extract::{State, WebSocketUpgrade},
     response::Response,
     routing::{any, get},
 };
 
-use crate::protocol::Client;
+use crate::server::Server;
 
 #[tokio::main]
-async fn main() {
+async fn main() -> anyhow::Result<()> {
+    let server = Server::new().await?;
+
     let app = Router::new()
         .route("/meta", get(|| async { "Hello, World!" }))
-        .route("/", any(ws_handler));
+        .route("/", any(ws_handler))
+        .with_state(server);
 
-    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
-    axum::serve(listener, app).await.unwrap();
+    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await?;
+
+    axum::serve(listener, app).await?;
+
+    Ok(())
 }
 
-async fn ws_handler(ws: WebSocketUpgrade) -> Response {
-    ws.on_upgrade(|socket: WebSocket| async {
-        match Client::initialize(socket).await {
-            Ok(mut client) => {
-                if let Err(e) = client.read_loop().await {
-                    eprintln!("Failed to handle client: {e}");
-                } else {
-                    println!("Client connection closed")
-                }
-            }
-
-            Err(e) => {
-                eprintln!("Failed to initialize client: {e}")
-            }
-        }
-    })
+async fn ws_handler(State(server): State<Arc<Server>>, ws: WebSocketUpgrade) -> Response {
+    server.ws_handler(ws).await
 }

src/server.rs

@@ -0,0 +1,41 @@
+use std::sync::Arc;
+
+use axum::{
+    extract::{WebSocketUpgrade, ws::WebSocket},
+    response::Response,
+};
+use ed25519_dalek::SigningKey;
+
+use crate::protocol::Client;
+
+pub struct Server {
+    pub key: SigningKey,
+}
+
+impl Server {
+    pub async fn new() -> anyhow::Result<Arc<Self>> {
+        Ok(Arc::new(Self {
+            key: crate::signature::get().await?,
+        }))
+    }
+}
+
+impl Server {
+    pub async fn ws_handler(self: &Arc<Self>, ws: WebSocketUpgrade) -> Response {
+        ws.on_upgrade(|socket: WebSocket| async {
+            match Client::initialize(socket).await {
+                Ok(mut client) => {
+                    if let Err(e) = client.read_loop().await {
+                        eprintln!("Failed to handle client: {e}");
+                    } else {
+                        println!("Client connection closed")
+                    }
+                }
+
+                Err(e) => {
+                    eprintln!("Failed to initialize client: {e}")
+                }
+            }
+        })
+    }
+}

src/signature.rs

@@ -0,0 +1,21 @@
+use std::path::PathBuf;
+
+use ed25519_dalek::SigningKey;
+use rand::rngs::OsRng;
+
+pub async fn get() -> anyhow::Result<SigningKey> {
+    let private_key_path = PathBuf::from("private.key");
+
+    if !private_key_path.exists() {
+        let key = SigningKey::generate(&mut OsRng);
+        tokio::fs::write(private_key_path, &key.to_bytes()).await?;
+        Ok(key)
+    } else {
+        Ok(SigningKey::from_bytes(
+            &tokio::fs::read(private_key_path)
+                .await?
+                .try_into()
+                .map_err(|_| anyhow::anyhow!("Invalid private key"))?,
+        ))
+    }
+}