add reorg test

feat: #575

Author: moisesPompilio

tests/common/mod.rs

@@ -29,9 +29,12 @@ use lightning_types::payment::{PaymentHash, PaymentPreimage};
 
 use lightning_persister::fs_store::FilesystemStore;
 
+use bitcoin::hashes::hex::FromHex;
 use bitcoin::hashes::sha256::Hash as Sha256;
 use bitcoin::hashes::Hash;
-use bitcoin::{Address, Amount, Network, OutPoint, Txid};
+use bitcoin::{
+	Address, Amount, Network, OutPoint, ScriptBuf, Sequence, Transaction, Txid, Witness,
+};
 
 use electrsd::corepc_node::Client as BitcoindClient;
 use electrsd::corepc_node::Node as BitcoinD;
@@ -40,7 +43,10 @@ use electrum_client::ElectrumApi;
 
 use rand::distributions::Alphanumeric;
 use rand::{thread_rng, Rng};
+use serde_json::{json, Value};
 
+use std::collections::HashMap;
+use std::collections::HashSet;
 use std::env;
 use std::path::PathBuf;
 use std::sync::{Arc, RwLock};
@@ -395,6 +401,21 @@ pub(crate) fn generate_blocks_and_wait<E: ElectrumApi>(
 	println!("\n");
 }
 
+pub(crate) fn invalidate_blocks(bitcoind: &BitcoindClient, num_blocks: usize) {
+	let blockchain_info = bitcoind.get_blockchain_info().expect("failed to get blockchain info");
+	let cur_height = blockchain_info.blocks as usize;
+	let target_height = cur_height - num_blocks + 1;
+	let block_hash = bitcoind
+		.get_block_hash(target_height as u64)
+		.expect("failed to get block hash")
+		.block_hash()
+		.expect("block hash should be present");
+	bitcoind.invalidate_block(block_hash).expect("failed to invalidate block");
+	let blockchain_info = bitcoind.get_blockchain_info().expect("failed to get blockchain info");
+	let new_cur_height = blockchain_info.blocks as usize;
+	assert!(new_cur_height + num_blocks == cur_height);
+}
+
 pub(crate) fn wait_for_block<E: ElectrumApi>(electrs: &E, min_height: usize) {
 	let mut header = match electrs.block_headers_subscribe() {
 		Ok(header) => header,
@@ -467,25 +488,47 @@ where
 	}
 }
 
-pub(crate) fn premine_and_distribute_funds<E: ElectrumApi>(
-	bitcoind: &BitcoindClient, electrs: &E, addrs: Vec<Address>, amount: Amount,
-) {
+pub(crate) fn premine_blocks<E: ElectrumApi>(bitcoind: &BitcoindClient, electrs: &E) {
 	let _ = bitcoind.create_wallet("ldk_node_test");
 	let _ = bitcoind.load_wallet("ldk_node_test");
 	generate_blocks_and_wait(bitcoind, electrs, 101);
-
-	for addr in addrs {
-		let txid = bitcoind.send_to_address(&addr, amount).unwrap().0.parse().unwrap();
-		wait_for_tx(electrs, txid);
+}
+pub(crate) fn distribute_funds<E: ElectrumApi>(
+	bitcoind: &BitcoindClient, electrs: &E, addrs: Vec<Address>, amount: Amount,
+) -> Txid {
+	let mut amounts = HashMap::<String, f64>::new();
+	for addr in &addrs {
+		amounts.insert(addr.to_string(), amount.to_btc());
 	}
 
+	let empty_account = json!("");
+	let amounts_json = json!(amounts);
+	let txid = bitcoind
+		.call::<Value>("sendmany", &[empty_account, amounts_json])
+		.unwrap()
+		.as_str()
+		.unwrap()
+		.parse()
+		.unwrap();
+
+	wait_for_tx(electrs, txid);
 	generate_blocks_and_wait(bitcoind, electrs, 1);
+
+	txid
+}
+
+pub(crate) fn premine_and_distribute_funds<E: ElectrumApi>(
+	bitcoind: &BitcoindClient, electrs: &E, addrs: Vec<Address>, amount: Amount,
+) {
+	premine_blocks(bitcoind, electrs);
+
+	distribute_funds(bitcoind, electrs, addrs, amount);
 }
 
 pub fn open_channel(
 	node_a: &TestNode, node_b: &TestNode, funding_amount_sat: u64, should_announce: bool,
 	electrsd: &ElectrsD,
-) {
+) -> OutPoint {
 	if should_announce {
 		node_a
 			.open_announced_channel(
@@ -513,6 +556,8 @@ pub fn open_channel(
 	let funding_txo_b = expect_channel_pending_event!(node_b, node_a.node_id());
 	assert_eq!(funding_txo_a, funding_txo_b);
 	wait_for_tx(&electrsd.client, funding_txo_a.txid);
+
+	funding_txo_a
 }
 
 pub(crate) fn do_channel_full_cycle<E: ElectrumApi>(

tests/prop_test.rs

@@ -0,0 +1,218 @@
+mod common;
+use bitcoin::Amount;
+use electrsd::corepc_node::Node as BitcoinD;
+use electrsd::ElectrsD;
+use ldk_node::payment::{PaymentDirection, PaymentKind};
+use ldk_node::{Event, LightningBalance, PendingSweepBalance};
+use proptest::prelude::*;
+use std::collections::HashMap;
+use std::sync::OnceLock;
+
+use crate::common::open_channel;
+use crate::common::setup_bitcoind_and_electrsd;
+use crate::common::{
+	distribute_funds, expect_event, generate_blocks_and_wait, invalidate_blocks, premine_blocks,
+	random_config, setup_node, wait_for_outpoint_spend, TestChainSource,
+};
+
+static SETUP_REORG: OnceLock<(BitcoinD, ElectrsD)> = OnceLock::new();
+
+fn init_setup_test_reorg() -> &'static (BitcoinD, ElectrsD) {
+	SETUP_REORG.get_or_init(|| {
+		let (bitcoind, electrsd) = setup_bitcoind_and_electrsd();
+		let (bitcoind_client, electrsd_client) = (&bitcoind.client, &electrsd.client);
+		premine_blocks(bitcoind_client, electrsd_client);
+		(bitcoind, electrsd)
+	})
+}
+
+proptest! {
+	#![proptest_config(ProptestConfig::with_cases(5))]
+	#[test]
+	fn test_reorg(
+		reorg_depth in 1..=6usize,
+		force_close in prop::bool::ANY,
+	) {
+
+	let (bitcoind, electrsd) = init_setup_test_reorg();
+	let chain_source_bitcoind = TestChainSource::BitcoindRpcSync(&bitcoind);
+	let chain_source_electrsd = TestChainSource::Electrum(&electrsd);
+	let chain_source_esplora = TestChainSource::Esplora(&electrsd);
+
+	macro_rules! config_node {
+		($chain_source: expr, $anchor_channels: expr) => {{
+			let config_a = random_config($anchor_channels);
+			let node = setup_node(&$chain_source, config_a, None);
+			node
+		}};
+	}
+	let anchor_channels = true;
+	let nodes = vec![
+		config_node!(chain_source_electrsd, anchor_channels),
+		config_node!(chain_source_bitcoind, anchor_channels),
+		config_node!(chain_source_esplora, anchor_channels),
+	];
+
+	let (bitcoind, electrs) = (&bitcoind.client, &electrsd.client);
+	macro_rules! reorg {
+		($reorg_depth: expr) => {{
+			invalidate_blocks(bitcoind, $reorg_depth);
+			generate_blocks_and_wait(bitcoind, electrs, $reorg_depth);
+		}};
+	}
+
+	let amount = 2_100_000;
+	let addr_nodes =
+		nodes.iter().map(|node| node.onchain_payment().new_address().unwrap()).collect::<Vec<_>>();
+	distribute_funds(bitcoind, electrs, addr_nodes, Amount::from_sat(amount));
+
+	macro_rules! sync_wallets {
+		() => {
+			nodes.iter().for_each(|node| node.sync_wallets().unwrap())
+		};
+	}
+	sync_wallets!();
+	nodes.iter().for_each(|node| {
+		assert_eq!(node.list_balances().spendable_onchain_balance_sats, amount);
+		assert_eq!(node.list_balances().total_onchain_balance_sats, amount);
+	});
+
+
+	let mut nodes_funding_tx = HashMap::new();
+	let funding_amount_sat = 2_000_000;
+	for index in 0..nodes.len() {
+		let node = &nodes[index];
+		let next_node = if index >= nodes.len() - 1 { &nodes[0] } else { &nodes[index + 1] };
+
+		let funding_txo_b = open_channel(node, next_node, funding_amount_sat, true, electrsd);
+		nodes_funding_tx.insert(node.node_id(), funding_txo_b);
+	}
+
+	generate_blocks_and_wait(bitcoind, electrs, 6);
+	sync_wallets!();
+	reorg!(reorg_depth);
+	sync_wallets!();
+
+	macro_rules! collect_channel_ready_events {
+		($node:expr, $expected:expr) => {{
+			let mut user_channels = HashMap::new();
+			for _ in 0..$expected {
+				match $node.wait_next_event() {
+					Event::ChannelReady { user_channel_id, counterparty_node_id, .. } => {
+						$node.event_handled().unwrap();
+						user_channels.insert(counterparty_node_id, user_channel_id);
+					},
+					other => panic!("Unexpected event: {:?}", other),
+				}
+			}
+			user_channels
+		}};
+	}
+
+	let mut node_channels_id = HashMap::new();
+	for index in 0..nodes.len() {
+		let node = &nodes[index];
+
+		assert_eq!(
+			node
+				.list_payments_with_filter(|p| p.direction == PaymentDirection::Outbound
+					&& matches!(p.kind, PaymentKind::Onchain { .. }))
+				.len(),
+			1
+		);
+
+		let user_channels = collect_channel_ready_events!(node, 2);
+		let next_node = if index >= nodes.len() - 1 { &nodes[0] } else { &nodes[index + 1] };
+		let prev_node = if index == 0 { &nodes[nodes.len() - 1] } else { &nodes[index - 1] };
+
+		assert!(user_channels.get(&Some(next_node.node_id())) != None);
+		assert!(user_channels.get(&Some(prev_node.node_id())) != None);
+
+		let user_channel_id =
+			user_channels.get(&Some(next_node.node_id())).expect("Missing user channel for node");
+		node_channels_id.insert(node.node_id(), *user_channel_id);
+	}
+
+
+	for index in 0..nodes.len() {
+		let node_a = &nodes[index];
+		let node_b = if index >= nodes.len() - 1 { &nodes[0] } else { &nodes[index + 1] };
+
+		let user_channel_id = node_channels_id.get(&node_a.node_id()).expect("user channel id not exist");
+		let funding = nodes_funding_tx.get(&node_a.node_id()).expect("Funding tx not exist");
+
+		if force_close {
+			node_a.force_close_channel(&user_channel_id, node_b.node_id(), None).unwrap();
+		} else {
+			node_a.close_channel(&user_channel_id, node_b.node_id()).unwrap();
+		}
+
+		expect_event!(node_a, ChannelClosed);
+		expect_event!(node_b, ChannelClosed);
+
+		wait_for_outpoint_spend(electrs, *funding);
+	}
+
+	reorg!(reorg_depth);
+	sync_wallets!();
+
+	generate_blocks_and_wait(bitcoind, electrs, 1);
+	sync_wallets!();
+
+	if force_close {
+		nodes.iter().for_each(|node| {
+			node.sync_wallets().unwrap();
+			// If there is no more balance, there is nothing to process here.
+			if node.list_balances().lightning_balances.len() < 1 {
+				return;
+			}
+			match node.list_balances().lightning_balances[0] {
+			LightningBalance::ClaimableAwaitingConfirmations {
+				confirmation_height,
+				..
+			} => {
+				let cur_height = node.status().current_best_block.height;
+				let blocks_to_go = confirmation_height - cur_height;
+				generate_blocks_and_wait(bitcoind, electrs, blocks_to_go as usize);
+				node.sync_wallets().unwrap();
+			},
+			_ => panic!("Unexpected balance state for node_hub!"),
+			}
+
+			assert!(node.list_balances().lightning_balances.len() < 2);
+
+			assert!(node.list_balances().pending_balances_from_channel_closures.len() > 0);
+			match node.list_balances().pending_balances_from_channel_closures[0] {
+				PendingSweepBalance::BroadcastAwaitingConfirmation { .. } => {},
+				_ => println!("Unexpected balance state!"),
+			}
+
+			generate_blocks_and_wait(&bitcoind, electrs, 1);
+			node.sync_wallets().unwrap();
+			assert!(node.list_balances().lightning_balances.len() < 2);
+			assert!(node.list_balances().pending_balances_from_channel_closures.len() > 0);
+			match node.list_balances().pending_balances_from_channel_closures[0] {
+				PendingSweepBalance::AwaitingThresholdConfirmations { .. } => {},
+				_ => println!("Unexpected balance state!"),
+			}
+		});
+	}
+
+	generate_blocks_and_wait(bitcoind, electrs, 6);
+	sync_wallets!();
+
+	reorg!(reorg_depth);
+	sync_wallets!();
+
+	// Check balance after close channel
+	nodes.iter().for_each(|node| {
+			assert!(node.list_balances().spendable_onchain_balance_sats > amount - 7000);
+			assert!(node.list_balances().spendable_onchain_balance_sats < amount);
+
+			assert_eq!(node.list_balances().total_anchor_channels_reserve_sats, 0);
+			assert!(node.list_balances().lightning_balances.is_empty());
+
+			assert_eq!(node.next_event(), None);
+		});
+	}
+}