Next

examples/wallet/src/main.rs

@@ -1,4 +1,4 @@
-use std::{collections::HashSet, str::FromStr, sync::Arc};
+use std::{collections::HashSet, str::FromStr};
 
 use bip39::Mnemonic;
 use chia_bls::{
@@ -11,17 +11,18 @@ use chia_client::Peer;
 use chia_protocol::{Bytes32, CoinState, NodeType};
 use chia_wallet::standard::DEFAULT_HIDDEN_PUZZLE_HASH;
 use chia_wallet_sdk::{
-    connect_peer, create_tls_connector, load_ssl_cert, migrate, CoinStore, HardenedKeyStore,
-    PuzzleStore, UnhardenedKeyStore,
+    connect_peer, create_tls_connector, load_ssl_cert, migrate, unused_indices, CoinStore,
+    HardenedKeyStore, KeyStore, PuzzleStore, SqliteCoinStore, UnhardenedKeyStore,
 };
 use sqlx::SqlitePool;
 use thiserror::Error;
 
 struct Wallet {
-    peer: Arc<Peer>,
-    hardened_keys: HardenedKeyStore,
+    peer: Peer,
     unhardened_keys: UnhardenedKeyStore,
-    standard_coins: CoinStore,
+    hardened_keys: HardenedKeyStore,
+    standard_coins: SqliteCoinStore,
+    derivation_size: u32,
 }
 
 #[derive(Debug, Error)]
@@ -32,49 +33,115 @@ enum WalletError {
 
 impl Wallet {
     async fn initial_sync(&self) -> Result<(), WalletError> {
-        let mut hardened_phs = HashSet::new();
-        let mut unhardened_phs = HashSet::new();
-        let mut puzzle_hashes = Vec::new();
-
-        for puzzle_hash in self.hardened_keys.puzzle_hashes().await {
-            hardened_phs.insert(Bytes32::new(puzzle_hash));
-            puzzle_hashes.push(Bytes32::new(puzzle_hash));
-        }
+        self.unhardened_keys
+            .derive_to_index(self.derivation_size)
+            .await;
 
-        for puzzle_hash in self.unhardened_keys.puzzle_hashes().await {
-            unhardened_phs.insert(Bytes32::new(puzzle_hash));
-            puzzle_hashes.push(Bytes32::new(puzzle_hash));
-        }
-
-        for puzzle_hashes in puzzle_hashes.chunks(10000) {
+        for puzzle_hashes in self.unhardened_keys.puzzle_hashes().await.chunks(10000) {
             let coin_states = self
                 .peer
-                .register_for_ph_updates(puzzle_hashes.to_vec(), 0)
+                .register_for_ph_updates(
+                    puzzle_hashes.iter().map(|ph| Bytes32::new(*ph)).collect(),
+                    0,
+                )
                 .await?;
 
             self.apply_updates(coin_states).await?;
         }
 
+        loop {
+            match unused_indices(&self.unhardened_keys, &self.standard_coins).await {
+                Ok(indices) => {
+                    if indices.len() < self.derivation_size as usize {
+                        if !self
+                            .derive_unhardened_to(indices.end + self.derivation_size)
+                            .await?
+                        {
+                            break;
+                        }
+                        continue;
+                    }
+                    break;
+                }
+                Err(index) => {
+                    if !self.derive_unhardened_to(index).await? {
+                        break;
+                    }
+                }
+            }
+        }
+
+        loop {
+            match unused_indices(&self.hardened_keys, &self.standard_coins).await {
+                Ok(indices) => {
+                    if indices.len() < self.derivation_size as usize {
+                        if !self
+                            .derive_hardened_to(indices.end + self.derivation_size)
+                            .await?
+                        {
+                            break;
+                        }
+                        continue;
+                    }
+                    break;
+                }
+                Err(index) => {
+                    if !self.derive_hardened_to(index).await? {
+                        break;
+                    }
+                }
+            }
+        }
+
         Ok(())
     }
 
+    async fn derive_unhardened_to(&self, index: u32) -> Result<bool, WalletError> {
+        self.unhardened_keys.derive_to_index(index).await;
+        let puzzle_hashes = self.unhardened_keys.puzzle_hashes().await;
+
+        let coin_states = self
+            .peer
+            .register_for_ph_updates(
+                puzzle_hashes.iter().map(|ph| Bytes32::new(*ph)).collect(),
+                0,
+            )
+            .await?;
+        let found = !coin_states.is_empty();
+
+        self.apply_updates(coin_states).await?;
+
+        Ok(found)
+    }
+
+    async fn derive_hardened_to(&self, index: u32) -> Result<bool, WalletError> {
+        self.hardened_keys.derive_to_index(index).await;
+        let puzzle_hashes = self.hardened_keys.puzzle_hashes().await;
+
+        let coin_states = self
+            .peer
+            .register_for_ph_updates(
+                puzzle_hashes.iter().map(|ph| Bytes32::new(*ph)).collect(),
+                0,
+            )
+            .await?;
+        let found = !coin_states.is_empty();
+
+        self.apply_updates(coin_states).await?;
+
+        Ok(found)
+    }
+
     async fn apply_updates(&self, coin_states: Vec<CoinState>) -> Result<(), WalletError> {
-        let mut p2_puzzle_hashes: HashSet<Bytes32> = self
-            .hardened_keys
+        let p2_puzzle_hashes: HashSet<Bytes32> = self
+            .unhardened_keys
             .puzzle_hashes()
             .await
             .into_iter()
+            .chain(self.hardened_keys.puzzle_hashes().await)
             .map(Bytes32::new)
             .collect();
 
-        p2_puzzle_hashes.extend(
-            self.unhardened_keys
-                .puzzle_hashes()
-                .await
-                .into_iter()
-                .map(Bytes32::new),
-        );
-
         let mut standard_coins = Vec::new();
 
         for coin_state in coin_states {
@@ -97,9 +164,9 @@ impl Wallet {
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let cert = load_ssl_cert("wallet.key", "wallet.crt");
-    let tls_connector = create_tls_connector(&cert);
-    let peer = connect_peer("localhost:56342", tls_connector).await?;
+    let cert = load_ssl_cert("wallet.key", "wallet.crt")?;
+    let tls_connector = create_tls_connector(&cert)?;
+    let peer = connect_peer("localhost:38926", tls_connector).await?;
 
     peer.send_handshake("simulator0".to_string(), NodeType::Wallet)
         .await?;
@@ -117,20 +184,22 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let seed = mnemonic.to_seed("");
     let sk = SecretKey::from_seed(&seed);
 
-    let root_sk = master_to_wallet_hardened_intermediate(&sk);
-    let hardened_keys = HardenedKeyStore::new(pool.clone(), root_sk, DEFAULT_HIDDEN_PUZZLE_HASH);
-
-    let root_pk = master_to_wallet_unhardened_intermediate(&sk.public_key());
+    let intermediate_pk = master_to_wallet_unhardened_intermediate(&sk.public_key());
     let unhardened_keys =
-        UnhardenedKeyStore::new(pool.clone(), root_pk, DEFAULT_HIDDEN_PUZZLE_HASH);
+        UnhardenedKeyStore::new(pool.clone(), intermediate_pk, DEFAULT_HIDDEN_PUZZLE_HASH);
 
-    let standard_coins = CoinStore::new(pool.clone());
+    let intermediate_sk = master_to_wallet_hardened_intermediate(&sk);
+    let hardened_keys =
+        HardenedKeyStore::new(pool.clone(), intermediate_sk, DEFAULT_HIDDEN_PUZZLE_HASH);
+
+    let standard_coins = SqliteCoinStore::new(pool.clone());
 
     let wallet = Wallet {
         peer,
-        hardened_keys,
         unhardened_keys,
+        hardened_keys,
         standard_coins,
+        derivation_size: 500,
     };
 
     wallet.initial_sync().await?;

src/sqlite.rs

@@ -1,11 +1,11 @@
 use sqlx::{migrate::MigrateError, SqlitePool};
 
-mod coin_store;
 mod hardened_key_store;
+mod sqlite_coin_store;
 mod unhardened_key_store;
 
-pub use coin_store::*;
 pub use hardened_key_store::*;
+pub use sqlite_coin_store::*;
 pub use unhardened_key_store::*;
 
 pub async fn migrate(pool: &SqlitePool) -> Result<(), MigrateError> {

src/sqlite/coin_store.rs → src/sqlite/sqlite_coin_store.rs

@@ -1,16 +1,56 @@
 use chia_protocol::{Coin, CoinState};
 use sqlx::SqlitePool;
 
-pub struct CoinStore {
+use crate::CoinStore;
+
+pub struct SqliteCoinStore {
     pool: SqlitePool,
 }
 
-impl CoinStore {
+impl SqliteCoinStore {
     pub fn new(pool: SqlitePool) -> Self {
         Self { pool }
     }
+}
+
+impl CoinStore for SqliteCoinStore {
+    async fn apply_updates(&self, coin_states: Vec<CoinState>) {
+        let mut tx = self.pool.begin().await.unwrap();
+
+        for coin_state in coin_states {
+            let coin_id = coin_state.coin.coin_id().to_vec();
+            let parent_coin_info = coin_state.coin.parent_coin_info.to_bytes().to_vec();
+            let puzzle_hash = coin_state.coin.puzzle_hash.to_bytes().to_vec();
+            let amount = coin_state.coin.amount as i64;
+
+            sqlx::query!(
+                "
+                REPLACE INTO `coin_states` (
+                    `coin_id`,
+                    `parent_coin_info`,
+                    `puzzle_hash`,
+                    `amount`,
+                    `created_height`,
+                    `spent_height`
+                )
+                VALUES (?, ?, ?, ?, ?, ?)
+                ",
+                coin_id,
+                parent_coin_info,
+                puzzle_hash,
+                amount,
+                coin_state.created_height,
+                coin_state.spent_height
+            )
+            .execute(&mut *tx)
+            .await
+            .unwrap();
+        }
+
+        tx.commit().await.unwrap();
+    }
 
-    pub async fn unspent_coins(&self) -> Vec<Coin> {
+    async fn unspent_coins(&self) -> Vec<Coin> {
         let rows = sqlx::query!(
             "
             SELECT `parent_coin_info`, `puzzle_hash`, `amount`
@@ -37,7 +77,7 @@ impl CoinStore {
             .collect()
     }
 
-    pub async fn coin_state(&self, coin_id: [u8; 32]) -> CoinState {
+    async fn coin_state(&self, coin_id: [u8; 32]) -> Option<CoinState> {
         let coin_id = coin_id.to_vec();
 
         let row = sqlx::query!(
@@ -48,54 +88,36 @@ impl CoinStore {
             ",
             coin_id
         )
-        .fetch_one(&self.pool)
+        .fetch_optional(&self.pool)
         .await
-        .unwrap();
+        .unwrap()?;
 
-        CoinState {
+        Some(CoinState {
             coin: Coin {
                 parent_coin_info: row.parent_coin_info.try_into().unwrap(),
                 puzzle_hash: row.puzzle_hash.try_into().unwrap(),
                 amount: row.amount as u64,
             },
             created_height: row.created_height.map(|height| height as u32),
             spent_height: row.spent_height.map(|height| height as u32),
-        }
+        })
     }
 
-    pub async fn apply_updates(&self, coin_states: Vec<CoinState>) {
-        let mut tx = self.pool.begin().await.unwrap();
-
-        for coin_state in coin_states {
-            let coin_id = coin_state.coin.coin_id().to_vec();
-            let parent_coin_info = coin_state.coin.parent_coin_info.to_bytes().to_vec();
-            let puzzle_hash = coin_state.coin.puzzle_hash.to_bytes().to_vec();
-            let amount = coin_state.coin.amount as i64;
+    async fn is_used(&self, puzzle_hash: [u8; 32]) -> bool {
+        let puzzle_hash = puzzle_hash.to_vec();
 
-            sqlx::query!(
-                "
-                REPLACE INTO `coin_states` (
-                    `coin_id`,
-                    `parent_coin_info`,
-                    `puzzle_hash`,
-                    `amount`,
-                    `created_height`,
-                    `spent_height`
-                )
-                VALUES (?, ?, ?, ?, ?, ?)
-                ",
-                coin_id,
-                parent_coin_info,
-                puzzle_hash,
-                amount,
-                coin_state.created_height,
-                coin_state.spent_height
-            )
-            .execute(&mut *tx)
-            .await
-            .unwrap();
-        }
+        let row = sqlx::query!(
+            "
+            SELECT COUNT(*) AS `count`
+            FROM `coin_states`
+            WHERE `puzzle_hash` = ?
+            ",
+            puzzle_hash
+        )
+        .fetch_one(&self.pool)
+        .await
+        .unwrap();
 
-        tx.commit().await.unwrap();
+        row.count > 0
     }
 }