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

@@ -2122,7 +2122,7 @@ where
 	entropy_source: ES,
 	node_signer: NS,
 	signer_provider: SP,
-
+	peer_storage: Mutex<HashMap<PublicKey, Vec<u8>>>,
 	logger: L,
 }
 
@@ -2901,7 +2901,7 @@ where
 			entropy_source,
 			node_signer,
 			signer_provider,
-
+			peer_storage: Mutex::new(new_hash_map()),
 			logger,
 		}
 	}
@@ -7172,6 +7172,25 @@ where
 		}
 	}
 
+	fn internal_peer_storage(&self, counterparty_node_id: &PublicKey, msg: &msgs::PeerStorageMessage) {
+		let per_peer_state = self.per_peer_state.write().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.total_channel_count() == 0 {
+			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();
@@ -9672,6 +9691,8 @@ where
 	}
 
 	fn handle_peer_storage(&self, counterparty_node_id: &PublicKey, msg: &msgs::PeerStorageMessage) {
+		let _persistence_guard = PersistenceNotifierGuard::notify_on_drop(self);
+		self.internal_peer_storage(counterparty_node_id, msg);
 	}
 
 	fn handle_your_peer_storage(&self, counterparty_node_id: &PublicKey, msg: &msgs::YourPeerStorageMessage) {
@@ -10025,6 +10046,17 @@ where
 				let mut peer_state_lock = peer_state_mutex.lock().unwrap();
 				let peer_state = &mut *peer_state_lock;
 				let pending_msg_events = &mut peer_state.pending_msg_events;
+				let peer_storage = self.peer_storage.lock().unwrap().get(counterparty_node_id).unwrap_or(&Vec::<u8>::new()).clone();
+
+				if peer_storage.len() > 0 {
+					pending_msg_events.push(events::MessageSendEvent::SendYourPeerStorageMessage { 
+						node_id: counterparty_node_id.clone(),
+						msg: msgs::YourPeerStorageMessage {
+							data: peer_storage
+						},
+					});
+				}
+
 
 				for (_, phase) in peer_state.channel_by_id.iter_mut() {
 					match phase {
@@ -11034,6 +11066,13 @@ where
 			pending_payment.write(writer)?;
 		}
 
+		let peer_storage = self.peer_storage.lock().unwrap();
+		(peer_storage.len() as u64).write(writer)?;
+		for (node_id, peer_data) in peer_storage.iter() {
+			node_id.write(writer)?;
+			peer_data.write(writer)?;
+		}
+
 		// 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() {
@@ -11539,6 +11578,14 @@ where
 			}
 		}
 
+		let peer_storage_count: u64 = Readable::read(reader)?;
+		let mut peer_storage: HashMap<PublicKey, Vec<u8>> = hash_map_with_capacity(cmp::min(peer_storage_count as usize, MAX_ALLOC_SIZE/(3*32)));
+		for _ in 0..peer_storage_count {
+			if peer_storage.insert(Readable::read(reader)?, Readable::read(reader)?).is_some() {
+				return Err(DecodeError::InvalidValue);
+			}
+		}
+
 		let pending_outbound_payments_count_compat: u64 = Readable::read(reader)?;
 		let mut pending_outbound_payments_compat: HashMap<PaymentId, PendingOutboundPayment> =
 			hash_map_with_capacity(cmp::min(pending_outbound_payments_count_compat as usize, MAX_ALLOC_SIZE/32));
@@ -12168,6 +12215,7 @@ where
 
 			last_days_feerates: Mutex::new(VecDeque::new()),
 
+			peer_storage: Mutex::new(peer_storage),
 			logger: args.logger,
 			default_configuration: args.default_config,
 		};