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

@@ -544,6 +544,9 @@ pub(crate) enum ChannelMonitorUpdateStep {
 	ShutdownScript {
 		scriptpubkey: ScriptBuf,
 	},
+	LatestPeerStorage {
+		data: Vec<u8>,
+	},
 }
 
 impl ChannelMonitorUpdateStep {
@@ -555,6 +558,7 @@ impl ChannelMonitorUpdateStep {
 			ChannelMonitorUpdateStep::CommitmentSecret { .. } => "CommitmentSecret",
 			ChannelMonitorUpdateStep::ChannelForceClosed { .. } => "ChannelForceClosed",
 			ChannelMonitorUpdateStep::ShutdownScript { .. } => "ShutdownScript",
+			ChannelMonitorUpdateStep::LatestPeerStorage { .. } => "LatestPeerStorage",
 		}
 	}
 }
@@ -588,6 +592,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.
@@ -835,6 +842,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
@@ -1110,6 +1119,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(())
@@ -1289,7 +1299,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,
@@ -1406,6 +1416,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)>)> {
@@ -2700,6 +2715,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) -> (Vec<PackageTemplate>, Vec<TransactionOutputs>) {
 		let funding_outp = HolderFundingOutput::build(
 			self.funding_redeemscript.clone(),
@@ -2869,6 +2888,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());
+				}
 			}
 		}
 
@@ -2900,6 +2923,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
 	}
@@ -4601,6 +4628,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),
@@ -4612,6 +4640,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),
 		});
 
 		// Monitors for anchor outputs channels opened in v0.0.116 suffered from a bug in which the
@@ -4676,7 +4705,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/ln/channelmanager.rs

@@ -1401,6 +1401,7 @@ where
 	node_signer: NS,
 	signer_provider: SP,
 
+	peer_storage: RwLock<HashMap<PublicKey, Vec<u8>>>,
 	logger: L,
 }
 
@@ -2483,7 +2484,7 @@ where
 			entropy_source,
 			node_signer,
 			signer_provider,
-
+			peer_storage: RwLock::new(HashMap::new()),
 			logger,
 		}
 	}
@@ -6432,6 +6433,73 @@ where
 		}
 	}
 
+	fn internal_peer_storage(&self, counterparty_node_id: &PublicKey, msg: &msgs::PeerStorageMessage) -> Result<(), ()> {
+		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 Err(()),
+		};
+
+		// Check if we have any channels with the peer (Currently we only provide the servie to peers we have a channel with).
+		let chan_info = self.list_channels_with_counterparty(counterparty_node_id);
+		if chan_info.is_empty() {
+			return Err(());
+		}
+
+		// Send ChannelMonitor Update.
+		let sorted_chan_info: Vec<ChannelDetails> = chan_info.iter().cloned().collect();
+		let sorted_chan_info = {
+			let mut sorted_vec: Vec<ChannelDetails> = sorted_chan_info;
+			sorted_vec.sort_by_cached_key(|s| s.funding_txo.unwrap().get_txid());
+			sorted_vec
+		};
+
+		let mut found_funded_chan = false;
+		for chan in &sorted_chan_info {
+			if let Some(funding_txo) = chan.funding_txo {
+				found_funded_chan = true;
+				let peer_storage_update = ChannelMonitorUpdate {
+					update_id: CLOSED_CHANNEL_UPDATE_ID,
+					counterparty_node_id: None,
+					updates: vec![ChannelMonitorUpdateStep::LatestPeerStorage { 
+						data: msg.data.clone(),
+					}],
+					channel_id: Some(chan.channel_id),
+				};
+				// Since channel monitor is already loaded so we do not need to push this onto background event.
+				let update_res: ChannelMonitorUpdateStatus = self.chain_monitor.update_channel(funding_txo, &peer_storage_update);
+				if update_res != ChannelMonitorUpdateStatus::Completed {
+					log_error!(WithContext::from(&self.logger, None, Some(chan.channel_id)),
+						"Critical error: failed to update channel monitor with latest peer storage {:?}: {:?}",
+						msg.data, update_res);
+				}
+				break;
+			}
+		}
+
+		if !found_funded_chan {
+			return Err(());
+		}
+		
+
+		// Update the store.
+		let mut peer_storage = self.peer_storage.write().unwrap();
+		peer_storage.insert(*counterparty_node_id, msg.data.clone());
+
+		let mut peer_state_lock = peer_state_mutex.lock().unwrap();
+		let peer_state = &mut *peer_state_lock;
+
+		// Send ACK.
+		peer_state.pending_msg_events.push(events::MessageSendEvent::SendYourPeerStorageMessage {
+			node_id: counterparty_node_id.clone(),
+			msg: msgs::YourPeerStorageMessage {
+				data: msg.data.clone()
+			},
+		});
+
+		Ok(())
+	}
+
 	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();
@@ -8743,6 +8811,13 @@ where
 	}
 
 	fn handle_peer_storage(&self, counterparty_node_id: &PublicKey, msg: &msgs::PeerStorageMessage) {
+		match self.internal_peer_storage(counterparty_node_id, msg) {
+			Ok(_) => {},
+			Err(_err) => {
+				let logger = WithContext::from(&self.logger, Some(*counterparty_node_id), None);
+				log_debug!(logger, "Could not store Peer Storage for peer {}", log_pubkey!(counterparty_node_id));
+			},
+		}
 	}
 
 	fn handle_your_peer_storage(&self, counterparty_node_id: &PublicKey, msg: &msgs::YourPeerStorageMessage) {
@@ -9104,6 +9179,19 @@ where
 						}
 					}
 				}
+
+				let peer_storage = self.peer_storage.read().unwrap();
+
+				if let Some(value) = peer_storage.get(counterparty_node_id) {
+					log_debug!(logger, "Generating peerstorage for {}", log_pubkey!(counterparty_node_id));
+
+					pending_msg_events.push(events::MessageSendEvent::SendYourPeerStorageMessage { 
+						node_id: counterparty_node_id.clone(),
+						msg: msgs::YourPeerStorageMessage {
+							data: value.clone()
+						},
+					});
+				}
 			}
 
 			return NotifyOption::SkipPersistHandleEvents;
@@ -10391,6 +10479,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)
@@ -10400,6 +10489,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() ||
@@ -11209,7 +11301,7 @@ where
 			entropy_source: args.entropy_source,
 			node_signer: args.node_signer,
 			signer_provider: args.signer_provider,
-
+			peer_storage: RwLock::new(peer_storage_dir),
 			logger: args.logger,
 			default_configuration: args.default_config,
 		};