Merge pull request #698 from xch-dev/arbor

Add Arbor Wallet compatibility

Author: Rigidity

.sqlx/query-3d10c346c9307a0fb2f71cadcf04881e163e2bf96d5c54b475dedcf8c981dcbc.json

@@ -1,5 +1,8 @@
 use chia::{bls::PublicKey, clvm_utils::ToTreeHash, protocol::Bytes32};
-use chia_wallet_sdk::driver::{ClawbackV2, OptionType, OptionUnderlying};
+use chia_wallet_sdk::{
+    driver::{ClawbackV2, OptionType, OptionUnderlying},
+    types::{puzzles::P2DelegatedConditionsArgs, Mod},
+};
 use sqlx::{query, SqliteExecutor};
 
 use crate::{Convert, Database, DatabaseError, DatabaseTx, Result};
@@ -9,13 +12,15 @@ pub enum P2PuzzleKind {
     PublicKey,
     Clawback,
     Option,
+    Arbor,
 }
 
 #[derive(Debug, Clone, Copy)]
 pub enum P2Puzzle {
     PublicKey(PublicKey),
     Clawback(Clawback),
     Option(Underlying),
+    Arbor(PublicKey),
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -100,6 +105,13 @@ impl Database {
                     strike_type: underlying.strike_type,
                 }))
             }
+            P2PuzzleKind::Arbor => {
+                let Some(key) = arbor_key(&self.pool, puzzle_hash).await? else {
+                    return Err(DatabaseError::PublicKeyNotFound);
+                };
+
+                Ok(P2Puzzle::Arbor(key))
+            }
         }
     }
 
@@ -162,10 +174,14 @@ impl DatabaseTx<'_> {
     pub async fn insert_option_p2_puzzle(&mut self, underlying: OptionUnderlying) -> Result<()> {
         insert_option_p2_puzzle(&mut *self.tx, underlying).await
     }
+
+    pub async fn insert_arbor_p2_puzzle(&mut self, key: PublicKey) -> Result<()> {
+        insert_arbor_p2_puzzle(&mut *self.tx, key).await
+    }
 }
 
 async fn custody_p2_puzzle_hashes(conn: impl SqliteExecutor<'_>) -> Result<Vec<Bytes32>> {
-    query!("SELECT hash FROM p2_puzzles WHERE kind = 0")
+    query!("SELECT hash FROM p2_puzzles WHERE kind IN (0, 3)")
         .fetch_all(conn)
         .await?
         .into_iter()
@@ -200,7 +216,7 @@ async fn is_custody_p2_puzzle_hash(
     let puzzle_hash = puzzle_hash.as_ref();
 
     Ok(query!(
-        "SELECT COUNT(*) AS count FROM p2_puzzles WHERE hash = ? AND kind = 0",
+        "SELECT COUNT(*) AS count FROM p2_puzzles WHERE hash = ? AND kind IN (0, 3)",
         puzzle_hash
     )
     .fetch_one(conn)
@@ -405,6 +421,30 @@ async fn insert_option_p2_puzzle(
     Ok(())
 }
 
+async fn insert_arbor_p2_puzzle(conn: impl SqliteExecutor<'_>, key: PublicKey) -> Result<()> {
+    let p2_puzzle_hash = P2DelegatedConditionsArgs::new(key)
+        .curry_tree_hash()
+        .to_vec();
+    let key = key.to_bytes();
+    let key = key.as_ref();
+
+    query!(
+        "
+        INSERT OR IGNORE INTO p2_puzzles (hash, kind) VALUES (?, 3);
+
+        INSERT OR IGNORE INTO p2_arbor (p2_puzzle_id, key)
+        VALUES ((SELECT id FROM p2_puzzles WHERE hash = ?), ?);
+        ",
+        p2_puzzle_hash,
+        p2_puzzle_hash,
+        key,
+    )
+    .execute(conn)
+    .await?;
+
+    Ok(())
+}
+
 async fn p2_puzzle_kind(
     conn: impl SqliteExecutor<'_>,
     p2_puzzle_hash: Bytes32,
@@ -419,6 +459,7 @@ async fn p2_puzzle_kind(
         0 => P2PuzzleKind::PublicKey,
         1 => P2PuzzleKind::Clawback,
         2 => P2PuzzleKind::Option,
+        3 => P2PuzzleKind::Arbor,
         _ => return Err(DatabaseError::InvalidEnumVariant),
     })
 }
@@ -496,6 +537,27 @@ async fn underlying_launcher_id(
     .convert()
 }
 
+async fn arbor_key(
+    conn: impl SqliteExecutor<'_>,
+    p2_puzzle_hash: Bytes32,
+) -> Result<Option<PublicKey>> {
+    let p2_puzzle_hash = p2_puzzle_hash.as_ref();
+
+    let row = query!(
+        "
+        SELECT key
+        FROM p2_puzzles
+        INNER JOIN p2_arbor ON p2_arbor.p2_puzzle_id = p2_puzzles.id
+        WHERE p2_puzzles.hash = ?
+        ",
+        p2_puzzle_hash
+    )
+    .fetch_optional(conn)
+    .await?;
+
+    row.map(|row| row.key.convert()).transpose()
+}
+
 async fn derivation(
     conn: impl SqliteExecutor<'_>,
     public_key: PublicKey,

…0e41f5e7b40a1434df439f0ccca126ee8ac.json …d9c136cb5c929e4b14ffa0f0f4422339b99.json.sqlx/query-91054392bb590bd729d257989bee20e41f5e7b40a1434df439f0ccca126ee8ac.json renamed to .sqlx/query-881596d83c520b3f27ed2e82f09d0d9c136cb5c929e4b14ffa0f0f4422339b99.json .sqlx/query-91054392bb590bd729d257989bee20e41f5e7b40a1434df439f0ccca126ee8ac.json renamed to .sqlx/query-881596d83c520b3f27ed2e82f09d0d9c136cb5c929e4b14ffa0f0f4422339b99.json

@@ -13,10 +13,11 @@ use chia_puzzles::SETTLEMENT_PAYMENT_HASH;
 use chia_wallet_sdk::{
     driver::{
         Action, Cat, ClawbackV2, Deltas, DriverError, Id, Layer, OptionUnderlying, Outputs,
-        Relation, SettlementLayer, SpendContext, SpendKind, SpendWithConditions, SpendableAsset,
-        Spends, StandardLayer,
+        P2DelegatedConditionsLayer, Relation, SettlementLayer, SpendContext, SpendKind,
+        SpendWithConditions, SpendableAsset, Spends, StandardLayer,
     },
     signer::AggSigConstants,
+    types::puzzles::P2DelegatedConditionsSolution,
     utils::select_coins,
 };
 use indexmap::IndexMap;
@@ -354,6 +355,13 @@ impl Wallet {
 
                             underlying.clawback_spend(ctx, spend)
                         }
+                        P2Puzzle::Arbor(key) => P2DelegatedConditionsLayer::new(*key)
+                            .construct_spend(
+                                ctx,
+                                P2DelegatedConditionsSolution::new(
+                                    spend.finish().into_iter().collect(),
+                                ),
+                            ),
                     }
                 }
                 SpendKind::Settlement(spend) => SettlementLayer

…5eb866a5176fb5df024ac8b0897f82d3008.json …6f4018f81f4b5d471c0d2cb329bc9a77e86.json.sqlx/query-23a3a5e21db7039a455a7fae4ad0a5eb866a5176fb5df024ac8b0897f82d3008.json renamed to .sqlx/query-8d85ede79ebf0aadc828252ef02b66f4018f81f4b5d471c0d2cb329bc9a77e86.json .sqlx/query-23a3a5e21db7039a455a7fae4ad0a5eb866a5176fb5df024ac8b0897f82d3008.json renamed to .sqlx/query-8d85ede79ebf0aadc828252ef02b66f4018f81f4b5d471c0d2cb329bc9a77e86.json

@@ -44,10 +44,7 @@ impl Wallet {
                 continue;
             };
             let Some(derivation) = self.db.derivation(required.public_key).await? else {
-                if partial {
-                    continue;
-                }
-                return Err(WalletError::UnknownPublicKey);
+                continue;
             };
             derivations.insert(required.public_key, derivation);
         }
@@ -75,9 +72,18 @@ impl Wallet {
             let RequiredSignature::Bls(required) = required else {
                 continue;
             };
-            let Some(sk) = secret_keys.get(&required.public_key).cloned() else {
-                continue;
+
+            let sk = if required.public_key == master_sk.public_key() {
+                master_sk.clone()
+            } else if let Some(sk) = secret_keys.get(&required.public_key).cloned() {
+                sk
+            } else {
+                if partial {
+                    continue;
+                }
+                return Err(WalletError::UnknownPublicKey);
             };
+
             aggregated_signature += &sign(&sk, required.message());
         }
 

.sqlx/query-bad04dbba3a2f7b238f59b1badae2b9c2098511951341f83e01db561430733ff.json

@@ -216,6 +216,8 @@ impl Sage {
             }
         }
 
+        tx.insert_arbor_p2_puzzle(master_pk).await?;
+
         tx.commit().await?;
 
         if req.login {

crates/sage-database/src/tables/p2_puzzles.rs

@@ -6,7 +6,9 @@ use chia::{
     protocol::{Bytes32, Coin, CoinSpend, SpendBundle},
     puzzles::{DeriveSynthetic, Proof},
 };
-use chia_wallet_sdk::driver::{ClawbackV2, Layer, OptionUnderlying, SpendContext, StandardLayer};
+use chia_wallet_sdk::driver::{
+    ClawbackV2, Layer, OptionUnderlying, P2DelegatedConditionsLayer, SpendContext, StandardLayer,
+};
 use sage_api::wallet_connect::{
     self, AssetCoinType, FilterUnlockedCoins, FilterUnlockedCoinsResponse, GetAssetCoins,
     GetAssetCoinsResponse, LineageProof, SendTransactionImmediately,
@@ -104,6 +106,9 @@ impl Sage {
                     );
                     underlying.into_1_of_n().construct_puzzle(&mut ctx)?
                 }
+                P2Puzzle::Arbor(key) => {
+                    P2DelegatedConditionsLayer::new(key).construct_puzzle(&mut ctx)?
+                }
             };
 
             let (puzzle, proof) = match req.kind {

crates/sage-wallet/src/wallet.rs

@@ -0,0 +1,9 @@
+/*
+ * P2 arbor is a p2 puzzle with kind = 3
+ */
+ CREATE TABLE p2_arbor (
+  id INTEGER NOT NULL PRIMARY KEY,
+  p2_puzzle_id INTEGER NOT NULL UNIQUE,
+  key BLOB NOT NULL,
+  FOREIGN KEY (p2_puzzle_id) REFERENCES p2_puzzles(id) ON DELETE CASCADE
+);