1  use std::collections::BTreeMap;
  2  use std::iter::repeat;
  3  use std::sync::Arc;
  4  use std::time::Duration;
  5  
  6  use anyhow::format_err;
  7  use async_trait::async_trait;
  8  use bitcoin::absolute::LockTime;
  9  use bitcoin::block::{Header as BlockHeader, Version};
 10  use bitcoin::blockdata::constants::genesis_block;
 11  use bitcoin::hash_types::Txid;
 12  use bitcoin::hashes::Hash;
 13  use bitcoin::merkle_tree::PartialMerkleTree;
 14  use bitcoin::{
 15      Address, Block, BlockHash, CompactTarget, Network, OutPoint, ScriptBuf, Transaction, TxOut,
 16  };
 17  use fedimint_bitcoind::{
 18      register_bitcoind, DynBitcoindRpc, IBitcoindRpc, IBitcoindRpcFactory,
 19      Result as BitcoinRpcResult,
 20  };
 21  use fedimint_core::envs::BitcoinRpcConfig;
 22  use fedimint_core::task::{sleep_in_test, TaskHandle};
 23  use fedimint_core::txoproof::TxOutProof;
 24  use fedimint_core::util::SafeUrl;
 25  use fedimint_core::{Amount, Feerate};
 26  use rand::rngs::OsRng;
 27  use tracing::debug;
 28  
 29  use super::BitcoinTest;
 30  
 31  #[derive(Debug, Clone)]
 32  pub struct FakeBitcoinFactory {
 33      pub bitcoin: FakeBitcoinTest,
 34      pub config: BitcoinRpcConfig,
 35  }
 36  
 37  impl FakeBitcoinFactory {
 38      /// Registers a fake bitcoin rpc factory for testing
 39      pub fn register_new() -> FakeBitcoinFactory {
 40          let kind = format!("test_btc-{}", rand::random::<u64>());
 41          let factory = FakeBitcoinFactory {
 42              bitcoin: FakeBitcoinTest::new(),
 43              config: BitcoinRpcConfig {
 44                  kind: kind.clone(),
 45                  url: "http://ignored".parse().unwrap(),
 46              },
 47          };
 48          register_bitcoind(kind, factory.clone().into());
 49          factory
 50      }
 51  }
 52  
 53  impl IBitcoindRpcFactory for FakeBitcoinFactory {
 54      fn create_connection(
 55          &self,
 56          _url: &SafeUrl,
 57          _handle: TaskHandle,
 58      ) -> anyhow::Result<DynBitcoindRpc> {
 59          Ok(self.bitcoin.clone().into())
 60      }
 61  }
 62  
 63  #[derive(Debug)]
 64  struct FakeBitcoinTestInner {
 65      /// Simulates mined bitcoin blocks
 66      blocks: Vec<Block>,
 67      /// Simulates pending transactions in the mempool
 68      pending: Vec<Transaction>,
 69      /// Tracks how much bitcoin was sent to an address (doesn't track sending
 70      /// out of it)
 71      addresses: BTreeMap<Txid, Amount>,
 72      /// Simulates the merkle tree proofs
 73      proofs: BTreeMap<Txid, TxOutProof>,
 74      /// Simulates the script history
 75      scripts: BTreeMap<ScriptBuf, Vec<Transaction>>,
 76  }
 77  
 78  #[derive(Clone, Debug)]
 79  pub struct FakeBitcoinTest {
 80      inner: Arc<std::sync::RwLock<FakeBitcoinTestInner>>,
 81  }
 82  
 83  impl Default for FakeBitcoinTest {
 84      fn default() -> Self {
 85          Self::new()
 86      }
 87  }
 88  
 89  impl FakeBitcoinTest {
 90      pub fn new() -> Self {
 91          let inner = FakeBitcoinTestInner {
 92              blocks: vec![genesis_block(Network::Regtest)],
 93              pending: vec![],
 94              addresses: Default::default(),
 95              proofs: Default::default(),
 96              scripts: Default::default(),
 97          };
 98          FakeBitcoinTest {
 99              inner: std::sync::RwLock::new(inner).into(),
100          }
101      }
102  
103      fn pending_merkle_tree(pending: &[Transaction]) -> PartialMerkleTree {
104          let txs = pending.iter().map(|tx| tx.txid()).collect::<Vec<Txid>>();
105          let matches = repeat(true).take(txs.len()).collect::<Vec<bool>>();
106          PartialMerkleTree::from_txids(txs.as_slice(), matches.as_slice())
107      }
108  
109      fn new_transaction(out: Vec<TxOut>) -> Transaction {
110          Transaction {
111              version: 0,
112              lock_time: LockTime::ZERO,
113              input: vec![],
114              output: out,
115          }
116      }
117  
118      fn mine_block(
119          addresses: &mut BTreeMap<Txid, Amount>,
120          blocks: &mut Vec<Block>,
121          pending: &mut Vec<Transaction>,
122      ) {
123          debug!(
124              "Mining block: {} transactions, {} blocks",
125              pending.len(),
126              blocks.len()
127          );
128          let root = BlockHash::hash(&[0]);
129          for tx in pending.iter() {
130              addresses.insert(tx.txid(), Amount::from_sats(output_sum(tx)));
131          }
132          // all blocks need at least one transaction
133          if pending.is_empty() {
134              pending.push(Self::new_transaction(vec![]));
135          }
136          let merkle_root = Self::pending_merkle_tree(pending)
137              .extract_matches(&mut vec![], &mut vec![])
138              .unwrap();
139          let block = Block {
140              header: BlockHeader {
141                  version: Version::from_consensus(0),
142                  prev_blockhash: blocks.last().map(|b| b.header.block_hash()).unwrap_or(root),
143                  merkle_root,
144                  time: 0,
145                  bits: CompactTarget::from_consensus(0),
146                  nonce: 0,
147              },
148              txdata: pending.clone(),
149          };
150          pending.clear();
151          blocks.push(block);
152      }
153  }
154  
155  #[async_trait]
156  impl BitcoinTest for FakeBitcoinTest {
157      async fn lock_exclusive(&self) -> Box<dyn BitcoinTest + Send + Sync> {
158          // With  FakeBitcoinTest, every test spawns their own instance,
159          // so not need to lock anything
160          Box::new(self.clone())
161      }
162  
163      async fn mine_blocks(&self, block_num: u64) {
164          let mut inner = self.inner.write().unwrap();
165  
166          let FakeBitcoinTestInner {
167              ref mut blocks,
168              ref mut pending,
169              ref mut addresses,
170              ..
171          } = *inner;
172  
173          for _ in 1..=block_num {
174              FakeBitcoinTest::mine_block(addresses, blocks, pending);
175          }
176      }
177  
178      async fn prepare_funding_wallet(&self) {
179          // In fake wallet this might not be technically necessary,
180          // but it makes it behave more like the `RealBitcoinTest`.
181          let block_count = self.inner.write().unwrap().blocks.len() as u64;
182          if block_count < 100 {
183              self.mine_blocks(100 - block_count).await;
184          }
185      }
186  
187      async fn send_and_mine_block(
188          &self,
189          address: &Address,
190          amount: bitcoin::Amount,
191      ) -> (TxOutProof, Transaction) {
192          let mut inner = self.inner.write().unwrap();
193  
194          let transaction = FakeBitcoinTest::new_transaction(vec![TxOut {
195              value: amount.to_sat(),
196              script_pubkey: address.payload.script_pubkey(),
197          }]);
198          inner.addresses.insert(transaction.txid(), amount.into());
199  
200          inner.pending.push(transaction.clone());
201          let merkle_proof = FakeBitcoinTest::pending_merkle_tree(&inner.pending);
202  
203          let FakeBitcoinTestInner {
204              ref mut blocks,
205              ref mut pending,
206              ref mut addresses,
207              ..
208          } = *inner;
209          FakeBitcoinTest::mine_block(addresses, blocks, pending);
210          let block_header = inner.blocks.last().unwrap().header;
211          let proof = TxOutProof {
212              block_header,
213              merkle_proof,
214          };
215          inner.proofs.insert(transaction.txid(), proof.clone());
216          inner
217              .scripts
218              .insert(address.payload.script_pubkey(), vec![transaction.clone()]);
219  
220          (proof, transaction)
221      }
222  
223      async fn get_new_address(&self) -> Address {
224          let ctx = bitcoin::secp256k1::Secp256k1::new();
225          let (_, public_key) = ctx.generate_keypair(&mut OsRng);
226  
227          Address::p2wpkh(&bitcoin::PublicKey::new(public_key), Network::Regtest).unwrap()
228      }
229  
230      async fn mine_block_and_get_received(&self, address: &Address) -> Amount {
231          self.mine_blocks(1).await;
232          let sats = self
233              .inner
234              .read()
235              .unwrap()
236              .blocks
237              .clone()
238              .into_iter()
239              .flat_map(|block| block.txdata.into_iter().flat_map(|tx| tx.output))
240              .find(|out| out.script_pubkey == address.payload.script_pubkey())
241              .map(|tx| tx.value)
242              .unwrap_or(0);
243          Amount::from_sats(sats)
244      }
245  
246      async fn get_mempool_tx_fee(&self, txid: &Txid) -> Amount {
247          loop {
248              let (pending, addresses) = {
249                  let inner = self.inner.read().unwrap();
250                  (inner.pending.clone(), inner.addresses.clone())
251              };
252  
253              let mut fee = Amount::ZERO;
254              let maybe_tx = pending.iter().find(|tx| tx.txid() == *txid);
255  
256              let tx = match maybe_tx {
257                  None => {
258                      sleep_in_test("no transaction found", Duration::from_millis(100)).await;
259                      continue;
260                  }
261                  Some(tx) => tx,
262              };
263  
264              for input in tx.input.iter() {
265                  fee += *addresses
266                      .get(&input.previous_output.txid)
267                      .expect("previous transaction should be known");
268              }
269  
270              for output in tx.output.iter() {
271                  fee -= Amount::from_sats(output.value);
272              }
273  
274              return fee;
275          }
276      }
277  }
278  
279  #[async_trait]
280  impl IBitcoindRpc for FakeBitcoinTest {
281      async fn get_network(&self) -> BitcoinRpcResult<bitcoin::Network> {
282          Ok(bitcoin::Network::Regtest)
283      }
284  
285      async fn get_block_count(&self) -> BitcoinRpcResult<u64> {
286          Ok(self.inner.read().unwrap().blocks.len() as u64)
287      }
288  
289      async fn get_block_hash(&self, height: u64) -> BitcoinRpcResult<bitcoin::BlockHash> {
290          Ok(self.inner.read().unwrap().blocks[height as usize]
291              .header
292              .block_hash())
293      }
294  
295      async fn get_fee_rate(&self, _confirmation_target: u16) -> BitcoinRpcResult<Option<Feerate>> {
296          Ok(Some(Feerate { sats_per_kvb: 2000 }))
297      }
298  
299      async fn submit_transaction(&self, transaction: bitcoin::Transaction) {
300          let mut inner = self.inner.write().unwrap();
301          inner.pending.push(transaction);
302  
303          let mut filtered = BTreeMap::<Vec<OutPoint>, bitcoin::Transaction>::new();
304  
305          // Simulate the mempool keeping txs with higher fees (less output)
306          // TODO: This looks borked, should remove from `filtered` on higher fee or
307          // something, and check per-input anyway. Probably doesn't matter, and I
308          // don't want to touch it.
309          for tx in inner.pending.iter() {
310              match filtered.get(&inputs(tx)) {
311                  Some(found) if output_sum(tx) > output_sum(found) => {}
312                  _ => {
313                      filtered.insert(inputs(tx), tx.clone());
314                  }
315              }
316          }
317  
318          inner.pending = filtered.into_values().collect();
319      }
320  
321      async fn get_tx_block_height(&self, txid: &bitcoin::Txid) -> BitcoinRpcResult<Option<u64>> {
322          for (height, block) in self.inner.read().unwrap().blocks.iter().enumerate() {
323              if block.txdata.iter().any(|tx| tx.txid() == *txid) {
324                  return Ok(Some(height as u64));
325              }
326          }
327          Ok(None)
328      }
329  
330      async fn watch_script_history(&self, _: &ScriptBuf) -> BitcoinRpcResult<()> {
331          Ok(())
332      }
333  
334      async fn get_script_history(
335          &self,
336          script: &ScriptBuf,
337      ) -> BitcoinRpcResult<Vec<bitcoin::Transaction>> {
338          let inner = self.inner.read().unwrap();
339          let script = inner.scripts.get(script).cloned().unwrap_or_default();
340          Ok(script)
341      }
342  
343      async fn get_txout_proof(&self, txid: bitcoin::Txid) -> BitcoinRpcResult<TxOutProof> {
344          let inner = self.inner.read().unwrap();
345          let proof = inner.proofs.get(&txid);
346          Ok(proof.ok_or(format_err!("No proof stored"))?.clone())
347      }
348  }
349  
350  fn output_sum(tx: &Transaction) -> u64 {
351      tx.output.iter().map(|output| output.value).sum()
352  }
353  
354  fn inputs(tx: &Transaction) -> Vec<OutPoint> {
355      tx.input.iter().map(|input| input.previous_output).collect()
356  }