lightning: Handle peer storage message, it's persistance and send it to the respective peer upon reconnection.

Author: Aditya Sharma

lightning/src/ln/channelmanager.rs

@@ -2303,7 +2303,7 @@ where
 	entropy_source: ES,
 	node_signer: NS,
 	signer_provider: SP,
-
+	peer_storage: Mutex<HashMap<PublicKey, Vec<u8>>>,
 	logger: L,
 }
 
@@ -3108,7 +3108,7 @@ where
 			entropy_source,
 			node_signer,
 			signer_provider,
-
+			peer_storage: Mutex::new(new_hash_map()),
 			logger,
 		}
 	}
@@ -7621,6 +7621,26 @@ where
 		}
 	}
 
+	fn internal_peer_storage(&self, counterparty_node_id: &PublicKey, msg: &msgs::PeerStorageMessage) {
+		let per_peer_state = self.per_peer_state.read().unwrap();
+		let peer_state_mutex = match per_peer_state.get(counterparty_node_id) {
+			Some(peer_state_mutex) => peer_state_mutex,
+			None => return,
+		};
+		let mut peer_state_lock = peer_state_mutex.lock().unwrap();
+		let peer_state = &mut *peer_state_lock;
+		let logger = WithContext::from(&self.logger, Some(*counterparty_node_id), None);
+
+		// Check if we have any channels with the peer (Currently we only provide the servie to peers we have a channel with).
+		if !peer_state.channel_by_id.values().any(|phase| matches!(phase, ChannelPhase::Funded(_))) {
+			log_debug!(logger, "We do not have any channel with {}", log_pubkey!(counterparty_node_id));
+			return;
+		}
+
+		log_trace!(logger, "Received Peer Storage from {}", log_pubkey!(counterparty_node_id));
+		self.peer_storage.lock().unwrap().insert(*counterparty_node_id, msg.data.clone());
+	}
+
 	fn internal_funding_signed(&self, counterparty_node_id: &PublicKey, msg: &msgs::FundingSigned) -> Result<(), MsgHandleErrInternal> {
 		let best_block = *self.best_block.read().unwrap();
 		let per_peer_state = self.per_peer_state.read().unwrap();
@@ -10203,6 +10223,8 @@ where
 	}
 
 	fn handle_peer_storage(&self, counterparty_node_id: &PublicKey, msg: &msgs::PeerStorageMessage) {
+		let _persistence_guard = PersistenceNotifierGuard::optionally_notify(self, || NotifyOption::SkipPersistNoEvents);
+		self.internal_peer_storage(counterparty_node_id, msg);
 	}
 
 	fn handle_your_peer_storage(&self, counterparty_node_id: &PublicKey, msg: &msgs::YourPeerStorageMessage) {
@@ -10557,6 +10579,16 @@ where
 				let peer_state = &mut *peer_state_lock;
 				let pending_msg_events = &mut peer_state.pending_msg_events;
 
+				if let Some(peer_storage) = self.peer_storage.lock().unwrap().get(counterparty_node_id) {
+					pending_msg_events.push(events::MessageSendEvent::SendYourPeerStorageMessage { 
+						node_id: counterparty_node_id.clone(),
+						msg: msgs::YourPeerStorageMessage {
+							data: peer_storage.clone()
+						},
+					});
+				}
+
+
 				for (_, phase) in peer_state.channel_by_id.iter_mut() {
 					match phase {
 						ChannelPhase::Funded(chan) => {
@@ -11614,6 +11646,13 @@ where
 			pending_payment.write(writer)?;
 		}
 
+		let mut peer_storage = None;
+		let peer_storage_dir = self.peer_storage.lock().unwrap();
+
+		if !peer_storage_dir.is_empty() {
+			peer_storage = Some(peer_storage_dir);
+		}
+
 		// For backwards compat, write the session privs and their total length.
 		let mut num_pending_outbounds_compat: u64 = 0;
 		for (_, outbound) in pending_outbound_payments.iter() {
@@ -11686,6 +11725,7 @@ where
 			(11, self.probing_cookie_secret, required),
 			(13, htlc_onion_fields, optional_vec),
 			(14, decode_update_add_htlcs_opt, option),
+			(16, peer_storage, option),
 		});
 
 		Ok(())
@@ -12146,6 +12186,7 @@ where
 		let mut events_override = None;
 		let mut in_flight_monitor_updates: Option<HashMap<(PublicKey, OutPoint), Vec<ChannelMonitorUpdate>>> = None;
 		let mut decode_update_add_htlcs: Option<HashMap<u64, Vec<msgs::UpdateAddHTLC>>> = None;
+		let mut peer_storage: Option<HashMap<PublicKey, Vec<u8>>> = None;
 		read_tlv_fields!(reader, {
 			(1, pending_outbound_payments_no_retry, option),
 			(2, pending_intercepted_htlcs, option),
@@ -12160,7 +12201,9 @@ where
 			(11, probing_cookie_secret, option),
 			(13, claimable_htlc_onion_fields, optional_vec),
 			(14, decode_update_add_htlcs, option),
+			(16, peer_storage, option),
 		});
+		let peer_storage = peer_storage.unwrap_or_else(|| new_hash_map());
 		let mut decode_update_add_htlcs = decode_update_add_htlcs.unwrap_or_else(|| new_hash_map());
 		if fake_scid_rand_bytes.is_none() {
 			fake_scid_rand_bytes = Some(args.entropy_source.get_secure_random_bytes());
@@ -12753,6 +12796,7 @@ where
 
 			last_days_feerates: Mutex::new(VecDeque::new()),
 
+			peer_storage: Mutex::new(peer_storage),
 			logger: args.logger,
 			default_configuration: args.default_config,
 		};