Merge pull request #142 from ChrisSon15/connect-wallet-protocol

move MemWallet to crate wallet and let it implement TradeWallet and ProtocolWalletApi traits

Author: ChrisSon15

Cargo.lock

@@ -16,6 +16,7 @@ rand_chacha = { workspace = true }
 testenv = { path = "../testenv" }
 thiserror = { workspace = true }
 tracing = { workspace = true }
+wallet = { path = "../wallet" }
 
 [dev-dependencies]
 bdk_wallet = { workspace = true, features = ["test-utils"] }

protocol/Cargo.toml

@@ -1,62 +1,20 @@
-use std::io::Write as _;
 use std::str::FromStr as _;
 
-use bdk_electrum::BdkElectrumClient;
-use bdk_electrum::electrum_client::Client;
-use bdk_wallet::bitcoin::bip32::Xpriv;
 use bdk_wallet::bitcoin::{
     Address, Amount, FeeRate, Network, OutPoint, Psbt, Script, ScriptBuf, Transaction, TxOut, Txid,
     relative,
 };
-use bdk_wallet::template::{Bip86, DescriptorTemplate as _};
-use bdk_wallet::{AddressInfo, KeychainKind, SignOptions, Wallet};
 use musig2::secp::{MaybeScalar, Point};
 use musig2::{PartialSignature, PubNonce};
-use rand::RngCore as _;
-use testenv::TestEnv;
 
+use wallet::protocol_wallet_api::MemWallet;
 use crate::multisig::{KeyCtx, SigCtx};
 use crate::receiver::{Receiver, ReceiverList};
 use crate::transaction::{
     DepositTxBuilder, ForwardingTxBuilder, RedirectTxBuilder, WarningTxBuilder, WithWitnesses as _,
 };
 use crate::wallet_service::WalletService;
 
-pub struct MemWallet {
-    wallet: Wallet,
-    client: BdkElectrumClient<Client>,
-}
-
-impl MemWallet {
-    pub fn transaction_broadcast(&self, tx: &Transaction) -> anyhow::Result<Txid> {
-        let result = self.client.transaction_broadcast(tx);
-
-        if let Err(e) = result {
-            if e.to_string().contains("Transaction already in block chain") {
-                return Ok(tx.compute_txid());
-            }
-            return Err(e.into());
-        }
-
-        Ok(result?)
-    }
-
-    pub fn funded_wallet(env: &TestEnv) -> Self {
-        // TODO move this line to TestEnv
-        let client =
-            BdkElectrumClient::new(Client::new(&env.electrum_url()).unwrap());
-        let mut wallet = Self::new(client).unwrap();
-        let address = wallet.next_unused_address();
-        let txid = env
-            .fund_address(&address, Amount::from_btc(10f64).unwrap())
-            .unwrap();
-        env.mine_block().unwrap();
-        env.wait_for_tx(txid).unwrap();
-        wallet.sync().unwrap();
-        wallet
-    }
-}
-
 #[derive(PartialEq, Eq, Clone, Copy, Debug)]
 #[expect(clippy::exhaustive_enums)]
 pub enum ProtocolRole {
@@ -73,94 +31,6 @@ impl ProtocolRole {
     }
 }
 
-// TODO think about stop_gap and batch_size
-const STOP_GAP: usize = 50;
-const BATCH_SIZE: usize = 5;
-
-impl MemWallet {
-    pub fn new(client: BdkElectrumClient<Client>) -> anyhow::Result<Self> {
-        let mut seed: [u8; 32] = [0u8; 32];
-        rand::rng().fill_bytes(&mut seed);
-
-        let network: Network = Network::Regtest;
-        let xprv: Xpriv = Xpriv::new_master(network, &seed)?;
-        tracing::info!(
-            "Generated Master Private Key:\n{xprv}\nWarning: be very careful with private \
-            keys when using MainNet! We are logging these values for convenience only because this \
-            is an example on RegTest.\n"
-        );
-
-        let (descriptor, external_map, _) = Bip86(xprv, KeychainKind::External)
-            .build(network)
-            .expect("Failed to build external descriptor");
-
-        let (change_descriptor, internal_map, _) = Bip86(xprv, KeychainKind::Internal)
-            .build(network)
-            .expect("Failed to build internal descriptor");
-
-        let wallet = Wallet::create(descriptor, change_descriptor)
-            .network(network)
-            .keymap(KeychainKind::External, external_map)
-            .keymap(KeychainKind::Internal, internal_map)
-            .create_wallet_no_persist()?;
-
-        Ok(Self { wallet, client })
-    }
-
-    pub fn sync(&mut self) -> anyhow::Result<()> {
-        // Populate the electrum client's transaction cache so it doesn't re-download transaction we
-        // already have.
-        self.client
-            .populate_tx_cache(self.wallet.tx_graph().full_txs().map(|tx_node| tx_node.tx));
-
-        let request = self.wallet.start_full_scan().inspect({
-            let mut stdout = std::io::stdout();
-            // let mut once = HashSet::<KeychainKind>::new();
-            move |_k, _spk_i, _| {
-                // if once.insert(k) {
-                //     print!("\nScanning keychain [{:?}]", k);
-                // }
-                // print!(" {:<3}", spk_i);
-                stdout.flush().expect("must flush");
-            }
-        });
-        tracing::info!("requesting update...");
-        let update = self
-            .client
-            .full_scan(request, STOP_GAP, BATCH_SIZE, false)?;
-        self.wallet.apply_update(update)?;
-        Ok(())
-    }
-
-    pub fn balance(&self) -> Amount {
-        self.wallet.balance().trusted_spendable()
-    }
-
-    pub fn next_unused_address(&mut self) -> AddressInfo {
-        // FIXME: `next_unused_address` just returns the same unused address over and over. It has
-        //  to either be marked as used (which change isn't staged and therefore presumably never
-        //  persisted) or a fresh address requested with `reveal_next_address`.
-        self.wallet.next_unused_address(KeychainKind::External)
-    }
-
-    fn _transfer_to_address(
-        &mut self,
-        address: &AddressInfo,
-        amount: Amount,
-    ) -> anyhow::Result<Txid> {
-        let mut tx_builder = self.wallet.build_tx();
-        tx_builder.add_recipient(address.script_pubkey(), amount);
-
-        let mut psbt = tx_builder.finish()?;
-        let finalized = self.wallet.sign(&mut psbt, SignOptions::default())?;
-        assert!(finalized);
-
-        let tx = psbt.extract_tx()?;
-        self.client.transaction_broadcast(&tx)?;
-        Ok(tx.compute_txid())
-    }
-}
-
 #[derive(Debug)]
 pub struct Round1Parameter {
     // DepositTx --------
@@ -497,7 +367,7 @@ impl RedirectTx {
     }
 
     pub fn broadcast(&self, me: &BMPContext) -> anyhow::Result<Txid> {
-        Ok(me.funds.client.transaction_broadcast(self.builder.signed_tx()?)?)
+        me.funds.transaction_broadcast(self.builder.signed_tx()?)
     }
 
     /// sum of all f64 must be 1
@@ -559,7 +429,7 @@ impl ClaimTx {
     }
 
     pub fn broadcast(&self, me: &BMPContext) -> anyhow::Result<Txid> {
-        Ok(me.funds.client.transaction_broadcast(self.signed_tx()?)?)
+        me.funds.transaction_broadcast(self.signed_tx()?)
     }
 }
 
@@ -651,7 +521,7 @@ impl WarningTx {
     }
 
     pub fn broadcast(&self, me: &BMPContext) -> anyhow::Result<Txid> {
-        Ok(me.funds.client.transaction_broadcast(self.signed_tx()?)?)
+        me.funds.transaction_broadcast(self.signed_tx()?)
     }
 }
 
@@ -759,7 +629,7 @@ impl SwapTx {
     }
 
     pub fn broadcast(&self, me: &BMPContext) -> anyhow::Result<Txid> {
-        Ok(me.funds.client.transaction_broadcast(self.builder.signed_tx()?)?)
+        me.funds.transaction_broadcast(self.builder.signed_tx()?)
     }
 }
 
@@ -785,11 +655,11 @@ impl DepositTx {
 
         let psbt = if ctx.am_buyer() {
             self.builder
-                .init_buyers_half_psbt(&mut ctx.funds.wallet, &mut rand::rng())?
+                    .init_buyers_half_psbt(&mut ctx.funds, &mut rand::rng())?
                 .buyers_half_psbt()?
         } else {
             self.builder
-                .init_sellers_half_psbt(&mut ctx.funds.wallet, &mut rand::rng())?
+                    .init_sellers_half_psbt(&mut ctx.funds, &mut rand::rng())?
                 .sellers_half_psbt()?
         };
         Ok(psbt.clone())
@@ -810,9 +680,9 @@ impl DepositTx {
 
     fn sign(&mut self, ctx: &mut BMPContext) -> anyhow::Result<()> {
         if ctx.am_buyer() {
-            self.builder.sign_buyer_inputs(&ctx.funds.wallet)?;
+            self.builder.sign_buyer_inputs(&ctx.funds)?;
         } else {
-            self.builder.sign_seller_inputs(&ctx.funds.wallet)?;
+            self.builder.sign_seller_inputs(&ctx.funds)?;
         }
         Ok(())
     }

protocol/src/protocol_musig_adaptor.rs

@@ -14,6 +14,7 @@ use bdk_wallet::miniscript::{Descriptor, ToPublicKey as _};
 use bdk_wallet::{KeychainKind, SignOptions, TxOrdering, Wallet};
 use rand::{RngCore, SeedableRng as _};
 use rand_chacha::ChaCha20Rng;
+use wallet::protocol_wallet_api::MemWallet;
 
 use crate::receiver::Receiver;
 use crate::swap::Swap as _;
@@ -55,7 +56,9 @@ struct MockTradeWallet<Cs: Iterator<Item=TxOutput>, As: Iterator<Item=Address>>
     internal_key: Option<XOnlyPublicKey>,
 }
 
-impl<Cs: Iterator<Item=TxOutput>, As: Iterator<Item=Address>> TradeWallet for MockTradeWallet<Cs, As> {
+impl<Cs: Iterator<Item = TxOutput>, As: Iterator<Item = Address>> TradeWallet
+    for MockTradeWallet<Cs, As>
+{
     fn network(&self) -> Network { Network::Regtest }
 
     fn new_address(&mut self) -> Result<Address> {
@@ -235,19 +238,19 @@ impl TradeWallet for Wallet {
     ) -> Result<Psbt> {
         let mut builder = self.build_tx();
         builder
-            .ordering(TxOrdering::Untouched)
-            .nlocktime(absolute::LockTime::ZERO)
-            .fee_rate(fee_rate)
-            .add_recipient(half_deposit_placeholder_spk(rng), deposit_amount);
+                .ordering(TxOrdering::Untouched)
+                .nlocktime(absolute::LockTime::ZERO)
+                .fee_rate(fee_rate)
+                .add_recipient(half_deposit_placeholder_spk(rng), deposit_amount);
         for receiver in trade_fee_receivers {
             builder.add_recipient(receiver.address.script_pubkey(), receiver.amount);
         }
         let mut psbt = builder.finish()?;
         // Calculate tx fee overpay unconditionally, as this performs additional checks on the PSBT:
         let overpay_msat: u64 = half_psbt_fee_overpay_msat(&psbt, fee_rate)
-            .ok_or(TransactionErrorKind::Overflow)?
-            .try_into()
-            .map_err(|_| TransactionErrorKind::InvalidPsbt)?;
+                .ok_or(TransactionErrorKind::Overflow)?
+                .try_into()
+                .map_err(|_| TransactionErrorKind::InvalidPsbt)?;
         let change_output_index = 1 + trade_fee_receivers.len();
         if psbt.unsigned_tx.output.len() > change_output_index {
             // Correct any tx fee overpay due to overly conservative input witness size estimation
@@ -277,6 +280,59 @@ impl TradeWallet for Wallet {
     }
 }
 
+impl TradeWallet for MemWallet {
+    fn network(&self) -> Network { self.network() }
+
+    fn new_address(&mut self) -> Result<Address> {
+        // For privacy, always get fresh addresses for the trade protocol.
+        // FIXME: Need to find a way to prevent gaps of unused addresses from growing too large.
+        Ok(self.reveal_next_address())
+    }
+
+    fn new_internal_key(&mut self) -> Result<XOnlyPublicKey> {
+        Self::new_internal_key(self).map_err(TransactionErrorKind::from)
+    }
+
+    fn create_half_deposit_psbt(
+        &mut self,
+        deposit_amount: Amount,
+        fee_rate: FeeRate,
+        trade_fee_receivers: &[Receiver],
+        rng: &mut dyn RngCore,
+    ) -> Result<Psbt> {
+        let mut res: Vec<(ScriptBuf, Amount)> = trade_fee_receivers
+                .iter()
+                .map(|receiver| (receiver.address.script_pubkey(), deposit_amount))
+                .collect();
+        res.push((half_deposit_placeholder_spk(rng), deposit_amount));
+        let mut psbt = self.create_psbt(res, fee_rate)?;
+        // Calculate tx fee overpay unconditionally, as this performs additional checks on the PSBT:
+        let overpay_msat: u64 = half_psbt_fee_overpay_msat(&psbt, fee_rate)
+            .ok_or(TransactionErrorKind::Overflow)?
+            .try_into()
+            .map_err(|_| TransactionErrorKind::InvalidPsbt)?;
+        let change_output_index = 1 + trade_fee_receivers.len();
+        if psbt.unsigned_tx.output.len() > change_output_index {
+            // Correct any tx fee overpay due to overly conservative input witness size estimation
+            // by BDK (as each witness is assumed to potentially have a non-default sighash type in
+            // the case of p2tr and not grind for low R in the case of p2wpkh), and also due to the
+            // fact that each would-be-signed half-deposit PSBT is 1 wu bigger than ideal.
+            let overpay = Amount::from_sat(overpay_msat / 1000);
+            psbt.unsigned_tx.output[change_output_index].value += overpay;
+        }
+        psbt.redact_sensitive_fields();
+        Ok(psbt)
+    }
+
+    fn sign_selected_inputs(&self, psbt: &mut Psbt, is_selected: &dyn Fn(&OutPoint) -> bool) -> Result<()> {
+        if !is_well_formed(psbt) {
+            return Err(TransactionErrorKind::InvalidPsbt);
+        }
+        self.sign_the_selected_inputs(psbt, is_selected)?;
+        Ok(())
+    }
+}
+
 trait Redact {
     fn redact_sensitive_fields(&mut self);
 }

protocol/src/psbt.rs

@@ -568,6 +568,8 @@ pub enum TransactionErrorKind {
     CreateTx(#[from] bdk_wallet::error::CreateTxError),
     Signer(#[from] bdk_wallet::signer::SignerError),
     Conversion(#[from] bdk_wallet::miniscript::descriptor::ConversionError),
+    Wallet(#[from] wallet::protocol_wallet_api::WalletErrorKind),
+    Anyhow(#[from] anyhow::Error),
 }
 
 impl From<ExtractTxError> for TransactionErrorKind {

protocol/src/transaction.rs

@@ -3,7 +3,7 @@ This represents a service to get access to the wallet
 as of now its a fake implementation, you need to pass in the wallet which you
 want to retrieve later. In the real version this should discover and actually load the user's wallet.
  */
-use crate::protocol_musig_adaptor::MemWallet;
+use wallet::protocol_wallet_api::MemWallet;
 
 #[derive(Default)]
 pub struct WalletService {