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

Author: Aditya Sharma

lightning/src/chain/channelmonitor.rs

@@ -560,6 +560,9 @@ pub(crate) enum ChannelMonitorUpdateStep {
 	ShutdownScript {
 		scriptpubkey: ScriptBuf,
 	},
+	LatestPeerStorage {
+		data: Vec<u8>,
+	},
 }
 
 impl ChannelMonitorUpdateStep {
@@ -571,6 +574,7 @@ impl ChannelMonitorUpdateStep {
 			ChannelMonitorUpdateStep::CommitmentSecret { .. } => "CommitmentSecret",
 			ChannelMonitorUpdateStep::ChannelForceClosed { .. } => "ChannelForceClosed",
 			ChannelMonitorUpdateStep::ShutdownScript { .. } => "ShutdownScript",
+			ChannelMonitorUpdateStep::LatestPeerStorage { .. } => "LatestPeerStorage",
 		}
 	}
 }
@@ -604,6 +608,9 @@ impl_writeable_tlv_based_enum_upgradable!(ChannelMonitorUpdateStep,
 	(5, ShutdownScript) => {
 		(0, scriptpubkey, required),
 	},
+	(6, LatestPeerStorage) => {
+		(0, data, required_vec),
+	},
 );
 
 /// Details about the balance(s) available for spending once the channel appears on chain.
@@ -851,6 +858,8 @@ pub(crate) struct ChannelMonitorImpl<Signer: WriteableEcdsaChannelSigner> {
 	/// revoked.
 	payment_preimages: HashMap<PaymentHash, PaymentPreimage>,
 
+	peer_storage: Vec<u8>,
+
 	// Note that `MonitorEvent`s MUST NOT be generated during update processing, only generated
 	// during chain data processing. This prevents a race in `ChainMonitor::update_channel` (and
 	// presumably user implementations thereof as well) where we update the in-memory channel
@@ -1144,6 +1153,7 @@ impl<Signer: WriteableEcdsaChannelSigner> Writeable for ChannelMonitorImpl<Signe
 			(15, self.counterparty_fulfilled_htlcs, required),
 			(17, self.initial_counterparty_commitment_info, option),
 			(19, self.channel_id, required),
+			(21, self.peer_storage, required_vec),
 		});
 
 		Ok(())
@@ -1323,7 +1333,7 @@ impl<Signer: WriteableEcdsaChannelSigner> ChannelMonitor<Signer> {
 			confirmed_commitment_tx_counterparty_output: None,
 			htlcs_resolved_on_chain: Vec::new(),
 			spendable_txids_confirmed: Vec::new(),
-
+			peer_storage: Vec::new(),
 			best_block,
 			counterparty_node_id: Some(counterparty_node_id),
 			initial_counterparty_commitment_info: None,
@@ -1440,6 +1450,11 @@ impl<Signer: WriteableEcdsaChannelSigner> ChannelMonitor<Signer> {
 		self.inner.lock().unwrap().channel_id()
 	}
 
+	/// Fets peer_storage of this peer.
+	pub fn get_peer_storage(&self) -> Vec<u8> {
+		self.inner.lock().unwrap().peer_storage()
+	}
+
 	/// Gets a list of txids, with their output scripts (in the order they appear in the
 	/// transaction), which we must learn about spends of via block_connected().
 	pub fn get_outputs_to_watch(&self) -> Vec<(Txid, Vec<(u32, ScriptBuf)>)> {
@@ -2761,6 +2776,10 @@ impl<Signer: WriteableEcdsaChannelSigner> ChannelMonitorImpl<Signer> {
 		}
 	}
 
+	fn update_peer_storage(&mut self, new_data: Vec<u8>) {
+		self.peer_storage = new_data;
+	}
+
 	fn generate_claimable_outpoints_and_watch_outputs(&mut self, reason: ClosureReason) -> (Vec<PackageTemplate>, Vec<TransactionOutputs>) {
 		let funding_outp = HolderFundingOutput::build(
 			self.funding_redeemscript.clone(),
@@ -2936,6 +2955,10 @@ impl<Signer: WriteableEcdsaChannelSigner> ChannelMonitorImpl<Signer> {
 						panic!("Attempted to replace shutdown script {} with {}", shutdown_script, scriptpubkey);
 					}
 				},
+				ChannelMonitorUpdateStep::LatestPeerStorage { data } => {
+					log_trace!(logger, "Updating ChannelMonitor with latest recieved PeerStorage");
+					self.update_peer_storage(data.clone());
+				}
 			}
 		}
 
@@ -2967,6 +2990,10 @@ impl<Signer: WriteableEcdsaChannelSigner> ChannelMonitorImpl<Signer> {
 		&self.funding_info
 	}
 
+	pub fn peer_storage(&self) -> Vec<u8> {
+		self.peer_storage.clone()
+	}
+
 	pub fn channel_id(&self) -> ChannelId {
 		self.channel_id
 	}
@@ -4632,6 +4659,7 @@ impl<'a, 'b, ES: EntropySource, SP: SignerProvider> ReadableArgs<(&'a ES, &'b SP
 		let mut counterparty_fulfilled_htlcs = Some(new_hash_map());
 		let mut initial_counterparty_commitment_info = None;
 		let mut channel_id = None;
+		let mut peer_storage = Some(Vec::new());
 		read_tlv_fields!(reader, {
 			(1, funding_spend_confirmed, option),
 			(3, htlcs_resolved_on_chain, optional_vec),
@@ -4643,6 +4671,7 @@ impl<'a, 'b, ES: EntropySource, SP: SignerProvider> ReadableArgs<(&'a ES, &'b SP
 			(15, counterparty_fulfilled_htlcs, option),
 			(17, initial_counterparty_commitment_info, option),
 			(19, channel_id, option),
+			(21, peer_storage, optional_vec),
 		});
 
 		// `HolderForceClosedWithInfo` replaced `HolderForceClosed` in v0.0.122. If we have both
@@ -4717,7 +4746,7 @@ impl<'a, 'b, ES: EntropySource, SP: SignerProvider> ReadableArgs<(&'a ES, &'b SP
 			confirmed_commitment_tx_counterparty_output,
 			htlcs_resolved_on_chain: htlcs_resolved_on_chain.unwrap(),
 			spendable_txids_confirmed: spendable_txids_confirmed.unwrap(),
-
+			peer_storage: peer_storage.unwrap(),
 			best_block,
 			counterparty_node_id,
 			initial_counterparty_commitment_info,

lightning/src/chain/transaction.rs

@@ -66,6 +66,10 @@ impl OutPoint {
 			vout: self.index as u32,
 		}
 	}
+
+	pub fn get_txid(self) -> Txid {
+		self.txid
+	}
 }
 
 impl core::fmt::Display for OutPoint {

lightning/src/ln/channel.rs

@@ -2101,6 +2101,11 @@ impl<SP: Deref> ChannelContext<SP> where SP::Target: SignerProvider  {
 		self.update_time_counter
 	}
 
+	pub fn increment_and_fetch_monitor_update_id(&mut self) -> u64 {
+		self.latest_monitor_update_id +=1;
+		self.latest_monitor_update_id
+	}
+
 	pub fn get_latest_monitor_update_id(&self) -> u64 {
 		self.latest_monitor_update_id
 	}

lightning/src/ln/channelmanager.rs

@@ -904,6 +904,7 @@ pub(super) struct PeerState<SP: Deref> where SP::Target: SignerProvider {
 	/// [`ChannelMessageHandler::peer_connected`] and no corresponding
 	/// [`ChannelMessageHandler::peer_disconnected`].
 	is_connected: bool,
+	peer_storage: Vec<u8>,
 }
 
 impl <SP: Deref> PeerState<SP> where SP::Target: SignerProvider {
@@ -2526,7 +2527,6 @@ where
 			entropy_source,
 			node_signer,
 			signer_provider,
-
 			logger,
 		}
 	}
@@ -6685,6 +6685,53 @@ 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;
+		}
+
+		let mut funded_channels: Vec<&mut Channel<SP>> = peer_state.channel_by_id.values_mut()
+			.filter_map(|phase| {
+				if let ChannelPhase::Funded(channel) = phase {
+					Some(channel)
+				} else {
+					None
+				}
+			})
+			.collect();
+
+		let min_funded_chan = funded_channels
+			.iter_mut()
+			.min_by_key(|s| s.context.get_funding_txo().unwrap().get_txid())
+			.unwrap();
+		
+		// Send ChannelMonitor Update.
+		let peer_storage_update = ChannelMonitorUpdate {
+			update_id: min_funded_chan.context.increment_and_fetch_monitor_update_id(),
+			counterparty_node_id: None,
+			updates: vec![ChannelMonitorUpdateStep::LatestPeerStorage { 
+				data: msg.data.clone(),
+			}],
+			channel_id: Some(min_funded_chan.context.channel_id()),
+		};
+		
+		// Update the store.
+		peer_state.peer_storage = msg.data.clone();
+
+		handle_new_monitor_update!(self, min_funded_chan.context.get_funding_txo().unwrap(), peer_storage_update, peer_state_lock, peer_state, per_peer_state, min_funded_chan);
+	}
+
 	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();
@@ -9069,6 +9116,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) {
@@ -9384,6 +9433,7 @@ where
 							monitor_update_blocked_actions: BTreeMap::new(),
 							actions_blocking_raa_monitor_updates: BTreeMap::new(),
 							is_connected: true,
+							peer_storage: Vec::new(),
 						}));
 					},
 					hash_map::Entry::Occupied(e) => {
@@ -9456,6 +9506,15 @@ where
 						},
 					}
 				}
+
+				let peer_storage = peer_state.peer_storage.clone();
+
+				pending_msg_events.push(events::MessageSendEvent::SendYourPeerStorageMessage { 
+					node_id: counterparty_node_id.clone(),
+					msg: msgs::YourPeerStorageMessage {
+						data: peer_storage
+					},
+				});
 			}
 
 			return NotifyOption::SkipPersistHandleEvents;
@@ -10783,6 +10842,7 @@ where
 		let mut channel_closures = VecDeque::new();
 		let mut close_background_events = Vec::new();
 		let mut funding_txo_to_channel_id = hash_map_with_capacity(channel_count as usize);
+		let mut peer_storage_dir: HashMap<PublicKey, Vec<u8>> = HashMap::new();
 		for _ in 0..channel_count {
 			let mut channel: Channel<SP> = Channel::read(reader, (
 				&args.entropy_source, &args.signer_provider, best_block_height, &provided_channel_type_features(&args.default_config)
@@ -10792,6 +10852,9 @@ where
 			funding_txo_to_channel_id.insert(funding_txo, channel.context.channel_id());
 			funding_txo_set.insert(funding_txo.clone());
 			if let Some(ref mut monitor) = args.channel_monitors.get_mut(&funding_txo) {
+				// Load Peer_storage from ChannelMonitor to memory.
+				peer_storage_dir.insert(channel.context.get_counterparty_node_id(), monitor.get_peer_storage());
+
 				if channel.get_cur_holder_commitment_transaction_number() > monitor.get_cur_holder_commitment_number() ||
 						channel.get_revoked_counterparty_commitment_transaction_number() > monitor.get_min_seen_secret() ||
 						channel.get_cur_counterparty_commitment_transaction_number() > monitor.get_cur_counterparty_commitment_number() ||
@@ -10938,7 +11001,7 @@ where
 			claimable_htlcs_list.push((payment_hash, previous_hops));
 		}
 
-		let peer_state_from_chans = |channel_by_id| {
+		let peer_state_from_chans = |channel_by_id, peer_storage_blob| {
 			PeerState {
 				channel_by_id,
 				inbound_channel_request_by_id: new_hash_map(),
@@ -10948,6 +11011,7 @@ where
 				monitor_update_blocked_actions: BTreeMap::new(),
 				actions_blocking_raa_monitor_updates: BTreeMap::new(),
 				is_connected: false,
+				peer_storage: peer_storage_blob,
 			}
 		};
 
@@ -10956,7 +11020,7 @@ where
 		for _ in 0..peer_count {
 			let peer_pubkey = Readable::read(reader)?;
 			let peer_chans = funded_peer_channels.remove(&peer_pubkey).unwrap_or(new_hash_map());
-			let mut peer_state = peer_state_from_chans(peer_chans);
+			let mut peer_state = peer_state_from_chans(peer_chans, peer_storage_dir.get(&peer_pubkey).cloned().unwrap_or_default());
 			peer_state.latest_features = Readable::read(reader)?;
 			per_peer_state.insert(peer_pubkey, Mutex::new(peer_state));
 		}
@@ -11167,7 +11231,7 @@ where
 					// still open, we need to replay any monitor updates that are for closed channels,
 					// creating the neccessary peer_state entries as we go.
 					let peer_state_mutex = per_peer_state.entry(counterparty_id).or_insert_with(|| {
-						Mutex::new(peer_state_from_chans(new_hash_map()))
+						Mutex::new(peer_state_from_chans(new_hash_map(), peer_storage_dir.get(&counterparty_id).cloned().unwrap_or_default()))
 					});
 					let mut peer_state = peer_state_mutex.lock().unwrap();
 					handle_in_flight_updates!(counterparty_id, chan_in_flight_updates,
@@ -11617,7 +11681,6 @@ where
 			entropy_source: args.entropy_source,
 			node_signer: args.node_signer,
 			signer_provider: args.signer_provider,
-
 			logger: args.logger,
 			default_configuration: args.default_config,
 		};