From a968d5e3576d074863a7b81a88e55c832b9620bc Mon Sep 17 00:00:00 2001 From: zrll Date: Tue, 6 Jun 2023 20:05:58 +0800 Subject: [PATCH] add authenticate --- Cargo.lock | 486 ++++++++++++++++++++++++++++++++++++++ Cargo.toml | 2 + src/authenticate/black.rs | 84 +++++++ src/authenticate/mod.rs | 14 ++ src/authenticate/token.rs | 52 ++++ src/broadcast/fire.rs | 10 +- src/config.rs | 92 +++++++- src/main.rs | 71 +++--- src/monitor/mod.rs | 1 - 9 files changed, 769 insertions(+), 43 deletions(-) create mode 100644 src/authenticate/black.rs create mode 100644 src/authenticate/mod.rs create mode 100644 src/authenticate/token.rs diff --git a/Cargo.lock b/Cargo.lock index a29e8c3..fb128cb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -23,6 +23,12 @@ dependencies = [ "libc", ] +[[package]] +name = "anyhow" +version = "1.0.71" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c7d0618f0e0b7e8ff11427422b64564d5fb0be1940354bfe2e0529b18a9d9b8" + [[package]] name = "arc-swap" version = "0.4.8" @@ -35,6 +41,30 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" +[[package]] +name = "base16ct" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c7f02d4ea65f2c1853089ffd8d2787bdbc63de2f0d29dedbcf8ccdfa0ccd4cf" + +[[package]] +name = "base64" +version = "0.21.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "604178f6c5c21f02dc555784810edfb88d34ac2c73b2eae109655649ee73ce3d" + +[[package]] +name = "base64ct" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8c3c1a368f70d6cf7302d78f8f7093da241fb8e8807c05cc9e51a125895a6d5b" + +[[package]] +name = "binstring" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e0d60973d9320722cb1206f412740e162a33b8547ea8d6be75d7cff237c7a85" + [[package]] name = "bitflags" version = "1.3.2" @@ -95,6 +125,24 @@ dependencies = [ "winapi", ] +[[package]] +name = "coarsetime" +version = "0.1.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a90d114103adbc625300f346d4d09dfb4ab1c4a8df6868435dd903392ecf4354" +dependencies = [ + "libc", + "once_cell", + "wasi 0.11.0+wasi-snapshot-preview1", + "wasm-bindgen", +] + +[[package]] +name = "const-oid" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520fbf3c07483f94e3e3ca9d0cfd913d7718ef2483d2cfd91c0d9e91474ab913" + [[package]] name = "convert_case" version = "0.5.0" @@ -168,6 +216,18 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crypto-bigint" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf4c2f4e1afd912bc40bfd6fed5d9dc1f288e0ba01bfcc835cc5bc3eb13efe15" +dependencies = [ + "generic-array", + "rand_core", + "subtle", + "zeroize", +] + [[package]] name = "crypto-common" version = "0.1.6" @@ -178,12 +238,40 @@ dependencies = [ "typenum", ] +[[package]] +name = "ct-codecs" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3b7eb4404b8195a9abb6356f4ac07d8ba267045c8d6d220ac4dc992e6cc75df" + [[package]] name = "data-encoding" version = "2.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2e66c9d817f1720209181c316d28635c050fa304f9c79e47a520882661b7308" +[[package]] +name = "der" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1a467a65c5e759bce6e65eaf91cc29f466cdc57cb65777bd646872a8a1fd4de" +dependencies = [ + "const-oid", + "pem-rfc7468 0.6.0", + "zeroize", +] + +[[package]] +name = "der" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56acb310e15652100da43d130af8d97b509e95af61aab1c5a7939ef24337ee17" +dependencies = [ + "const-oid", + "pem-rfc7468 0.7.0", + "zeroize", +] + [[package]] name = "digest" version = "0.10.7" @@ -191,7 +279,33 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ "block-buffer", + "const-oid", "crypto-common", + "subtle", +] + +[[package]] +name = "ecdsa" +version = "0.16.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0997c976637b606099b9985693efa3581e84e41f5c11ba5255f88711058ad428" +dependencies = [ + "der 0.7.6", + "digest", + "elliptic-curve", + "rfc6979", + "signature 2.1.0", + "spki 0.7.2", +] + +[[package]] +name = "ed25519-compact" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a3d382e8464107391c8706b4c14b087808ecb909f6c15c34114bc42e53a9e4c" +dependencies = [ + "ct-codecs", + "getrandom", ] [[package]] @@ -200,6 +314,37 @@ version = "1.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7fcaabb2fef8c910e7f4c7ce9f67a1283a1715879a7c230ca9d6d1ae31f16d91" +[[package]] +name = "elliptic-curve" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "968405c8fdc9b3bf4df0a6638858cc0b52462836ab6b1c87377785dd09cf1c0b" +dependencies = [ + "base16ct", + "crypto-bigint", + "digest", + "ff", + "generic-array", + "group", + "hkdf", + "pem-rfc7468 0.7.0", + "pkcs8 0.10.2", + "rand_core", + "sec1", + "subtle", + "zeroize", +] + +[[package]] +name = "ff" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ded41244b729663b1e574f1b4fb731469f69f79c17667b5d776b16cda0479449" +dependencies = [ + "rand_core", + "subtle", +] + [[package]] name = "flate2" version = "1.0.26" @@ -266,6 +411,7 @@ checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" dependencies = [ "typenum", "version_check", + "zeroize", ] [[package]] @@ -279,6 +425,17 @@ dependencies = [ "wasi 0.11.0+wasi-snapshot-preview1", ] +[[package]] +name = "group" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0f9ef7462f7c099f518d754361858f86d8a07af53ba9af0fe635bbccb151a63" +dependencies = [ + "ff", + "rand_core", + "subtle", +] + [[package]] name = "hashbrown" version = "0.12.3" @@ -294,6 +451,48 @@ dependencies = [ "libc", ] +[[package]] +name = "hkdf" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "791a029f6b9fc27657f6f188ec6e5e43f6911f6f878e0dc5501396e09809d437" +dependencies = [ + "hmac", +] + +[[package]] +name = "hmac" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" +dependencies = [ + "digest", +] + +[[package]] +name = "hmac-sha1-compact" +version = "1.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05e2440a0078e20c3b68ca01234cea4219f23e64b0c0bdb1200c5550d54239bb" + +[[package]] +name = "hmac-sha256" +version = "1.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc736091aacb31ddaa4cd5f6988b3c21e99913ac846b41f32538c5fae5d71bfe" +dependencies = [ + "digest", +] + +[[package]] +name = "hmac-sha512" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520c9c3f6040661669bc5c91e551b605a520c8e0a63a766a91a65adef734d151" +dependencies = [ + "digest", +] + [[package]] name = "http" version = "0.2.9" @@ -432,11 +631,54 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "jwt-simple" +version = "0.11.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0537086995d782ba2fb6c120a88f0d66c5ee5f1208a3559826d4cf2264b170da" +dependencies = [ + "anyhow", + "binstring", + "coarsetime", + "ct-codecs", + "ed25519-compact", + "hmac-sha1-compact", + "hmac-sha256", + "hmac-sha512", + "k256", + "p256", + "p384", + "rand", + "rsa", + "serde", + "serde_json", + "spki 0.6.0", + "thiserror", + "zeroize", +] + +[[package]] +name = "k256" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cadb76004ed8e97623117f3df85b17aaa6626ab0b0831e6573f104df16cd1bcc" +dependencies = [ + "cfg-if", + "ecdsa", + "elliptic-curve", + "once_cell", + "sha2", + "signature 2.1.0", +] + [[package]] name = "lazy_static" version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" +dependencies = [ + "spin", +] [[package]] name = "libc" @@ -444,6 +686,12 @@ version = "0.2.145" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc86cde3ff845662b8f4ef6cb50ea0e20c524eb3d29ae048287e06a1b3fa6a81" +[[package]] +name = "libm" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7012b1bbb0719e1097c47611d3898568c546d597c2e74d66f6087edd5233ff4" + [[package]] name = "linked-hash-map" version = "0.5.6" @@ -503,8 +751,10 @@ dependencies = [ name = "machine_monitor" version = "0.1.0" dependencies = [ + "base64", "chrono", "hyper", + "jwt-simple", "lazy_static", "serde", "serde_json", @@ -559,6 +809,44 @@ dependencies = [ "winapi", ] +[[package]] +name = "num-bigint-dig" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2399c9463abc5f909349d8aa9ba080e0b88b3ce2885389b60b993f39b1a56905" +dependencies = [ + "byteorder", + "lazy_static", + "libm", + "num-integer", + "num-iter", + "num-traits", + "rand", + "smallvec", + "zeroize", +] + +[[package]] +name = "num-integer" +version = "0.1.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "225d3389fb3509a24c93f5c29eb6bde2586b98d9f016636dff58d7c6f7569cd9" +dependencies = [ + "autocfg", + "num-traits", +] + +[[package]] +name = "num-iter" +version = "0.1.43" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d03e6c028c5dc5cac6e2dec0efda81fc887605bb3d884578bb6d6bf7514e252" +dependencies = [ + "autocfg", + "num-integer", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.15" @@ -566,6 +854,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "578ede34cf02f8924ab9447f50c28075b4d3e5b269972345e7e0372b38c6cdcd" dependencies = [ "autocfg", + "libm", ] [[package]] @@ -593,6 +882,30 @@ dependencies = [ "num-traits", ] +[[package]] +name = "p256" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c9863ad85fa8f4460f9c48cb909d38a0d689dba1f6f6988a5e3e0d31071bcd4b" +dependencies = [ + "ecdsa", + "elliptic-curve", + "primeorder", + "sha2", +] + +[[package]] +name = "p384" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70786f51bcc69f6a4c0360e063a4cac5419ef7c5cd5b3c99ad70f3be5ba79209" +dependencies = [ + "ecdsa", + "elliptic-curve", + "primeorder", + "sha2", +] + [[package]] name = "parking_lot" version = "0.11.2" @@ -641,6 +954,24 @@ dependencies = [ "windows-sys 0.45.0", ] +[[package]] +name = "pem-rfc7468" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d159833a9105500e0398934e205e0773f0b27529557134ecfc51c27646adac" +dependencies = [ + "base64ct", +] + +[[package]] +name = "pem-rfc7468" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88b39c9bfcfc231068454382784bb460aae594343fb030d46e9f50a645418412" +dependencies = [ + "base64ct", +] + [[package]] name = "percent-encoding" version = "2.3.0" @@ -659,12 +990,53 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "pkcs1" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eff33bdbdfc54cc98a2eca766ebdec3e1b8fb7387523d5c9c9a2891da856f719" +dependencies = [ + "der 0.6.1", + "pkcs8 0.9.0", + "spki 0.6.0", + "zeroize", +] + +[[package]] +name = "pkcs8" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9eca2c590a5f85da82668fa685c09ce2888b9430e83299debf1f34b65fd4a4ba" +dependencies = [ + "der 0.6.1", + "spki 0.6.0", +] + +[[package]] +name = "pkcs8" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7" +dependencies = [ + "der 0.7.6", + "spki 0.7.2", +] + [[package]] name = "ppv-lite86" version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" +[[package]] +name = "primeorder" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c2fcef82c0ec6eefcc179b978446c399b3cdf73c392c35604e399eee6df1ee3" +dependencies = [ + "elliptic-curve", +] + [[package]] name = "proc-macro2" version = "1.0.59" @@ -756,6 +1128,37 @@ dependencies = [ "bitflags", ] +[[package]] +name = "rfc6979" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dd2a808d456c4a54e300a23e9f5a67e122c3024119acbfd73e3bf664491cb2" +dependencies = [ + "hmac", + "subtle", +] + +[[package]] +name = "rsa" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "094052d5470cbcef561cb848a7209968c9f12dfa6d668f4bca048ac5de51099c" +dependencies = [ + "byteorder", + "digest", + "num-bigint-dig", + "num-integer", + "num-iter", + "num-traits", + "pkcs1", + "pkcs8 0.9.0", + "rand_core", + "signature 1.6.4", + "smallvec", + "subtle", + "zeroize", +] + [[package]] name = "ryu" version = "1.0.13" @@ -768,6 +1171,20 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" +[[package]] +name = "sec1" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0aec48e813d6b90b15f0b8948af3c63483992dee44c03e9930b3eebdabe046e" +dependencies = [ + "base16ct", + "der 0.7.6", + "generic-array", + "pkcs8 0.10.2", + "subtle", + "zeroize", +] + [[package]] name = "serde" version = "1.0.163" @@ -841,6 +1258,17 @@ dependencies = [ "digest", ] +[[package]] +name = "sha2" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82e6b795fe2e3b1e845bafcb27aa35405c4d47cdfc92af5fc8d3002f76cebdc0" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "signal-hook-registry" version = "1.4.1" @@ -850,6 +1278,26 @@ dependencies = [ "libc", ] +[[package]] +name = "signature" +version = "1.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74233d3b3b2f6d4b006dc19dee745e73e2a6bfb6f93607cd3b02bd5b00797d7c" +dependencies = [ + "digest", + "rand_core", +] + +[[package]] +name = "signature" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e1788eed21689f9cf370582dfc467ef36ed9c707f073528ddafa8d83e3b8500" +dependencies = [ + "digest", + "rand_core", +] + [[package]] name = "simple-log" version = "1.6.0" @@ -880,6 +1328,38 @@ dependencies = [ "winapi", ] +[[package]] +name = "spin" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e63cff320ae2c57904679ba7cb63280a3dc4613885beafb148ee7bf9aa9042d" + +[[package]] +name = "spki" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67cf02bbac7a337dc36e4f5a693db6c21e7863f45070f7064577eb4367a3212b" +dependencies = [ + "base64ct", + "der 0.6.1", +] + +[[package]] +name = "spki" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1e996ef02c474957d681f1b05213dfb0abab947b446a62d37770b23500184a" +dependencies = [ + "base64ct", + "der 0.7.6", +] + +[[package]] +name = "subtle" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81cdd64d312baedb58e21336b31bc043b77e01cc99033ce76ef539f78e965ebc" + [[package]] name = "syn" version = "2.0.18" @@ -1384,3 +1864,9 @@ checksum = "56c1936c4cc7a1c9ab21a1ebb602eb942ba868cbd44a99cb7cdc5892335e1c85" dependencies = [ "linked-hash-map", ] + +[[package]] +name = "zeroize" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a0956f1ba7c7909bfb66c2e9e4124ab6f6482560f6628b5aaeba39207c9aad9" diff --git a/Cargo.toml b/Cargo.toml index c88240f..735689e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,3 +16,5 @@ simple-log = "1.6.0" chrono = "0.4.26" hyper = "0.14.26" tungstenite = "0.19.0" +base64 = "0.21.2" +jwt-simple = "0.11.5" diff --git a/src/authenticate/black.rs b/src/authenticate/black.rs new file mode 100644 index 0000000..a84aa1d --- /dev/null +++ b/src/authenticate/black.rs @@ -0,0 +1,84 @@ +use std::collections::HashMap; +use std::sync::Mutex; +use lazy_static::lazy_static; +use simple_log::{info, warn}; + +lazy_static! { + static ref BLACKLIST: Mutex> = Mutex::new(HashMap::new()); +} + +pub async fn is_in_blacklist(ip: String) -> bool { + let mut trys = 0; + loop { + trys += 1; + if trys > 100 { + return true; + } + + if let Ok(map) = BLACKLIST.try_lock() { + return match map.get(ip.as_str()) { + None => { + false + } + Some(a) => { + if *a > 50 { + true + } else { + false + } + } + } + } + } +} + +pub async fn clear_blacklist() { + loop { + tokio::time::sleep(std::time::Duration::from_secs(1200)).await; + + let mut trys = 0; + loop { + trys += 1; + if trys > 100 { + break; + } + + if let Ok(mut map) = BLACKLIST.try_lock() { + if map.len() > 0 { + map.clear(); + info!("Clearing black list."); + } + break; + } + } + + } +} + +pub async fn record(ip: String) { + warn!("Recording {} into blacklist.", ip.as_str()); + let mut trys = 0; + loop { + trys += 1; + if trys > 100 { + return; + } + + if let Ok(mut map) = BLACKLIST.try_lock() { + let num = match map.get(ip.as_str()) { + None => { + 0 + } + Some(a) => { + *a + } + }; + map.insert(ip.clone(), num + 1); + return; + } + } +} + +pub fn init() { + tokio::spawn(clear_blacklist()); +} \ No newline at end of file diff --git a/src/authenticate/mod.rs b/src/authenticate/mod.rs new file mode 100644 index 0000000..7410dd2 --- /dev/null +++ b/src/authenticate/mod.rs @@ -0,0 +1,14 @@ +pub mod token; +pub mod black; + +use serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct UserInfo {//_id: mc uuid + pub _id: String, + pub display_name: String, + pub enabled: bool, + pub group: Vec, + pub bind_qq: Option, + pub ban_reason: Option +} diff --git a/src/authenticate/token.rs b/src/authenticate/token.rs new file mode 100644 index 0000000..3656ee7 --- /dev/null +++ b/src/authenticate/token.rs @@ -0,0 +1,52 @@ +use base64::Engine; +use jwt_simple::prelude::{HS256Key, MACLike}; +use simple_log::debug; +use simple_log::log::error; +use crate::authenticate::black::{is_in_blacklist, record}; +use crate::authenticate::UserInfo; +use crate::CONFIG; + +impl UserInfo { + pub async fn from_token(token: String) -> Result { + let key = match &CONFIG.authenticate.secret { + None => { return Err("Key is not defined.".to_string())} + Some(a) => {a} + }; + let key = HS256Key::from_bytes(base64::engine::general_purpose::STANDARD.decode(key).unwrap().as_slice()); + + let claims = key.verify_token::(&token, None); + match claims { + Ok(info) => Ok(info.custom), + Err(err) => { + + error!("{}", err.to_string()); + Err("Cannot parse token or token is valid".to_string()) + } + } + } +} + +pub async fn authenticate_token(token: String, ip: String) -> Result<(), String> { + if is_in_blacklist(ip.clone()).await { + debug!("{} is in black list.", ip); + return Err("In blacklist.".to_string()); + } + + let info = match UserInfo::from_token(token).await { + Ok(a) => {a} + Err(err) => { + record(ip.clone()).await; + return Err(err); + } + }; + + let group = info.group; + for permission in group { + if permission == "admin" { + return Ok(()); + } + } + + record(ip).await; + Err("No permission".to_string()) +} \ No newline at end of file diff --git a/src/broadcast/fire.rs b/src/broadcast/fire.rs index babc822..ab55093 100644 --- a/src/broadcast/fire.rs +++ b/src/broadcast/fire.rs @@ -1,6 +1,6 @@ use std::ops::{DerefMut}; -use simple_log::{info, warn}; -use tungstenite::{Message, WebSocket}; +use simple_log::{debug, info, warn}; +use tungstenite::{Message}; use crate::broadcast::{CLIENTS}; use crate::monitor::MachineInfo; @@ -22,15 +22,15 @@ pub async fn fire(data: MachineInfo) { for client_info in clients.deref_mut() { let (ip, client) = client_info.as_mut().unwrap(); match client.write_message(Message::Text(data_str.clone())) { - Ok(_) => {}, + Ok(_) => {} Err(_) => { warn!("Removing client: {}.", ip); *client_info = None; } }; } - clients.retain(|x| x.is_some()); if clients.len() > 0 { - info!("Send to {} clients.", clients.len()); + clients.retain(|x| x.is_some()); + debug!("Send to {} clients.", clients.len()); } } \ No newline at end of file diff --git a/src/config.rs b/src/config.rs index f450314..306ab79 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,11 +1,15 @@ use std::fs::OpenOptions; use std::io::{Read, Write}; +use base64::Engine; +use jwt_simple::prelude::HS256Key; use serde::{Deserialize, Serialize}; +use simple_log::error; use toml::Value; #[derive(Serialize, Deserialize, Debug, Clone)] pub struct Config { pub connection: ConnectionSetting, + pub authenticate: AuthenticateSetting, } #[derive(Serialize, Deserialize, Debug, Clone)] @@ -14,16 +18,21 @@ pub struct ConnectionSetting { pub port: i32, } +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct AuthenticateSetting { + pub method: String, + pub secret: Option, + pub db_url: Option, + pub db_name: Option, +} + pub fn get_config() -> Config { let mut config: String = Default::default(); let mut file = OpenOptions::new().read(true).write(true).create(true).open("config.toml").expect("Can not open 'config.toml'"); file.read_to_string(&mut config).expect("Can not read 'config.toml'"); - if config == "" { - config = "[mongodb]\n[connect]\n[feature]".to_string(); - } - let mut config: Value = toml::from_str(config.as_str()).expect("Cannot parse config file, you should probably generate another."); + let config: Value = toml::from_str(config.as_str()).expect("Cannot parse config file, you should probably generate another."); let mut edit = false; @@ -32,6 +41,12 @@ pub fn get_config() -> Config { address: "".to_string(), port: 0, }, + authenticate: AuthenticateSetting { + method: "".to_string(), + secret: None, + db_url: None, + db_name: None, + }, }; result_config.connection = if let Some(a) = config.get("connection") { @@ -39,16 +54,22 @@ pub fn get_config() -> Config { if let Some(c) = b.as_str() { c.to_string() } else { + edit = true; "127.0.0.1".to_string() } } else { edit = true; "127.0.0.1".to_string() }; - let port = if let Some(b) = a.get("address") { + + + + + let port = if let Some(b) = a.get("port") { if let Some(c) = b.as_integer() { c as i32 } else { + edit = true; 7890 } } else { @@ -62,15 +83,74 @@ pub fn get_config() -> Config { port, } } else { + edit = true; ConnectionSetting { address: "127.0.0.1".to_string(), port: 7890, } }; + + result_config.authenticate = if let Some(a) = config.get("authenticate") { + let method = if let Some(b) = a.get("method") { + if let Some(c) = b.as_str() { + if c == "token" || c == "key" { + c.to_string() + } else { + edit = true; + "token".to_string() + } + } else { + edit = true; + "token".to_string() + } + } else { + edit = true; + "token".to_string() + }; + + let secret = if method == "token" { + let key = if let Some(b) = a.get("secret") { + if let Some(c) = b.as_str() { + c.to_string() + } else { + edit = true; + let key = HS256Key::generate(); + base64::engine::general_purpose::STANDARD.encode(key.to_bytes()) + } + } else { + edit = true; + let key = HS256Key::generate(); + base64::engine::general_purpose::STANDARD.encode(key.to_bytes()) + }; + + Some(key) + } else { + None + }; + + AuthenticateSetting { + method, + secret, + db_url: None, + db_name: None, + } + } else { + let key = HS256Key::generate(); + let key = base64::engine::general_purpose::STANDARD.encode(key.to_bytes()); + AuthenticateSetting { + method: "token".to_string(), + secret: Some(key), + db_url: None, + db_name: None, + } + }; + if edit { + error!("Saving config."); let mut file = OpenOptions::new().write(true).create(true).open("config.toml").expect("Can not open 'config.toml'"); - file.write(toml::to_string(&config).unwrap().as_bytes()).expect("Cannot save config."); + file.write(toml::to_string(&result_config).unwrap().as_bytes()).expect("Cannot save config."); + panic!("Config changed") } result_config diff --git a/src/main.rs b/src/main.rs index 745bd20..2b5c100 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,16 +1,17 @@ mod config; mod monitor; mod broadcast; +mod authenticate; use std::borrow::Cow; -use std::net::{TcpListener}; -use std::thread::spawn; +use std::net::{TcpListener, TcpStream}; use chrono::Local; use lazy_static::lazy_static; -use simple_log::{info, LogConfigBuilder}; +use simple_log::{info, LogConfigBuilder, warn}; use tungstenite::{accept}; use tungstenite::protocol::CloseFrame; use tungstenite::protocol::frame::coding::CloseCode; +use crate::authenticate::token::authenticate_token; use crate::broadcast::CLIENTS; use crate::config::{Config, get_config}; @@ -38,39 +39,47 @@ async fn main() { simple_log::new(config).expect("Cannot init logger"); monitor::init(); + authenticate::black::init(); let server = TcpListener::bind(format!("{}:{}", &CONFIG.connection.address, &CONFIG.connection.port)).unwrap(); info!("listening at {}:{}.", &CONFIG.connection.address, &CONFIG.connection.port); for stream in server.incoming() { - spawn(move || { - let stream = stream.unwrap(); - let addr = &stream.peer_addr().unwrap(); - let mut websocket = accept(stream).unwrap(); + tokio::spawn(handle_income(stream.unwrap())); + } +} + +async fn handle_income(stream: TcpStream) { + let addr = &stream.peer_addr().unwrap().ip(); + let mut websocket = accept(stream).unwrap(); - if let Ok(a) = websocket.read_message() { - info!("Received: {} from {}", a.to_string(), addr.to_string()); + if let Ok(a) = websocket.read_message() { + info!("Received: {} from {}", a.to_string(), addr.to_string()); - if a.to_string() == "Hello!password" { - let mut trys = 0; - loop { - trys += 1; - if trys > 100 { - break; - } + match authenticate_token(a.to_string(), addr.to_string()).await { + Ok(_) => { + let mut trys = 0; + loop { + trys += 1; + if trys > 100 { + websocket.close(Some(CloseFrame { + code: CloseCode::Again, + reason: Cow::from("Server error"), + })).unwrap(); + warn!("Failed to connect with {}.", addr.to_string()); + break; + } - if let Ok(mut clients) = CLIENTS.try_lock() { - clients.push(Some((addr.to_string(), websocket))); - break; - } else { - break; - } - }; - } else { - websocket.close(Some(CloseFrame { - code: CloseCode::Policy, - reason: Cow::from("Wrong password"), - })).unwrap(); - } + if let Ok(mut clients) = CLIENTS.try_lock() { + clients.push(Some((addr.to_string(), websocket))); + break; + } + }; } - }); + Err(err) => { + websocket.close(Some(CloseFrame { + code: CloseCode::Policy, + reason: Cow::from(err), + })).unwrap(); + } + } } -} +} \ No newline at end of file diff --git a/src/monitor/mod.rs b/src/monitor/mod.rs index 30a18e7..17f42d0 100644 --- a/src/monitor/mod.rs +++ b/src/monitor/mod.rs @@ -1,7 +1,6 @@ pub mod refresh; mod query; -use lazy_static::lazy_static; use serde::{Deserialize, Serialize}; use crate::monitor::query::query;