Allow intercepting HTLCs based on the source channel

It may be useful in some situations to select HTLCs for
interception based on the source channel in addition to the sink.
Here we add the ability to do so by adding new flags to
`HTLCInterceptionFlags`.

Author: TheBlueMatt

lightning/src/ln/channelmanager.rs

@@ -4729,7 +4729,9 @@ impl<
 		}
 	}
 
-	fn forward_needs_intercept_to_known_chan(&self, outbound_chan: &FundedChannel<SP>) -> bool {
+	fn forward_needs_intercept_to_known_chan(
+		&self, prev_chan_public: bool, outbound_chan: &FundedChannel<SP>,
+	) -> bool {
 		let intercept_flags = self.config.read().unwrap().htlc_interception_flags;
 		if !outbound_chan.context.should_announce() {
 			if outbound_chan.context.is_connected() {
@@ -4746,6 +4748,23 @@ impl<
 				return true;
 			}
 		}
+		if prev_chan_public {
+			if outbound_chan.context.should_announce() {
+				if intercept_flags & (HTLCInterceptionFlags::FromPublicToPublicChannels as u8) != 0
+				{
+					return true;
+				}
+			} else {
+				if intercept_flags & (HTLCInterceptionFlags::FromPublicToPrivateChannels as u8) != 0
+				{
+					return true;
+				}
+			}
+		} else {
+			if intercept_flags & (HTLCInterceptionFlags::FromPrivateChannels as u8) != 0 {
+				return true;
+			}
+		}
 		false
 	}
 
@@ -4839,7 +4858,7 @@ impl<
 	}
 
 	fn can_forward_htlc_should_intercept(
-		&self, msg: &msgs::UpdateAddHTLC, next_hop: &NextPacketDetails,
+		&self, msg: &msgs::UpdateAddHTLC, prev_chan_public: bool, next_hop: &NextPacketDetails,
 	) -> Result<bool, LocalHTLCFailureReason> {
 		let outgoing_scid = match next_hop.outgoing_connector {
 			HopConnector::ShortChannelId(scid) => scid,
@@ -4858,7 +4877,7 @@ impl<
 		// times we do it.
 		let intercept =
 			match self.do_funded_channel_callback(outgoing_scid, |chan: &mut FundedChannel<SP>| {
-				let intercept = self.forward_needs_intercept_to_known_chan(chan);
+				let intercept = self.forward_needs_intercept_to_known_chan(prev_chan_public, chan);
 				self.can_forward_htlc_to_outgoing_channel(chan, msg, next_hop, intercept)?;
 				Ok(intercept)
 			}) {
@@ -6869,34 +6888,29 @@ impl<
 		'outer_loop: for (incoming_scid_alias, update_add_htlcs) in decode_update_add_htlcs {
 			// If any decoded update_add_htlcs were processed, we need to persist.
 			should_persist = true;
-			let incoming_channel_details_opt = self.do_funded_channel_callback(
-				incoming_scid_alias,
-				|chan: &mut FundedChannel<SP>| {
-					let counterparty_node_id = chan.context.get_counterparty_node_id();
-					let channel_id = chan.context.channel_id();
-					let funding_txo = chan.funding.get_funding_txo().unwrap();
-					let user_channel_id = chan.context.get_user_id();
-					let accept_underpaying_htlcs = chan.context.config().accept_underpaying_htlcs;
-					(
-						counterparty_node_id,
-						channel_id,
-						funding_txo,
-						user_channel_id,
-						accept_underpaying_htlcs,
-					)
-				},
-			);
 			let (
 				incoming_counterparty_node_id,
 				incoming_channel_id,
 				incoming_funding_txo,
 				incoming_user_channel_id,
 				incoming_accept_underpaying_htlcs,
-			) = if let Some(incoming_channel_details) = incoming_channel_details_opt {
-				incoming_channel_details
-			} else {
+				incoming_chan_is_public,
+			) = match self.do_funded_channel_callback(
+				incoming_scid_alias,
+				|chan: &mut FundedChannel<SP>| {
+					(
+						chan.context.get_counterparty_node_id(),
+						chan.context.channel_id(),
+						chan.funding.get_funding_txo().unwrap(),
+						chan.context.get_user_id(),
+						chan.context.config().accept_underpaying_htlcs,
+						chan.context.should_announce(),
+					)
+				},
+			) {
+				Some(incoming_channel_details) => incoming_channel_details,
 				// The incoming channel no longer exists, HTLCs should be resolved onchain instead.
-				continue;
+				None => continue,
 			};
 
 			let mut htlc_forwards = Vec::new();
@@ -7016,9 +7030,11 @@ impl<
 				// Now process the HTLC on the outgoing channel if it's a forward.
 				let mut intercept_forward = false;
 				if let Some(next_packet_details) = next_packet_details_opt.as_ref() {
-					match self
-						.can_forward_htlc_should_intercept(&update_add_htlc, next_packet_details)
-					{
+					match self.can_forward_htlc_should_intercept(
+						&update_add_htlc,
+						incoming_chan_is_public,
+						next_packet_details,
+					) {
 						Err(reason) => {
 							fail_htlc_continue_to_next!(reason);
 						},
@@ -16492,9 +16508,29 @@ impl<
 				);
 				log_trace!(logger, "Releasing held htlc with intercept_id {}", intercept_id);
 
+				let prev_chan_public = {
+					let per_peer_state = self.per_peer_state.read().unwrap();
+					let peer_state = per_peer_state
+						.get(&htlc.prev_counterparty_node_id)
+						.map(|mtx| mtx.lock().unwrap());
+					let chan_state = peer_state
+						.as_ref()
+						.map(|state| state.channel_by_id.get(&htlc.prev_channel_id))
+						.flatten();
+					if let Some(chan_state) = chan_state {
+						chan_state.context().should_announce()
+					} else {
+						// If the inbound channel has closed since the HTLC was held, we really
+						// shouldn't forward it - forwarding it now would result in, at best,
+						// having to claim the HTLC on chain. Instead, drop the HTLC and let the
+						// counterparty claim their money on chain.
+						return;
+					}
+				};
+
 				let should_intercept = self
 					.do_funded_channel_callback(next_hop_scid, |chan| {
-						self.forward_needs_intercept_to_known_chan(chan)
+						self.forward_needs_intercept_to_known_chan(prev_chan_public, chan)
 					})
 					.unwrap_or_else(|| self.forward_needs_intercept_to_unknown_chan(next_hop_scid));
 

lightning/src/ln/interception_tests.rs

@@ -51,37 +51,48 @@ fn do_test_htlc_interception_flags(
 	let node_chanmgrs = create_node_chanmgrs(3, &node_cfgs, &[None, Some(intercept_config), None]);
 	let nodes = create_network(3, &node_cfgs, &node_chanmgrs);
 
-	create_announced_chan_between_nodes(&nodes, 0, 1);
+	let inbound_private = match flag {
+		Flag::FromPrivateChannels => {
+			create_unannounced_chan_between_nodes_with_value(&nodes, 0, 1, 100000, 0);
+			true
+		},
+		_ => {
+			create_announced_chan_between_nodes(&nodes, 0, 1);
+			false
+		},
+	};
 
 	let node_0_id = nodes[0].node.get_our_node_id();
 	let node_1_id = nodes[1].node.get_our_node_id();
 	let node_2_id = nodes[2].node.get_our_node_id();
 
 	// First open the right type of channel (and get it in the right state) for the bit we're
 	// testing.
-	let (target_scid, target_chan_id) = match flag {
-		Flag::ToOfflinePrivateChannels | Flag::ToOnlinePrivateChannels => {
+	let (target_scid, target_chan_id, outbound_private_for_known_scids) = match flag {
+		Flag::ToOfflinePrivateChannels
+		| Flag::ToOnlinePrivateChannels
+		| Flag::FromPublicToPrivateChannels => {
 			create_unannounced_chan_between_nodes_with_value(&nodes, 1, 2, 100000, 0);
 			let chan_id = nodes[2].node.list_channels()[0].channel_id;
 			let scid = nodes[2].node.list_channels()[0].short_channel_id.unwrap();
 			if flag == Flag::ToOfflinePrivateChannels {
 				nodes[1].node.peer_disconnected(node_2_id);
 				nodes[2].node.peer_disconnected(node_1_id);
-			} else {
-				assert_eq!(flag, Flag::ToOnlinePrivateChannels);
 			}
-			(scid, chan_id)
+			(scid, chan_id, Some(true))
 		},
-		Flag::ToInterceptSCIDs | Flag::ToPublicChannels | Flag::ToUnknownSCIDs => {
+		Flag::ToInterceptSCIDs
+		| Flag::ToPublicChannels
+		| Flag::FromPrivateChannels
+		| Flag::FromPublicToPublicChannels
+		| Flag::ToUnknownSCIDs => {
 			let (chan_upd, _, chan_id, _) = create_announced_chan_between_nodes(&nodes, 1, 2);
 			if flag == Flag::ToInterceptSCIDs {
-				(nodes[1].node.get_intercept_scid(), chan_id)
-			} else if flag == Flag::ToPublicChannels {
-				(chan_upd.contents.short_channel_id, chan_id)
+				(nodes[1].node.get_intercept_scid(), chan_id, None)
 			} else if flag == Flag::ToUnknownSCIDs {
-				(42424242, chan_id)
+				(42424242, chan_id, None)
 			} else {
-				panic!();
+				(chan_upd.contents.short_channel_id, chan_id, Some(false))
 			}
 		},
 		_ => panic!("Combined flags aren't allowed"),
@@ -101,21 +112,50 @@ fn do_test_htlc_interception_flags(
 		get_route_and_payment_hash!(nodes[0], nodes[2], pay_params, amt_msat);
 	route.paths[0].hops[1].short_channel_id = target_scid;
 
-	let interception_bit_match = (flags_bitmask & (flag as u8)) != 0;
+	let mut should_intercept = false;
+	for a_flag in ALL_FLAGS {
+		if flags_bitmask & (a_flag as u8) != 0 {
+			match a_flag {
+				Flag::ToInterceptSCIDs => {
+					should_intercept |= flag == Flag::ToInterceptSCIDs;
+				},
+				Flag::ToOfflinePrivateChannels => {
+					should_intercept |= flag == Flag::ToOfflinePrivateChannels;
+				},
+				Flag::ToOnlinePrivateChannels => {
+					should_intercept |= flag != Flag::ToOfflinePrivateChannels
+						&& outbound_private_for_known_scids == Some(true);
+				},
+				Flag::ToPublicChannels => {
+					should_intercept |= outbound_private_for_known_scids == Some(false);
+				},
+				Flag::ToUnknownSCIDs => {
+					should_intercept |= flag == Flag::ToUnknownSCIDs;
+				},
+				Flag::FromPrivateChannels => {
+					should_intercept |= inbound_private;
+				},
+				Flag::FromPublicToPrivateChannels => {
+					should_intercept |=
+						!inbound_private && outbound_private_for_known_scids == Some(true);
+				},
+				Flag::FromPublicToPublicChannels => {
+					should_intercept |=
+						!inbound_private && outbound_private_for_known_scids == Some(false);
+				},
+				_ => panic!("Combined flags aren't allowed"),
+			}
+		}
+	}
+
 	match modification {
 		Some(ForwardingMod::FeeTooLow) => {
-			assert!(
-				interception_bit_match,
-				"No reason to test failing if we aren't trying to intercept",
-			);
+			assert!(should_intercept, "No reason to test failing if we aren't trying to intercept");
 			route.paths[0].hops[0].fee_msat = 500;
 		},
 		Some(ForwardingMod::CLTVBelowConfig) => {
 			route.paths[0].hops[0].cltv_expiry_delta = 6 * 12;
-			assert!(
-				interception_bit_match,
-				"No reason to test failing if we aren't trying to intercept",
-			);
+			assert!(should_intercept, "No reason to test failing if we aren't trying to intercept");
 		},
 		Some(ForwardingMod::CLTVBelowMin) => {
 			route.paths[0].hops[0].cltv_expiry_delta = 6;
@@ -133,7 +173,7 @@ fn do_test_htlc_interception_flags(
 	do_commitment_signed_dance(&nodes[1], &nodes[0], &payment_event.commitment_msg, false, true);
 	expect_and_process_pending_htlcs(&nodes[1], false);
 
-	if interception_bit_match && modification.is_none() {
+	if should_intercept && modification.is_none() {
 		// If we were set to intercept, check that we got an interception event then
 		// forward the HTLC on to nodes[2] and claim the payment.
 		let intercept_id;
@@ -172,7 +212,14 @@ fn do_test_htlc_interception_flags(
 		// If we were not set to intercept, check that the HTLC either failed or was
 		// automatically forwarded as appropriate.
 		match (modification, flag) {
-			(None, Flag::ToOnlinePrivateChannels | Flag::ToPublicChannels) => {
+			(
+				None,
+				Flag::ToOnlinePrivateChannels
+				| Flag::ToPublicChannels
+				| Flag::FromPrivateChannels
+				| Flag::FromPublicToPrivateChannels
+				| Flag::FromPublicToPublicChannels,
+			) => {
 				check_added_monitors(&nodes[1], 1);
 
 				let forward_ev = SendEvent::from_node(&nodes[1]);
@@ -241,31 +288,55 @@ fn do_test_htlc_interception_flags(
 }
 
 const MAX_BITMASK: u8 = HTLCInterceptionFlags::AllValidHTLCs as u8;
-const ALL_FLAGS: [HTLCInterceptionFlags; 5] = [
+const ALL_FLAGS: [HTLCInterceptionFlags; 8] = [
 	HTLCInterceptionFlags::ToInterceptSCIDs,
 	HTLCInterceptionFlags::ToOfflinePrivateChannels,
 	HTLCInterceptionFlags::ToOnlinePrivateChannels,
 	HTLCInterceptionFlags::ToPublicChannels,
 	HTLCInterceptionFlags::ToUnknownSCIDs,
+	HTLCInterceptionFlags::FromPrivateChannels,
+	HTLCInterceptionFlags::FromPublicToPrivateChannels,
+	HTLCInterceptionFlags::FromPublicToPublicChannels,
 ];
-
 #[test]
-fn test_htlc_interception_flags() {
+fn check_all_flags() {
 	let mut all_flag_bits = 0;
 	for flag in ALL_FLAGS {
 		all_flag_bits |= flag as isize;
 	}
 	assert_eq!(all_flag_bits, MAX_BITMASK as isize, "all flags must test all bits");
+}
 
+fn test_htlc_interception_flags_subrange<I: Iterator<Item = u8>>(r: I) {
 	// Test all 2^5 = 32 combinations of the HTLCInterceptionFlags bitmask
 	// For each combination, test 5 different HTLC forwards and verify correct interception behavior
-	for flags_bitmask in 0..=MAX_BITMASK {
+	for flags_bitmask in r {
 		for flag in ALL_FLAGS {
 			do_test_htlc_interception_flags(flags_bitmask, flag, None);
 		}
 	}
 }
 
+#[test]
+fn test_htlc_interception_flags_a() {
+	test_htlc_interception_flags_subrange(0..MAX_BITMASK / 4);
+}
+
+#[test]
+fn test_htlc_interception_flags_b() {
+	test_htlc_interception_flags_subrange(MAX_BITMASK / 4..MAX_BITMASK / 2);
+}
+
+#[test]
+fn test_htlc_interception_flags_c() {
+	test_htlc_interception_flags_subrange(MAX_BITMASK / 2..MAX_BITMASK / 4 * 3);
+}
+
+#[test]
+fn test_htlc_interception_flags_d() {
+	test_htlc_interception_flags_subrange(MAX_BITMASK / 4 * 3..=MAX_BITMASK);
+}
+
 #[test]
 fn test_htlc_bad_for_chan_config() {
 	// Test that interception won't be done if an HTLC fails to meet the target channel's channel
@@ -274,6 +345,9 @@ fn test_htlc_bad_for_chan_config() {
 		HTLCInterceptionFlags::ToOfflinePrivateChannels,
 		HTLCInterceptionFlags::ToOnlinePrivateChannels,
 		HTLCInterceptionFlags::ToPublicChannels,
+		HTLCInterceptionFlags::FromPrivateChannels,
+		HTLCInterceptionFlags::FromPublicToPrivateChannels,
+		HTLCInterceptionFlags::FromPublicToPublicChannels,
 	];
 	for flag in have_chan_flags {
 		do_test_htlc_interception_flags(flag as u8, flag, Some(ForwardingMod::FeeTooLow));