Merge pull request lightningdevkit#3302 from TheBlueMatt/2024-09-atomic-cleanups

Simplify and fix AtomicCounter

Author: tnull

lightning/src/ln/functional_tests.rs

@@ -7670,8 +7670,8 @@ fn test_bump_penalty_txn_on_revoked_htlcs() {
 		assert_ne!(node_txn[0].input[0].previous_output, node_txn[2].input[0].previous_output);
 		assert_ne!(node_txn[1].input[0].previous_output, node_txn[2].input[0].previous_output);
 
-		assert_eq!(node_txn[0].input[0].previous_output, revoked_htlc_txn[1].input[0].previous_output);
-		assert_eq!(node_txn[1].input[0].previous_output, revoked_htlc_txn[0].input[0].previous_output);
+		assert_eq!(node_txn[1].input[0].previous_output, revoked_htlc_txn[1].input[0].previous_output);
+		assert_eq!(node_txn[0].input[0].previous_output, revoked_htlc_txn[0].input[0].previous_output);
 
 		// node_txn[3] spends the revoked outputs from the revoked_htlc_txn (which only have one
 		// output, checked above).

lightning/src/ln/interactivetxs.rs

@@ -1328,7 +1328,7 @@ mod tests {
 	impl EntropySource for TestEntropySource {
 		fn get_secure_random_bytes(&self) -> [u8; 32] {
 			let mut res = [0u8; 32];
-			let increment = self.0.get_increment();
+			let increment = self.0.next();
 			for (i, byte) in res.iter_mut().enumerate() {
 				// Rotate the increment value by 'i' bits to the right, to avoid clashes
 				// when `generate_local_serial_id` does a parity flip on consecutive calls for the

lightning/src/ln/max_payment_path_len_tests.rs

@@ -255,6 +255,13 @@ fn blinded_path_with_custom_tlv() {
 	create_announced_chan_between_nodes(&nodes, 1, 2);
 	let chan_upd_2_3 = create_announced_chan_between_nodes_with_value(&nodes, 2, 3, 1_000_000, 0).0.contents;
 
+	// Ensure all nodes are at the same height
+	let node_max_height = nodes.iter().map(|node| node.blocks.lock().unwrap().len()).max().unwrap() as u32;
+	connect_blocks(&nodes[0], node_max_height - nodes[0].best_block_info().1);
+	connect_blocks(&nodes[1], node_max_height - nodes[1].best_block_info().1);
+	connect_blocks(&nodes[2], node_max_height - nodes[2].best_block_info().1);
+	connect_blocks(&nodes[3], node_max_height - nodes[3].best_block_info().1);
+
 	// Construct the route parameters for sending to nodes[3]'s blinded path.
 	let amt_msat = 100_000;
 	let (payment_preimage, payment_hash, payment_secret) = get_payment_preimage_hash(&nodes[3], Some(amt_msat), None);

lightning/src/ln/monitor_tests.rs

@@ -2251,6 +2251,10 @@ 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.entropy_source.set_counter(0x1_0000_0004);
+	nodes[1].keys_manager.backing.inner.entropy_source.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/ln/peer_handler.rs

@@ -1055,7 +1055,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, OM: Deref, L: Deref, CM
 
 	fn get_ephemeral_key(&self) -> SecretKey {
 		let mut ephemeral_hash = self.ephemeral_key_midstate.clone();
-		let counter = self.peer_counter.get_increment();
+		let counter = self.peer_counter.next();
 		ephemeral_hash.input(&counter.to_le_bytes());
 		SecretKey::from_slice(&Sha256::from_engine(ephemeral_hash).to_byte_array()).expect("You broke SHA-256!")
 	}

lightning/src/sign/mod.rs

@@ -1027,7 +1027,6 @@ pub trait ChangeDestinationSource {
 ///
 /// This implementation performs no policy checks and is insufficient by itself as
 /// a secure external signer.
-#[derive(Debug)]
 pub struct InMemorySigner {
 	/// Holder secret key in the 2-of-2 multisig script of a channel. This key also backs the
 	/// holder's anchor output in a commitment transaction, if one is present.
@@ -1854,6 +1853,9 @@ pub struct KeysManager {
 	channel_master_key: Xpriv,
 	channel_child_index: AtomicUsize,
 
+	#[cfg(test)]
+	pub(crate) entropy_source: RandomBytes,
+	#[cfg(not(test))]
 	entropy_source: RandomBytes,
 
 	seed: [u8; 32],
@@ -2310,6 +2312,9 @@ impl SignerProvider for KeysManager {
 /// Switching between this struct and [`KeysManager`] will invalidate any previously issued
 /// invoices and attempts to pay previous invoices will fail.
 pub struct PhantomKeysManager {
+	#[cfg(test)]
+	pub(crate) inner: KeysManager,
+	#[cfg(not(test))]
 	inner: KeysManager,
 	inbound_payment_key: KeyMaterial,
 	phantom_secret: SecretKey,
@@ -2475,7 +2480,6 @@ impl PhantomKeysManager {
 }
 
 /// An implementation of [`EntropySource`] using ChaCha20.
-#[derive(Debug)]
 pub struct RandomBytes {
 	/// Seed from which all randomness produced is derived from.
 	seed: [u8; 32],
@@ -2489,11 +2493,18 @@ 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 {
 	fn get_secure_random_bytes(&self) -> [u8; 32] {
-		let index = self.index.get_increment();
+		let index = self.index.next();
 		let mut nonce = [0u8; 16];
 		nonce[..8].copy_from_slice(&index.to_be_bytes());
 		ChaCha20::get_single_block(&self.seed, &nonce)

lightning/src/util/atomic_counter.rs

@@ -1,32 +1,44 @@
-//! A simple atomic counter that uses AtomicUsize to give a u64 counter.
+//! A simple atomic counter that uses mutexes if the platform doesn't support atomic u64s.
 
-#[cfg(not(any(target_pointer_width = "32", target_pointer_width = "64")))]
-compile_error!("We need at least 32-bit pointers for atomic counter (and to have enough memory to run LDK)");
+#[cfg(target_has_atomic = "64")]
+use core::sync::atomic::{AtomicU64, Ordering};
+#[cfg(not(target_has_atomic = "64"))]
+use crate::sync::Mutex;
 
-use core::sync::atomic::{AtomicUsize, Ordering};
-
-#[derive(Debug)]
 pub(crate) struct AtomicCounter {
-	// Usize needs to be at least 32 bits to avoid overflowing both low and high. If usize is 64
-	// bits we will never realistically count into high:
-	counter_low: AtomicUsize,
-	counter_high: AtomicUsize,
+	#[cfg(target_has_atomic = "64")]
+	counter: AtomicU64,
+	#[cfg(not(target_has_atomic = "64"))]
+	counter: Mutex<u64>,
 }
 
 impl AtomicCounter {
 	pub(crate) fn new() -> Self {
 		Self {
-			counter_low: AtomicUsize::new(0),
-			counter_high: AtomicUsize::new(0),
+			#[cfg(target_has_atomic = "64")]
+			counter: AtomicU64::new(0),
+			#[cfg(not(target_has_atomic = "64"))]
+			counter: Mutex::new(0),
+		}
+	}
+	pub(crate) fn next(&self) -> u64 {
+		#[cfg(target_has_atomic = "64")] {
+			self.counter.fetch_add(1, Ordering::AcqRel)
+		}
+		#[cfg(not(target_has_atomic = "64"))] {
+			let mut mtx = self.counter.lock().unwrap();
+			*mtx += 1;
+			*mtx - 1
 		}
 	}
-	pub(crate) fn get_increment(&self) -> u64 {
-		let low = self.counter_low.fetch_add(1, Ordering::AcqRel) as u64;
-		let high = if low == 0 {
-			self.counter_high.fetch_add(1, Ordering::AcqRel) as u64
-		} else {
-			self.counter_high.load(Ordering::Acquire) as u64
-		};
-		(high << 32) | low
+	#[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;
+		}
 	}
 }