Add deregistration cleanup for OnionMessageInterceptor in LSPS2

When intercept SCIDs are removed during cleanup, also clean up the
handler-level `peer_by_intercept_scid` map and deregister the peer
from onion message interception if they have no remaining active SCIDs.

Cleanup is added in all relevant paths:
- `prune_expired_request_state()` now returns pruned SCIDs
- `peer_disconnected()` cleans up after pruning
- `htlc_intercepted()` error path (fixes existing TODO)
- `channel_open_abandoned()` cleans up after removing the SCID
- `persist()` cleans up both `peer_by_intercept_scid` and
  `peer_by_channel_id` when removing a peer entry entirely

Co-Authored-By: HAL 9000

Author: tnull

lightning-liquidity/src/lsps2/service.rs

@@ -632,17 +632,20 @@ impl PeerState {
 		});
 	}
 
-	fn prune_expired_request_state(&mut self) {
+	fn prune_expired_request_state(&mut self) -> Vec<u64> {
+		let mut pruned_scids = Vec::new();
 		self.outbound_channels_by_intercept_scid.retain(|intercept_scid, entry| {
 			if entry.is_prunable() {
 				// We abort the flow, and prune any data kept.
 				self.intercept_scid_by_channel_id.retain(|_, iscid| intercept_scid != iscid);
 				self.intercept_scid_by_user_channel_id.retain(|_, iscid| intercept_scid != iscid);
 				self.needs_persist |= true;
+				pruned_scids.push(*intercept_scid);
 				return false;
 			}
 			true
 		});
+		pruned_scids
 	}
 
 	fn pending_requests_and_channels(&self) -> usize {
@@ -788,6 +791,29 @@ where
 		&self.config
 	}
 
+	/// Cleans up `peer_by_intercept_scid` entries for the given SCIDs, and deregisters the peer
+	/// from onion message interception if they have no remaining active intercept SCIDs.
+	fn cleanup_intercept_scids(
+		&self, counterparty_node_id: &PublicKey, pruned_scids: &[u64], has_remaining_channels: bool,
+	) {
+		if pruned_scids.is_empty() {
+			return;
+		}
+
+		{
+			let mut peer_by_intercept_scid = self.peer_by_intercept_scid.write().unwrap();
+			for scid in pruned_scids {
+				peer_by_intercept_scid.remove(scid);
+			}
+		}
+
+		if !has_remaining_channels {
+			if let Some(ref interceptor) = self.onion_message_interceptor {
+				interceptor.deregister_peer_for_interception(counterparty_node_id);
+			}
+		}
+	}
+
 	/// Returns whether the peer has any active LSPS2 requests.
 	pub(crate) fn has_active_requests(&self, counterparty_node_id: &PublicKey) -> bool {
 		let outer_state_lock = self.per_peer_state.read().unwrap();
@@ -1067,7 +1093,15 @@ where
 								peer_state
 									.outbound_channels_by_intercept_scid
 									.remove(&intercept_scid);
-								// TODO: cleanup peer_by_intercept_scid
+								let has_remaining =
+									!peer_state.outbound_channels_by_intercept_scid.is_empty();
+								drop(peer_state);
+								drop(outer_state_lock);
+								self.cleanup_intercept_scids(
+									counterparty_node_id,
+									&[intercept_scid],
+									has_remaining,
+								);
 								return Err(APIError::APIMisuseError { err: e.err });
 							},
 						}
@@ -1286,7 +1320,7 @@ where
 	pub async fn channel_open_abandoned(
 		&self, counterparty_node_id: &PublicKey, user_channel_id: u128,
 	) -> Result<(), APIError> {
-		{
+		let (intercept_scid, has_remaining) = {
 			let outer_state_lock = self.per_peer_state.read().unwrap();
 			let inner_state_lock = outer_state_lock.get(counterparty_node_id).ok_or_else(|| {
 				APIError::APIMisuseError {
@@ -1333,7 +1367,11 @@ where
 			peer_state.outbound_channels_by_intercept_scid.remove(&intercept_scid);
 			peer_state.intercept_scid_by_channel_id.retain(|_, &mut scid| scid != intercept_scid);
 			peer_state.needs_persist |= true;
-		}
+			let has_remaining = !peer_state.outbound_channels_by_intercept_scid.is_empty();
+			(intercept_scid, has_remaining)
+		};
+
+		self.cleanup_intercept_scids(counterparty_node_id, &[intercept_scid], has_remaining);
 
 		self.persist_peer_state(*counterparty_node_id).await.map_err(|e| {
 			APIError::APIMisuseError {
@@ -1817,17 +1855,32 @@ where
 			{
 				// First build a list of peers to persist and prune with the read lock. This allows
 				// us to avoid the write lock unless we actually need to remove a node.
+				let mut all_pruned_scids = Vec::new();
 				let outer_state_lock = self.per_peer_state.read().unwrap();
 				for (counterparty_node_id, inner_state_lock) in outer_state_lock.iter() {
 					let mut peer_state_lock = inner_state_lock.lock().unwrap();
-					peer_state_lock.prune_expired_request_state();
+					let pruned_scids = peer_state_lock.prune_expired_request_state();
+					if !pruned_scids.is_empty() {
+						let has_remaining =
+							!peer_state_lock.outbound_channels_by_intercept_scid.is_empty();
+						all_pruned_scids.push((*counterparty_node_id, pruned_scids, has_remaining));
+					}
 					let is_prunable = peer_state_lock.is_prunable();
 					if is_prunable {
 						need_remove.push(*counterparty_node_id);
 					} else if peer_state_lock.needs_persist {
 						need_persist.push(*counterparty_node_id);
 					}
 				}
+				drop(outer_state_lock);
+
+				for (counterparty_node_id, pruned_scids, has_remaining) in all_pruned_scids {
+					self.cleanup_intercept_scids(
+						&counterparty_node_id,
+						&pruned_scids,
+						has_remaining,
+					);
+				}
 			}
 
 			for counterparty_node_id in need_persist.into_iter() {
@@ -1838,6 +1891,7 @@ where
 
 			for counterparty_node_id in need_remove {
 				let mut future_opt = None;
+				let mut was_removed = false;
 				{
 					// We need to take the `per_peer_state` write lock to remove an entry, but also
 					// have to hold it until after the `remove` call returns (but not through
@@ -1849,6 +1903,7 @@ where
 						let state = entry.get_mut().get_mut().unwrap();
 						if state.is_prunable() {
 							entry.remove();
+							was_removed = true;
 							let key = counterparty_node_id.to_string();
 							future_opt = Some(self.kv_store.remove(
 								LIQUIDITY_MANAGER_PERSISTENCE_PRIMARY_NAMESPACE,
@@ -1866,6 +1921,20 @@ where
 						debug_assert!(false);
 					}
 				}
+				if was_removed {
+					// Clean up handler-level maps for the removed peer.
+					self.peer_by_intercept_scid
+						.write()
+						.unwrap()
+						.retain(|_, node_id| *node_id != counterparty_node_id);
+					self.peer_by_channel_id
+						.write()
+						.unwrap()
+						.retain(|_, node_id| *node_id != counterparty_node_id);
+					if let Some(ref interceptor) = self.onion_message_interceptor {
+						interceptor.deregister_peer_for_interception(&counterparty_node_id);
+					}
+				}
 				if let Some(future) = future_opt {
 					future.await?;
 					did_persist = true;
@@ -1893,7 +1962,11 @@ where
 			// We clean up the peer state, but leave removing the peer entry to the prune logic in
 			// `persist` which removes it from the store.
 			peer_state_lock.prune_pending_requests();
-			peer_state_lock.prune_expired_request_state();
+			let pruned_scids = peer_state_lock.prune_expired_request_state();
+			let has_remaining = !peer_state_lock.outbound_channels_by_intercept_scid.is_empty();
+			drop(peer_state_lock);
+			drop(outer_state_lock);
+			self.cleanup_intercept_scids(&counterparty_node_id, &pruned_scids, has_remaining);
 		}
 	}