From f5abc93978ed2279f315ca36e8e639e9245bcbee Mon Sep 17 00:00:00 2001
From: Jeffrey Czyz <jkczyz@gmail.com>
Date: Tue, 14 Sep 2021 11:44:19 -0500
Subject: [PATCH] Pass hop index in construct_onion_keys_callback

This simplifies failing route hop calculation, which will be useful for
later identifying the failing hop for PaymentFailed events.
---
 lightning/src/ln/onion_utils.rs | 25 +++++++++++++------------
 1 file changed, 13 insertions(+), 12 deletions(-)

diff --git a/lightning/src/ln/onion_utils.rs b/lightning/src/ln/onion_utils.rs
index f1b617fd..4509ee77 100644
--- a/lightning/src/ln/onion_utils.rs
+++ b/lightning/src/ln/onion_utils.rs
@@ -75,11 +75,11 @@ pub(super) fn gen_ammag_from_shared_secret(shared_secret: &[u8]) -> [u8; 32] {
 
 // can only fail if an intermediary hop has an invalid public key or session_priv is invalid
 #[inline]
-pub(super) fn construct_onion_keys_callback<T: secp256k1::Signing, FType: FnMut(SharedSecret, [u8; 32], PublicKey, &RouteHop)> (secp_ctx: &Secp256k1<T>, path: &Vec<RouteHop>, session_priv: &SecretKey, mut callback: FType) -> Result<(), secp256k1::Error> {
+pub(super) fn construct_onion_keys_callback<T: secp256k1::Signing, FType: FnMut(SharedSecret, [u8; 32], PublicKey, &RouteHop, usize)> (secp_ctx: &Secp256k1<T>, path: &Vec<RouteHop>, session_priv: &SecretKey, mut callback: FType) -> Result<(), secp256k1::Error> {
 	let mut blinded_priv = session_priv.clone();
 	let mut blinded_pub = PublicKey::from_secret_key(secp_ctx, &blinded_priv);
 
-	for hop in path.iter() {
+	for (idx, hop) in path.iter().enumerate() {
 		let shared_secret = SharedSecret::new(&hop.pubkey, &blinded_priv);
 
 		let mut sha = Sha256::engine();
@@ -92,7 +92,7 @@ pub(super) fn construct_onion_keys_callback<T: secp256k1::Signing, FType: FnMut(
 		blinded_priv.mul_assign(&blinding_factor)?;
 		blinded_pub = PublicKey::from_secret_key(secp_ctx, &blinded_priv);
 
-		callback(shared_secret, blinding_factor, ephemeral_pubkey, hop);
+		callback(shared_secret, blinding_factor, ephemeral_pubkey, hop, idx);
 	}
 
 	Ok(())
@@ -102,7 +102,7 @@ pub(super) fn construct_onion_keys_callback<T: secp256k1::Signing, FType: FnMut(
 pub(super) fn construct_onion_keys<T: secp256k1::Signing>(secp_ctx: &Secp256k1<T>, path: &Vec<RouteHop>, session_priv: &SecretKey) -> Result<Vec<OnionKeys>, secp256k1::Error> {
 	let mut res = Vec::with_capacity(path.len());
 
-	construct_onion_keys_callback(secp_ctx, path, session_priv, |shared_secret, _blinding_factor, ephemeral_pubkey, _| {
+	construct_onion_keys_callback(secp_ctx, path, session_priv, |shared_secret, _blinding_factor, ephemeral_pubkey, _, _| {
 		let (rho, mu) = gen_rho_mu_from_shared_secret(&shared_secret[..]);
 
 		res.push(OnionKeys {
@@ -337,12 +337,10 @@ pub(super) fn process_onion_failure<T: secp256k1::Signing, L: Deref>(secp_ctx: &
 		let mut htlc_msat = *first_hop_htlc_msat;
 		let mut error_code_ret = None;
 		let mut error_packet_ret = None;
-		let mut next_route_hop_ix = 0;
 		let mut is_from_final_node = false;
 
 		// Handle packed channel/node updates for passing back for the route handler
-		construct_onion_keys_callback(secp_ctx, path, session_priv, |shared_secret, _, _, route_hop| {
-			next_route_hop_ix += 1;
+		construct_onion_keys_callback(secp_ctx, path, session_priv, |shared_secret, _, _, route_hop, route_hop_idx| {
 			if res.is_some() { return; }
 
 			let amt_to_forward = htlc_msat - route_hop.fee_msat;
@@ -356,7 +354,7 @@ pub(super) fn process_onion_failure<T: secp256k1::Signing, L: Deref>(secp_ctx: &
 			chacha.process(&packet_decrypted, &mut decryption_tmp[..]);
 			packet_decrypted = decryption_tmp;
 
-			is_from_final_node = path.last().unwrap().pubkey == route_hop.pubkey;
+			is_from_final_node = route_hop_idx + 1 == path.len();
 
 			if let Ok(err_packet) = msgs::DecodedOnionErrorPacket::read(&mut Cursor::new(&packet_decrypted)) {
 				let um = gen_um_from_shared_secret(&shared_secret[..]);
@@ -388,10 +386,13 @@ pub(super) fn process_onion_failure<T: secp256k1::Signing, L: Deref>(secp_ctx: &
 							network_update = Some(NetworkUpdate::NodeFailure { node_id: route_hop.pubkey, is_permanent: error_code & PERM == PERM });
 						}
 						else if error_code & PERM == PERM {
-							network_update = if payment_failed { None } else { Some(NetworkUpdate::ChannelClosed {
-								short_channel_id: path[next_route_hop_ix - if next_route_hop_ix == path.len() { 1 } else { 0 }].short_channel_id,
-								is_permanent: true,
-							})};
+							network_update = if payment_failed { None } else {
+								let failing_route_hop = if is_from_final_node { route_hop } else { &path[route_hop_idx + 1] };
+								Some(NetworkUpdate::ChannelClosed {
+									short_channel_id: failing_route_hop.short_channel_id,
+									is_permanent: true,
+								})
+							};
 						}
 						else if error_code & UPDATE == UPDATE {
 							if let Some(update_len_slice) = err_packet.failuremsg.get(debug_field_size+2..debug_field_size+4) {
-- 
2.30.2