diff --git a/rewrite/Cargo.lock b/rewrite/Cargo.lock index 8030362..24a7635 100644 --- a/rewrite/Cargo.lock +++ b/rewrite/Cargo.lock @@ -1738,6 +1738,7 @@ dependencies = [ "prost-types", "rand 0.8.5", "serde", + "serde_json", "solana-client", "solana-program", "solana-sdk", diff --git a/rewrite/Cargo.toml b/rewrite/Cargo.toml index db7e5c4..66a73e0 100644 --- a/rewrite/Cargo.toml +++ b/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" diff --git a/rewrite/rustfmt.toml b/rewrite/rustfmt.toml new file mode 100644 index 0000000..01959db --- /dev/null +++ b/rewrite/rustfmt.toml @@ -0,0 +1,3 @@ +reorder_impl_items = true +use_small_heuristics = "Max" +merge_imports = true diff --git a/rewrite/src/datastore.rs b/rewrite/src/datastore.rs index 3c07bcc..2f62e7d 100644 --- a/rewrite/src/datastore.rs +++ b/rewrite/src/datastore.rs @@ -41,8 +41,6 @@ pub struct Store { sol_client: SolClient, nodes: DashMap, conns: DashSet, - // 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> { 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(); diff --git a/rewrite/src/grpc/server.rs b/rewrite/src/grpc/server.rs index fb5b52d..5dcb7be 100644 --- a/rewrite/src/grpc/server.rs +++ b/rewrite/src/grpc/server.rs @@ -40,6 +40,7 @@ impl Update for MyServer { &self, req: Request>, ) -> Result, 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(); diff --git a/rewrite/src/main.rs b/rewrite/src/main.rs index a672de4..5f5822e 100644 --- a/rewrite/src/main.rs +++ b/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::grpc::challenge::NodeUpdate; +use crate::{datastore::LOCALHOST, grpc::challenge::NodeUpdate, solana::Client as SolClient}; use datastore::Store; use solana_sdk::signer::Signer; -use std::fs::File; -use std::io::{BufRead, BufReader}; -use std::sync::Arc; -use tokio::sync::broadcast; -use tokio::sync::broadcast::Sender; -use tokio::time::{sleep, Duration}; -use tokio::task::JoinSet; -use crate::solana::Client as SolClient; +use std::{ + fs::File, + io::{BufRead, BufReader}, + sync::Arc, +}; +use tokio::{ + sync::{broadcast, broadcast::Sender}, + 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, tx: Sender) { loop { @@ -26,6 +29,16 @@ pub async fn localhost_cron(ds: Arc, tx: Sender) { } 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) => { diff --git a/rewrite/src/persistence.rs b/rewrite/src/persistence.rs new file mode 100644 index 0000000..985f3dc --- /dev/null +++ b/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> { + 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> { + 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) + } +}