From 331788fb6e7795ea05432000b656dfff7abca34c Mon Sep 17 00:00:00 2001 From: Valentyn Faychuk Date: Sun, 22 Dec 2024 16:39:15 +0200 Subject: [PATCH] switch to external ips Signed-off-by: Valentyn Faychuk --- Cargo.lock | 271 ++++++++++++++++++++++++++++++++++-- Cargo.toml | 3 +- docker/challenge.Dockerfile | 2 +- src/datastore.rs | 82 ++++++----- src/grpc/client.rs | 30 ++-- src/grpc/server.rs | 29 ++-- src/main.rs | 37 +++-- 7 files changed, 365 insertions(+), 89 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 46abf2f..a25f625 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -115,7 +115,7 @@ dependencies = [ "futures-core", "futures-util", "mio", - "socket2", + "socket2 0.5.7", "tokio", "tracing", ] @@ -178,7 +178,7 @@ dependencies = [ "serde_json", "serde_urlencoded", "smallvec", - "socket2", + "socket2 0.5.7", "time", "url", ] @@ -1387,14 +1387,38 @@ dependencies = [ "zeroize", ] +[[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.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6f63b86c8a8826a49b8c21f08a2d07338eec8d900540f8630dc76284be802989" dependencies = [ - "darling_core", - "darling_macro", + "darling_core 0.20.10", + "darling_macro 0.20.10", +] + +[[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]] @@ -1411,13 +1435,24 @@ dependencies = [ "syn 2.0.79", ] +[[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.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d336a2a514f6ccccaa3e09b02d41d35330c07ddf03a62165fcec10bb561c7806" dependencies = [ - "darling_core", + "darling_core 0.20.10", "quote", "syn 2.0.79", ] @@ -1510,6 +1545,31 @@ dependencies = [ "syn 1.0.109", ] +[[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 = "derive_more" version = "0.99.18" @@ -1614,6 +1674,18 @@ dependencies = [ "syn 2.0.79", ] +[[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 = "dunce" version = "1.0.5" @@ -1682,6 +1754,24 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "endian-type" +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 = "enum-iterator" version = "1.5.0" @@ -2032,6 +2122,7 @@ dependencies = [ "once_cell", "prost", "prost-types", + "public-ip", "rand 0.8.5", "rustls 0.23.20", "serde", @@ -2097,6 +2188,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "heck" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8" + [[package]] name = "heck" version = "0.5.0" @@ -2266,7 +2363,7 @@ dependencies = [ "httpdate", "itoa", "pin-project-lite", - "socket2", + "socket2 0.5.7", "tokio", "tower-service", "tracing", @@ -2327,6 +2424,20 @@ dependencies = [ "tower-service", ] +[[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.30", + "tokio", + "tower-service", + "tracing", +] + [[package]] name = "hyper-timeout" version = "0.5.1" @@ -2353,7 +2464,7 @@ dependencies = [ "http-body 1.0.1", "hyper 1.4.1", "pin-project-lite", - "socket2", + "socket2 0.5.7", "tokio", "tower-service", "tracing", @@ -2388,6 +2499,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 = "0.5.0" @@ -2675,6 +2797,12 @@ version = "0.4.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24" +[[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" @@ -2773,6 +2901,15 @@ version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "defc4c55412d89136f966bbb339008b474350e5e6e78d2714439c386b3137a03" +[[package]] +name = "nibble_vec" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a5d83df9f36fe23f0c3648c6bbb8b0298bb5f1939c8f2704431371f4b84d43" +dependencies = [ + "smallvec", +] + [[package]] name = "nix" version = "0.28.0" @@ -3242,7 +3379,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c1318b19085f08681016926435853bbf7858f9c082d0999b80550ff5d9abe15" dependencies = [ "bytes", - "heck", + "heck 0.5.0", "itertools 0.13.0", "log", "multimap", @@ -3278,6 +3415,27 @@ dependencies = [ "prost", ] +[[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.30", + "hyper-system-resolver", + "pin-project-lite", + "thiserror", + "tokio", + "tracing", + "tracing-futures", + "trust-dns-client", + "trust-dns-proto", +] + [[package]] name = "qstring" version = "0.7.2" @@ -3330,7 +3488,7 @@ checksum = "055b4e778e8feb9f93c4e439f71dc2156ef13360b432b799e179a8c4cdf0b1d7" dependencies = [ "bytes", "libc", - "socket2", + "socket2 0.5.7", "tracing", "windows-sys 0.48.0", ] @@ -3344,6 +3502,16 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "radix_trie" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c069c179fcdc6a2fe24d8d18305cf085fdbd4f922c041943e203685d6a1c58fd" +dependencies = [ + "endian-type", + "nibble_vec", +] + [[package]] name = "rand" version = "0.7.3" @@ -3916,7 +4084,7 @@ version = "2.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "881b6f881b17d13214e5d494c939ebab463d01264ce1811e9d4ac3a882e7695f" dependencies = [ - "darling", + "darling 0.20.10", "proc-macro2", "quote", "syn 2.0.79", @@ -3928,7 +4096,7 @@ version = "3.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9d846214a9854ef724f3da161b426242d8de7c1fc7de2f89bb1efcb154dca79d" dependencies = [ - "darling", + "darling 0.20.10", "proc-macro2", "quote", "syn 2.0.79", @@ -4039,6 +4207,16 @@ version = "1.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" +[[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.7" @@ -4243,7 +4421,7 @@ dependencies = [ "rand 0.8.5", "serde", "serde_derive", - "socket2", + "socket2 0.5.7", "solana-logger", "solana-sdk", "solana-version", @@ -5238,6 +5416,12 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ea5119cdb4c55b55d432abb513a0429384878c15dde60cc77b1c99de1a95a6a" +[[package]] +name = "strsim" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6446ced80d6c486436db5c078dde11a9f73d42b57fb273121e160b84f63d894c" + [[package]] name = "strsim" version = "0.10.0" @@ -5490,7 +5674,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "signal-hook-registry", - "socket2", + "socket2 0.5.7", "tokio-macros", "windows-sys 0.52.0", ] @@ -5613,7 +5797,7 @@ dependencies = [ "percent-encoding", "pin-project", "prost", - "socket2", + "socket2 0.5.7", "tokio", "tokio-stream", "tower 0.4.13", @@ -5782,6 +5966,63 @@ 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 = "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", + "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", + "tinyvec", + "tokio", + "url", +] + [[package]] name = "try-lock" version = "0.2.5" @@ -5905,7 +6146,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22784dbdf76fdde8af1aeda5622b546b422b6fc585325248a2bf9f5e41e94d6c" dependencies = [ "form_urlencoded", - "idna", + "idna 0.5.0", "percent-encoding", ] diff --git a/Cargo.toml b/Cargo.toml index 4de45e0..98ed587 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,8 +11,10 @@ actix-web = "4.9" async-stream = "0.3" chrono = "0.4" dashmap = "6.1" +tonic = "0.12" prost = "0.13" prost-types = "0.13" +public-ip = "0.2" rand = "0.8" serde = { version = "1.0", features = ["derive"] } serde_with = { version = "3.11", features = ["macros", "base64"] } @@ -24,7 +26,6 @@ spl-associated-token-account = "5.0" spl-token = "6.0" tokio = { version = "1.40", features = ["macros", "rt-multi-thread", "fs"] } # this can be "full" tokio-stream = { version = "0.1", features = ["sync"] } -tonic = "0.12" # sgx poc dependencies rustls = "0.23" tokio-rustls = "0.26" diff --git a/docker/challenge.Dockerfile b/docker/challenge.Dockerfile index e53a821..0551a04 100644 --- a/docker/challenge.Dockerfile +++ b/docker/challenge.Dockerfile @@ -7,7 +7,7 @@ RUN mkdir sgx_libs &&\ FROM ubuntu:20.04 -RUN apt update -y && apt install -y jq curl +RUN apt update -y && apt install -y jq curl && rm -rf /var/lib/apt/lists/* COPY challenge.sh /challenge.sh COPY challenge.tar.gz /challenge.tar.gz diff --git a/src/datastore.rs b/src/datastore.rs index f7caab4..dbc1a18 100644 --- a/src/datastore.rs +++ b/src/datastore.rs @@ -5,7 +5,6 @@ use serde_with::{serde_as, TimestampSeconds}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; type IP = String; -pub const LOCALHOST: &str = "localhost"; const LOG_PATH: &str = "/host/main/logs"; const LOCAL_INFO_FILE: &str = "/host/main/node_info"; @@ -103,16 +102,17 @@ impl NodeInfo { /// Multithreaded state, designed to be /// shared everywhere in the code pub struct State { + my_ip: String, nodes: DashMap, conns: DashSet, } impl State { - pub fn new_with_localhost() -> Self { - let localhost_info = NodeInfo::load(); - let state = Self { nodes: DashMap::new(), conns: DashSet::new() }; - state.nodes.insert(LOCALHOST.to_string(), localhost_info); - state + pub fn new(my_ip: String) -> Self { + let nodes = DashMap::new(); + let my_info = NodeInfo::load(); + nodes.insert(my_ip.clone(), my_info); + Self { my_ip, nodes, conns: DashSet::new() } } pub fn add_conn(&self, ip: &str) { @@ -124,49 +124,67 @@ impl State { } pub fn increase_mint_requests(&self) { - if let Some(mut localhost_info) = self.nodes.get_mut(LOCALHOST) { - localhost_info.mint_requests += 1; - localhost_info.log(localhost_info.key()); - localhost_info.save(); + if let Some(mut my_info) = self.nodes.get_mut(&self.my_ip) { + my_info.mint_requests += 1; + my_info.log(my_info.key()); + my_info.save(); } } pub fn increase_mints(&self) { - if let Some(mut localhost_info) = self.nodes.get_mut(LOCALHOST) { - localhost_info.mints += 1; - localhost_info.log(localhost_info.key()); - localhost_info.save(); + if let Some(mut my_info) = self.nodes.get_mut(&self.my_ip) { + my_info.mints += 1; + my_info.log(my_info.key()); + my_info.save(); } } pub fn increase_mratls_conns(&self) { - if let Some(mut localhost_info) = self.nodes.get_mut(LOCALHOST) { - localhost_info.mratls_conns += 1; - localhost_info.log(localhost_info.key()); - localhost_info.save(); + if let Some(mut my_info) = self.nodes.get_mut(&self.my_ip) { + my_info.mratls_conns += 1; + my_info.log(my_info.key()); + my_info.save(); + } + } + + pub fn decrease_mratls_conns(&self) { + if let Some(mut my_info) = self.nodes.get_mut(&self.my_ip) { + if my_info.mratls_conns > 0 { + my_info.mratls_conns -= 1; + my_info.log(my_info.key()); + my_info.save(); + } } } pub fn increase_disk_attacks(&self) { - if let Some(mut localhost_info) = self.nodes.get_mut(LOCALHOST) { - localhost_info.disk_attacks += 1; - localhost_info.log(localhost_info.key()); - localhost_info.save(); + if let Some(mut my_info) = self.nodes.get_mut(&self.my_ip) { + my_info.disk_attacks += 1; + my_info.log(my_info.key()); + my_info.save(); } } pub fn increase_net_attacks(&self) { - if let Some(mut localhost_info) = self.nodes.get_mut(LOCALHOST) { - localhost_info.net_attacks += 1; - localhost_info.log(localhost_info.key()); - localhost_info.save(); + if let Some(mut my_info) = self.nodes.get_mut(&self.my_ip) { + my_info.net_attacks += 1; + my_info.log(my_info.key()); + my_info.save(); } } - pub fn get_localhost(&self) -> NodeInfo { - let mut localhost = self.nodes.get_mut(LOCALHOST).expect("no localhost node"); - localhost.keepalive = SystemTime::now(); - localhost.clone() + pub fn declare_myself_public(&self) { + if let Some(mut my_info) = self.nodes.get_mut(&self.my_ip) { + my_info.public = true; + my_info.log(my_info.key()); + my_info.save(); + } + } + + pub fn get_my_info(&self) -> NodeInfo { + let mut my_info = self.nodes.get_mut(&self.my_ip).expect("no info for this node"); + my_info.keepalive = SystemTime::now(); + my_info.clone() } /// This returns true if NodeInfo got modified. @@ -188,7 +206,7 @@ impl State { } // returns a random node that does not have an active connection - pub fn get_random_node(&self) -> Option { + pub fn get_random_disconnected_node(&self) -> Option { use rand::{rngs::OsRng, RngCore}; let len = self.nodes.len(); if len == 0 { @@ -200,7 +218,7 @@ impl State { .map(|n| n.key().clone()) .cycle() .skip(skip) - .find(|k| !self.conns.contains(k)) + .find(|ip| ip != &self.my_ip && !self.conns.contains(ip)) } pub fn remove_inactive_nodes(&self) { diff --git a/src/grpc/client.rs b/src/grpc/client.rs index 418ad5b..31a4355 100644 --- a/src/grpc/client.rs +++ b/src/grpc/client.rs @@ -1,6 +1,6 @@ use super::challenge::{Keys, NodeUpdate}; use crate::{ - datastore::{State, LOCALHOST}, + datastore::State, grpc::challenge::{update_client::UpdateClient, Empty}, }; use detee_sgx::RaTlsConfig; @@ -13,14 +13,20 @@ use tokio_stream::{wrappers::BroadcastStream, StreamExt}; #[derive(Clone)] pub struct ConnManager { + my_ip: String, state: Arc, tx: Sender, ratls_config: RaTlsConfig, } impl ConnManager { - pub fn init(state: Arc, ratls_config: RaTlsConfig, tx: Sender) -> Self { - Self { state, ratls_config, tx } + pub fn init( + my_ip: String, + state: Arc, + ratls_config: RaTlsConfig, + tx: Sender, + ) -> Self { + Self { my_ip, state, ratls_config, tx } } pub async fn start_with_node(self, node_ip: String) { @@ -29,10 +35,8 @@ impl ConnManager { pub async fn start(self) { loop { - if let Some(node) = self.state.get_random_node() { - if node != LOCALHOST { - self.connect_wrapper(node.clone()).await; - } + if let Some(ip) = self.state.get_random_disconnected_node() { + self.connect_wrapper(ip.clone()).await; } sleep(Duration::from_secs(3)).await; } @@ -96,17 +100,9 @@ impl ConnManager { })?; let mut resp_stream = response.into_inner(); - let _ = self.tx.send((LOCALHOST.to_string(), self.state.get_localhost()).into()); - - while let Some(mut update) = resp_stream.message().await? { - // "localhost" IPs need to be changed to the real IP of the counterpart - if update.ip == LOCALHOST { - update.ip = node_ip.clone(); - // since we are connecting TO this server, we have a guarantee that this - // server is not behind NAT, so we can set it public - update.public = true; - } + let _ = self.tx.send((self.my_ip.clone(), self.state.get_my_info()).into()); + while let Some(update) = resp_stream.message().await? { // update the entire network in case the information is new if self.state.process_node_update(update.clone().into()) && self.tx.send(update.clone()).is_err() diff --git a/src/grpc/server.rs b/src/grpc/server.rs index a233d54..8d148f9 100644 --- a/src/grpc/server.rs +++ b/src/grpc/server.rs @@ -129,36 +129,38 @@ impl Update for MyServer { &self, req: Request>, ) -> Result, Status> { + // connection info is added in the tower for tls connection and must be present let conn_info = req.extensions().get::>().unwrap(); - self.state.increase_mratls_conns(); let remote_ip = conn_info.addr.ip().to_string(); let tx = self.tx.clone(); let mut rx = self.tx.subscribe(); let mut inbound = req.into_inner(); - let ds = self.state.clone(); + let state = self.state.clone(); let stream = async_stream::stream! { - let full_update_list: Vec = ds.get_node_list().into_iter().map(Into::::into).collect(); + state.declare_myself_public(); + state.increase_mratls_conns(); + + let full_update_list: Vec = state.get_node_list().into_iter().map(Into::::into).collect(); for update in full_update_list { yield Ok(update); } + let error_status: Status; loop { tokio::select! { Some(msg) = inbound.next() => { match msg { - Ok(mut update) => { - if update.ip == "localhost" { - update.ip = remote_ip.clone(); - // note that we don't set this node online, - // as it can be behind NAT + Ok(update) => { + if update.ip != remote_ip { + println!("node {remote_ip} is forwarding us the update for {}", update.ip); } - if update.ip != "127.0.0.1" && ds.process_node_update(update.clone().into()) && tx.send(update.clone()).is_err() { + if state.process_node_update(update.clone().into()) && tx.send(update.clone()).is_err() { println!("tokio broadcast receivers had an issue consuming the channel"); }; } Err(e) => { - yield Err(Status::internal(format!("Error receiving client stream: {}", e))); + error_status = Status::internal(format!("Error receiving client stream: {}", e)); break; } } @@ -167,13 +169,14 @@ impl Update for MyServer { yield Ok(update); // disconnect client if too many connections are active if tx.receiver_count() > 9 { - yield Err(Status::internal("Already have too many clients. Connect to another server.")); - return; + error_status = Status::internal("Already have too many clients. Connect to another server."); + break; } } - } } + yield Err(error_status); + state.decrease_mratls_conns(); }; Ok(Response::new(Box::pin(stream) as Self::GetUpdatesStream)) diff --git a/src/main.rs b/src/main.rs index 36a487a..76bfc8c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,14 +6,14 @@ mod solana; use crate::persistence::SealError; use crate::{ - datastore::LOCALHOST, grpc::challenge::NodeUpdate, persistence::KeysFile, - persistence::SealedFile, solana::SolClient, + grpc::challenge::NodeUpdate, persistence::KeysFile, persistence::SealedFile, solana::SolClient, }; use datastore::State; use detee_sgx::{InstanceMeasurement, RaTlsConfig}; use std::{ fs::File, io::{BufRead, BufReader}, + io::{Error, ErrorKind}, sync::Arc, }; use tokio::{ @@ -26,10 +26,16 @@ const INIT_NODES_FILE: &str = "/host/detee_challenge_nodes"; const KEYS_FILE: &str = "/host/main/TRY_TO_HACK_THIS"; const MAX_CONNECTIONS: usize = 3; -pub async fn localhost_cron(state: Arc, tx: Sender) { +async fn resolve_my_ip() -> Result { + let err = "Can't resolve my external IP, try again"; + let ip = public_ip::addr_v4().await.ok_or(Error::new(ErrorKind::Other, err))?; + Ok(format!("{}", ip)) +} + +pub async fn heartbeat_cron(my_ip: String, state: Arc, tx: Sender) { loop { sleep(Duration::from_secs(60)).await; - let _ = tx.send((LOCALHOST.to_string(), state.get_localhost()).into()); + let _ = tx.send((my_ip.clone(), state.get_my_info()).into()); state.remove_inactive_nodes(); } } @@ -103,14 +109,15 @@ async fn main() { let ratls_config = RaTlsConfig::new() .allow_instance_measurement(InstanceMeasurement::new().with_current_mrenclave().unwrap()); - let state = Arc::new(State::new_with_localhost()); + let my_ip = resolve_my_ip().await.unwrap(); + let state = Arc::new(State::new(my_ip.clone())); let sol_client = Arc::new(get_sol_client(state.clone(), ratls_config.clone()).await); let (tx, _) = broadcast::channel(500); let mut tasks = JoinSet::new(); - tasks.spawn(localhost_cron(state.clone(), tx.clone())); + tasks.spawn(heartbeat_cron(my_ip.clone(), state.clone(), tx.clone())); tasks.spawn(http_server::init(state.clone(), sol_client.clone())); tasks.spawn( grpc::server::MyServer::init( @@ -125,16 +132,26 @@ async fn main() { if let Ok(input) = std::fs::read_to_string(INIT_NODES_FILE) { for line in input.lines() { tasks.spawn( - grpc::client::ConnManager::init(state.clone(), ratls_config.clone(), tx.clone()) - .start_with_node(line.to_string()), + grpc::client::ConnManager::init( + my_ip.clone(), + state.clone(), + ratls_config.clone(), + tx.clone(), + ) + .start_with_node(line.to_string()), ); } } for _ in 0..MAX_CONNECTIONS { tasks.spawn( - grpc::client::ConnManager::init(state.clone(), ratls_config.clone(), tx.clone()) - .start(), + grpc::client::ConnManager::init( + my_ip.clone(), + state.clone(), + ratls_config.clone(), + tx.clone(), + ) + .start(), ); }