From cf38693810ad2bb462003a468ebcf1c58407f570 Mon Sep 17 00:00:00 2001 From: Alex Emmet <111742636+Alex-Emmet@users.noreply.github.com> Date: Tue, 3 Mar 2026 16:30:35 +0100 Subject: [PATCH] [WIP] QUIC connection --- Cargo.lock | 599 +------------ Cargo.toml | 43 +- src/main.rs | 6 +- src/server/api.rs | 22 +- src/server/omikron_connection.rs | 1370 +++++++++++++++++++++++------- src/server/omikron_manager.rs | 8 +- 6 files changed, 1123 insertions(+), 925 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 40c314f..833d221 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -14,9 +14,7 @@ dependencies = [ "ansi_term", "anyhow", "async-trait", - "async-tungstenite", - "axum", - "base64 0.22.1", + "base64", "bytes", "color-eyre", "dashmap", @@ -28,9 +26,6 @@ dependencies = [ "hex", "hkdf", "http-body-util", - "hyper", - "hyper-rustls", - "hyper-util", "json", "once_cell", "pnet", @@ -38,14 +33,15 @@ dependencies = [ "rand 0.8.5", "rand_core 0.6.4", "reqwest", - "rustls 0.23.37", - "rustls-pemfile 2.2.0", + "rustls", + "rustls-pemfile", "sha1", "sha2", "sqlx", "strum 0.27.2", "strum_macros 0.27.2", "sysinfo", + "thiserror 2.0.18", "tokio", "tokio-rustls", "tokio-tungstenite", @@ -111,7 +107,7 @@ dependencies = [ "actix-service", "actix-tls", "actix-utils", - "base64 0.22.1", + "base64", "bitflags 2.11.0", "brotli", "bytes", @@ -531,18 +527,6 @@ dependencies = [ "pin-project-lite", ] -[[package]] -name = "async-native-tls" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9343dc5acf07e79ff82d0c37899f079db3534d99f189a1837c8e549c99405bec" -dependencies = [ - "futures-util", - "native-tls", - "thiserror 1.0.69", - "url", -] - [[package]] name = "async-std" version = "1.13.2" @@ -575,19 +559,6 @@ version = "4.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b75356056920673b02621b35afd0f7dda9306d03c79a30f5c56c44cf256e3de" -[[package]] -name = "async-tls" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b2ae3c9eba89d472a0e4fe1dea433df78fbbe63d2b764addaf2ba3a6bde89a5e" -dependencies = [ - "futures-core", - "futures-io", - "rustls 0.21.12", - "rustls-pemfile 1.0.4", - "webpki-roots 0.22.6", -] - [[package]] name = "async-trait" version = "0.1.89" @@ -599,36 +570,6 @@ dependencies = [ "syn", ] -[[package]] -name = "async-tungstenite" -version = "0.32.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8acc405d38be14342132609f06f02acaf825ddccfe76c4824a69281e0458ebd4" -dependencies = [ - "async-native-tls", - "async-std", - "async-tls", - "atomic-waker", - "futures-core", - "futures-io", - "futures-task", - "futures-util", - "gio", - "glib", - "log", - "native-tls", - "openssl", - "pin-project-lite", - "rustls-native-certs", - "rustls-pki-types", - "tokio", - "tokio-native-tls", - "tokio-openssl", - "tokio-rustls", - "tungstenite", - "webpki-roots 1.0.6", -] - [[package]] name = "atoi" version = "2.0.0" @@ -652,9 +593,9 @@ checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" [[package]] name = "aws-lc-rs" -version = "1.16.0" +version = "1.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d9a7b350e3bb1767102698302bc37256cbd48422809984b98d292c40e2579aa9" +checksum = "94bffc006df10ac2a68c83692d734a465f8ee6c5b384d8545a636f81d858f4bf" dependencies = [ "aws-lc-sys", "zeroize", @@ -662,9 +603,9 @@ dependencies = [ [[package]] name = "aws-lc-sys" -version = "0.37.1" +version = "0.38.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b092fe214090261288111db7a2b2c2118e5a7f30dc2569f1732c4069a6840549" +checksum = "4321e568ed89bb5a7d291a7f37997c2c0df89809d7b6d12062c81ddb54aa782e" dependencies = [ "cc", "cmake", @@ -672,61 +613,6 @@ dependencies = [ "fs_extra", ] -[[package]] -name = "axum" -version = "0.8.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b52af3cb4058c895d37317bb27508dccc8e5f2d39454016b297bf4a400597b8" -dependencies = [ - "axum-core", - "base64 0.22.1", - "bytes", - "form_urlencoded", - "futures-util", - "http 1.4.0", - "http-body", - "http-body-util", - "hyper", - "hyper-util", - "itoa", - "matchit", - "memchr", - "mime", - "percent-encoding", - "pin-project-lite", - "serde_core", - "serde_json", - "serde_path_to_error", - "serde_urlencoded", - "sha1", - "sync_wrapper", - "tokio", - "tokio-tungstenite", - "tower", - "tower-layer", - "tower-service", - "tracing", -] - -[[package]] -name = "axum-core" -version = "0.5.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" -dependencies = [ - "bytes", - "futures-core", - "http 1.4.0", - "http-body", - "http-body-util", - "mime", - "pin-project-lite", - "sync_wrapper", - "tower-layer", - "tower-service", - "tracing", -] - [[package]] name = "backtrace" version = "0.3.76" @@ -742,12 +628,6 @@ dependencies = [ "windows-link 0.2.1", ] -[[package]] -name = "base64" -version = "0.21.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567" - [[package]] name = "base64" version = "0.22.1" @@ -872,16 +752,6 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6d43a04d8753f35258c91f8ec639f792891f748a1edbd759cf1dcea3382ad83c" -[[package]] -name = "cfg-expr" -version = "0.20.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78cef5b5a1a6827c7322ae2a636368a573006b27cfa76c7ebd53e834daeaab6a" -dependencies = [ - "smallvec", - "target-lexicon", -] - [[package]] name = "cfg-if" version = "1.0.4" @@ -1247,9 +1117,10 @@ dependencies = [ [[package]] name = "epsilon-core" version = "0.1.0" -source = "git+https://github.com/Tensamin/Epsilon.git#dc17684bc9308c79559e0d4648c6e730cc567c89" +source = "git+https://github.com/Tensamin/Epsilon.git#eb91e812e559270ff22572542491f8215db27a76" dependencies = [ "byteorder", + "quinn", "rand 0.8.5", "strum 0.28.0", "strum_macros 0.28.0", @@ -1258,14 +1129,14 @@ dependencies = [ [[package]] name = "epsilon-native" version = "0.1.0" -source = "git+https://github.com/Tensamin/Epsilon.git#dc17684bc9308c79559e0d4648c6e730cc567c89" +source = "git+https://github.com/Tensamin/Epsilon.git#eb91e812e559270ff22572542491f8215db27a76" dependencies = [ "anyhow", "async-trait", "bytes", "epsilon-core", "quinn", - "rustls 0.23.37", + "rustls", "thiserror 2.0.18", "tokio", ] @@ -1407,21 +1278,6 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" -[[package]] -name = "foreign-types" -version = "0.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1" -dependencies = [ - "foreign-types-shared", -] - -[[package]] -name = "foreign-types-shared" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b" - [[package]] name = "form_urlencoded" version = "1.2.2" @@ -1630,80 +1486,6 @@ version = "0.32.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e629b9b98ef3dd8afe6ca2bd0f89306cec16d43d907889945bc5d6687f2f13c7" -[[package]] -name = "gio" -version = "0.21.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c5ff48bf600c68b476e61dc6b7c762f2f4eb91deef66583ba8bb815c30b5811a" -dependencies = [ - "futures-channel", - "futures-core", - "futures-io", - "futures-util", - "gio-sys", - "glib", - "libc", - "pin-project-lite", - "smallvec", -] - -[[package]] -name = "gio-sys" -version = "0.21.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0071fe88dba8e40086c8ff9bbb62622999f49628344b1d1bf490a48a29d80f22" -dependencies = [ - "glib-sys", - "gobject-sys", - "libc", - "system-deps", - "windows-sys 0.61.2", -] - -[[package]] -name = "glib" -version = "0.21.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16de123c2e6c90ce3b573b7330de19be649080ec612033d397d72da265f1bd8b" -dependencies = [ - "bitflags 2.11.0", - "futures-channel", - "futures-core", - "futures-executor", - "futures-task", - "futures-util", - "gio-sys", - "glib-macros", - "glib-sys", - "gobject-sys", - "libc", - "memchr", - "smallvec", -] - -[[package]] -name = "glib-macros" -version = "0.21.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf59b675301228a696fe01c3073974643365080a76cc3ed5bc2cbc466ad87f17" -dependencies = [ - "heck", - "proc-macro-crate", - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "glib-sys" -version = "0.21.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d95e1a3a19ae464a7286e14af9a90683c64d70c02532d88d87ce95056af3e6c" -dependencies = [ - "libc", - "system-deps", -] - [[package]] name = "glob" version = "0.3.3" @@ -1722,17 +1504,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "gobject-sys" -version = "0.21.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2dca35da0d19a18f4575f3cb99fe1c9e029a2941af5662f326f738a21edaf294" -dependencies = [ - "glib-sys", - "libc", - "system-deps", -] - [[package]] name = "h2" version = "0.3.27" @@ -1924,7 +1695,6 @@ dependencies = [ "http 1.4.0", "http-body", "httparse", - "httpdate", "itoa", "pin-project-lite", "pin-utils", @@ -1942,9 +1712,7 @@ dependencies = [ "http 1.4.0", "hyper", "hyper-util", - "log", - "rustls 0.23.37", - "rustls-native-certs", + "rustls", "rustls-pki-types", "tokio", "tokio-rustls", @@ -1957,7 +1725,7 @@ version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ - "base64 0.22.1", + "base64", "bytes", "futures-channel", "futures-util", @@ -1971,7 +1739,6 @@ dependencies = [ "socket2 0.6.2", "system-configuration", "tokio", - "tower-layer", "tower-service", "tracing", "windows-registry", @@ -2140,9 +1907,9 @@ dependencies = [ [[package]] name = "ipnet" -version = "2.11.0" +version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "469fb0b9cefa57e3ef31275ee7cacb78f2fdca44e4765491884a2b119d4eb130" +checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" [[package]] name = "ipnetwork" @@ -2267,12 +2034,13 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "libredox" -version = "0.1.12" +version = "0.1.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d0b95e02c851351f877147b7deea7b1afb1df71b63aa5f8270716e0c5720616" +checksum = "1744e39d1d6a9948f4f388969627434e31128196de472883b39f148769bfe30a" dependencies = [ "bitflags 2.11.0", "libc", + "plain", "redox_syscall 0.7.3", ] @@ -2355,12 +2123,6 @@ dependencies = [ "sha2", ] -[[package]] -name = "matchit" -version = "0.8.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" - [[package]] name = "md-5" version = "0.10.6" @@ -2405,23 +2167,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "native-tls" -version = "0.2.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "465500e14ea162429d264d44189adc38b199b62b1c21eea9f69e4b73cb03bbf2" -dependencies = [ - "libc", - "log", - "openssl", - "openssl-probe", - "openssl-sys", - "schannel", - "security-framework", - "security-framework-sys", - "tempfile", -] - [[package]] name = "no-std-net" version = "0.6.0" @@ -2529,50 +2274,12 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" -[[package]] -name = "openssl" -version = "0.10.75" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08838db121398ad17ab8531ce9de97b244589089e290a384c900cb9ff7434328" -dependencies = [ - "bitflags 2.11.0", - "cfg-if", - "foreign-types", - "libc", - "once_cell", - "openssl-macros", - "openssl-sys", -] - -[[package]] -name = "openssl-macros" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - [[package]] name = "openssl-probe" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" -[[package]] -name = "openssl-sys" -version = "0.9.111" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "82cab2d520aa75e3c58898289429321eb788c3106963d0dc886ec7a5f4adc321" -dependencies = [ - "cc", - "libc", - "pkg-config", - "vcpkg", -] - [[package]] name = "owo-colors" version = "4.3.0" @@ -2647,9 +2354,9 @@ checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" [[package]] name = "piper" -version = "0.2.4" +version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96c8c490f422ef9a4efd2cb5b42b76c8613d7e7dfc1caf667b8a3350a5acc066" +checksum = "c835479a4443ded371d6c535cbfd8d31ad92c5d23ae9770a61bc155e4992a3c1" dependencies = [ "atomic-waker", "fastrand 2.3.0", @@ -2683,6 +2390,12 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" +[[package]] +name = "plain" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4596b6d070b27117e987119b4dac604f3c58cfb0b191112e24771b2faeac1a6" + [[package]] name = "pnet" version = "0.35.0" @@ -2856,15 +2569,6 @@ dependencies = [ "syn", ] -[[package]] -name = "proc-macro-crate" -version = "3.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "219cb19e96be00ab2e37d6e299658a0cfa83e52429179969b0f0121b4ac46983" -dependencies = [ - "toml_edit", -] - [[package]] name = "proc-macro2" version = "1.0.106" @@ -2886,7 +2590,7 @@ dependencies = [ "quinn-proto", "quinn-udp", "rustc-hash", - "rustls 0.23.37", + "rustls", "socket2 0.6.2", "thiserror 2.0.18", "tokio", @@ -2908,7 +2612,7 @@ dependencies = [ "rand 0.9.2", "ring", "rustc-hash", - "rustls 0.23.37", + "rustls", "rustls-pki-types", "rustls-platform-verifier", "slab", @@ -3071,7 +2775,7 @@ version = "0.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ab3f43e3283ab1488b624b44b0e988d0acea0b3214e694730a055cb6b2efa801" dependencies = [ - "base64 0.22.1", + "base64", "bytes", "encoding_rs", "futures-core", @@ -3088,7 +2792,7 @@ dependencies = [ "percent-encoding", "pin-project-lite", "quinn", - "rustls 0.23.37", + "rustls", "rustls-pki-types", "rustls-platform-verifier", "sync_wrapper", @@ -3185,18 +2889,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "rustls" -version = "0.21.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f56a14d1f48b391359b22f731fd4bd7e43c97f3c50eee276f3aa09c94784d3e" -dependencies = [ - "log", - "ring", - "rustls-webpki 0.101.7", - "sct", -] - [[package]] name = "rustls" version = "0.23.37" @@ -3208,7 +2900,7 @@ dependencies = [ "once_cell", "ring", "rustls-pki-types", - "rustls-webpki 0.103.9", + "rustls-webpki", "subtle", "zeroize", ] @@ -3225,15 +2917,6 @@ dependencies = [ "security-framework", ] -[[package]] -name = "rustls-pemfile" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1c74cae0a4cf6ccbbf5f359f08efdf8ee7e1dc532573bf0db71968cb56b1448c" -dependencies = [ - "base64 0.21.7", -] - [[package]] name = "rustls-pemfile" version = "2.2.0" @@ -3264,10 +2947,10 @@ dependencies = [ "jni", "log", "once_cell", - "rustls 0.23.37", + "rustls", "rustls-native-certs", "rustls-platform-verifier-android", - "rustls-webpki 0.103.9", + "rustls-webpki", "security-framework", "security-framework-sys", "webpki-root-certs", @@ -3280,16 +2963,6 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" -[[package]] -name = "rustls-webpki" -version = "0.101.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b6275d1ee7a1cd780b64aca7726599a1dbc893b1e64144529e55c3c2f745765" -dependencies = [ - "ring", - "untrusted", -] - [[package]] name = "rustls-webpki" version = "0.103.9" @@ -3338,16 +3011,6 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" -[[package]] -name = "sct" -version = "0.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da046153aa2352493d6cb7da4b6e5c0c057d8a1d0a9aa8560baffdd945acd414" -dependencies = [ - "ring", - "untrusted", -] - [[package]] name = "security-framework" version = "3.7.0" @@ -3420,26 +3083,6 @@ dependencies = [ "zmij", ] -[[package]] -name = "serde_path_to_error" -version = "0.1.20" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" -dependencies = [ - "itoa", - "serde", - "serde_core", -] - -[[package]] -name = "serde_spanned" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8bbf91e5a4d6315eee45e704372590b30e260ee83af6639d64557f51b067776" -dependencies = [ - "serde_core", -] - [[package]] name = "serde_urlencoded" version = "0.7.1" @@ -3606,7 +3249,7 @@ checksum = "ee6798b1838b6a0f69c007c133b8df5866302197e404e8b6ee8ed3e3a5e68dc6" dependencies = [ "async-io 1.13.0", "async-std", - "base64 0.22.1", + "base64", "bytes", "crc", "crossbeam-queue", @@ -3677,7 +3320,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "aa003f0038df784eb8fecbbac13affe3da23b45194bd57dba231c8f48199c526" dependencies = [ "atoi", - "base64 0.22.1", + "base64", "bitflags 2.11.0", "byteorder", "bytes", @@ -3719,7 +3362,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db58fcd5a53cf07c184b154801ff91347e4c30d17a3562a635ff028ad5deda46" dependencies = [ "atoi", - "base64 0.22.1", + "base64", "bitflags 2.11.0", "byteorder", "crc", @@ -3898,38 +3541,6 @@ dependencies = [ "libc", ] -[[package]] -name = "system-deps" -version = "7.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48c8f33736f986f16d69b6cb8b03f55ddcad5c41acc4ccc39dd88e84aa805e7f" -dependencies = [ - "cfg-expr", - "heck", - "pkg-config", - "toml", - "version-compare", -] - -[[package]] -name = "target-lexicon" -version = "0.13.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df7f62577c25e07834649fc3b39fafdc597c0a3527dc1c60129201ccfcbaa50c" - -[[package]] -name = "tempfile" -version = "3.26.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "82a72c767771b47409d2345987fda8628641887d5466101319899796367354a0" -dependencies = [ - "fastrand 2.3.0", - "getrandom 0.4.1", - "once_cell", - "rustix 1.1.4", - "windows-sys 0.61.2", -] - [[package]] name = "thiserror" version = "1.0.69" @@ -4037,9 +3648,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.49.0" +version = "1.50.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72a2903cd7736441aac9df9d7688bd0ce48edccaadf181c3b90be801e81d3d86" +checksum = "27ad5e34374e03cfffefc301becb44e9dc3c17584f414349ebe29ed26661822d" dependencies = [ "bytes", "libc", @@ -4054,43 +3665,22 @@ dependencies = [ [[package]] name = "tokio-macros" -version = "2.6.0" +version = "2.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af407857209536a95c8e56f8231ef2c2e2aff839b22e07a1ffcbc617e9db9fa5" +checksum = "5c55a2eff8b69ce66c84f85e1da1c233edc36ceb85a2058d11b0d6a3c7e7569c" dependencies = [ "proc-macro2", "quote", "syn", ] -[[package]] -name = "tokio-native-tls" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbae76ab933c85776efabc971569dd6119c580d8f5d448769dec1764bf796ef2" -dependencies = [ - "native-tls", - "tokio", -] - -[[package]] -name = "tokio-openssl" -version = "0.6.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "59df6849caa43bb7567f9a36f863c447d95a11d5903c9cc334ba32576a27eadd" -dependencies = [ - "openssl", - "openssl-sys", - "tokio", -] - [[package]] name = "tokio-rustls" version = "0.26.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" dependencies = [ - "rustls 0.23.37", + "rustls", "tokio", ] @@ -4123,57 +3713,6 @@ dependencies = [ "tokio", ] -[[package]] -name = "toml" -version = "0.9.12+spec-1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf92845e79fc2e2def6a5d828f0801e29a2f8acc037becc5ab08595c7d5e9863" -dependencies = [ - "indexmap", - "serde_core", - "serde_spanned", - "toml_datetime", - "toml_parser", - "toml_writer", - "winnow", -] - -[[package]] -name = "toml_datetime" -version = "0.7.5+spec-1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92e1cfed4a3038bc5a127e35a2d360f145e1f4b971b551a2ba5fd7aedf7e1347" -dependencies = [ - "serde_core", -] - -[[package]] -name = "toml_edit" -version = "0.23.10+spec-1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "84c8b9f757e028cee9fa244aea147aab2a9ec09d5325a9b01e0a49730c2b5269" -dependencies = [ - "indexmap", - "toml_datetime", - "toml_parser", - "winnow", -] - -[[package]] -name = "toml_parser" -version = "1.0.9+spec-1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "702d4415e08923e7e1ef96cd5727c0dfed80b4d2fa25db9647fe5eb6f7c5a4c4" -dependencies = [ - "winnow", -] - -[[package]] -name = "toml_writer" -version = "1.0.6+spec-1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ab16f14aed21ee8bfd8ec22513f7287cd4a91aa92e44edfe2c17ddd004e92607" - [[package]] name = "tower" version = "0.5.3" @@ -4187,7 +3726,6 @@ dependencies = [ "tokio", "tower-layer", "tower-service", - "tracing", ] [[package]] @@ -4291,13 +3829,9 @@ dependencies = [ "http 1.4.0", "httparse", "log", - "native-tls", "rand 0.9.2", - "rustls 0.23.37", - "rustls-pki-types", "sha1", "thiserror 2.0.18", - "url", "utf-8", ] @@ -4415,12 +3949,6 @@ version = "0.2.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" -[[package]] -name = "version-compare" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "03c2856837ef78f57382f06b2b8563a2f512f7185d732608fd9176cb3b8edf0e" - [[package]] name = "version_check" version = "0.9.5" @@ -4595,16 +4123,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "webpki" -version = "0.22.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed63aea5ce73d0ff405984102c42de94fc55a6b75765d621c65262469b3c9b53" -dependencies = [ - "ring", - "untrusted", -] - [[package]] name = "webpki-root-certs" version = "1.0.6" @@ -4614,24 +4132,6 @@ dependencies = [ "rustls-pki-types", ] -[[package]] -name = "webpki-roots" -version = "0.22.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6c71e40d7d2c34a5106301fb632274ca37242cd0c9d3e64dbece371a40a2d87" -dependencies = [ - "webpki", -] - -[[package]] -name = "webpki-roots" -version = "1.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22cfaf3c063993ff62e73cb4311efde4db1efb31ab78a3e5c457939ad5cc0bed" -dependencies = [ - "rustls-pki-types", -] - [[package]] name = "whoami" version = "1.6.1" @@ -5107,15 +4607,6 @@ version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" -[[package]] -name = "winnow" -version = "0.7.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a5364e9d77fcdeeaa6062ced926ee3381faa2ee02d3eb83a5c27a8825540829" -dependencies = [ - "memchr", -] - [[package]] name = "wit-bindgen" version = "0.51.0" @@ -5367,9 +4858,9 @@ dependencies = [ [[package]] name = "zlib-rs" -version = "0.6.2" +version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c745c48e1007337ed136dc99df34128b9faa6ed542d80a1c673cf55a6d7236c8" +checksum = "3be3d40e40a133f9c916ee3f9f4fa2d9d63435b5fbe1bfc6d9dae0aa0ada1513" [[package]] name = "zmij" diff --git a/Cargo.toml b/Cargo.toml index 8432d26..a5cb77f 100755 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,39 +15,6 @@ aes-gcm = "*" ansi_term = "0.12.1" anyhow = "1.0.101" async-trait = "0.1.89" -async-tungstenite = { version = "0.32.0", features = [ - "futures-03-sink", - "futures-util", - "handshake", - "__rustls-tls", - "async-native-tls", - "async-std", - "async-std-runtime", - "async-tls", - "gio", - "gio-runtime", - "glib", - "openssl", - "real-async-native-tls", - "real-async-tls", - "real-native-tls", - "real-tokio-native-tls", - "real-tokio-openssl", - "real-tokio-rustls", - "rustls-native-certs", - "rustls-pki-types", - "tokio", - "tokio-native-tls", - "tokio-openssl", - "tokio-runtime", - "tokio-rustls-manual-roots", - "tokio-rustls-native-certs", - "tokio-rustls-webpki-roots", - "url", - "verbose-logging", - "webpki-roots", -] } -axum = { version = "0.8.8", features = ["ws", "http2"] } base64 = "0.22.1" bytes = "1.11.1" color-eyre = "0.6.5" @@ -58,16 +25,14 @@ futures-util = "0.3.31" hex = "0.4.3" hkdf = "0.12.4" http-body-util = "0.1.3" -hyper = { version = "1.8.1", features = ["http2", "full"] } -hyper-rustls = { version = "0.27.7", features = ["http2"] } -hyper-util = { version = "0.1.20", features = ["full"] } json = "0.12.4" once_cell = "1.21.3" pnet = "0.35.0" rand = "0.8" rand_core = { version = "0.6", features = ["getrandom", "std"] } -reqwest = "0.13.2" -rustls = "0.23.37" +reqwest = { version = "0.13.2" } +rustls = { version = "0.23.37", default-features = false, features = ["std", "tls12", "aws-lc-rs", "prefer-post-quantum"] } +quinn = { version = "0.11.9", default-features = false, features = ["runtime-tokio", "rustls-aws-lc-rs"] } rustls-pemfile = "2.2.0" sha1 = "0.10.6" sha2 = "0.10.9" @@ -85,4 +50,4 @@ uuid = { version = "1.19.0", features = ["v4"] } walkdir = "2.5.0" x448 = "0.6.0" zip = "6.0.0" -quinn = "0.11.9" +thiserror = "2.0.18" diff --git a/src/main.rs b/src/main.rs index 043700c..72c42c5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -27,6 +27,9 @@ async fn main() { startup(); log_in!("Incoming messages"); log_out!("Outgoing messages"); + + let _ = omikron_connection::start(9187).await; + log!("Started"); log!(" .env"); if let Err(e) = initialize_db().await { @@ -43,7 +46,8 @@ async fn main() { } else { log!(" Users"); } + let _ = server::server::start(9188).await; - let _ = omikron_connection::OmikronServer::start(9187).await; + tokio::signal::ctrl_c().await.unwrap(); } diff --git a/src/server/api.rs b/src/server/api.rs index 7f3c728..b216112 100644 --- a/src/server/api.rs +++ b/src/server/api.rs @@ -65,16 +65,20 @@ pub async fn handle(path: &str, body_string: Option) -> HttpResponse { // ================================================== ["api", "get", "omikron"] => { if let Ok(omikron_conn) = get_random_omikron().await { - let id = omikron_conn.get_omikron_id().await; + if let Some(id) = omikron_conn.get_omikron_id().await { + if let Ok((public_key, ip_address)) = sql::get_omikron_by_id(id).await { + let mut res = JsonValue::new_object(); + res["status"] = "success".into(); + res["id"] = id.into(); + res["public_key"] = public_key.into(); + res["ip_address"] = ip_address.into(); - if let Ok((public_key, ip_address)) = sql::get_omikron_by_id(id).await { - let mut res = JsonValue::new_object(); - res["status"] = "success".into(); - res["id"] = id.into(); - res["public_key"] = public_key.into(); - res["ip_address"] = ip_address.into(); - - (StatusCode::OK, res.dump()) + (StatusCode::OK, res.dump()) + } else { + let mut res = JsonValue::new_object(); + res["status"] = "error".into(); + (StatusCode::INTERNAL_SERVER_ERROR, res.dump()) + } } else { let mut res = JsonValue::new_object(); res["status"] = "error".into(); diff --git a/src/server/omikron_connection.rs b/src/server/omikron_connection.rs index ade9746..b509787 100644 --- a/src/server/omikron_connection.rs +++ b/src/server/omikron_connection.rs @@ -1,412 +1,1140 @@ -use crate::server::{omikron_manager, short_link::add_short_link}; -use crate::sql::{sql, sql::get_omikron_by_id, user_online_tracker}; -use crate::util::file_util::load_file_buf; -use crate::util::{crypto_helper::encrypt, logger::PrintType}; -use crate::{get_private_key, get_public_key, log_cv_in, log_cv_out, log_in}; +use crate::{ + get_private_key, get_public_key, log, log_cv_in, log_cv_out, log_err, log_in, + server::{omikron_manager, short_link::add_short_link}, + sql::{ + connection_status::UserStatus, + sql::{self, get_by_user_id, get_by_username, get_iota_by_id, get_omikron_by_id}, + user_online_tracker::{self}, + }, + util::{ + crypto_helper::encrypt, + file_util::{load_file_buf, load_file_vec}, + logger::PrintType, + }, +}; use base64::{Engine as _, engine::general_purpose::STANDARD}; use dashmap::DashMap; use epsilon_core::{CommunicationType, CommunicationValue, DataTypes, DataValue}; -use epsilon_native::{Receiver, Sender, host}; +use epsilon_native::{Host, Receiver, Sender}; use quinn::ServerConfig; use rand::{Rng, distributions::Alphanumeric}; -use rustls::pki_types::PrivateKeyDer; -use std::sync::Arc; -use tokio::sync::RwLock; +use rustls::{ServerConfig as CryptoConfig, crypto::aws_lc_rs}; +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; +use tokio::{ + sync::{Mutex, RwLock}, + time::interval, +}; use x448::PublicKey; -pub struct OmikronServer; +// ============================================================================ +// Configuration +// ============================================================================ -impl OmikronServer { - pub async fn start(port: u16) -> Result<(), Box> { - let tls_cfg = OmikronServer::load_tls().expect("TLS config failed"); - let server_crypto = quinn::crypto::rustls::QuicServerConfig::try_from(tls_cfg) - .expect("Failed to convert to QuicServerConfig"); - let mut host = host(port, ServerConfig::with_crypto(Arc::new(server_crypto))).await?; - tokio::spawn(async move { - while let Some((sender, receiver)) = host.next().await { - let connection = OmikronConnection::new(sender); - tokio::spawn(Self::connection_loop(connection, receiver)); +const CLEANUP_INTERVAL: Duration = Duration::from_secs(30); +const MAX_WAITING_AGE: Duration = Duration::from_secs(60); + +// ============================================================================ +// Error Types +// ============================================================================ + +#[derive(Debug, thiserror::Error)] +pub enum OmikronError { + #[error("Not connected")] + NotConnected, + #[error("Not authenticated")] + NotAuthenticated, + #[error("Invalid response")] + InvalidResponse, + #[error("Authentication failed")] + AuthenticationFailed, + #[error("SQL error: {0}")] + Sql(String), + #[error("Send error: {0}")] + Send(String), +} + +pub type OmikronResult = Result; + +// ============================================================================ +// Waiting Task System (Preserved from original) +// ============================================================================ + +pub struct WaitingTask { + pub task: Box, CommunicationValue) -> bool + Send + Sync>, + pub inserted_at: Instant, +} + +// ============================================================================ +// Connection State +// ============================================================================ + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum AuthState { + Unauthenticated, + Identified { omikron_id: i64 }, + Authenticated { omikron_id: i64 }, +} + +impl AuthState { + fn is_authenticated(&self) -> bool { + match self { + AuthState::Authenticated { omikron_id } => true, + _ => false, + } + } + + fn omikron_id(&self) -> Option { + match self { + AuthState::Identified { omikron_id } | AuthState::Authenticated { omikron_id } => { + Some(*omikron_id) } - }); - - Ok(()) - } - fn load_tls() -> Option { - let mut cert_file_buf = load_file_buf("certs", "cert.pem").ok()?; - let mut key_file_buf = load_file_buf("certs", "cert.key").ok()?; - - let cert_chain = rustls_pemfile::certs(&mut cert_file_buf) - .collect::, _>>() - .ok()?; - - let mut keys: Vec = rustls_pemfile::pkcs8_private_keys(&mut key_file_buf) - .map(|k| k.map(Into::into)) - .collect::, _>>() - .ok()?; - - if keys.is_empty() { - let mut key_file_buf = load_file_buf("certs", "cert.key").ok()?; - keys = rustls_pemfile::rsa_private_keys(&mut key_file_buf) - .map(|k| k.map(Into::into)) - .collect::, _>>() - .ok()?; + _ => None, } - - if keys.is_empty() { - return None; - } - - let cfg = rustls::ServerConfig::builder() - .with_no_client_auth() - .with_single_cert(cert_chain, keys.remove(0)) - .ok()?; - - Some(cfg) - } - - async fn connection_loop(conn: Arc, receiver: Receiver) { - while let Ok(cv) = receiver.receive().await { - log_cv_in!(PrintType::Omikron, cv); - conn.clone().handle_message(cv).await; - } - - conn.handle_close().await; } } +// ============================================================================ +// Omikron Connection (Epsilon/QUIC-based) +// ============================================================================ + pub struct OmikronConnection { - sender: Arc>>, - omikron_id: Arc>, - pub_key: Arc>>>, + id: u64, + sender: Mutex>, + state: RwLock, - identified: Arc>, - challenged: Arc>, - challenge: Arc>, + // Authentication state (preserved from original) + challenge: RwLock, + pub_key: RwLock>>, - ping: Arc>, + // Ping tracking + pub ping: RwLock, - waiting_tasks: - DashMap, CommunicationValue) -> bool + Send + Sync>>, + // Waiting tasks for request/response pattern + waiting_tasks: DashMap, + + // Cleanup handle + cleanup_handle: Mutex>>, } impl OmikronConnection { + // ------------------------------------------------------------------------- + // Construction + // ------------------------------------------------------------------------- + pub fn new(sender: Sender) -> Arc { - Arc::new(Self { - sender: Arc::new(RwLock::new(Some(sender))), - omikron_id: Arc::new(RwLock::new(0)), - pub_key: Arc::new(RwLock::new(None)), - identified: Arc::new(RwLock::new(false)), - challenged: Arc::new(RwLock::new(false)), - challenge: Arc::new(RwLock::new(String::new())), - ping: Arc::new(RwLock::new(-1)), + let conn = Arc::new(Self { + id: rand::random(), + sender: Mutex::new(Some(sender)), + state: RwLock::new(AuthState::Unauthenticated), + challenge: RwLock::new(String::new()), + pub_key: RwLock::new(None), + ping: RwLock::new(-1), waiting_tasks: DashMap::new(), - }) + cleanup_handle: Mutex::new(None), + }); + + // Start cleanup task for waiting tasks + let cleanup_conn = conn.clone(); + let handle = tokio::spawn(async move { + let mut ticker = interval(CLEANUP_INTERVAL); + loop { + ticker.tick().await; + cleanup_conn + .waiting_tasks + .retain(|_, v| v.inserted_at.elapsed() < MAX_WAITING_AGE); + } + }); + + // Store handle (would need block_in_place or similar to set this immediately) + // For now, we'll handle this differently in handle() + + conn } - async fn send(&self, cv: &CommunicationValue) { - log_cv_out!(PrintType::Omikron, cv); + // ------------------------------------------------------------------------- + // Main Handler Loop + // ------------------------------------------------------------------------- - if let Some(sender) = self.sender.read().await.as_ref() { - if sender.send(cv).await.is_err() { - self.handle_close().await; + pub async fn handle(self: Arc, mut receiver: Receiver) { + log_in!( + self.id as i64, + PrintType::Omega, + "Omikron connection started" + ); + + // Start cleanup task + let cleanup_conn = self.clone(); + let _cleanup_handle = tokio::spawn(async move { + let mut ticker = interval(CLEANUP_INTERVAL); + loop { + ticker.tick().await; + cleanup_conn + .waiting_tasks + .retain(|_, v| v.inserted_at.elapsed() < MAX_WAITING_AGE); + } + }); + + while let Ok(cv) = receiver.receive().await { + if let Err(e) = self.clone().process_message(cv).await { + log_err!(0, PrintType::Omega, "Error processing message: {}", e); + // Don't break on error unless critical - match original WebSocket behavior + if matches!(e, OmikronError::NotConnected) { + break; + } + } + } + + // Connection closed + self.cleanup().await; + log_in!( + self.id as i64, + PrintType::Omega, + "Omikron connection closed" + ); + } + + // ------------------------------------------------------------------------- + // Message Processing + // ------------------------------------------------------------------------- + + async fn process_message(self: Arc, cv: CommunicationValue) -> OmikronResult<()> { + // Log incoming + log_cv_in!(PrintType::Omikron, &cv); + + let msg_id = cv.get_id(); + + // Check waiting tasks first (response to previous request) + if let Some((_, task)) = self.waiting_tasks.remove(&msg_id) { + let _ = (task.task)(self.clone(), cv); + return Ok(()); + } + + // Handle ping regardless of auth state + if cv.is_type(CommunicationType::ping) { + return self.handle_ping(cv).await; + } + + // Route based on authentication state + match *self.state.read().await { + AuthState::Unauthenticated => self.handle_unauthenticated(cv).await, + AuthState::Identified { .. } => self.handle_identified(cv).await, + AuthState::Authenticated { omikron_id } => { + self.handle_authenticated(cv, omikron_id).await } } } - async fn send_error_response(&self, id: u32, comm_type: CommunicationType) { - let response = CommunicationValue::new(comm_type).with_id(id); - self.send(&response).await; - } + // ------------------------------------------------------------------------- + // Authentication Handlers + // ------------------------------------------------------------------------- - pub async fn get_omikron_id(&self) -> i64 { - *self.omikron_id.read().await - } - - async fn is_identified(&self) -> bool { - *self.identified.read().await && *self.challenged.read().await - } - - pub async fn handle_message(self: Arc, cv: CommunicationValue) { - if cv.is_type(CommunicationType::ping) { - self.handle_ping(cv).await; - return; + async fn handle_unauthenticated(&self, cv: CommunicationValue) -> OmikronResult<()> { + if !cv.is_type(CommunicationType::identification) { + self.send_error_response(cv.get_id(), CommunicationType::error_not_authenticated) + .await; + return Err(OmikronError::NotAuthenticated); } - if let Some((_, task)) = self.waiting_tasks.remove(&cv.get_id()) { - let _ = task(self.clone(), cv.clone()); - return; - } + // Extract omikron ID + let omikron_id = cv + .get_data(DataTypes::omikron) + .as_number() + .ok_or(OmikronError::InvalidResponse)?; - if !self.is_identified().await { - self.handle_identification(cv).await; - return; - } + // Lookup omikron in database + let (public_key, _) = get_omikron_by_id(omikron_id) + .await + .map_err(|e| OmikronError::Sql(e.to_string()))?; - self.handle_authenticated(cv).await; + let pub_key_bytes = STANDARD + .decode(&public_key) + .map_err(|_| OmikronError::AuthenticationFailed)?; + + let omikron_pub_key = + PublicKey::from_bytes(&pub_key_bytes).ok_or(OmikronError::AuthenticationFailed)?; + + // Generate challenge + let challenge: String = rand::thread_rng() + .sample_iter(&Alphanumeric) + .take(32) + .map(char::from) + .collect(); + + // Store state + *self.challenge.write().await = challenge.clone(); + *self.pub_key.write().await = Some(pub_key_bytes); + *self.state.write().await = AuthState::Identified { omikron_id }; + + // Encrypt challenge + let encrypted = encrypt(get_private_key(), omikron_pub_key, &challenge) + .map_err(|_| OmikronError::AuthenticationFailed)?; + + // Send challenge response + let response = CommunicationValue::new(CommunicationType::challenge) + .with_id(cv.get_id()) + .add_data( + DataTypes::public_key, + DataValue::Str(STANDARD.encode(get_public_key().as_bytes())), + ) + .add_data(DataTypes::challenge, DataValue::Str(encrypted)); + + self.send(&response).await } - async fn handle_identification(self: &Arc, cv: CommunicationValue) { - let identified = *self.identified.read().await; - let challenged = *self.challenged.read().await; + async fn handle_identified(&self, cv: CommunicationValue) -> OmikronResult<()> { + if !cv.is_type(CommunicationType::challenge_response) { + self.send_error_response(cv.get_id(), CommunicationType::error_not_authenticated) + .await; + return Err(OmikronError::NotAuthenticated); + } - if !identified && cv.is_type(CommunicationType::identification) { - let omikron_id = cv.get_data(DataTypes::omikron).as_number().unwrap_or(0); + let client_response = cv + .get_data(DataTypes::challenge) + .as_str() + .ok_or(OmikronError::InvalidResponse)?; - let (public_key, _) = match get_omikron_by_id(omikron_id).await { - Ok(v) => v, - Err(_) => { - let _ = self - .send( - &CommunicationValue::new(CommunicationType::error_not_authenticated) - .with_id(cv.get_id()), - ) - .await; - return; - } - }; + let expected_challenge = self.challenge.read().await.clone(); - let pub_key_bytes = match STANDARD.decode(&public_key) { - Ok(b) => b, - Err(_) => { - let _ = self - .send( - &CommunicationValue::new(CommunicationType::error_invalid_omikron_id) - .with_id(cv.get_id()), - ) - .await; - return; - } - }; + 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 }; - let omikron_pub_key = match PublicKey::from_bytes(&pub_key_bytes) { - Some(k) => k, - _ => { - let _ = self - .send( - &CommunicationValue::new(CommunicationType::error_invalid_public_key) - .with_id(cv.get_id()), - ) - .await; - return; - } - }; + // Register with manager + omikron_manager::add_omikron(self.arc_self()).await; - let challenge: String = rand::thread_rng() - .sample_iter(&Alphanumeric) - .take(32) - .map(char::from) - .collect(); - - *self.omikron_id.write().await = omikron_id; - *self.challenge.write().await = challenge.clone(); - *self.pub_key.write().await = Some(pub_key_bytes); - *self.identified.write().await = true; - - let encrypted = - encrypt(get_private_key(), omikron_pub_key, &challenge).unwrap_or_default(); - - let response = CommunicationValue::new(CommunicationType::challenge) + // Send success response + let response = CommunicationValue::new(CommunicationType::identification_response) .with_id(cv.get_id()) - .add_data( - DataTypes::public_key, - DataValue::Str(STANDARD.encode(get_public_key().as_bytes())), - ) - .add_data(DataTypes::challenge, DataValue::Str(encrypted)); + .add_data(DataTypes::accepted, DataValue::Bool(true)); + + self.send(&response).await?; + log_in!(omikron_id, PrintType::Omega, "Omikron authenticated"); + Ok(()) + } else { + self.send_error_response(cv.get_id(), CommunicationType::error_invalid_challenge) + .await; + Err(OmikronError::AuthenticationFailed) + } + } + + // ------------------------------------------------------------------------- + // Authenticated Message Handlers + // ------------------------------------------------------------------------- + + async fn handle_authenticated( + &self, + cv: CommunicationValue, + omikron_id: i64, + ) -> OmikronResult<()> { + match cv.get_type() { + // Link shortening + CommunicationType::shorten_link => self.handle_shorten_link(cv).await, + + // Online status tracking + CommunicationType::user_connected => { + self.handle_user_connected(cv, omikron_id).await; + Ok(()) + } + CommunicationType::user_disconnected => { + self.handle_user_disconnected(cv, omikron_id).await; + Ok(()) + } + CommunicationType::iota_connected => { + self.handle_iota_connected(cv, omikron_id).await; + Ok(()) + } + CommunicationType::iota_disconnected => { + self.handle_iota_disconnected(cv, omikron_id).await; + Ok(()) + } + CommunicationType::sync_client_iota_status => { + self.handle_sync_status(cv, omikron_id).await; + 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 + } + CommunicationType::complete_register_user => { + 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, + + _ => { + log_err!( + 0, + PrintType::Omega, + "Unknown message type: {:?}", + cv.get_type() + ); + Ok(()) + } + } + } + + // ------------------------------------------------------------------------- + // Specific Handlers (ported from original WebSocket implementation) + // ------------------------------------------------------------------------- + + async fn handle_shorten_link(&self, cv: CommunicationValue) -> OmikronResult<()> { + let link = cv + .get_data(DataTypes::link) + .as_str() + .ok_or(OmikronError::InvalidResponse)?; + + let short = add_short_link(link) + .await + .map_err(|_| OmikronError::Sql("Shortend link Error".to_string()))?; + + let response = CommunicationValue::new(CommunicationType::shorten_link) + .with_id(cv.get_id()) + .add_data(DataTypes::link, DataValue::Str(short)); + + self.send(&response).await + } + + async fn handle_user_connected(&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( + user_id as i64, + UserStatus::user_online, + omikron_id, + ); + } + } + + async fn handle_user_disconnected(&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) { + user_online_tracker::track_user_status( + user_id as i64, + UserStatus::user_offline, + status.omikron_id, + ); + } + } + } + + async fn handle_iota_connected(&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; + user_online_tracker::track_iota_connection(iota_id, omikron_id, true); + + let mut user_ids = Vec::new(); + if let Ok(users) = sql::get_users_by_iota_id(iota_id).await { + for (user_id, _, _, _, _, _, _, _, _, _, _, _) in users { + user_ids.push(DataValue::Number(user_id)); + user_online_tracker::track_user_status( + user_id, + UserStatus::user_offline, + omikron_id, + ); + } + } else { + log_in!(PrintType::General, "SQL error loading users for IOTA"); + } + + let response = CommunicationValue::new(CommunicationType::iota_user_data) + .with_id(cv.get_id()) + .add_data(DataTypes::user_ids, DataValue::Array(user_ids)); let _ = self.send(&response).await; - return; + } else { + log_in!(PrintType::General, "No IOTA ID found"); } + } - if identified && !challenged && cv.is_type(CommunicationType::challenge_response) { - let client_response = cv.get_data(DataTypes::challenge).as_str().unwrap_or(""); - - if client_response == *self.challenge.read().await { - *self.challenged.write().await = true; - - omikron_manager::add_omikron(self.clone()).await; - - let _ = self - .send( - &CommunicationValue::new(CommunicationType::identification_response) - .with_id(cv.get_id()) - .add_data(DataTypes::accepted, DataValue::BoolTrue), - ) - .await; - - log_in!(PrintType::Omega, "Omikron Connected"); - } else { - let _ = self - .send( - &CommunicationValue::new(CommunicationType::error_invalid_challenge) - .with_id(cv.get_id()), - ) - .await; + async fn handle_iota_disconnected(&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; + let iota_offline = user_online_tracker::untrack_iota_connection(iota_id, omikron_id); + if iota_offline { + if let Ok(users) = sql::get_users_by_iota_id(iota_id).await { + let user_ids: Vec = users.iter().map(|u| u.0).collect(); + user_online_tracker::untrack_many_users(&user_ids); + } } } } - async fn handle_authenticated(self: &Arc, cv: CommunicationValue) { - if cv.is_type(CommunicationType::shorten_link) { - if let Some(link) = cv.get_data(DataTypes::link).as_str() { - if let Ok(short) = add_short_link(link).await { - let _ = self - .send( - &CommunicationValue::new(CommunicationType::shorten_link) - .with_id(cv.get_id()) - .add_data(DataTypes::link, DataValue::Str(short)), - ) - .await; + async fn handle_sync_status(&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 { + user_online_tracker::track_user_status( + *user_id, + UserStatus::user_online, + omikron_id, + ); } } - return; } - if cv.is_type(CommunicationType::get_register) { - let register_id = sql::get_register_id().await; - let _ = self - .send( - &CommunicationValue::new(CommunicationType::get_register) - .with_id(cv.get_id()) - .add_data(DataTypes::user_id, DataValue::Number(register_id as i64)), - ) - .await; - return; - } - - if cv.is_type(CommunicationType::delete_user) { - let user_id = cv.get_sender(); - match sql::delete_user(user_id as i64).await { - Ok(_) => { - let _ = self - .send( - &CommunicationValue::new(CommunicationType::success) - .with_id(cv.get_id()), - ) - .await; - } - Err(e) => { - let _ = self - .send( - &CommunicationValue::new(CommunicationType::error) - .with_id(cv.get_id()) - .add_data(DataTypes::error_type, DataValue::Str(e.to_string())), - ) - .await; + if let DataValue::Array(iota_ids) = cv.get_data(DataTypes::iota_ids) { + for iota_id_val in iota_ids { + if let DataValue::Number(iota_id) = iota_id_val { + user_online_tracker::track_iota_connection(*iota_id, omikron_id, true); } } - return; + } + } + + async fn handle_get_user_data(&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; + return self.send(&response).await; + } } - if cv.is_type(CommunicationType::delete_iota) { - if let DataValue::Number(iota_id) = cv.get_data(DataTypes::iota_id) { - match sql::delete_iota(*iota_id).await { + // 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; + return self.send(&response).await; + } + } + + // Not found + let response = + CommunicationValue::new(CommunicationType::error_not_found).with_id(cv.get_id()); + self.send(&response).await + } + + async fn build_user_data_response( + &self, + msg_id: u32, + user: ( + i64, + i64, + String, + Option, + Option, + Option, + Option>, + i32, + i64, + String, + String, + String, + ), + ) -> CommunicationValue { + let ( + id, + iota_id, + username, + display, + status, + about, + avatar, + sub_level, + sub_end, + public_key, + _, + _, + ) = user; + + let mut response = CommunicationValue::new(CommunicationType::get_user_data) + .with_id(msg_id) + .add_data(DataTypes::username, DataValue::Str(username.clone())) + .add_data(DataTypes::public_key, DataValue::Str(public_key)) + .add_data(DataTypes::user_id, DataValue::Number(id)) + .add_data(DataTypes::iota_id, DataValue::Number(iota_id)) + .add_data(DataTypes::sub_level, DataValue::Number(sub_level as i64)) + .add_data(DataTypes::sub_end, DataValue::Number(sub_end)); + + // Display name (fallback to username) + let display_name = display.filter(|d| !d.is_empty()).unwrap_or(username); + response = response.add_data(DataTypes::display, DataValue::Str(display_name)); + + // Optional fields + if let Some(s) = status.filter(|s| !s.is_empty()) { + response = response.add_data(DataTypes::status, DataValue::Str(s)); + } + if let Some(a) = about.filter(|a| !a.is_empty()) { + response = response.add_data(DataTypes::about, DataValue::Str(a)); + } + if let Some(av) = avatar { + response = response.add_data(DataTypes::avatar, DataValue::Str(STANDARD.encode(av))); + } + + // Online status + let user_status = user_online_tracker::get_user_status(id); + let iota_connections = + user_online_tracker::get_iota_omikron_connections(iota_id).unwrap_or_default(); + + if let Some(us) = user_status { + response = response.add_data( + DataTypes::online_status, + DataValue::Str(us.connection_type.to_string()), + ); + response = response.add_data(DataTypes::omikron_id, DataValue::Number(us.omikron_id)); + } else { + response = response.add_data( + DataTypes::online_status, + DataValue::Str(UserStatus::iota_offline.to_string()), + ); + } + + response = response.add_data( + DataTypes::omikron_connections, + DataValue::Array( + iota_connections + .into_iter() + .map(DataValue::Number) + .collect(), + ), + ); + + response + } + + async fn handle_get_iota_data(&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 + .build_iota_data_response(cv.get_id(), iota_id, public_key, None, None) + .await; + return self.send(&response).await; + } + } + + // Try by user_id + if let Some(user_id) = cv.get_data(DataTypes::user_id).as_number() { + if let Ok((_, iota_id, _, _, _, _, _, _, _, _, _, _)) = + get_by_user_id(user_id as i64).await + { + if let Ok((iota_id, public_key)) = get_iota_by_id(iota_id).await { + let response = self + .build_iota_data_response( + cv.get_id(), + iota_id, + public_key, + Some(user_id as i64), + None, + ) + .await; + return self.send(&response).await; + } + } + } + + // Try by username + if let Some(username) = cv.get_data(DataTypes::username).as_str() { + if let Ok((user_id, iota_id, _, _, _, _, _, _, _, _, _, _)) = + get_by_username(username).await + { + if let Ok((iota_id, public_key)) = get_iota_by_id(iota_id).await { + let response = self + .build_iota_data_response( + cv.get_id(), + iota_id, + public_key, + Some(user_id), + Some(username.to_string()), + ) + .await; + return self.send(&response).await; + } + } + } + + let response = + CommunicationValue::new(CommunicationType::error_not_found).with_id(cv.get_id()); + self.send(&response).await + } + + async fn build_iota_data_response( + &self, + msg_id: u32, + iota_id: i64, + public_key: String, + user_id: Option, + username: Option, + ) -> CommunicationValue { + let mut response = CommunicationValue::new(CommunicationType::get_iota_data) + .with_id(msg_id) + .add_data(DataTypes::public_key, DataValue::Str(public_key)) + .add_data(DataTypes::iota_id, DataValue::Number(iota_id)); + + if let Some(uid) = user_id { + response = response.add_data(DataTypes::user_id, DataValue::Number(uid)); + } + if let Some(uname) = username { + response = response.add_data(DataTypes::username, DataValue::Str(uname)); + } + + let iota_connections = + user_online_tracker::get_iota_omikron_connections(iota_id).unwrap_or_default(); + + response.add_data( + DataTypes::omikron_connections, + DataValue::Array( + iota_connections + .into_iter() + .map(DataValue::Number) + .collect(), + ), + ) + } + + async fn handle_get_register(&self, cv: CommunicationValue) -> OmikronResult<()> { + let register_id = sql::get_register_id().await; + let response = CommunicationValue::new(CommunicationType::get_register) + .with_id(cv.get_id()) + .add_data(DataTypes::user_id, DataValue::Number(register_id as i64)); + self.send(&response).await + } + + async fn handle_complete_register_iota(&self, cv: CommunicationValue) -> OmikronResult<()> { + let iota_id_opt = cv + .get_data(DataTypes::iota_id) + .as_number() + .map(|n| n as i64); + + if let Some(public_key) = cv.get_data(DataTypes::public_key).as_str() { + if let Some(iota_id) = iota_id_opt { + // Register existing IOTA + match sql::register_complete_iota(iota_id, public_key.to_string()).await { Ok(_) => { let response = CommunicationValue::new(CommunicationType::success) .with_id(cv.get_id()); - self.send(&response).await; + self.send(&response).await } Err(e) => { - self.send( - &CommunicationValue::new(CommunicationType::error) - .with_id(cv.get_id()) - .add_data(DataTypes::error_type, DataValue::Str(e.to_string())), - ) - .await; + let response = CommunicationValue::new(CommunicationType::error) + .with_id(cv.get_id()) + .add_data(DataTypes::error_type, DataValue::Str(e.to_string())); + self.send(&response).await } } } else { - self.send_error_response(cv.get_id(), CommunicationType::error_invalid_data) - .await; + // Create new IOTA + match sql::create_new_iota(public_key.to_string()).await { + Ok(new_iota_id) => { + let response = + CommunicationValue::new(CommunicationType::complete_register_iota) + .with_id(cv.get_id()) + .add_data(DataTypes::iota_id, DataValue::Number(new_iota_id)); + self.send(&response).await + } + Err(e) => { + let response = CommunicationValue::new(CommunicationType::error) + .with_id(cv.get_id()) + .add_data(DataTypes::error_type, DataValue::Str(e.to_string())); + self.send(&response).await + } + } + } + } else { + self.send_error_response(cv.get_id(), CommunicationType::error_invalid_data) + .await + } + } + + async fn handle_complete_register_user(&self, cv: CommunicationValue) -> OmikronResult<()> { + let user_id = cv + .get_data(DataTypes::user_id) + .as_number() + .map(|n| n as i64); + let username = cv + .get_data(DataTypes::username) + .as_str() + .map(|s| s.to_string()); + let public_key = cv + .get_data(DataTypes::public_key) + .as_str() + .map(|s| s.to_string()); + let iota_id = cv + .get_data(DataTypes::iota_id) + .as_number() + .map(|n| n as i64); + let reset_token = cv + .get_data(DataTypes::reset_token) + .as_str() + .map(|s| s.to_string()); + + if let (Some(uid), Some(uname), Some(pk), Some(iid), Some(rt)) = + (user_id, username, public_key, iota_id, reset_token) + { + match sql::register_complete_user(uid, uname, pk, iid, rt).await { + Ok(_) => { + let response = + CommunicationValue::new(CommunicationType::success).with_id(cv.get_id()); + self.send(&response).await + } + Err(e) => { + let response = CommunicationValue::new(CommunicationType::error) + .with_id(cv.get_id()) + .add_data(DataTypes::error_type, DataValue::Str(e.to_string())); + self.send(&response).await + } + } + } else { + self.send_error_response(cv.get_id(), CommunicationType::error_invalid_data) + .await + } + } + + async fn handle_change_user_data(&self, cv: CommunicationValue) -> OmikronResult<()> { + let user_id = cv.get_sender() as i64; + let mut success = true; + let mut error_message = String::new(); + + // Process each field + if let Some(username) = cv.get_data(DataTypes::username).as_str() { + if let Err(e) = sql::change_username(user_id, username.to_string()).await { + success = false; + error_message = e.to_string(); + } + } + if let Some(display) = cv.get_data(DataTypes::display).as_str() { + if let Err(e) = sql::change_display_name(user_id, display.to_string()).await { + success = false; + error_message = e.to_string(); + } + } + if let Some(avatar) = cv.get_data(DataTypes::avatar).as_str() { + if let Err(e) = sql::change_avatar(user_id, avatar.to_string()).await { + success = false; + error_message = e.to_string(); + } + } + if let Some(about) = cv.get_data(DataTypes::about).as_str() { + if let Err(e) = sql::change_about(user_id, about.to_string()).await { + success = false; + error_message = e.to_string(); + } + } + if let Some(status) = cv.get_data(DataTypes::status).as_str() { + if let Err(e) = sql::change_status(user_id, status.to_string()).await { + success = false; + error_message = e.to_string(); + } + } + if let (Some(public_key), Some(private_key_hash)) = ( + cv.get_data(DataTypes::public_key).as_str(), + cv.get_data(DataTypes::private_key_hash).as_str(), + ) { + if let Err(e) = sql::change_keys( + user_id, + public_key.to_string(), + private_key_hash.to_string(), + ) + .await + { + success = false; + error_message = e.to_string(); } - return; } - // NOTIFICATIONS - if cv.is_type(CommunicationType::get_notifications) { - let user_id = cv.get_sender(); - if let Ok(notifications) = sql::get_notifications(user_id as i64).await { - let mut json_array = Vec::new(); - for (sender, amount) in notifications { - let mut obj = Vec::new(); - let _ = obj.push((DataTypes::sender_id, DataValue::Number(sender))); - let _ = obj.push((DataTypes::amount, DataValue::Number(amount))); - json_array.push(DataValue::Container(obj)); + if success { + let response = CommunicationValue::new(CommunicationType::success).with_id(cv.get_id()); + self.send(&response).await + } else { + let response = CommunicationValue::new(CommunicationType::error) + .with_id(cv.get_id()) + .add_data(DataTypes::error_type, DataValue::Str(error_message)); + self.send(&response).await + } + } + + async fn handle_change_iota_data(&self, cv: CommunicationValue) -> OmikronResult<()> { + let user_id = cv.get_sender() as i64; + + if let (Some(iota_id), Some(reset_token), Some(new_token)) = ( + cv.get_data(DataTypes::iota_id) + .as_number() + .map(|n| n as i64), + cv.get_data(DataTypes::reset_token).as_str(), + cv.get_data(DataTypes::new_token).as_str(), + ) { + match sql::get_by_user_id(user_id).await { + Ok(user) => { + let current_token = user.11; // reset_token field + if current_token == reset_token { + let mut success = true; + let mut error_message = String::new(); + + if let Err(e) = sql::change_iota_id(user_id, iota_id).await { + success = false; + error_message = e.to_string(); + } + if success { + if let Err(e) = sql::change_token(user_id, new_token.to_string()).await + { + success = false; + error_message = e.to_string(); + } + } + + if success { + let response = CommunicationValue::new(CommunicationType::success) + .with_id(cv.get_id()); + self.send(&response).await + } else { + let response = CommunicationValue::new(CommunicationType::error) + .with_id(cv.get_id()) + .add_data(DataTypes::error_type, DataValue::Str(error_message)); + self.send(&response).await + } + } else { + self.send_error_response( + cv.get_id(), + CommunicationType::error_invalid_challenge, + ) + .await + } } - let response = CommunicationValue::new(CommunicationType::get_notifications) + Err(_) => { + self.send_error_response(cv.get_id(), CommunicationType::error_not_found) + .await + } + } + } else { + self.send_error_response(cv.get_id(), CommunicationType::error_invalid_data) + .await + } + } + + async fn handle_delete_user(&self, cv: CommunicationValue) -> OmikronResult<()> { + let user_id = cv.get_sender() as i64; + match sql::delete_user(user_id).await { + Ok(_) => { + let response = + CommunicationValue::new(CommunicationType::success).with_id(cv.get_id()); + self.send(&response).await + } + Err(e) => { + let response = CommunicationValue::new(CommunicationType::error) .with_id(cv.get_id()) - .add_data(DataTypes::notifications, DataValue::Array(json_array)); - self.send(&response).await; - } - } - if cv.is_type(CommunicationType::read_notification) { - if let (user_id, Some(other_id)) = ( - cv.get_sender(), - cv.get_data(DataTypes::sender_id).as_number(), - ) { - if let Ok(_) = sql::read_notification(user_id as i64, other_id).await { - let response = CommunicationValue::new(CommunicationType::read_notification) - .with_id(cv.get_id()); - self.send(&response).await; - } - } - } - if cv.is_type(CommunicationType::push_notification) { - if let (user_id, Some(other_id)) = ( - cv.get_sender(), - cv.get_data(DataTypes::sender_id).as_number(), - ) { - if let Ok(_) = sql::add_notification(user_id as i64, other_id).await { - let response = CommunicationValue::new(CommunicationType::push_notification) - .with_id(cv.get_id()); - self.send(&response).await; - } + .add_data(DataTypes::error_type, DataValue::Str(e.to_string())); + self.send(&response).await } } } - async fn send_error_response(&self, message_id: &Uuid, error_type: CommunicationType) { - let error = CommunicationValue::new(error_type).with_id(*message_id); - self.send_message(&error).await; - } - pub async fn close(&self) { - let omikron_id = self.get_omikron_id().await; - if omikron_id != 0 { - log_in!(omikron_id, PrintType::Omega, "Omikron Disconnected"); - omikron_manager::remove_omikron(omikron_id).await; - user_online_tracker::untrack_omikron(omikron_id).await; - } - } - pub async fn handle_close(self: Arc) { - let omikron_id = self.get_omikron_id().await; - if omikron_id != 0 { - log_in!(omikron_id, PrintType::Omega, "Omikron Disconnected"); - omikron_manager::remove_omikron(omikron_id).await; - user_online_tracker::untrack_omikron(omikron_id).await; + async fn handle_delete_iota(&self, cv: CommunicationValue) -> OmikronResult<()> { + if let Some(iota_id) = cv + .get_data(DataTypes::iota_id) + .as_number() + .map(|n| n as i64) + { + match sql::delete_iota(iota_id).await { + Ok(_) => { + let response = + CommunicationValue::new(CommunicationType::success).with_id(cv.get_id()); + self.send(&response).await + } + Err(e) => { + let response = CommunicationValue::new(CommunicationType::error) + .with_id(cv.get_id()) + .add_data(DataTypes::error_type, DataValue::Str(e.to_string())); + self.send(&response).await + } + } + } else { + self.send_error_response(cv.get_id(), CommunicationType::error_invalid_data) + .await } } - async fn handle_ping(&self, cv: CommunicationValue) { + async fn handle_get_notifications(&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 = notifications + .into_iter() + .map(|(sender, amount)| { + DataValue::Container(vec![ + (DataTypes::sender_id, DataValue::Number(sender)), + (DataTypes::amount, DataValue::Number(amount)), + ]) + }) + .collect(); + + let response = CommunicationValue::new(CommunicationType::get_notifications) + .with_id(cv.get_id()) + .add_data(DataTypes::notifications, DataValue::Array(json_array)); + self.send(&response).await + } else { + Ok(()) + } + } + + async fn handle_read_notification(&self, cv: CommunicationValue) -> OmikronResult<()> { + let user_id = cv.get_sender() as i64; + if let Some(other_id) = cv + .get_data(DataTypes::sender_id) + .as_number() + .map(|n| n as i64) + { + if sql::read_notification(user_id, other_id).await.is_ok() { + let response = CommunicationValue::new(CommunicationType::read_notification) + .with_id(cv.get_id()); + self.send(&response).await + } else { + Ok(()) + } + } else { + Ok(()) + } + } + + async fn handle_push_notification(&self, cv: CommunicationValue) -> OmikronResult<()> { + let user_id = cv.get_sender() as i64; + if let Some(other_id) = cv + .get_data(DataTypes::sender_id) + .as_number() + .map(|n| n as i64) + { + if sql::add_notification(user_id, other_id).await.is_ok() { + let response = CommunicationValue::new(CommunicationType::push_notification) + .with_id(cv.get_id()); + self.send(&response).await + } else { + Ok(()) + } + } else { + Ok(()) + } + } + + async fn handle_ping(&self, cv: CommunicationValue) -> OmikronResult<()> { if let DataValue::Number(last_ping) = cv.get_data(DataTypes::last_ping) { - if let Ok(val) = last_ping.to_string().parse::() { - *self.ping.write().await = val; - } + *self.ping.write().await = *last_ping; } - let _ = self - .send(&CommunicationValue::new(CommunicationType::pong).with_id(cv.get_id())) - .await; + let response = CommunicationValue::new(CommunicationType::pong).with_id(cv.get_id()); + self.send(&response).await } - pub async fn handle_close(&self) { - if self.is_identified().await { - let omikron_id = self.get_omikron_id().await; + // ------------------------------------------------------------------------- + // Utilities + // ------------------------------------------------------------------------- + async fn send(&self, cv: &CommunicationValue) -> OmikronResult<()> { + log_cv_out!(PrintType::Omikron, cv); + + let guard = self.sender.lock().await; + let sender = guard.as_ref().ok_or(OmikronError::NotConnected)?; + + sender + .send(cv) + .await + .map_err(|e| OmikronError::Send(e.to_string())) + } + + async fn send_error_response( + &self, + message_id: u32, + error_type: CommunicationType, + ) -> OmikronResult<()> { + let error = CommunicationValue::new(error_type).with_id(message_id); + self.send(&error).await + } + + pub async fn close(&self) {} + + async fn cleanup(&self) { + if let Some(omikron_id) = self.state.read().await.omikron_id() { if omikron_id != 0 { - log_in!(PrintType::Omega, "Omikron Disconnected"); - + log_in!(omikron_id, PrintType::Omega, "Omikron disconnected"); omikron_manager::remove_omikron(omikron_id).await; user_online_tracker::untrack_omikron(omikron_id).await; } } + + // Cancel cleanup task + if let Some(handle) = self.cleanup_handle.lock().await.take() { + handle.abort(); + } + } + + fn arc_self(&self) -> Arc { + // 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 directly") + } + + // Public API for external use + pub async fn is_authenticated(&self) -> bool { + self.state.read().await.is_authenticated() + } + + pub async fn get_omikron_id(&self) -> Option { + self.state.read().await.omikron_id() + } + + pub async fn send_message(&self, cv: &CommunicationValue) -> OmikronResult<()> { + self.send(cv).await } } + +// ============================================================================ +// Server Startup +// ============================================================================ + +pub async fn start(port: u16) -> Result<(), Box> { + let _ = aws_lc_rs::default_provider().install_default(); + + let tls_cfg = load_tls().expect("TLS config failed"); + let server_crypto = quinn::crypto::rustls::QuicServerConfig::try_from(tls_cfg)?; + let server_cfg = ServerConfig::with_crypto(Arc::new(server_crypto)); + + let mut host: Host = epsilon_native::host(port, server_cfg).await?; + log!("OmikronServer listening on port {}", port); + + while let Some((sender, receiver)) = host.next().await { + tokio::spawn(async move { + let conn = OmikronConnection::new(sender); + conn.handle(receiver).await; + }); + } + + Ok(()) +} + +fn load_tls() -> Option { + let _ = aws_lc_rs::default_provider().install_default(); + + let mut cert_pem = load_file_buf("certs", "cert.pem").ok()?; + let cert_chain = rustls_pemfile::certs(&mut cert_pem) + .collect::, _>>() + .ok()?; + + let key_pem = load_file_vec("certs", "key.pem").ok()?; + let key_der = rustls_pemfile::private_key(&mut &*key_pem).ok()??; + + let cfg = CryptoConfig::builder() + .with_no_client_auth() + .with_single_cert(cert_chain, key_der) + .ok()?; + + Some(cfg) +} diff --git a/src/server/omikron_manager.rs b/src/server/omikron_manager.rs index f481702..37bad20 100644 --- a/src/server/omikron_manager.rs +++ b/src/server/omikron_manager.rs @@ -8,7 +8,13 @@ pub static OMIKRON_CONNECTIONS: Lazy>> = Lazy::new(|| DashMap::new()); pub async fn add_omikron(conn: Arc) { - let id = conn.get_omikron_id().await; + let id = match conn.get_omikron_id().await { + Some(id) => id, + _ => { + conn.close().await; + return; + } + }; if let Some(old) = OMIKRON_CONNECTIONS.insert(id, conn.clone()) { old.close().await;