Merge pull request #31 from getAlby/feat/custom-tlvs-storage

feat: custom TLVs storage
getAlby · Jun 14, 2024 · 63f0d41 · 63f0d41
2 parents fe5cc20 + b124287
commit 63f0d41
Show file tree

Hide file tree

Showing 5 changed files with 86 additions and 27 deletions.
diff --git a/bindings/ldk_node.udl b/bindings/ldk_node.udl
@@ -237,7 +237,7 @@ interface PaymentKind {
 	Onchain();
 	Bolt11(PaymentHash hash, PaymentPreimage? preimage, PaymentSecret? secret, string? bolt11_invoice);
 	Bolt11Jit(PaymentHash hash, PaymentPreimage? preimage, PaymentSecret? secret, LSPFeeLimits lsp_fee_limits);
-	Spontaneous(PaymentHash hash, PaymentPreimage? preimage);
+	Spontaneous(PaymentHash hash, PaymentPreimage? preimage, sequence<TlvEntry> custom_tlvs);
 };
 
 enum PaymentDirection {

diff --git a/src/event.rs b/src/event.rs
@@ -1,7 +1,7 @@
 use crate::types::{DynStore, Sweeper, Wallet};
 use crate::{
 	hex_utils, BumpTransactionEventHandler, ChannelManager, Config, Error, Graph, PeerInfo,
-	PeerStore, UserChannelId,
+	PeerStore, TlvEntry, UserChannelId,
 };
 
 use crate::payment::store::{
@@ -431,7 +431,7 @@ where
 				via_channel_id: _,
 				via_user_channel_id: _,
 				claim_deadline: _,
-				onion_fields: _,
+				onion_fields,
 				counterparty_skimmed_fee_msat,
 			} => {
 				let payment_id = PaymentId(payment_hash.0);
@@ -553,6 +553,62 @@ where
 						panic!("Failed to access payment store");
 					});
 				}
+
+				if let PaymentPurpose::SpontaneousPayment(preimage) = purpose {
+					let custom_tlvs = onion_fields
+						.map(|of| {
+							of.custom_tlvs()
+								.iter()
+								.map(|(t, v)| TlvEntry { r#type: *t, value: v.clone() })
+								.collect()
+						})
+						.unwrap_or_default();
+					log_info!(
+						self.logger,
+						"Saving spontaneous payment with custom TLVs {:?} for payment hash {} of {}msat",
+						custom_tlvs,
+						hex_utils::to_string(&payment_hash.0),
+						amount_msat,
+					);
+
+					let payment = PaymentDetails {
+						id: payment_id,
+						kind: PaymentKind::Spontaneous {
+							hash: payment_hash,
+							preimage: Some(preimage),
+							custom_tlvs,
+						},
+						amount_msat: Some(amount_msat),
+						direction: PaymentDirection::Inbound,
+						status: PaymentStatus::Pending,
+						last_update: time::SystemTime::now()
+							.duration_since(time::UNIX_EPOCH)
+							.unwrap_or(time::Duration::ZERO)
+							.as_secs(),
+						fee_msat: None,
+					};
+
+					match self.payment_store.insert(payment) {
+						Ok(false) => (),
+						Ok(true) => {
+							log_error!(
+								self.logger,
+								"Spontaneous payment with hash {} was previously known",
+								hex_utils::to_string(&payment_hash.0)
+							);
+							debug_assert!(false);
+						},
+						Err(e) => {
+							log_error!(
+								self.logger,
+								"Failed to insert spontaneous payment with hash {}: {}",
+								hex_utils::to_string(&payment_hash.0),
+								e
+							);
+							debug_assert!(false);
+						},
+					};
+				}
 			},
 			LdkEvent::PaymentClaimed {
 				payment_hash,
@@ -621,37 +677,26 @@ where
 						);
 						return;
 					},
-					PaymentPurpose::SpontaneousPayment(preimage) => {
-						let payment = PaymentDetails {
-							id: payment_id,
-							kind: PaymentKind::Spontaneous {
-								hash: payment_hash,
-								preimage: Some(preimage),
-							},
-							amount_msat: Some(amount_msat),
-							direction: PaymentDirection::Inbound,
-							status: PaymentStatus::Succeeded,
-							last_update: time::SystemTime::now()
-								.duration_since(time::UNIX_EPOCH)
-								.unwrap_or(time::Duration::ZERO)
-								.as_secs(),
-							fee_msat: None,
+					PaymentPurpose::SpontaneousPayment(_) => {
+						let update = PaymentDetailsUpdate {
+							status: Some(PaymentStatus::Succeeded),
+							..PaymentDetailsUpdate::new(payment_id)
 						};
 
-						match self.payment_store.insert(payment) {
-							Ok(false) => (),
-							Ok(true) => {
+						match self.payment_store.update(&update) {
+							Ok(true) => (),
+							Ok(false) => {
 								log_error!(
 									self.logger,
-									"Spontaneous payment with hash {} was previously known",
+									"Spontaneous payment with hash {} couldn't be found in store",
 									hex_utils::to_string(&payment_hash.0)
 								);
 								debug_assert!(false);
 							},
 							Err(e) => {
 								log_error!(
 									self.logger,
-									"Failed to insert payment with hash {}: {}",
+									"Failed to update payment with hash {}: {}",
 									hex_utils::to_string(&payment_hash.0),
 									e
 								);

diff --git a/src/payment/spontaneous.rs b/src/payment/spontaneous.rs
@@ -68,7 +68,9 @@ impl SpontaneousPayment {
 			amount_msat,
 		);
 		let recipient_fields = RecipientOnionFields::spontaneous_empty()
-			.with_custom_tlvs(custom_tlvs.into_iter().map(|tlv| (tlv.r#type, tlv.value)).collect())
+			.with_custom_tlvs(
+				custom_tlvs.iter().map(|tlv| (tlv.r#type, tlv.value.clone())).collect(),
+			)
 			.map_err(|_| {
 				log_error!(self.logger, "Payment error: invalid custom TLVs.");
 				Error::InvalidCustomTlv
@@ -89,6 +91,7 @@ impl SpontaneousPayment {
 					kind: PaymentKind::Spontaneous {
 						hash: payment_hash,
 						preimage: Some(payment_preimage),
+						custom_tlvs,
 					},
 					status: PaymentStatus::Pending,
 					direction: PaymentDirection::Outbound,
@@ -111,6 +114,7 @@ impl SpontaneousPayment {
 							kind: PaymentKind::Spontaneous {
 								hash: payment_hash,
 								preimage: Some(payment_preimage),
+								custom_tlvs,
 							},
 
 							status: PaymentStatus::Failed,

diff --git a/src/payment/store.rs b/src/payment/store.rs
@@ -3,7 +3,7 @@ use crate::io::{
 	PAYMENT_INFO_PERSISTENCE_PRIMARY_NAMESPACE, PAYMENT_INFO_PERSISTENCE_SECONDARY_NAMESPACE,
 };
 use crate::logger::{log_error, Logger};
-use crate::types::DynStore;
+use crate::types::{DynStore, TlvEntry};
 use crate::Error;
 
 use lightning::ln::channelmanager::PaymentId;
@@ -108,7 +108,7 @@ impl Readable for PaymentDetails {
 					PaymentKind::Bolt11 { hash, preimage, secret, bolt11_invoice: None }
 				}
 			} else {
-				PaymentKind::Spontaneous { hash, preimage }
+				PaymentKind::Spontaneous { hash, preimage, custom_tlvs: Vec::new() }
 			}
 		};
 
@@ -194,6 +194,8 @@ pub enum PaymentKind {
 		hash: PaymentHash,
 		/// The pre-image used by the payment.
 		preimage: Option<PaymentPreimage>,
+		/// Custom TLVs.
+		custom_tlvs: Vec<TlvEntry>,
 	},
 }
 
@@ -214,6 +216,7 @@ impl_writeable_tlv_based_enum!(PaymentKind,
 	(8, Spontaneous) => {
 		(0, hash, required),
 		(2, preimage, option),
+		(131072, custom_tlvs, optional_vec),
 	};
 );
 
@@ -602,7 +605,7 @@ mod tests {
 			);
 
 			match spontaneous_decoded.kind {
-				PaymentKind::Spontaneous { hash: h, preimage: p } => {
+				PaymentKind::Spontaneous { hash: h, preimage: p, custom_tlvs: _ } => {
 					assert_eq!(hash, h);
 					assert_eq!(preimage, p);
 				},

diff --git a/src/types.rs b/src/types.rs
@@ -2,6 +2,7 @@ use crate::logger::FilesystemLogger;
 use crate::message_handler::NodeCustomMessageHandler;
 
 use lightning::chain::chainmonitor;
+use lightning::impl_writeable_tlv_based;
 use lightning::ln::channelmanager::ChannelDetails as LdkChannelDetails;
 use lightning::ln::msgs::RoutingMessageHandler;
 use lightning::ln::msgs::SocketAddress;
@@ -478,10 +479,16 @@ impl Default for ChannelConfig {
 }
 
 /// Custom TLV entry.
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct TlvEntry {
 	/// Type number.
 	pub r#type: u64,
 
 	/// Serialized value.
 	pub value: Vec<u8>,
 }
+
+impl_writeable_tlv_based!(TlvEntry, {
+	(0, r#type, required),
+	(1, value, required),
+});