Merge pull request #396 from tnull/2024-11-onchain-recovery-test

Author: tnull

tests/common/mod.rs

@@ -264,7 +264,7 @@ pub(crate) fn setup_two_nodes(
 ) -> (TestNode, TestNode) {
 	println!("== Node A ==");
 	let config_a = random_config(anchor_channels);
-	let node_a = setup_node(chain_source, config_a);
+	let node_a = setup_node(chain_source, config_a, None);
 
 	println!("\n== Node B ==");
 	let mut config_b = random_config(anchor_channels);
@@ -279,11 +279,13 @@ pub(crate) fn setup_two_nodes(
 			.trusted_peers_no_reserve
 			.push(node_a.node_id());
 	}
-	let node_b = setup_node(chain_source, config_b);
+	let node_b = setup_node(chain_source, config_b, None);
 	(node_a, node_b)
 }
 
-pub(crate) fn setup_node(chain_source: &TestChainSource, config: Config) -> TestNode {
+pub(crate) fn setup_node(
+	chain_source: &TestChainSource, config: Config, seed_bytes: Option<Vec<u8>>,
+) -> TestNode {
 	setup_builder!(builder, config);
 	match chain_source {
 		TestChainSource::Esplora(electrsd) => {
@@ -302,6 +304,11 @@ pub(crate) fn setup_node(chain_source: &TestChainSource, config: Config) -> Test
 			builder.set_chain_source_bitcoind_rpc(rpc_host, rpc_port, rpc_user, rpc_password);
 		},
 	}
+
+	if let Some(seed) = seed_bytes {
+		builder.set_entropy_seed_bytes(seed).unwrap();
+	}
+
 	let test_sync_store = Arc::new(TestSyncStore::new(config.storage_dir_path.into()));
 	let node = builder.build_with_store(test_sync_store).unwrap();
 	node.start().unwrap();

tests/integration_tests_rust.rs

@@ -21,6 +21,8 @@ use ldk_node::{Builder, Event, NodeError};
 use lightning::ln::channelmanager::PaymentId;
 use lightning::util::persist::KVStore;
 
+use bitcoincore_rpc::RpcApi;
+
 use bitcoin::Amount;
 
 use std::sync::Arc;
@@ -315,12 +317,104 @@ fn onchain_spend_receive() {
 	assert!(node_b.list_balances().spendable_onchain_balance_sats < 100000);
 }
 
+#[test]
+fn onchain_wallet_recovery() {
+	let (bitcoind, electrsd) = setup_bitcoind_and_electrsd();
+
+	let chain_source = TestChainSource::Esplora(&electrsd);
+
+	let seed_bytes = vec![42u8; 64];
+
+	let original_config = random_config(true);
+	let original_node = setup_node(&chain_source, original_config, Some(seed_bytes.clone()));
+
+	let premine_amount_sat = 100_000;
+
+	let addr_1 = original_node.onchain_payment().new_address().unwrap();
+
+	premine_and_distribute_funds(
+		&bitcoind.client,
+		&electrsd.client,
+		vec![addr_1],
+		Amount::from_sat(premine_amount_sat),
+	);
+	original_node.sync_wallets().unwrap();
+	assert_eq!(original_node.list_balances().spendable_onchain_balance_sats, premine_amount_sat);
+
+	let addr_2 = original_node.onchain_payment().new_address().unwrap();
+
+	let txid = bitcoind
+		.client
+		.send_to_address(
+			&addr_2,
+			Amount::from_sat(premine_amount_sat),
+			None,
+			None,
+			None,
+			None,
+			None,
+			None,
+		)
+		.unwrap();
+	wait_for_tx(&electrsd.client, txid);
+
+	generate_blocks_and_wait(&bitcoind.client, &electrsd.client, 1);
+
+	original_node.sync_wallets().unwrap();
+	assert_eq!(
+		original_node.list_balances().spendable_onchain_balance_sats,
+		premine_amount_sat * 2
+	);
+
+	original_node.stop().unwrap();
+	drop(original_node);
+
+	// Now we start from scratch, only the seed remains the same.
+	let recovered_config = random_config(true);
+	let recovered_node = setup_node(&chain_source, recovered_config, Some(seed_bytes));
+
+	recovered_node.sync_wallets().unwrap();
+	assert_eq!(
+		recovered_node.list_balances().spendable_onchain_balance_sats,
+		premine_amount_sat * 2
+	);
+
+	// Check we sync even when skipping some addresses.
+	let _addr_3 = recovered_node.onchain_payment().new_address().unwrap();
+	let _addr_4 = recovered_node.onchain_payment().new_address().unwrap();
+	let _addr_5 = recovered_node.onchain_payment().new_address().unwrap();
+	let addr_6 = recovered_node.onchain_payment().new_address().unwrap();
+
+	let txid = bitcoind
+		.client
+		.send_to_address(
+			&addr_6,
+			Amount::from_sat(premine_amount_sat),
+			None,
+			None,
+			None,
+			None,
+			None,
+			None,
+		)
+		.unwrap();
+	wait_for_tx(&electrsd.client, txid);
+
+	generate_blocks_and_wait(&bitcoind.client, &electrsd.client, 1);
+
+	recovered_node.sync_wallets().unwrap();
+	assert_eq!(
+		recovered_node.list_balances().spendable_onchain_balance_sats,
+		premine_amount_sat * 3
+	);
+}
+
 #[test]
 fn sign_verify_msg() {
 	let (_bitcoind, electrsd) = setup_bitcoind_and_electrsd();
 	let config = random_config(true);
 	let chain_source = TestChainSource::Esplora(&electrsd);
-	let node = setup_node(&chain_source, config);
+	let node = setup_node(&chain_source, config, None);
 
 	// Tests arbitrary message signing and later verification
 	let msg = "OK computer".as_bytes();