switch to external ips

Signed-off-by: Valentyn Faychuk <valy@detee.ltd>

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_core 0.20.10",
- "darling_macro",
+ "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",
 ]
 

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"

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

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<IP, NodeInfo>,
     conns: DashSet<IP>,
 }
 
 impl State {
-    pub fn new_with_localhost() -> Self {
+    pub fn new(my_ip: String) -> Self {
-        let localhost_info = NodeInfo::load();
+        let nodes = DashMap::new();
-        let state = Self { nodes: DashMap::new(), conns: DashSet::new() };
+        let my_info = NodeInfo::load();
-        state.nodes.insert(LOCALHOST.to_string(), localhost_info);
+        nodes.insert(my_ip.clone(), my_info);
-        state
+        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) {
+        if let Some(mut my_info) = self.nodes.get_mut(&self.my_ip) {
-            localhost_info.mint_requests += 1;
+            my_info.mint_requests += 1;
-            localhost_info.log(localhost_info.key());
+            my_info.log(my_info.key());
-            localhost_info.save();
+            my_info.save();
         }
     }
 
     pub fn increase_mints(&self) {
-        if let Some(mut localhost_info) = self.nodes.get_mut(LOCALHOST) {
+        if let Some(mut my_info) = self.nodes.get_mut(&self.my_ip) {
-            localhost_info.mints += 1;
+            my_info.mints += 1;
-            localhost_info.log(localhost_info.key());
+            my_info.log(my_info.key());
-            localhost_info.save();
+            my_info.save();
         }
     }
 
     pub fn increase_mratls_conns(&self) {
-        if let Some(mut localhost_info) = self.nodes.get_mut(LOCALHOST) {
+        if let Some(mut my_info) = self.nodes.get_mut(&self.my_ip) {
-            localhost_info.mratls_conns += 1;
+            my_info.mratls_conns += 1;
-            localhost_info.log(localhost_info.key());
+            my_info.log(my_info.key());
-            localhost_info.save();
+            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) {
+        if let Some(mut my_info) = self.nodes.get_mut(&self.my_ip) {
-            localhost_info.disk_attacks += 1;
+            my_info.disk_attacks += 1;
-            localhost_info.log(localhost_info.key());
+            my_info.log(my_info.key());
-            localhost_info.save();
+            my_info.save();
         }
     }
 
     pub fn increase_net_attacks(&self) {
-        if let Some(mut localhost_info) = self.nodes.get_mut(LOCALHOST) {
+        if let Some(mut my_info) = self.nodes.get_mut(&self.my_ip) {
-            localhost_info.net_attacks += 1;
+            my_info.net_attacks += 1;
-            localhost_info.log(localhost_info.key());
+            my_info.log(my_info.key());
-            localhost_info.save();
+            my_info.save();
         }
     }
 
-    pub fn get_localhost(&self) -> NodeInfo {
+    pub fn declare_myself_public(&self) {
-        let mut localhost = self.nodes.get_mut(LOCALHOST).expect("no localhost node");
+        if let Some(mut my_info) = self.nodes.get_mut(&self.my_ip) {
-        localhost.keepalive = SystemTime::now();
+            my_info.public = true;
-        localhost.clone()
+            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<String> {
+    pub fn get_random_disconnected_node(&self) -> Option<String> {
         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) {

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<State>,
     tx: Sender<NodeUpdate>,
     ratls_config: RaTlsConfig,
 }
 
 impl ConnManager {
-    pub fn init(state: Arc<State>, ratls_config: RaTlsConfig, tx: Sender<NodeUpdate>) -> Self {
+    pub fn init(
-        Self { state, ratls_config, tx }
+        my_ip: String,
+        state: Arc<State>,
+        ratls_config: RaTlsConfig,
+        tx: Sender<NodeUpdate>,
+    ) -> 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 let Some(ip) = self.state.get_random_disconnected_node() {
-                if node != LOCALHOST {
+                self.connect_wrapper(ip.clone()).await;
-                    self.connect_wrapper(node.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());
+        let _ = self.tx.send((self.my_ip.clone(), self.state.get_my_info()).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;
-            }
 
+        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()

src/grpc/server.rs

@@ -129,36 +129,38 @@ impl Update for MyServer {
         &self,
         req: Request<Streaming<NodeUpdate>>,
     ) -> Result<Response<Self::GetUpdatesStream>, Status> {
+        // connection info is added in the tower for tls connection and must be present
         let conn_info = req.extensions().get::<Arc<ConnInfo>>().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<NodeUpdate> = ds.get_node_list().into_iter().map(Into::<NodeUpdate>::into).collect();
+            state.declare_myself_public();
+            state.increase_mratls_conns();
+
+            let full_update_list: Vec<NodeUpdate> = state.get_node_list().into_iter().map(Into::<NodeUpdate>::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) => {
+                            Ok(update) => {
-                                if update.ip == "localhost" {
+                                if update.ip != remote_ip {
-                                    update.ip = remote_ip.clone();
+                                    println!("node {remote_ip} is forwarding us the update for {}", update.ip);
-                                    // note that we don't set this node online,
-                                    // as it can be behind NAT
                                 }
-                                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."));
+                            error_status = Status::internal("Already have too many clients. Connect to another server.");
-                            return;
+                            break;
                         }
                     }
-
                 }
             }
+            yield Err(error_status);
+            state.decrease_mratls_conns();
         };
 
         Ok(Response::new(Box::pin(stream) as Self::GetUpdatesStream))

src/main.rs

@@ -6,14 +6,14 @@ mod solana;
 
 use crate::persistence::SealError;
 use crate::{
-    datastore::LOCALHOST, grpc::challenge::NodeUpdate, persistence::KeysFile,
+    grpc::challenge::NodeUpdate, persistence::KeysFile, persistence::SealedFile, solana::SolClient,
-    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<State>, tx: Sender<NodeUpdate>) {
+async fn resolve_my_ip() -> Result<String, Error> {
+    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<State>, tx: Sender<NodeUpdate>) {
     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())
+                grpc::client::ConnManager::init(
-                    .start_with_node(line.to_string()),
+                    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())
+            grpc::client::ConnManager::init(
-                .start(),
+                my_ip.clone(),
+                state.clone(),
+                ratls_config.clone(),
+                tx.clone(),
+            )
+            .start(),
         );
     }