Small refactors in onion_utils

Author: joostjager

lightning/src/ln/onion_utils.rs

@@ -877,10 +877,7 @@ fn crypt_failure_packet(shared_secret: &[u8], packet: &mut OnionErrorPacket) {
 	chacha.process_in_place(&mut packet.data);
 
 	if let Some(ref mut attribution_data) = packet.attribution_data {
-		let ammagext = gen_ammagext_from_shared_secret(&shared_secret);
-		let mut chacha = ChaCha20::new(&ammagext, &[0u8; 8]);
-		chacha.process_in_place(&mut attribution_data.hold_times);
-		chacha.process_in_place(&mut attribution_data.hmacs);
+		attribution_data.crypt(shared_secret);
 	}
 }
 
@@ -942,10 +939,7 @@ fn update_attribution_data(
 	let attribution_data =
 		onion_error_packet.attribution_data.get_or_insert(AttributionData::new());
 
-	let hold_time_bytes: [u8; 4] = hold_time.to_be_bytes();
-	attribution_data.hold_times[..HOLD_TIME_LEN].copy_from_slice(&hold_time_bytes);
-
-	attribution_data.add_hmacs(shared_secret, &onion_error_packet.data);
+	attribution_data.update(&onion_error_packet.data, shared_secret, hold_time);
 }
 
 pub(super) fn build_failure_packet(
@@ -2636,6 +2630,14 @@ impl_writeable!(AttributionData, {
 });
 
 impl AttributionData {
+	/// Encrypts or decrypts the attribution data using the provided shared secret.
+	pub(crate) fn crypt(&mut self, shared_secret: &[u8]) {
+		let ammagext = gen_ammagext_from_shared_secret(&shared_secret);
+		let mut chacha = ChaCha20::new(&ammagext, &[0u8; 8]);
+		chacha.process_in_place(&mut self.hold_times);
+		chacha.process_in_place(&mut self.hmacs);
+	}
+
 	/// Adds the current node's HMACs for all possible positions to this packet.
 	pub(crate) fn add_hmacs(&mut self, shared_secret: &[u8], message: &[u8]) {
 		let um: [u8; 32] = gen_um_from_shared_secret(&shared_secret);
@@ -2685,7 +2687,7 @@ impl AttributionData {
 
 	/// Verifies the attribution data of a failure packet for the given position in the path. If the HMAC checks out, the
 	/// reported hold time is returned. If the HMAC does not match, None is returned.
-	fn verify(&self, message: &Vec<u8>, shared_secret: &[u8], position: usize) -> Option<u32> {
+	fn verify(&self, message: &[u8], shared_secret: &[u8], position: usize) -> Option<u32> {
 		// Calculate the expected HMAC.
 		let um = gen_um_from_shared_secret(shared_secret);
 		let mut hmac = HmacEngine::<Sha256>::new(&um);
@@ -2770,6 +2772,12 @@ impl AttributionData {
 	fn get_hold_time_bytes(&self, idx: usize) -> &[u8] {
 		&self.hold_times[idx * HOLD_TIME_LEN..(idx + 1) * HOLD_TIME_LEN]
 	}
+
+	fn update(&mut self, message: &[u8], shared_secret: &[u8], hold_time: u32) {
+		let hold_time_bytes: [u8; 4] = hold_time.to_be_bytes();
+		self.hold_times[..HOLD_TIME_LEN].copy_from_slice(&hold_time_bytes);
+		self.add_hmacs(shared_secret, message);
+	}
 }
 
 /// Updates the attribution data for an intermediate node.