added persistence

rewrite/Cargo.lock

@@ -1738,6 +1738,7 @@ dependencies = [
  "prost-types",
  "rand 0.8.5",
  "serde",
+ "serde_json",
  "solana-client",
  "solana-program",
  "solana-sdk",

rewrite/Cargo.toml

@@ -12,6 +12,7 @@ prost = "0.13.2"
 prost-types = "0.13.2"
 rand = "0.8.5"
 serde = { version = "1.0.210", features = ["derive"] }
+serde_json = "1.0.128"
 solana-client = "2.0.11"
 solana-program = "2.0.11"
 solana-sdk = "2.0.10"

rewrite/rustfmt.toml

@@ -0,0 +1,3 @@
+reorder_impl_items = true
+use_small_heuristics = "Max"
+merge_imports = true

rewrite/src/datastore.rs

@@ -41,8 +41,6 @@ pub struct Store {
     sol_client: SolClient,
     nodes: DashMap<IP, NodeInfo>,
     conns: DashSet<IP>,
-    // TODO: write persistence
-    // persistence: FileManager,
 }
 
 impl Store {
@@ -89,6 +87,18 @@ impl Store {
         }
     }
 
+    pub fn increase_mints(&self) {
+        if let Some(mut localhost_info) = self.nodes.get_mut(LOCALHOST) {
+            localhost_info.mints += 1;
+        }
+    }
+
+    pub fn increase_ratls_conns(&self) {
+        if let Some(mut localhost_info) = self.nodes.get_mut(LOCALHOST) {
+            localhost_info.ratls_conns += 1;
+        }
+    }
+
     pub fn mint(&self, recipient: &str) -> Result<String, Box<dyn std::error::Error>> {
         use std::str::FromStr;
         let recipient = solana_sdk::pubkey::Pubkey::from_str(recipient)?;
@@ -97,12 +107,6 @@ impl Store {
         Ok(sig)
     }
 
-    pub fn increase_mints(&self) {
-        if let Some(mut localhost_info) = self.nodes.get_mut(LOCALHOST) {
-            localhost_info.mints += 1;
-        }
-    }
-
     pub fn get_localhost(&self) -> NodeInfo {
         let mut localhost = self.nodes.get_mut(LOCALHOST).expect("no localhost node");
         localhost.keepalive = SystemTime::now();

rewrite/src/grpc/server.rs

@@ -40,6 +40,7 @@ impl Update for MyServer {
         &self,
         req: Request<Streaming<NodeUpdate>>,
     ) -> Result<Response<Self::GetUpdatesStream>, Status> {
+        self.ds.increase_ratls_conns();
         let remote_ip = req.remote_addr().unwrap().ip().to_string();
         let tx = self.tx.clone();
         let mut rx = self.tx.subscribe();

rewrite/src/main.rs

@@ -1,21 +1,24 @@
 mod datastore;
 mod grpc;
 mod http_server;
+mod persistence;
 mod solana;
-use crate::datastore::LOCALHOST;
+use crate::{datastore::LOCALHOST, grpc::challenge::NodeUpdate, solana::Client as SolClient};
-use crate::grpc::challenge::NodeUpdate;
 use datastore::Store;
 use solana_sdk::signer::Signer;
-use std::fs::File;
+use std::{
-use std::io::{BufRead, BufReader};
+    fs::File,
-use std::sync::Arc;
+    io::{BufRead, BufReader},
-use tokio::sync::broadcast;
+    sync::Arc,
-use tokio::sync::broadcast::Sender;
+};
-use tokio::time::{sleep, Duration};
+use tokio::{
-use tokio::task::JoinSet;
+    sync::{broadcast, broadcast::Sender},
-use crate::solana::Client as SolClient;
+    task::JoinSet,
+    time::{sleep, Duration},
+};
 
 const INIT_NODES: &str = "detee_challenge_nodes";
+const DISK_PERSISTENCE: &str = "TRY_TO_HACK_THIS";
 
 pub async fn localhost_cron(ds: Arc<Store>, tx: Sender<NodeUpdate>) {
     loop {
@@ -26,6 +29,16 @@ pub async fn localhost_cron(ds: Arc<Store>, tx: Sender<NodeUpdate>) {
 }
 
 async fn get_sol_client() -> SolClient {
+    match crate::persistence::Data::read(DISK_PERSISTENCE).await {
+        Ok(data) => {
+            let (keypair, token) = data.parse();
+            println!("Found the following wallet saved to disk: {}", keypair.pubkey());
+            println!("Loading token mint address {}", token);
+            return SolClient::from(keypair, token);
+        }
+        Err(e) => println!("Did not find old pubkeys saved to disk: {e}"),
+    };
+
     let input = match File::open(INIT_NODES) {
         Ok(i) => i,
         Err(_) => {
@@ -43,6 +56,11 @@ async fn get_sol_client() -> SolClient {
                     bundle.0.pubkey()
                 );
                 println!("The address of the Token is {}", bundle.1.to_string());
+                println!("Saving this data to disk in the file {DISK_PERSISTENCE}");
+                let disk_data = crate::persistence::Data::init_from(&bundle.0, &bundle.1).await;
+                if let Err(e) = disk_data.write(DISK_PERSISTENCE).await {
+                    println!("Could not save data to disk due to: {e}");
+                }
                 return SolClient::from(bundle.0, bundle.1);
             }
             Err(e) => {

rewrite/src/persistence.rs

@@ -0,0 +1,49 @@
+use serde::{Deserialize, Serialize};
+use solana_sdk::{pubkey::Pubkey, signature::keypair::Keypair};
+use std::str::FromStr;
+use tokio::{
+    fs::File,
+    io::{AsyncReadExt, AsyncWriteExt},
+};
+
+#[derive(Serialize, Deserialize)]
+pub struct Data {
+    random: String,
+    keypair: String,
+    token: String,
+}
+
+impl Data {
+    pub async fn init_from(keypair: &Keypair, token: &Pubkey) -> Self {
+        use rand::{distributions::Alphanumeric, Rng};
+        let random_string: String =
+            rand::thread_rng().sample_iter(&Alphanumeric).take(128).map(char::from).collect();
+
+        Self {
+            random: random_string,
+            keypair: keypair.to_base58_string(),
+            token: token.to_string(),
+        }
+    }
+
+    pub async fn write(self, path: &str) -> Result<(), Box<dyn std::error::Error>> {
+        let serialized = serde_json::to_string(&self)?;
+        let mut file = File::create(path).await?;
+        file.write_all(serialized.as_bytes()).await?;
+        file.flush().await?;
+        Ok(())
+    }
+
+    pub async fn read(path: &str) -> Result<Self, Box<dyn std::error::Error>> {
+        let mut file = File::open(path).await?;
+        let mut contents = String::new();
+        file.read_to_string(&mut contents).await?;
+        Ok(serde_json::from_str(&contents)?)
+    }
+
+    pub fn parse(self) -> (Keypair, Pubkey) {
+        let keypair = Keypair::from_base58_string(&self.keypair);
+        let pubkey = Pubkey::from_str(&self.token).unwrap();
+        (keypair, pubkey)
+    }
+}