Merge pull request #3647 from moisesPompilio/2827_issue

Add deterministic random bytes in `override_random_bytes`

Author: TheBlueMatt

lightning/src/ln/functional_test_utils.rs

@@ -3265,6 +3265,10 @@ pub fn fail_payment<'a, 'b, 'c>(origin_node: &Node<'a, 'b, 'c>, expected_path: &
 }
 
 pub fn create_chanmon_cfgs(node_count: usize) -> Vec<TestChanMonCfg> {
+	create_chanmon_cfgs_with_keys(node_count, None)
+}
+
+pub fn create_chanmon_cfgs_with_keys(node_count: usize, predefined_keys_ids: Option<Vec<[u8; 32]>>) -> Vec<TestChanMonCfg> {
 	let mut chan_mon_cfgs = Vec::new();
 	for i in 0..node_count {
 		let tx_broadcaster = test_utils::TestBroadcaster::new(Network::Testnet);
@@ -3276,6 +3280,13 @@ pub fn create_chanmon_cfgs(node_count: usize) -> Vec<TestChanMonCfg> {
 		let keys_manager = test_utils::TestKeysInterface::new(&seed, Network::Testnet);
 		let scorer = RwLock::new(test_utils::TestScorer::new());
 
+		// Set predefined keys_id if provided
+		if let Some(keys_ids) = &predefined_keys_ids {
+			if let Some(keys_id) = keys_ids.get(i) {
+				keys_manager.set_next_keys_id(*keys_id);
+			}
+		}
+
 		chan_mon_cfgs.push(TestChanMonCfg { tx_broadcaster, fee_estimator, chain_source, logger, persister, keys_manager, scorer });
 	}
 

lightning/src/ln/monitor_tests.rs

@@ -2313,7 +2313,11 @@ fn test_claimable_balance_correct_while_payment_pending() {
 
 fn do_test_restored_packages_retry(check_old_monitor_retries_after_upgrade: bool) {
 	// Tests that we'll retry packages that were previously timelocked after we've restored them.
-	let chanmon_cfgs = create_chanmon_cfgs(2);
+	let node0_key_id = <[u8; 32]>::from_hex("0000000000000000000000004D49E5DA0000000000000000000000000000002A").unwrap();
+	let node1_key_id = <[u8; 32]>::from_hex("0000000000000000000000004D49E5DAD000D6201F116BAFD379F1D61DF161B9").unwrap();
+	let predefined_keys_ids = Some(vec![node0_key_id, node1_key_id]);
+
+	let chanmon_cfgs = create_chanmon_cfgs_with_keys(2, predefined_keys_ids);
 	let node_cfgs = create_node_cfgs(2, &chanmon_cfgs);
 	let persister;
 	let new_chain_monitor;
@@ -2323,10 +2327,6 @@ fn do_test_restored_packages_retry(check_old_monitor_retries_after_upgrade: bool
 
 	let mut nodes = create_network(2, &node_cfgs, &node_chanmgrs);
 
-	// Reset our RNG counters to mirror the RNG output from when this test was written.
-	nodes[0].keys_manager.backing.inner.set_counter(0x1_0000_0004);
-	nodes[1].keys_manager.backing.inner.set_counter(0x1_0000_0004);
-
 	// Open a channel, lock in an HTLC, and immediately broadcast the commitment transaction. This
 	// ensures that the HTLC timeout package is held until we reach its expiration height.
 	let (_, _, chan_id, funding_tx) = create_announced_chan_between_nodes_with_value(&nodes, 0, 1, 100_000, 50_000_000);

lightning/src/sign/mod.rs

@@ -2365,13 +2365,6 @@ impl RandomBytes {
 	pub fn new(seed: [u8; 32]) -> Self {
 		Self { seed, index: AtomicCounter::new() }
 	}
-
-	#[cfg(test)]
-	/// Force the counter to a value to produce the same output again. Mostly useful in tests where
-	/// we need to maintain behavior with a previous version which didn't use as much RNG output.
-	pub(crate) fn set_counter(&self, count: u64) {
-		self.index.set_counter(count);
-	}
 }
 
 impl EntropySource for RandomBytes {

lightning/src/util/atomic_counter.rs

@@ -33,16 +33,4 @@ impl AtomicCounter {
 			*mtx - 1
 		}
 	}
-	#[cfg(test)]
-	pub(crate) fn set_counter(&self, count: u64) {
-		#[cfg(target_has_atomic = "64")]
-		{
-			self.counter.store(count, Ordering::Release);
-		}
-		#[cfg(not(target_has_atomic = "64"))]
-		{
-			let mut mtx = self.counter.lock().unwrap();
-			*mtx = count;
-		}
-	}
 }

lightning/src/util/dyn_signer.rs

@@ -244,8 +244,6 @@ delegate!(DynKeysInterface, OutputSpender, inner,
 pub trait DynKeysInterfaceTrait:
 	NodeSigner + OutputSpender + SignerProvider<EcdsaSigner = DynSigner> + EntropySource + Send + Sync
 {
-	#[cfg(test)]
-	fn set_counter(&self, _count: u64) {}
 }
 
 #[cfg(taproot)]
@@ -258,8 +256,6 @@ pub trait DynKeysInterfaceTrait:
 	+ Send
 	+ Sync
 {
-	#[cfg(test)]
-	fn set_counter(&self, _count: u64) {}
 }
 
 /// A dyn wrapper for PhantomKeysManager
@@ -320,9 +316,4 @@ delegate!(DynPhantomKeysInterface, OutputSpender, inner,
 	) -> Result<Transaction, ()>
 );
 
-impl DynKeysInterfaceTrait for DynPhantomKeysInterface {
-	#[cfg(test)]
-	fn set_counter(&self, count: u64) {
-		self.inner.inner.entropy_source.set_counter(count);
-	}
-}
+impl DynKeysInterfaceTrait for DynPhantomKeysInterface {}

lightning/src/util/test_utils.rs

@@ -1496,6 +1496,7 @@ pub struct TestKeysInterface {
 	expectations: Mutex<Option<VecDeque<OnGetShutdownScriptpubkey>>>,
 	pub unavailable_signers_ops: Mutex<HashMap<[u8; 32], HashSet<SignerOp>>>,
 	pub next_signer_disabled_ops: Mutex<HashSet<SignerOp>>,
+	pub override_next_keys_id: Mutex<Option<[u8; 32]>>,
 }
 
 impl EntropySource for TestKeysInterface {
@@ -1546,6 +1547,13 @@ impl SignerProvider for TestKeysInterface {
 	type TaprootSigner = TestChannelSigner;
 
 	fn generate_channel_keys_id(&self, inbound: bool, user_channel_id: u128) -> [u8; 32] {
+		let mut override_keys = self.override_next_keys_id.lock().unwrap();
+
+		if let Some(keys_id) = *override_keys {
+			// Reset after use
+			*override_keys = None;
+			return keys_id;
+		}
 		self.backing.generate_channel_keys_id(inbound, user_channel_id)
 	}
 
@@ -1625,6 +1633,7 @@ impl TestKeysInterface {
 			expectations: Mutex::new(None),
 			unavailable_signers_ops: Mutex::new(new_hash_map()),
 			next_signer_disabled_ops: Mutex::new(new_hash_set()),
+			override_next_keys_id: Mutex::new(None),
 		}
 	}
 
@@ -1652,6 +1661,11 @@ impl TestKeysInterface {
 		let cell = states.get(&keys_id).unwrap();
 		Arc::clone(cell)
 	}
+
+	pub fn set_next_keys_id(&self, keys_id: [u8; 32]) -> &Self {
+		*self.override_next_keys_id.lock().unwrap() = Some(keys_id);
+		self
+	}
 }
 
 impl Drop for TestKeysInterface {