Merge pull request #105 from xch-dev/fixes

Fixes

crates/sage-api/Cargo.toml

@@ -15,6 +15,7 @@ categories = { workspace = true }
 workspace = true
 
 [dependencies]
+sage-config = { workspace = true }
 serde = { workspace = true }
 tauri-specta = { workspace = true, features = ["derive"] }
 specta = { workspace = true, features = ["derive", "bigdecimal"] }

crates/sage-api/src/requests/settings.rs

@@ -1,3 +1,5 @@
+use indexmap::IndexMap;
+use sage_config::Network;
 use serde::{Deserialize, Serialize};
 use specta::Type;
 
@@ -70,3 +72,11 @@ pub struct SetDerivationBatchSize {
 
 #[derive(Debug, Clone, Copy, Serialize, Deserialize, Type)]
 pub struct SetDerivationBatchSizeResponse {}
+
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, Type)]
+pub struct GetNetworks {}
+
+#[derive(Debug, Clone, Serialize, Deserialize, Type)]
+pub struct GetNetworksResponse {
+    pub networks: IndexMap<String, Network>,
+}

crates/sage-cli/src/router.rs

@@ -113,6 +113,7 @@ routes!(
     set_network_id await: SetNetworkId = "/set_network_id",
     set_derive_automatically: SetDeriveAutomatically = "/set_derive_automatically",
     set_derivation_batch_size: SetDerivationBatchSize = "/set_derivation_batch_size",
+    get_networks: GetNetworks = "/get_networks",
 
     remove_cat await: RemoveCat = "/remove_cat",
     update_cat await: UpdateCat = "/update_cat",

crates/sage-wallet/src/sync_manager/peer_state.rs

@@ -134,4 +134,20 @@ impl PeerState {
     pub(super) fn add_peer(&mut self, state: PeerInfo) {
         self.peers.insert(state.peer.socket_addr().ip(), state);
     }
+
+    pub fn trusted_peers(&self) -> &HashSet<IpAddr> {
+        &self.trusted_peers
+    }
+
+    pub fn banned_peers(&mut self) -> &HashMap<IpAddr, u64> {
+        self.banned_peers.retain(|_, ban_until| {
+            let start = SystemTime::now();
+            let since_the_epoch = start
+                .duration_since(UNIX_EPOCH)
+                .expect("Time went backwards");
+
+            *ban_until > since_the_epoch.as_secs()
+        });
+        &self.banned_peers
+    }
 }

crates/sage/Cargo.toml

@@ -41,3 +41,6 @@ hex-literal = { workspace = true }
 base64 = { workspace = true }
 bigdecimal = { workspace = true }
 clvmr = { workspace = true }
+serde = { workspace = true, features = ["derive"] }
+bincode = { workspace = true }
+serde_json = { workspace = true }

crates/sage/src/endpoints/offers.rs

@@ -7,9 +7,11 @@ use sage_api::{
 use sage_wallet::{
     fetch_nft_offer_details, insert_transaction, MakerSide, SyncCommand, TakerSide, Transaction,
 };
+use tracing::debug;
 
 use crate::{
-    parse_asset_id, parse_cat_amount, parse_genesis_challenge, parse_nft_id, Error, Result, Sage,
+    json_bundle, parse_asset_id, parse_cat_amount, parse_genesis_challenge, parse_nft_id, Error,
+    Result, Sage,
 };
 
 impl Sage {
@@ -118,6 +120,11 @@ impl Sage {
             )
             .await?;
 
+        debug!(
+            "{}",
+            serde_json::to_string(&json_bundle(&spend_bundle)).expect("msg")
+        );
+
         let mut tx = wallet.db.tx().await?;
 
         let subscriptions = insert_transaction(

crates/sage/src/endpoints/settings.rs

@@ -2,10 +2,10 @@ use std::time::Duration;
 
 use itertools::Itertools;
 use sage_api::{
-    AddPeer, AddPeerResponse, GetPeers, GetPeersResponse, PeerRecord, RemovePeer,
-    RemovePeerResponse, SetDerivationBatchSize, SetDerivationBatchSizeResponse,
-    SetDeriveAutomatically, SetDeriveAutomaticallyResponse, SetDiscoverPeers,
-    SetDiscoverPeersResponse, SetNetworkId, SetNetworkIdResponse, SetTargetPeers,
+    AddPeer, AddPeerResponse, GetNetworks, GetNetworksResponse, GetPeers, GetPeersResponse,
+    PeerRecord, RemovePeer, RemovePeerResponse, SetDerivationBatchSize,
+    SetDerivationBatchSizeResponse, SetDeriveAutomatically, SetDeriveAutomaticallyResponse,
+    SetDiscoverPeers, SetDiscoverPeersResponse, SetNetworkId, SetNetworkIdResponse, SetTargetPeers,
     SetTargetPeersResponse,
 };
 use sage_wallet::SyncCommand;
@@ -106,6 +106,7 @@ impl Sage {
             .await?;
 
         self.switch_wallet().await?;
+        self.setup_peers().await?;
 
         Ok(SetNetworkIdResponse {})
     }
@@ -136,4 +137,10 @@ impl Sage {
 
         Ok(SetDerivationBatchSizeResponse {})
     }
+
+    pub fn get_networks(&mut self, _req: GetNetworks) -> Result<GetNetworksResponse> {
+        Ok(GetNetworksResponse {
+            networks: self.networks.clone(),
+        })
+    }
 }

crates/sage/src/error.rs

@@ -74,6 +74,9 @@ pub enum Error {
     #[error("TOML serialization error: {0}")]
     TomlSer(#[from] toml::ser::Error),
 
+    #[error("Bincode error: {0}")]
+    Bincode(#[from] bincode::Error),
+
     #[error("Logging initialization error: {0}")]
     LogSubscriber(#[from] TryInitError),
 
@@ -205,7 +208,8 @@ impl Error {
             | Self::Database(..)
             | Self::Bech32m(..)
             | Self::ToClvm(..)
-            | Self::FromClvm(..) => ErrorKind::Internal,
+            | Self::FromClvm(..)
+            | Self::Bincode(..) => ErrorKind::Internal,
             Self::UnknownFingerprint
             | Self::UnknownNetwork
             | Self::MissingCoin(..)

crates/sage/src/lib.rs

@@ -2,6 +2,7 @@
 
 mod endpoints;
 mod error;
+mod peers;
 mod sage;
 mod utils;
 

crates/sage/src/peers.rs

@@ -0,0 +1,23 @@
+use std::{
+    collections::{HashMap, HashSet},
+    net::IpAddr,
+};
+
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Default, Clone, Serialize, Deserialize)]
+pub struct Peers {
+    pub connections: HashSet<IpAddr>,
+    pub trusted: HashSet<IpAddr>,
+    pub banned: HashMap<IpAddr, u64>,
+}
+
+impl Peers {
+    pub fn to_bytes(&self) -> bincode::Result<Vec<u8>> {
+        bincode::serialize(self)
+    }
+
+    pub fn from_bytes(bytes: &[u8]) -> bincode::Result<Self> {
+        bincode::deserialize(bytes)
+    }
+}

crates/sage/src/sage.rs

@@ -3,6 +3,7 @@ use std::{
     path::{Path, PathBuf},
     str::FromStr,
     sync::Arc,
+    time::{Duration, SystemTime, UNIX_EPOCH},
 };
 
 use chia::{bls::master_to_wallet_unhardened_intermediate, protocol::Bytes32};
@@ -22,7 +23,7 @@ use tracing::{error, info, Level};
 use tracing_appender::rolling::{Builder, Rotation};
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer};
 
-use crate::{Error, Result};
+use crate::{peers::Peers, Error, Result};
 
 #[derive(Debug)]
 pub struct Sage {
@@ -63,6 +64,7 @@ impl Sage {
 
         let receiver = self.setup_sync_manager()?;
         self.switch_wallet().await?;
+        self.setup_peers().await?;
 
         info!("Sage wallet initialized");
 
@@ -280,6 +282,83 @@ impl Sage {
         Ok(())
     }
 
+    pub async fn setup_peers(&mut self) -> Result<()> {
+        let peer_dir = self.path.join("peers");
+
+        if !peer_dir.exists() {
+            fs::create_dir_all(&peer_dir)?;
+        }
+
+        let peer_path = peer_dir.join(format!("{}.bin", self.config.network.network_id));
+
+        let peers = if peer_path.try_exists()? {
+            Peers::from_bytes(&fs::read(&peer_path)?).unwrap_or_else(|error| {
+                error!("Failed to load peers, reverting to default: {error}");
+                Peers::default()
+            })
+        } else {
+            Peers::default()
+        };
+
+        let mut state = self.peer_state.lock().await;
+
+        for (&ip, &timestamp) in &peers.banned {
+            let now = SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .expect("system time before epoch")
+                .as_secs();
+
+            if now >= timestamp {
+                continue;
+            }
+
+            state.ban(ip, Duration::from_secs(timestamp - now), "already banned");
+        }
+
+        for &ip in &peers.connections {
+            if state.peer(ip).is_some() {
+                continue;
+            }
+
+            self.command_sender
+                .send(SyncCommand::ConnectPeer {
+                    ip,
+                    trusted: peers.trusted.contains(&ip),
+                })
+                .await?;
+        }
+
+        Ok(())
+    }
+
+    pub async fn save_peers(&self) -> Result<()> {
+        let peer_dir = self.path.join("peers");
+
+        if !peer_dir.exists() {
+            fs::create_dir_all(&peer_dir)?;
+        }
+
+        let mut peers = Peers::default();
+        let mut state = self.peer_state.lock().await;
+
+        for peer in state.peers() {
+            peers.connections.insert(peer.socket_addr().ip());
+        }
+
+        for (&ip, &ban) in state.banned_peers() {
+            peers.banned.insert(ip, ban);
+        }
+
+        for &ip in state.trusted_peers() {
+            peers.trusted.insert(ip);
+        }
+
+        let peer_path = peer_dir.join(format!("{}.bin", self.config.network.network_id));
+        fs::write(&peer_path, peers.to_bytes()?)?;
+
+        Ok(())
+    }
+
     #[allow(clippy::needless_pass_by_value)]
     pub fn parse_address(&self, input: String) -> Result<Bytes32> {
         let (puzzle_hash, prefix) = decode_address(&input)?;

src-tauri/Cargo.toml

@@ -29,6 +29,7 @@ tokio = { workspace = true }
 chia-wallet-sdk = { workspace = true, features = ["rustls"] }
 tauri-plugin-shell = { workspace = true }
 serde_json = { workspace = true, features = ["arbitrary_precision"] }
+tracing = { workspace = true }
 
 # This is to ensure that the bindgen feature is enabled for the aws-lc-rs crate.
 # https://aws.github.io/aws-lc-rs/platform_support.html#tested-platforms

src-tauri/src/app_state.rs

@@ -1,6 +1,7 @@
 use std::{
     ops::{Deref, DerefMut},
     path::Path,
+    sync::Arc,
 };
 
 use sage::Sage;
@@ -11,7 +12,7 @@ use tokio::sync::Mutex;
 
 use crate::error::Result;
 
-pub type AppState = Mutex<AppStateInner>;
+pub type AppState = Arc<Mutex<AppStateInner>>;
 
 pub struct AppStateInner {
     pub app_handle: AppHandle,
@@ -42,9 +43,9 @@ impl AppStateInner {
         }
     }
 
-    pub async fn initialize(&mut self) -> Result<()> {
+    pub async fn initialize(&mut self) -> Result<bool> {
         if self.initialized {
-            return Ok(());
+            return Ok(true);
         }
 
         self.initialized = true;
@@ -74,6 +75,6 @@ impl AppStateInner {
             Result::Ok(())
         });
 
-        Ok(())
+        Ok(false)
     }
 }