[WIP] change to webtransport

Cargo.lock

@@ -417,6 +417,45 @@ dependencies = [
  "derive_arbitrary",
 ]
 
+[[package]]
+name = "asn1-rs"
+version = "0.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "56624a96882bb8c26d61312ae18cb45868e5a9992ea73c58e45c3101e56a1e60"
+dependencies = [
+ "asn1-rs-derive",
+ "asn1-rs-impl",
+ "displaydoc",
+ "nom",
+ "num-traits",
+ "rusticata-macros",
+ "thiserror 2.0.18",
+ "time",
+]
+
+[[package]]
+name = "asn1-rs-derive"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3109e49b1e4909e9db6515a30c633684d68cdeaa252f215214cb4fa1a5bfee2c"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+ "synstructure",
+]
+
+[[package]]
+name = "asn1-rs-impl"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7b18050c2cd6fe86c3a76584ef5e0baf286d038cda203eb6223df2cc413565f7"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "async-channel"
 version = "1.9.0"
@@ -598,6 +637,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "94bffc006df10ac2a68c83692d734a465f8ee6c5b384d8545a636f81d858f4bf"
 dependencies = [
  "aws-lc-sys",
+ "untrusted 0.7.1",
  "zeroize",
 ]
 
@@ -1001,6 +1041,20 @@ dependencies = [
  "zeroize",
 ]
 
+[[package]]
+name = "der-parser"
+version = "10.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "07da5016415d5a3c4dd39b11ed26f915f52fc4e0dc197d87908bc916e51bc1a6"
+dependencies = [
+ "asn1-rs",
+ "displaydoc",
+ "nom",
+ "num-bigint",
+ "num-traits",
+ "rusticata-macros",
+]
+
 [[package]]
 name = "deranged"
 version = "0.5.8"
@@ -1117,7 +1171,7 @@ dependencies = [
 [[package]]
 name = "epsilon-core"
 version = "0.1.0"
-source = "git+https://github.com/Tensamin/Epsilon.git#eb91e812e559270ff22572542491f8215db27a76"
+source = "git+https://github.com/Tensamin/Epsilon.git#4bdd32b44a77ac980cede20a1d401004f52e4a29"
 dependencies = [
  "byteorder",
  "quinn",
@@ -1129,16 +1183,19 @@ dependencies = [
 [[package]]
 name = "epsilon-native"
 version = "0.1.0"
-source = "git+https://github.com/Tensamin/Epsilon.git#eb91e812e559270ff22572542491f8215db27a76"
+source = "git+https://github.com/Tensamin/Epsilon.git#4bdd32b44a77ac980cede20a1d401004f52e4a29"
 dependencies = [
  "anyhow",
  "async-trait",
+ "aws-lc-rs",
  "bytes",
  "epsilon-core",
  "quinn",
+ "rcgen",
  "rustls",
  "thiserror 2.0.18",
  "tokio",
+ "wtransport",
 ]
 
 [[package]]
@@ -1452,20 +1509,20 @@ dependencies = [
  "cfg-if",
  "js-sys",
  "libc",
- "r-efi",
+ "r-efi 5.3.0",
  "wasip2",
  "wasm-bindgen",
 ]
 
 [[package]]
 name = "getrandom"
-version = "0.4.1"
+version = "0.4.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "139ef39800118c7683f2fd3c98c1b23c09ae076556b435f8e9064ae108aaeeec"
+checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555"
 dependencies = [
  "cfg-if",
  "libc",
- "r-efi",
+ "r-efi 6.0.0",
  "wasip2",
  "wasip3",
 ]
@@ -1625,6 +1682,12 @@ dependencies = [
  "windows-sys 0.61.2",
 ]
 
+[[package]]
+name = "httlib-huffman"
+version = "0.3.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1a9fcbcc408c5526c3ab80d534e5c86e7967c1fb7aa0a8c76abd1edc27deb877"
+
 [[package]]
 name = "http"
 version = "0.2.12"
@@ -2145,6 +2208,12 @@ version = "0.3.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
 
+[[package]]
+name = "minimal-lexical"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a"
+
 [[package]]
 name = "miniz_oxide"
 version = "0.8.9"
@@ -2173,6 +2242,16 @@ version = "0.6.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "43794a0ace135be66a25d3ae77d41b91615fb68ae937f904090203e81f755b65"
 
+[[package]]
+name = "nom"
+version = "7.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a"
+dependencies = [
+ "memchr",
+ "minimal-lexical",
+]
+
 [[package]]
 name = "ntapi"
 version = "0.4.3"
@@ -2182,6 +2261,16 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "num-bigint"
+version = "0.4.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9"
+dependencies = [
+ "num-integer",
+ "num-traits",
+]
+
 [[package]]
 name = "num-bigint-dig"
 version = "0.8.6"
@@ -2262,6 +2351,21 @@ dependencies = [
  "memchr",
 ]
 
+[[package]]
+name = "octets"
+version = "0.3.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8311fa8ab7a57759b4ff1f851a3048d9ef0effaa0130726426b742d26d8a88e7"
+
+[[package]]
+name = "oid-registry"
+version = "0.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "12f40cff3dde1b6087cc5d5f5d4d65712f34016a03ed60e9c08dcc392736b5b7"
+dependencies = [
+ "asn1-rs",
+]
+
 [[package]]
 name = "once_cell"
 version = "1.21.3"
@@ -2325,6 +2429,16 @@ dependencies = [
  "hmac",
 ]
 
+[[package]]
+name = "pem"
+version = "3.0.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be"
+dependencies = [
+ "base64",
+ "serde_core",
+]
+
 [[package]]
 name = "pem-rfc7468"
 version = "0.7.0"
@@ -2651,6 +2765,12 @@ version = "5.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
 
+[[package]]
+name = "r-efi"
+version = "6.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf"
+
 [[package]]
 name = "rand"
 version = "0.8.5"
@@ -2716,6 +2836,20 @@ dependencies = [
  "getrandom 0.3.4",
 ]
 
+[[package]]
+name = "rcgen"
+version = "0.14.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "10b99e0098aa4082912d4c649628623db6aba77335e4f4569ff5083a6448b32e"
+dependencies = [
+ "pem",
+ "ring",
+ "rustls-pki-types",
+ "time",
+ "x509-parser 0.18.1",
+ "yasna",
+]
+
 [[package]]
 name = "redox_syscall"
 version = "0.5.18"
@@ -2817,7 +2951,7 @@ dependencies = [
  "cfg-if",
  "getrandom 0.2.17",
  "libc",
- "untrusted",
+ "untrusted 0.9.0",
  "windows-sys 0.52.0",
 ]
 
@@ -2862,6 +2996,15 @@ dependencies = [
  "semver",
 ]
 
+[[package]]
+name = "rusticata-macros"
+version = "4.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "faf0c4a6ece9950b9abdb62b1cfcf2a68b3b67a10ba445b3bb85be2a293d0632"
+dependencies = [
+ "nom",
+]
+
 [[package]]
 name = "rustix"
 version = "0.37.28"
@@ -2972,7 +3115,7 @@ dependencies = [
  "aws-lc-rs",
  "ring",
  "rustls-pki-types",
- "untrusted",
+ "untrusted 0.9.0",
 ]
 
 [[package]]
@@ -3890,6 +4033,12 @@ dependencies = [
  "subtle",
 ]
 
+[[package]]
+name = "untrusted"
+version = "0.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a156c684c91ea7d62626509bce3cb4e1d9ed5c4d978f7b4352658f96a4c26b4a"
+
 [[package]]
 name = "untrusted"
 version = "0.9.0"
@@ -3926,7 +4075,7 @@ version = "1.21.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b672338555252d43fd2240c714dc444b8c6fb0a5c5335e65a07bba7742735ddb"
 dependencies = [
- "getrandom 0.4.1",
+ "getrandom 0.4.2",
  "js-sys",
  "wasm-bindgen",
 ]
@@ -4701,6 +4850,42 @@ version = "0.6.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9"
 
+[[package]]
+name = "wtransport"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e56b611f195638f3790e4e5a41e9d777643a6c324ae4ffd1f0f53f2738e7678c"
+dependencies = [
+ "bytes",
+ "pem",
+ "quinn",
+ "rcgen",
+ "rustls",
+ "rustls-native-certs",
+ "rustls-pki-types",
+ "sha2",
+ "socket2 0.5.10",
+ "thiserror 2.0.18",
+ "time",
+ "tokio",
+ "tracing",
+ "url",
+ "wtransport-proto",
+ "x509-parser 0.17.0",
+]
+
+[[package]]
+name = "wtransport-proto"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1627c5b59450278e9771aab35275d72bfa2788128c197c5af1e4a820c8737ef4"
+dependencies = [
+ "httlib-huffman",
+ "octets",
+ "thiserror 2.0.18",
+ "url",
+]
+
 [[package]]
 name = "x448"
 version = "0.6.0"
@@ -4712,6 +4897,50 @@ dependencies = [
  "rand_core 0.5.1",
 ]
 
+[[package]]
+name = "x509-parser"
+version = "0.17.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4569f339c0c402346d4a75a9e39cf8dad310e287eef1ff56d4c68e5067f53460"
+dependencies = [
+ "asn1-rs",
+ "data-encoding",
+ "der-parser",
+ "lazy_static",
+ "nom",
+ "oid-registry",
+ "rusticata-macros",
+ "thiserror 2.0.18",
+ "time",
+]
+
+[[package]]
+name = "x509-parser"
+version = "0.18.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d43b0f71ce057da06bc0851b23ee24f3f86190b07203dd8f567d0b706a185202"
+dependencies = [
+ "asn1-rs",
+ "data-encoding",
+ "der-parser",
+ "lazy_static",
+ "nom",
+ "oid-registry",
+ "ring",
+ "rusticata-macros",
+ "thiserror 2.0.18",
+ "time",
+]
+
+[[package]]
+name = "yasna"
+version = "0.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e17bb3549cc1321ae1296b9cdc2698e2b6cb1992adfa19a8c72e5b7a738f44cd"
+dependencies = [
+ "time",
+]
+
 [[package]]
 name = "yoke"
 version = "0.8.1"

src/main.rs

@@ -28,7 +28,9 @@ async fn main() {
     log_in!("Incoming messages");
     log_out!("Outgoing messages");
 
-    let _ = omikron_connection::start(9187).await;
+    tokio::spawn(async move {
+        let _ = omikron_connection::start(9187).await;
+    });
 
     log!("Started");
     log!("  .env");

src/server/omikron_connection.rs

@@ -16,11 +16,12 @@ use base64::{Engine as _, engine::general_purpose::STANDARD};
 use dashmap::DashMap;
 use epsilon_core::{CommunicationType, CommunicationValue, DataTypes, DataValue};
 use epsilon_native::{Host, Receiver, Sender};
-use quinn::ServerConfig;
 use quinn::crypto::rustls::QuicServerConfig;
+use quinn::{Endpoint, ServerConfig};
 use rand::{Rng, distributions::Alphanumeric};
 use rustls::{ServerConfig as CryptoConfig, crypto::aws_lc_rs};
 use std::{
+    net::SocketAddr,
     sync::Arc,
     time::{Duration, Instant},
 };
@@ -105,18 +106,10 @@ pub struct OmikronConnection {
     id: u64,
     sender: Mutex<Option<Sender>>,
     state: RwLock<AuthState>,
-
-    // Authentication state (preserved from original)
     challenge: RwLock<String>,
     pub_key: RwLock<Option<Vec<u8>>>,
-
-    // Ping tracking
     pub ping: RwLock<i64>,
-
-    // Waiting tasks for request/response pattern
     waiting_tasks: DashMap<u32, WaitingTask>,
-
-    // Cleanup handle
     cleanup_handle: Mutex<Option<tokio::task::JoinHandle<()>>>,
 }
 
@@ -189,7 +182,7 @@ impl OmikronConnection {
         }
 
         // Connection closed
-        self.cleanup().await;
+        self.clone().cleanup().await;
         log_in!(
             self.id as i64,
             PrintType::Omega,
@@ -220,10 +213,10 @@ impl OmikronConnection {
 
         // Route based on authentication state
         match *self.state.read().await {
-            AuthState::Unauthenticated => self.handle_unauthenticated(cv).await,
+            AuthState::Unauthenticated => self.clone().handle_unauthenticated(cv).await,
-            AuthState::Identified { .. } => self.handle_identified(cv).await,
+            AuthState::Identified { .. } => self.clone().handle_identified(cv).await,
             AuthState::Authenticated { omikron_id } => {
-                self.handle_authenticated(cv, omikron_id).await
+                self.clone().handle_authenticated(cv, omikron_id).await
             }
         }
     }
@@ -232,7 +225,7 @@ impl OmikronConnection {
     // Authentication Handlers
     // -------------------------------------------------------------------------
 
-    async fn handle_unauthenticated(&self, cv: CommunicationValue) -> OmikronResult<()> {
+    async fn handle_unauthenticated(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
         if !cv.is_type(CommunicationType::identification) {
             self.send_error_response(cv.get_id(), CommunicationType::error_not_authenticated)
                 .await;
@@ -285,7 +278,7 @@ impl OmikronConnection {
         self.send(&response).await
     }
 
-    async fn handle_identified(&self, cv: CommunicationValue) -> OmikronResult<()> {
+    async fn handle_identified(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
         if !cv.is_type(CommunicationType::challenge_response) {
             self.send_error_response(cv.get_id(), CommunicationType::error_not_authenticated)
                 .await;
@@ -300,19 +293,16 @@ impl OmikronConnection {
         let expected_challenge = self.challenge.read().await.clone();
 
         if client_response == expected_challenge {
-            // Challenge passed - mark as authenticated
             let omikron_id = self.state.read().await.omikron_id().unwrap_or(0);
             *self.state.write().await = AuthState::Authenticated { omikron_id };
 
-            // Register with manager
+            omikron_manager::add_omikron(self.clone()).await;
-            omikron_manager::add_omikron(self.arc_self()).await;
 
-            // Send success response
             let response = CommunicationValue::new(CommunicationType::identification_response)
                 .with_id(cv.get_id())
                 .add_data(DataTypes::accepted, DataValue::Bool(true));
 
-            self.send(&response).await?;
+            self.clone().send(&response).await?;
             log_in!(omikron_id, PrintType::Omega, "Omikron authenticated");
             Ok(())
         } else {
@@ -327,7 +317,7 @@ impl OmikronConnection {
     // -------------------------------------------------------------------------
 
     async fn handle_authenticated(
-        &self,
+        self: Arc<Self>,
         cv: CommunicationValue,
         omikron_id: i64,
     ) -> OmikronResult<()> {
@@ -357,11 +347,9 @@ impl OmikronConnection {
                 Ok(())
             }
 
-            // Data queries
             CommunicationType::get_user_data => self.handle_get_user_data(cv).await,
             CommunicationType::get_iota_data => self.handle_get_iota_data(cv).await,
 
-            // Registration
             CommunicationType::get_register => self.handle_get_register(cv).await,
             CommunicationType::complete_register_iota => {
                 self.handle_complete_register_iota(cv).await
@@ -370,13 +358,11 @@ impl OmikronConnection {
                 self.handle_complete_register_user(cv).await
             }
 
-            // Data modification
             CommunicationType::change_user_data => self.handle_change_user_data(cv).await,
             CommunicationType::change_iota_data => self.handle_change_iota_data(cv).await,
             CommunicationType::delete_user => self.handle_delete_user(cv).await,
             CommunicationType::delete_iota => self.handle_delete_iota(cv).await,
 
-            // Notifications
             CommunicationType::get_notifications => self.handle_get_notifications(cv).await,
             CommunicationType::read_notification => self.handle_read_notification(cv).await,
             CommunicationType::push_notification => self.handle_push_notification(cv).await,
@@ -397,7 +383,7 @@ impl OmikronConnection {
     // Specific Handlers (ported from original WebSocket implementation)
     // -------------------------------------------------------------------------
 
-    async fn handle_shorten_link(&self, cv: CommunicationValue) -> OmikronResult<()> {
+    async fn handle_shorten_link(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
         let link = cv
             .get_data(DataTypes::link)
             .as_str()
@@ -414,7 +400,7 @@ impl OmikronConnection {
         self.send(&response).await
     }
 
-    async fn handle_user_connected(&self, cv: CommunicationValue, omikron_id: i64) {
+    async fn handle_user_connected(self: Arc<Self>, cv: CommunicationValue, omikron_id: i64) {
         log_in!(PrintType::Omega, "User connected");
         if let Some(user_id) = cv.get_data(DataTypes::user_id).as_number() {
             user_online_tracker::track_user_status(
@@ -425,7 +411,7 @@ impl OmikronConnection {
         }
     }
 
-    async fn handle_user_disconnected(&self, cv: CommunicationValue, omikron_id: i64) {
+    async fn handle_user_disconnected(self: Arc<Self>, cv: CommunicationValue, _omikron_id: i64) {
         log_in!(PrintType::Omega, "User disconnected");
         if let Some(user_id) = cv.get_data(DataTypes::user_id).as_number() {
             if let Some(status) = user_online_tracker::get_user_status(user_id as i64) {
@@ -438,7 +424,7 @@ impl OmikronConnection {
         }
     }
 
-    async fn handle_iota_connected(&self, cv: CommunicationValue, omikron_id: i64) {
+    async fn handle_iota_connected(self: Arc<Self>, cv: CommunicationValue, omikron_id: i64) {
         log_in!(PrintType::Omega, "IOTA connected");
         if let Some(iota_id) = cv.get_data(DataTypes::iota_id).as_number() {
             let iota_id = iota_id as i64;
@@ -468,7 +454,7 @@ impl OmikronConnection {
         }
     }
 
-    async fn handle_iota_disconnected(&self, cv: CommunicationValue, omikron_id: i64) {
+    async fn handle_iota_disconnected(self: Arc<Self>, cv: CommunicationValue, omikron_id: i64) {
         log_in!(PrintType::Omega, "IOTA disconnected");
         if let Some(iota_id) = cv.get_data(DataTypes::iota_id).as_number() {
             let iota_id = iota_id as i64;
@@ -482,7 +468,7 @@ impl OmikronConnection {
         }
     }
 
-    async fn handle_sync_status(&self, cv: CommunicationValue, omikron_id: i64) {
+    async fn handle_sync_status(self: Arc<Self>, cv: CommunicationValue, omikron_id: i64) {
         if let DataValue::Array(user_ids) = cv.get_data(DataTypes::user_ids) {
             for user_id_val in user_ids {
                 if let DataValue::Number(user_id) = user_id_val {
@@ -504,11 +490,14 @@ impl OmikronConnection {
         }
     }
 
-    async fn handle_get_user_data(&self, cv: CommunicationValue) -> OmikronResult<()> {
+    async fn handle_get_user_data(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
         // Try by user_id first
         if let Some(user_id) = cv.get_data(DataTypes::user_id).as_number() {
             if let Ok(user_data) = get_by_user_id(user_id as i64).await {
-                let response = self.build_user_data_response(cv.get_id(), user_data).await;
+                let response = self
+                    .clone()
+                    .build_user_data_response(cv.get_id(), user_data)
+                    .await;
                 return self.send(&response).await;
             }
         }
@@ -516,7 +505,10 @@ impl OmikronConnection {
         // Try by username
         if let Some(username) = cv.get_data(DataTypes::username).as_str() {
             if let Ok(user_data) = get_by_username(username).await {
-                let response = self.build_user_data_response(cv.get_id(), user_data).await;
+                let response = self
+                    .clone()
+                    .build_user_data_response(cv.get_id(), user_data)
+                    .await;
                 return self.send(&response).await;
             }
         }
@@ -528,7 +520,7 @@ impl OmikronConnection {
     }
 
     async fn build_user_data_response(
-        &self,
+        self: Arc<Self>,
         msg_id: u32,
         user: (
             i64,
@@ -615,11 +607,12 @@ impl OmikronConnection {
         response
     }
 
-    async fn handle_get_iota_data(&self, cv: CommunicationValue) -> OmikronResult<()> {
+    async fn handle_get_iota_data(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
         // Try by iota_id
         if let Some(iota_id) = cv.get_data(DataTypes::iota_id).as_number() {
             if let Ok((iota_id, public_key)) = get_iota_by_id(iota_id as i64).await {
                 let response = self
+                    .clone()
                     .build_iota_data_response(cv.get_id(), iota_id, public_key, None, None)
                     .await;
                 return self.send(&response).await;
@@ -633,6 +626,7 @@ impl OmikronConnection {
             {
                 if let Ok((iota_id, public_key)) = get_iota_by_id(iota_id).await {
                     let response = self
+                        .clone()
                         .build_iota_data_response(
                             cv.get_id(),
                             iota_id,
@@ -653,6 +647,7 @@ impl OmikronConnection {
             {
                 if let Ok((iota_id, public_key)) = get_iota_by_id(iota_id).await {
                     let response = self
+                        .clone()
                         .build_iota_data_response(
                             cv.get_id(),
                             iota_id,
@@ -672,7 +667,7 @@ impl OmikronConnection {
     }
 
     async fn build_iota_data_response(
-        &self,
+        self: Arc<Self>,
         msg_id: u32,
         iota_id: i64,
         public_key: String,
@@ -705,7 +700,7 @@ impl OmikronConnection {
         )
     }
 
-    async fn handle_get_register(&self, cv: CommunicationValue) -> OmikronResult<()> {
+    async fn handle_get_register(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
         let register_id = sql::get_register_id().await;
         let response = CommunicationValue::new(CommunicationType::get_register)
             .with_id(cv.get_id())
@@ -713,7 +708,10 @@ impl OmikronConnection {
         self.send(&response).await
     }
 
-    async fn handle_complete_register_iota(&self, cv: CommunicationValue) -> OmikronResult<()> {
+    async fn handle_complete_register_iota(
+        self: Arc<Self>,
+        cv: CommunicationValue,
+    ) -> OmikronResult<()> {
         let iota_id_opt = cv
             .get_data(DataTypes::iota_id)
             .as_number()
@@ -759,7 +757,10 @@ impl OmikronConnection {
         }
     }
 
-    async fn handle_complete_register_user(&self, cv: CommunicationValue) -> OmikronResult<()> {
+    async fn handle_complete_register_user(
+        self: Arc<Self>,
+        cv: CommunicationValue,
+    ) -> OmikronResult<()> {
         let user_id = cv
             .get_data(DataTypes::user_id)
             .as_number()
@@ -803,7 +804,7 @@ impl OmikronConnection {
         }
     }
 
-    async fn handle_change_user_data(&self, cv: CommunicationValue) -> OmikronResult<()> {
+    async fn handle_change_user_data(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
         let user_id = cv.get_sender() as i64;
         let mut success = true;
         let mut error_message = String::new();
@@ -866,7 +867,7 @@ impl OmikronConnection {
         }
     }
 
-    async fn handle_change_iota_data(&self, cv: CommunicationValue) -> OmikronResult<()> {
+    async fn handle_change_iota_data(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
         let user_id = cv.get_sender() as i64;
 
         if let (Some(iota_id), Some(reset_token), Some(new_token)) = (
@@ -924,7 +925,7 @@ impl OmikronConnection {
         }
     }
 
-    async fn handle_delete_user(&self, cv: CommunicationValue) -> OmikronResult<()> {
+    async fn handle_delete_user(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
         let user_id = cv.get_sender() as i64;
         match sql::delete_user(user_id).await {
             Ok(_) => {
@@ -941,7 +942,7 @@ impl OmikronConnection {
         }
     }
 
-    async fn handle_delete_iota(&self, cv: CommunicationValue) -> OmikronResult<()> {
+    async fn handle_delete_iota(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
         if let Some(iota_id) = cv
             .get_data(DataTypes::iota_id)
             .as_number()
@@ -966,7 +967,10 @@ impl OmikronConnection {
         }
     }
 
-    async fn handle_get_notifications(&self, cv: CommunicationValue) -> OmikronResult<()> {
+    async fn handle_get_notifications(
+        self: Arc<Self>,
+        cv: CommunicationValue,
+    ) -> OmikronResult<()> {
         let user_id = cv.get_sender() as i64;
         if let Ok(notifications) = sql::get_notifications(user_id).await {
             let json_array: Vec<DataValue> = notifications
@@ -988,7 +992,10 @@ impl OmikronConnection {
         }
     }
 
-    async fn handle_read_notification(&self, cv: CommunicationValue) -> OmikronResult<()> {
+    async fn handle_read_notification(
+        self: Arc<Self>,
+        cv: CommunicationValue,
+    ) -> OmikronResult<()> {
         let user_id = cv.get_sender() as i64;
         if let Some(other_id) = cv
             .get_data(DataTypes::sender_id)
@@ -1007,7 +1014,10 @@ impl OmikronConnection {
         }
     }
 
-    async fn handle_push_notification(&self, cv: CommunicationValue) -> OmikronResult<()> {
+    async fn handle_push_notification(
+        self: Arc<Self>,
+        cv: CommunicationValue,
+    ) -> OmikronResult<()> {
         let user_id = cv.get_sender() as i64;
         if let Some(other_id) = cv
             .get_data(DataTypes::sender_id)
@@ -1026,7 +1036,7 @@ impl OmikronConnection {
         }
     }
 
-    async fn handle_ping(&self, cv: CommunicationValue) -> OmikronResult<()> {
+    async fn handle_ping(self: Arc<Self>, cv: CommunicationValue) -> OmikronResult<()> {
         if let DataValue::Number(last_ping) = cv.get_data(DataTypes::last_ping) {
             *self.ping.write().await = *last_ping;
         }
@@ -1039,7 +1049,7 @@ impl OmikronConnection {
     // Utilities
     // -------------------------------------------------------------------------
 
-    async fn send(&self, cv: &CommunicationValue) -> OmikronResult<()> {
+    async fn send(self: Arc<Self>, cv: &CommunicationValue) -> OmikronResult<()> {
         log_cv_out!(PrintType::Omikron, cv);
 
         let guard = self.sender.lock().await;
@@ -1052,7 +1062,7 @@ impl OmikronConnection {
     }
 
     async fn send_error_response(
-        &self,
+        self: Arc<Self>,
         message_id: u32,
         error_type: CommunicationType,
     ) -> OmikronResult<()> {
@@ -1060,9 +1070,9 @@ impl OmikronConnection {
         self.send(&error).await
     }
 
-    pub async fn close(&self) {}
+    pub async fn close(self: Arc<Self>) {}
 
-    async fn cleanup(&self) {
+    async fn cleanup(self: Arc<Self>) {
         if let Some(omikron_id) = self.state.read().await.omikron_id() {
             if omikron_id != 0 {
                 log_in!(omikron_id, PrintType::Omega, "Omikron disconnected");
@@ -1077,22 +1087,22 @@ impl OmikronConnection {
         }
     }
 
-    fn arc_self(&self) -> Arc<Self> {
+    fn arc_self(self: Arc<Self>) -> Arc<Self> {
         // This is a bit of a hack - in practice you'd store the Arc in the struct
         // or use weak references. For now, we rely on the caller having the Arc.
         panic!("Use the Arc<OmikronConnection> directly")
     }
 
     // Public API for external use
-    pub async fn is_authenticated(&self) -> bool {
+    pub async fn is_authenticated(self: Arc<Self>) -> bool {
         self.state.read().await.is_authenticated()
     }
 
-    pub async fn get_omikron_id(&self) -> Option<i64> {
+    pub async fn get_omikron_id(self: Arc<Self>) -> Option<i64> {
         self.state.read().await.omikron_id()
     }
 
-    pub async fn send_message(&self, cv: &CommunicationValue) -> OmikronResult<()> {
+    pub async fn send_message(self: Arc<Self>, cv: &CommunicationValue) -> OmikronResult<()> {
         self.send(cv).await
     }
 }

src/server/omikron_manager.rs

@@ -8,7 +8,7 @@ pub static OMIKRON_CONNECTIONS: Lazy<DashMap<i64, Arc<OmikronConnection>>> =
     Lazy::new(|| DashMap::new());
 
 pub async fn add_omikron(conn: Arc<OmikronConnection>) {
-    let id = match conn.get_omikron_id().await {
+    let id = match conn.clone().get_omikron_id().await {
         Some(id) => id,
         _ => {
             conn.close().await;