From af7d4157eee79aa4ed362476a1ebd74265fe35e5 Mon Sep 17 00:00:00 2001 From: Noor Date: Tue, 10 Jun 2025 16:45:16 +0530 Subject: [PATCH] WIP on implementing redirect implement public IP handling Add pub_sub_node field to vm_node table in SQL schema Refactor BRAIN_GRPC_ADDR to use LazyLock for dynamic IP assignment Implement current_brain_endpoint function to retrieve public IP --- Cargo.lock | 343 ++++++++++++++++++++++++++++++++----- Cargo.toml | 1 + src/bin/brain.rs | 4 +- src/constants.rs | 28 ++- src/db/mod.rs | 2 + src/grpc/mod.rs | 52 +++++- src/grpc/vm.rs | 9 +- surql/tables.sql | 2 + tests/grpc_general_test.rs | 1 + tests/grpc_vm_cli_test.rs | 53 +++++- 10 files changed, 444 insertions(+), 51 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5e912e8..d93a0f2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -263,7 +263,7 @@ dependencies = [ "fnv", "futures-timer", "futures-util", - "http", + "http 1.3.1", "indexmap 2.9.0", "mime", "multer", @@ -285,7 +285,7 @@ checksum = "29db05b624fb6352fc11bfe30c54ab1b16a1fe937d7c05a783f4e88ef1292b3b" dependencies = [ "Inflector", "async-graphql-parser", - "darling", + "darling 0.20.11", "proc-macro-crate", "proc-macro2", "quote", @@ -390,8 +390,8 @@ dependencies = [ "axum-core", "bytes", "futures-util", - "http", - "http-body", + "http 1.3.1", + "http-body 1.0.1", "http-body-util", "itoa", "matchit", @@ -416,8 +416,8 @@ dependencies = [ "async-trait", "bytes", "futures-util", - "http", - "http-body", + "http 1.3.1", + "http-body 1.0.1", "http-body-util", "mime", "pin-project-lite", @@ -908,14 +908,38 @@ dependencies = [ "syn 2.0.100", ] +[[package]] +name = "darling" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d706e75d87e35569db781a9b5e2416cff1236a47ed380831f959382ccd5f858" +dependencies = [ + "darling_core 0.10.2", + "darling_macro 0.10.2", +] + [[package]] name = "darling" version = "0.20.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" dependencies = [ - "darling_core", - "darling_macro", + "darling_core 0.20.11", + "darling_macro 0.20.11", +] + +[[package]] +name = "darling_core" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0c960ae2da4de88a91b2d920c2a7233b400bc33cb28453a2987822d8392519b" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim 0.9.3", + "syn 1.0.109", ] [[package]] @@ -928,17 +952,28 @@ dependencies = [ "ident_case", "proc-macro2", "quote", - "strsim", + "strsim 0.11.1", "syn 2.0.100", ] +[[package]] +name = "darling_macro" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9b5a2f4ac4969822c62224815d069952656cadc7084fdca9751e6d959189b72" +dependencies = [ + "darling_core 0.10.2", + "quote", + "syn 1.0.109", +] + [[package]] name = "darling_macro" version = "0.20.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" dependencies = [ - "darling_core", + "darling_core 0.20.11", "quote", "syn 2.0.100", ] @@ -997,6 +1032,31 @@ dependencies = [ "serde", ] +[[package]] +name = "derive_builder" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2658621297f2cf68762a6f7dc0bb7e1ff2cfd6583daef8ee0fed6f7ec468ec0" +dependencies = [ + "darling 0.10.2", + "derive_builder_core", + "proc-macro2", + "quote", + "syn 1.0.109", +] + +[[package]] +name = "derive_builder_core" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2791ea3e372c8495c0bc2033991d76b512cd799d07491fbd6890124db9458bef" +dependencies = [ + "darling 0.10.2", + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "detee-shared" version = "0.1.0" @@ -1072,6 +1132,18 @@ dependencies = [ "urlencoding", ] +[[package]] +name = "dns-lookup" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53ecafc952c4528d9b51a458d1a8904b81783feff9fde08ab6ed2545ff396872" +dependencies = [ + "cfg-if", + "libc", + "socket2 0.4.10", + "winapi", +] + [[package]] name = "doc-comment" version = "0.3.3" @@ -1149,6 +1221,18 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c34f04666d835ff5d62e058c3995147c06f42fe86ff053337632bca83e42702d" +[[package]] +name = "enum-as-inner" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "570d109b813e904becc80d8d5da38376818a143348413f7149f1340fe04754d4" +dependencies = [ + "heck 0.4.1", + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "env_filter" version = "0.1.3" @@ -1499,7 +1583,7 @@ dependencies = [ "fnv", "futures-core", "futures-sink", - "http", + "http 1.3.1", "indexmap 2.9.0", "slab", "tokio", @@ -1609,6 +1693,17 @@ dependencies = [ "syn 2.0.100", ] +[[package]] +name = "http" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "601cbb57e577e2f5ef5be8e7b83f0f63994f25aa94d673e54a92d5c516d101f1" +dependencies = [ + "bytes", + "fnv", + "itoa", +] + [[package]] name = "http" version = "1.3.1" @@ -1620,6 +1715,17 @@ dependencies = [ "itoa", ] +[[package]] +name = "http-body" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ceab25649e9960c0311ea418d17bee82c0dcec1bd053b5f9a66e265a693bed2" +dependencies = [ + "bytes", + "http 0.2.12", + "pin-project-lite", +] + [[package]] name = "http-body" version = "1.0.1" @@ -1627,7 +1733,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" dependencies = [ "bytes", - "http", + "http 1.3.1", ] [[package]] @@ -1638,8 +1744,8 @@ checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" dependencies = [ "bytes", "futures-core", - "http", - "http-body", + "http 1.3.1", + "http-body 1.0.1", "pin-project-lite", ] @@ -1661,6 +1767,29 @@ version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b112acc8b3adf4b107a8ec20977da0273a8c386765a3ec0229bd500a1443f9f" +[[package]] +name = "hyper" +version = "0.14.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41dfc780fdec9373c01bae43289ea34c972e40ee3c9f6b3c8801a35f35586ce7" +dependencies = [ + "bytes", + "futures-channel", + "futures-core", + "futures-util", + "http 0.2.12", + "http-body 0.4.6", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "socket2 0.5.9", + "tokio", + "tower-service", + "tracing", + "want", +] + [[package]] name = "hyper" version = "1.6.0" @@ -1671,8 +1800,8 @@ dependencies = [ "futures-channel", "futures-util", "h2", - "http", - "http-body", + "http 1.3.1", + "http-body 1.0.1", "httparse", "httpdate", "itoa", @@ -1689,8 +1818,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2d191583f3da1305256f22463b9bb0471acad48a4e534a5218b9963e9c1f59b2" dependencies = [ "futures-util", - "http", - "hyper", + "http 1.3.1", + "hyper 1.6.0", "hyper-util", "rustls", "rustls-pki-types", @@ -1700,13 +1829,27 @@ dependencies = [ "webpki-roots", ] +[[package]] +name = "hyper-system-resolver" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6eea26c5d0b6ab9d72219f65000af310f042a740926f7b2fa3553e774036e2e7" +dependencies = [ + "derive_builder", + "dns-lookup", + "hyper 0.14.32", + "tokio", + "tower-service", + "tracing", +] + [[package]] name = "hyper-timeout" version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" dependencies = [ - "hyper", + "hyper 1.6.0", "hyper-util", "pin-project-lite", "tokio", @@ -1722,12 +1865,12 @@ dependencies = [ "bytes", "futures-channel", "futures-util", - "http", - "http-body", - "hyper", + "http 1.3.1", + "http-body 1.0.1", + "hyper 1.6.0", "libc", "pin-project-lite", - "socket2", + "socket2 0.5.9", "tokio", "tower-service", "tracing", @@ -1881,6 +2024,17 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" +[[package]] +name = "idna" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "418a0a6fab821475f634efe3ccc45c013f742efe03d853e8d3355d5cb850ecf8" +dependencies = [ + "matches", + "unicode-bidi", + "unicode-normalization", +] + [[package]] name = "idna" version = "1.0.3" @@ -2181,6 +2335,12 @@ dependencies = [ "tendril", ] +[[package]] +name = "matches" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2532096657941c2fea9c289d370a250971c689d4f143798ff67113ec042024a5" + [[package]] name = "matchit" version = "0.7.3" @@ -2281,7 +2441,7 @@ dependencies = [ "bytes", "encoding_rs", "futures-util", - "http", + "http 1.3.1", "httparse", "memchr", "mime", @@ -2756,7 +2916,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf" dependencies = [ "heck 0.5.0", - "itertools 0.11.0", + "itertools 0.14.0", "log", "multimap", "once_cell", @@ -2776,7 +2936,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.11.0", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.100", @@ -2826,6 +2986,27 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "public-ip" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b4c40db5262d93298c363a299f8bc1b3a956a78eecddba3bc0e58b76e2f419a" +dependencies = [ + "dns-lookup", + "futures-core", + "futures-util", + "http 0.2.12", + "hyper 0.14.32", + "hyper-system-resolver", + "pin-project-lite", + "thiserror 1.0.69", + "tokio", + "tracing", + "tracing-futures", + "trust-dns-client", + "trust-dns-proto", +] + [[package]] name = "quick_cache" version = "0.5.2" @@ -2851,7 +3032,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls", - "socket2", + "socket2 0.5.9", "thiserror 2.0.12", "tokio", "tracing", @@ -2887,7 +3068,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2", + "socket2 0.5.9", "tracing", "windows-sys 0.59.0", ] @@ -3103,10 +3284,10 @@ dependencies = [ "bytes", "futures-core", "futures-util", - "http", - "http-body", + "http 1.3.1", + "http-body 1.0.1", "http-body-util", - "hyper", + "hyper 1.6.0", "hyper-rustls", "hyper-util", "ipnet", @@ -3517,7 +3698,7 @@ version = "3.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8d00caa5193a3c8362ac2b73be6b9e768aa5a4b2f721d8f4b339600c3cb51f8e" dependencies = [ - "darling", + "darling 0.20.11", "proc-macro2", "quote", "syn 2.0.100", @@ -3649,6 +3830,16 @@ version = "1.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b" +[[package]] +name = "socket2" +version = "0.4.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7916fc008ca5542385b89a3d3ce689953c143e9304a9bf8beec1de48994c0d" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "socket2" version = "0.5.9" @@ -3749,6 +3940,12 @@ dependencies = [ "quote", ] +[[package]] +name = "strsim" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6446ced80d6c486436db5c078dde11a9f73d42b57fb273121e160b84f63d894c" + [[package]] name = "strsim" version = "0.11.1" @@ -3801,6 +3998,7 @@ dependencies = [ "itertools 0.14.0", "log", "nanoid", + "public-ip", "rand 0.8.5", "serde", "serde_json", @@ -3886,7 +4084,7 @@ dependencies = [ "geo", "geo-types", "hex", - "http", + "http 1.3.1", "ipnet", "jsonwebtoken", "lexicmp", @@ -3924,7 +4122,7 @@ dependencies = [ "sha2", "snap", "storekey", - "strsim", + "strsim 0.11.1", "subtle", "sysinfo", "thiserror 1.0.69", @@ -4175,7 +4373,7 @@ dependencies = [ "libc", "mio", "pin-project-lite", - "socket2", + "socket2 0.5.9", "tokio-macros", "windows-sys 0.52.0", ] @@ -4271,17 +4469,17 @@ dependencies = [ "base64 0.22.1", "bytes", "h2", - "http", - "http-body", + "http 1.3.1", + "http-body 1.0.1", "http-body-util", - "hyper", + "hyper 1.6.0", "hyper-timeout", "hyper-util", "percent-encoding", "pin-project", "prost", "rustls-pemfile", - "socket2", + "socket2 0.5.9", "tokio", "tokio-rustls", "tokio-stream", @@ -4383,6 +4581,18 @@ dependencies = [ "once_cell", ] +[[package]] +name = "tracing-futures" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97d095ae15e245a057c8e8451bab9b3ee1e1f68e9ba2b4fbc18d0ac5237835f2" +dependencies = [ + "futures", + "futures-task", + "pin-project", + "tracing", +] + [[package]] name = "trice" version = "0.4.0" @@ -4394,6 +4604,51 @@ dependencies = [ "web-sys", ] +[[package]] +name = "trust-dns-client" +version = "0.20.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b4ef9b9bde0559b78a4abb00339143750085f05e5a453efb7b8bef1061f09dc" +dependencies = [ + "cfg-if", + "data-encoding", + "futures-channel", + "futures-util", + "lazy_static", + "log", + "radix_trie", + "rand 0.8.5", + "thiserror 1.0.69", + "time", + "tokio", + "trust-dns-proto", +] + +[[package]] +name = "trust-dns-proto" +version = "0.20.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca94d4e9feb6a181c690c4040d7a24ef34018d8313ac5044a61d21222ae24e31" +dependencies = [ + "async-trait", + "cfg-if", + "data-encoding", + "enum-as-inner", + "futures-channel", + "futures-io", + "futures-util", + "idna 0.2.3", + "ipnet", + "lazy_static", + "log", + "rand 0.8.5", + "smallvec", + "thiserror 1.0.69", + "tinyvec", + "tokio", + "url", +] + [[package]] name = "try-lock" version = "0.2.5" @@ -4409,7 +4664,7 @@ dependencies = [ "byteorder", "bytes", "data-encoding", - "http", + "http 1.3.1", "httparse", "log", "rand 0.8.5", @@ -4450,6 +4705,12 @@ version = "2.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "75b844d17643ee918803943289730bec8aac480150456169e647ed0b576ba539" +[[package]] +name = "unicode-bidi" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c1cb5db39152898a79168971543b1cb5020dff7fe43c8dc468b0885f5e29df5" + [[package]] name = "unicode-ident" version = "1.0.18" @@ -4518,7 +4779,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32f8b686cadd1473f4bd0117a5d28d36b1ade384ea9b5069a1c40aefed7fda60" dependencies = [ "form_urlencoded", - "idna", + "idna 1.0.3", "percent-encoding", ] diff --git a/Cargo.toml b/Cargo.toml index 9d96c45..048ee45 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,7 @@ thiserror = "2.0.12" nanoid = "0.4.0" dotenv = "0.15.0" hex = "0.4.3" +public-ip = { version = "0.2.2", features = ["dns-resolver"] } [profile.release] lto = true diff --git a/src/bin/brain.rs b/src/bin/brain.rs index 4798063..6020f1a 100644 --- a/src/bin/brain.rs +++ b/src/bin/brain.rs @@ -32,8 +32,6 @@ async fn main() { let db = db::db_connection(&db_url, &db_user, &db_pass, &db_ns, &db_name).await.unwrap(); let db_arc = Arc::new(db); - let addr = BRAIN_GRPC_ADDR.parse().unwrap(); - let snp_daemon_server = BrainVmDaemonServer::new(VmDaemonServer::new(db_arc.clone())); let snp_cli_server = BrainVmCliServer::new(VmCliServer::new(db_arc.clone())); let general_service_server = BrainGeneralCliServer::new(GeneralCliServer::new(db_arc.clone())); @@ -55,7 +53,7 @@ async fn main() { .add_service(general_service_server) .add_service(sgx_daemon_server) .add_service(sgx_cli_server) - .serve(addr) + .serve(*BRAIN_GRPC_ADDR) .await .unwrap(); } diff --git a/src/constants.rs b/src/constants.rs index c8aa89e..5ff7f88 100644 --- a/src/constants.rs +++ b/src/constants.rs @@ -1,6 +1,15 @@ -use std::sync::LazyLock; +use std::{ + net::{IpAddr, SocketAddr}, + sync::LazyLock, +}; +use tokio::sync::OnceCell; + +pub const BRAIN_GRPC_PORT: u16 = 31337; +pub static BRAIN_GRPC_ADDR: LazyLock = LazyLock::new(|| { + let ip = IpAddr::from([0, 0, 0, 0]); + SocketAddr::new(ip, BRAIN_GRPC_PORT) +}); -pub const BRAIN_GRPC_ADDR: &str = "0.0.0.0:31337"; pub const CERT_PATH: &str = "/etc/detee/brain/brain-crt.pem"; pub const CERT_KEY_PATH: &str = "/etc/detee/brain/brain-key.pem"; pub const CONFIG_PATH: &str = "/etc/detee/brain/config.ini"; @@ -20,6 +29,21 @@ pub static ADMIN_ACCOUNTS: LazyLock> = LazyLock::new(|| { .unwrap_or(default_admin_keys) }); +pub static CURRENT_BRAIN_ENDPOINT: OnceCell = OnceCell::const_new(); + +pub async fn current_brain_endpoint() -> Result { + if !CURRENT_BRAIN_ENDPOINT.initialized() { + let pub_ip = public_ip::addr().await.ok_or("Could not determine public IP")?; + let brain_endpoint = *CURRENT_BRAIN_ENDPOINT + .get_or_init(|| async { SocketAddr::new(pub_ip, BRAIN_GRPC_PORT) }) + .await; + Ok(brain_endpoint) + } else { + let brain_endpoint = *CURRENT_BRAIN_ENDPOINT.get().unwrap_or(&BRAIN_GRPC_ADDR); + Ok(brain_endpoint) + } +} + pub const OLD_BRAIN_DATA_PATH: &str = "./saved_data.yaml"; pub const ACCOUNT: &str = "account"; diff --git a/src/db/mod.rs b/src/db/mod.rs index 6af367b..7075287 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -51,6 +51,8 @@ pub enum Error { TooBigTransaction(String), #[error("Unknown: {0}")] Unknown(String), + #[error(transparent)] + Tonic(#[from] tonic::Status), } pub mod prelude { diff --git a/src/grpc/mod.rs b/src/grpc/mod.rs index c78c346..4f09718 100644 --- a/src/grpc/mod.rs +++ b/src/grpc/mod.rs @@ -3,15 +3,65 @@ pub mod general; pub mod types; pub mod vm; -use crate::constants::ADMIN_ACCOUNTS; +use crate::constants::{current_brain_endpoint, ADMIN_ACCOUNTS}; +use crate::db::prelude as db; use detee_shared::app_proto::*; use detee_shared::common_proto::{Empty, Pubkey}; use detee_shared::general_proto::{ AirdropReq, BanUserReq, KickReq, RegOperatorReq, ReportNodeReq, SlashReq, }; use detee_shared::vm_proto::*; +use serde::Deserialize; +use surrealdb::engine::remote::ws::Client; +use surrealdb::{RecordId, Surreal}; use tonic::{Request, Status}; +pub async fn set_pubsub_node(db: &Surreal, id: RecordId) { + match current_brain_endpoint().await { + Ok(public_endpoint) => { + let addr = public_endpoint.to_string(); + match db + .query(format!("UPDATE $id SET pub_sub_node = '{addr}'",)) + .bind(("id", id)) + .await + { + Ok(res) => log::info!("Updated pub_sub_node {:?}", res), + Err(e) => log::error!("Could not update pub_sub_node {:?}", e), + } + } + _ => { + log::error!("Could not determine public brain endpoint. Using local one."); + } + } +} + +pub async fn check_pubsub_node(db: &Surreal, id: RecordId) -> Result { + #[derive(Deserialize, Debug)] + struct PubSubNode { + pub pub_sub_node: String, + } + let mut query_resp = db.query(format!("select pub_sub_node from {id}")).await?; + let node_endpoint = query_resp + .take::>(0)? + .ok_or(Status::internal("Could not get current brain endpoint"))? + .pub_sub_node; + + let current_endpoint = crate::constants::current_brain_endpoint() + .await + .map_err(|_| Status::internal("Could not get current brain endpoint"))?; + + dbg!(¤t_endpoint, &node_endpoint); + + if current_endpoint.to_string() == node_endpoint { + return Ok(node_endpoint); + } + + let mut status = Status::new(tonic::Code::Unavailable, "moved"); + status.metadata_mut().insert("location", node_endpoint.parse().unwrap()); + + Err(status.into()) +} + pub trait PubkeyGetter { fn get_pubkey(&self) -> Option; } diff --git a/src/grpc/vm.rs b/src/grpc/vm.rs index 211f97d..88961a0 100644 --- a/src/grpc/vm.rs +++ b/src/grpc/vm.rs @@ -1,7 +1,7 @@ #![allow(dead_code)] use crate::constants::{ACCOUNT, NEW_VM_REQ, UPDATE_VM_REQ, VM_NODE}; use crate::db::prelude as db; -use crate::grpc::{check_sig_from_parts, check_sig_from_req}; +use crate::grpc::{check_pubsub_node, check_sig_from_parts, check_sig_from_req, set_pubsub_node}; use detee_shared::common_proto::Empty; use detee_shared::vm_proto::brain_vm_cli_server::BrainVmCli; use detee_shared::vm_proto::brain_vm_daemon_server::BrainVmDaemon; @@ -38,8 +38,9 @@ impl BrainVmDaemon for VmDaemonServer { ) -> Result, Status> { let req = check_sig_from_req(req)?; info!("Starting registration process for {:?}", req); + let id = surrealdb::RecordId::from((VM_NODE, req.node_pubkey.clone())); db::VmNode { - id: surrealdb::RecordId::from((VM_NODE, req.node_pubkey.clone())), + id: id.clone(), operator: surrealdb::RecordId::from((ACCOUNT, req.operator_wallet)), country: req.country, region: req.region, @@ -58,6 +59,7 @@ impl BrainVmDaemon for VmDaemonServer { } .register(&self.db) .await?; + set_pubsub_node(&self.db, id).await; info!("Sending deleted contracts to {}", req.node_pubkey); let deleted_vms = db::DeletedVm::list_by_node(&self.db, &req.node_pubkey).await?; @@ -222,6 +224,9 @@ impl BrainVmCli for VmCliServer { async fn new_vm(&self, req: Request) -> Result, Status> { let req = check_sig_from_req(req)?; + let id = surrealdb::RecordId::from((VM_NODE, req.node_pubkey.clone())); + check_pubsub_node(&self.db, id).await?; + // TODO: make it atleast 1 hour if req.locked_nano < 100 { log::error!("locking lessthan 100 nano lps: {}", req.locked_nano); diff --git a/surql/tables.sql b/surql/tables.sql index 65e9381..90012ff 100644 --- a/surql/tables.sql +++ b/surql/tables.sql @@ -6,6 +6,7 @@ DEFINE FIELD email ON TABLE account TYPE string DEFAULT ""; DEFINE TABLE vm_node SCHEMAFULL; DEFINE FIELD operator ON TABLE vm_node TYPE record; +DEFINE FIELD pub_sub_node ON TABLE vm_node TYPE string default "127.0.0.1:31337"; DEFINE FIELD country ON TABLE vm_node TYPE string; DEFINE FIELD region ON TABLE vm_node TYPE string; DEFINE FIELD city ON TABLE vm_node TYPE string; @@ -80,6 +81,7 @@ DEFINE FIELD price_per_unit ON TABLE deleted_vm TYPE int; DEFINE TABLE app_node SCHEMAFULL; DEFINE FIELD operator ON TABLE app_node TYPE record; +DEFINE FIELD pub_sub_node ON TABLE vm_node TYPE string default "127.0.0.1:31337"; DEFINE FIELD country ON TABLE app_node TYPE string; DEFINE FIELD region ON TABLE app_node TYPE string; DEFINE FIELD city ON TABLE app_node TYPE string; diff --git a/tests/grpc_general_test.rs b/tests/grpc_general_test.rs index f115392..7d42c59 100644 --- a/tests/grpc_general_test.rs +++ b/tests/grpc_general_test.rs @@ -343,6 +343,7 @@ async fn test_ban_user() { let operator_banned_you = create_new_vm(&db_conn, &user_key, node_pubkey, &brain_channel).await.err().unwrap(); + dbg!(&operator_banned_you); assert!(operator_banned_you.to_string().contains("This operator banned you")); } diff --git a/tests/grpc_vm_cli_test.rs b/tests/grpc_vm_cli_test.rs index dd5ebf6..820fbb1 100644 --- a/tests/grpc_vm_cli_test.rs +++ b/tests/grpc_vm_cli_test.rs @@ -4,12 +4,15 @@ use common::vm_cli_utils::{create_new_vm, user_list_vm_contracts}; use common::vm_daemon_utils::{mock_vm_daemon, register_vm_node}; use detee_shared::vm_proto::brain_vm_cli_client::BrainVmCliClient; use detee_shared::vm_proto::brain_vm_daemon_client::BrainVmDaemonClient; -use detee_shared::vm_proto::{self, DeleteVmReq}; +use detee_shared::vm_proto::{self, DeleteVmReq, RegisterVmNodeReq}; use detee_shared::vm_proto::{ExtendVmReq, ListVmContractsReq, NewVmReq}; use futures::StreamExt; use std::vec; -use surreal_brain::constants::{ACCOUNT, ACTIVE_VM, DELETED_VM, NEW_VM_REQ, TOKEN_DECIMAL}; +use surreal_brain::constants::{ + ACCOUNT, ACTIVE_VM, DELETED_VM, NEW_VM_REQ, TOKEN_DECIMAL, VM_NODE, +}; use surreal_brain::db::prelude as db; +use surreal_brain::grpc::check_pubsub_node; mod common; @@ -301,3 +304,49 @@ async fn test_extend_vm() { } // TODO: test register vm node, delete vm contract while node offline, kick, etc.. + +#[tokio::test] +async fn test_pub_sub_redirect() { + let db = prepare_test_db().await.unwrap(); + + let brain_channel = run_service_for_stream().await.unwrap(); + let mut vm_daemon_client = BrainVmDaemonClient::new(brain_channel.clone()); + + let node_key = Key::new(); + let operator_wallet = "Hv5q3enK249RUnLRLi9YNQMrPCRxvL2XnhznkzrtCmkG"; + + let req = RegisterVmNodeReq { + node_pubkey: node_key.pubkey.clone(), + operator_wallet: operator_wallet.to_string(), + main_ip: String::from("185.243.218.213"), + city: String::from("Oslo"), + country: String::from("Norway"), + region: String::from("EU"), + price: 1200, + }; + + vm_daemon_client + .register_vm_node(node_key.sign_request(req).unwrap()) + .await + .unwrap() + .into_inner(); + + // let client_key = Key::new(); + + // let new_vm_req = vm_proto::NewVmReq { + // admin_pubkey: client_key.pubkey.clone(), + // node_pubkey: node_key.pubkey.clone(), + // price_per_unit: 1200, + // extra_ports: vec![8080, 8081], + // locked_nano: 100, + // ..Default::default() + // }; + + let id = surrealdb::RecordId::from((VM_NODE, node_key.pubkey.clone())); + + let endpoint = check_pubsub_node(&db, id).await; + dbg!(&endpoint); + + // let mut client_vm_cli = BrainVmCliClient::new(brain_channel.clone()); + // let redirect = client_vm_cli.new_vm(client_key.sign_request(new_vm_req).unwrap()).await; +}