Implement KVStoreUnpersister

Author: tnull

Cargo.toml

@@ -60,6 +60,7 @@ futures = "0.3"
 serde_json = { version = "1.0" }
 tokio = { version = "1", default-features = false, features = [ "rt-multi-thread", "time", "sync" ] }
 esplora-client = { version = "=0.3", default-features = false }
+libc = "0.2"
 
 [dev-dependencies]
 electrsd = { version = "0.22.0", features = ["legacy", "esplora_a33e97e1", "bitcoind_23_0"] }

src/event.rs

@@ -3,6 +3,7 @@ use crate::{
 	PaymentInfo, PaymentInfoStorage, PaymentStatus, Wallet,
 };
 
+use crate::io_utils::KVStoreUnpersister;
 use crate::logger::{log_error, log_info, Logger};
 
 use lightning::chain::chaininterface::{BroadcasterInterface, ConfirmationTarget, FeeEstimator};
@@ -181,9 +182,9 @@ impl Writeable for EventQueueSerWrapper<'_> {
 	}
 }
 
-pub(crate) struct EventHandler<K: Deref, L: Deref>
+pub(crate) struct EventHandler<K: Deref + Clone, L: Deref>
 where
-	K::Target: KVStorePersister,
+	K::Target: KVStorePersister + KVStoreUnpersister,
 	L::Target: Logger,
 {
 	wallet: Arc<Wallet<bdk::database::SqliteDatabase>>,
@@ -197,9 +198,9 @@ where
 	_config: Arc<Config>,
 }
 
-impl<K: Deref, L: Deref> EventHandler<K, L>
+impl<K: Deref + Clone, L: Deref> EventHandler<K, L>
 where
-	K::Target: KVStorePersister,
+	K::Target: KVStorePersister + KVStoreUnpersister,
 	L::Target: Logger,
 {
 	pub fn new(
@@ -222,9 +223,9 @@ where
 	}
 }
 
-impl<K: Deref, L: Deref> LdkEventHandler for EventHandler<K, L>
+impl<K: Deref + Clone, L: Deref> LdkEventHandler for EventHandler<K, L>
 where
-	K::Target: KVStorePersister,
+	K::Target: KVStorePersister + KVStoreUnpersister,
 	L::Target: Logger,
 {
 	fn handle_event(&self, event: LdkEvent) {
@@ -363,25 +364,24 @@ where
 					PaymentPurpose::SpontaneousPayment(preimage) => (Some(preimage), None),
 				};
 
-				let payment_info =
-					if let Some(mut payment_info) = self.payment_store.get(&payment_hash) {
+				let mut locked_store = self.payment_store.lock().unwrap();
+				locked_store
+					.entry(payment_hash)
+					.and_modify(|payment_info| {
 						payment_info.status = PaymentStatus::Succeeded;
 						payment_info.preimage = payment_preimage;
 						payment_info.secret = payment_secret;
 						payment_info.amount_msat = Some(amount_msat);
-						payment_info
-					} else {
-						PaymentInfo {
-							preimage: payment_preimage,
-							payment_hash,
-							secret: payment_secret,
-							amount_msat: Some(amount_msat),
-							direction: PaymentDirection::Inbound,
-							status: PaymentStatus::Succeeded,
-						}
-					};
+					})
+					.or_insert(PaymentInfo {
+						preimage: payment_preimage,
+						payment_hash,
+						secret: payment_secret,
+						amount_msat: Some(amount_msat),
+						direction: PaymentDirection::Inbound,
+						status: PaymentStatus::Succeeded,
+					});
 
-				self.payment_store.insert(payment_info).expect("Failed to access payment store");
 				self.event_queue
 					.add_event(Event::PaymentReceived { payment_hash, amount_msat })
 					.expect("Failed to push to event queue");

src/io_utils.rs

@@ -3,12 +3,15 @@ use crate::{Config, FilesystemLogger, NetworkGraph, Scorer, WALLET_KEYS_SEED_LEN
 
 use lightning::routing::scoring::{ProbabilisticScorer, ProbabilisticScoringParameters};
 use lightning::util::ser::{Readable, ReadableArgs};
+use lightning_persister::FilesystemPersister;
 
 use rand::{thread_rng, RngCore};
 
 use std::fs;
 use std::io::{BufReader, Write};
+use std::os::unix::io::AsRawFd;
 use std::path::Path;
+use std::path::PathBuf;
 use std::sync::Arc;
 
 pub(crate) fn read_or_generate_seed_file(keys_seed_path: &str) -> [u8; WALLET_KEYS_SEED_LEN] {
@@ -83,3 +86,43 @@ pub(crate) fn read_payment_info(config: &Config) -> Vec<PaymentInfo> {
 
 	payments
 }
+
+/// Provides an interface that allows a previously persisted key to be unpersisted.
+pub trait KVStoreUnpersister {
+	/// Unpersist (i.e., remove) the writeable previously persisted under the provided key.
+	/// Returns `true` if the key was present, and `false` otherwise.
+	fn unpersist(&self, key: &str) -> std::io::Result<bool>;
+}
+
+impl KVStoreUnpersister for FilesystemPersister {
+	fn unpersist(&self, key: &str) -> std::io::Result<bool> {
+		let mut dest_file = PathBuf::from(self.get_data_dir());
+		dest_file.push(key);
+
+		if !dest_file.is_file() {
+			return Ok(false);
+		}
+
+		fs::remove_file(&dest_file)?;
+		#[cfg(not(target_os = "windows"))]
+		{
+			let parent_directory = dest_file.parent().unwrap();
+			let dir_file = fs::OpenOptions::new().read(true).open(parent_directory)?;
+			unsafe {
+				// The above call to `fs::remove_file` corresponds to POSIX `unlink`, whose changes
+				// to the inode might get cached (and hence possibly lost on crash), depending on
+				// the target platform and file system.
+				//
+				// In order to assert we permanently removed the file in question we therefore
+				// call `fsync` on the parent directory on platforms that support it,
+				libc::fsync(dir_file.as_raw_fd());
+			}
+		}
+
+		if dest_file.is_file() {
+			return Err(std::io::Error::new(std::io::ErrorKind::Other, "Unpersisting key failed"));
+		}
+
+		return Ok(true);
+	}
+}

src/payment_store.rs

@@ -1,14 +1,16 @@
 use crate::hex_utils;
+use crate::io_utils::KVStoreUnpersister;
 use crate::Error;
 
 use lightning::ln::{PaymentHash, PaymentPreimage, PaymentSecret};
 use lightning::util::persist::KVStorePersister;
 use lightning::{impl_writeable_tlv_based, impl_writeable_tlv_based_enum};
 
-use std::collections::HashMap;
+use std::collections::hash_map;
+use std::collections::{HashMap, HashSet};
 use std::iter::FromIterator;
 use std::ops::Deref;
-use std::sync::Mutex;
+use std::sync::{Mutex, MutexGuard};
 
 /// Represents a payment.
 #[derive(Clone, Debug, PartialEq, Eq)]
@@ -70,21 +72,26 @@ impl_writeable_tlv_based_enum!(PaymentStatus,
 /// The payment information will be persisted under this prefix.
 pub(crate) const PAYMENT_INFO_PERSISTENCE_PREFIX: &str = "payments";
 
-pub(crate) struct PaymentInfoStorage<K: Deref>
+pub(crate) struct PaymentInfoStorage<K: Deref + Clone>
 where
-	K::Target: KVStorePersister,
+	K::Target: KVStorePersister + KVStoreUnpersister,
 {
 	payments: Mutex<HashMap<PaymentHash, PaymentInfo>>,
 	persister: K,
 }
 
-impl<K: Deref> PaymentInfoStorage<K>
+impl<K: Deref + Clone> PaymentInfoStorage<K>
 where
-	K::Target: KVStorePersister,
+	K::Target: KVStorePersister + KVStoreUnpersister,
 {
-	pub(crate) fn from_payments(mut payments: Vec<PaymentInfo>, persister: K) -> Self {
+	pub(crate) fn new(persister: K) -> Self {
+		let payments = Mutex::new(HashMap::new());
+		Self { payments, persister }
+	}
+
+	pub(crate) fn from_payments(payments: Vec<PaymentInfo>, persister: K) -> Self {
 		let payments = Mutex::new(HashMap::from_iter(
-			payments.drain(..).map(|payment_info| (payment_info.payment_hash, payment_info)),
+			payments.into_iter().map(|payment_info| (payment_info.payment_hash, payment_info)),
 		));
 		Self { payments, persister }
 	}
@@ -106,9 +113,20 @@ where
 		return Ok(());
 	}
 
-	// TODO: Need an `unpersist` method for this?
-	//pub(crate) fn remove_payment(&self, payment_hash: &PaymentHash) -> Result<(), Error> {
-	//}
+	pub(crate) fn lock(&self) -> Result<PaymentInfoGuard<K>, ()> {
+		let locked_store = self.payments.lock().map_err(|_| ())?;
+		Ok(PaymentInfoGuard::new(locked_store, self.persister.clone()))
+	}
+
+	pub(crate) fn remove(&self, payment_hash: &PaymentHash) -> Result<(), Error> {
+		let key = format!(
+			"{}/{}",
+			PAYMENT_INFO_PERSISTENCE_PREFIX,
+			hex_utils::to_string(&payment_hash.0)
+		);
+		self.persister.unpersist(&key).map_err(|_| Error::PersistenceFailed)?;
+		Ok(())
+	}
 
 	pub(crate) fn get(&self, payment_hash: &PaymentHash) -> Option<PaymentInfo> {
 		self.payments.lock().unwrap().get(payment_hash).cloned()
@@ -136,3 +154,79 @@ where
 		Ok(())
 	}
 }
+
+pub(crate) struct PaymentInfoGuard<'a, K: Deref>
+where
+	K::Target: KVStorePersister + KVStoreUnpersister,
+{
+	inner: MutexGuard<'a, HashMap<PaymentHash, PaymentInfo>>,
+	touched_keys: HashSet<PaymentHash>,
+	persister: K,
+}
+
+impl<'a, K: Deref> PaymentInfoGuard<'a, K>
+where
+	K::Target: KVStorePersister + KVStoreUnpersister,
+{
+	pub fn new(inner: MutexGuard<'a, HashMap<PaymentHash, PaymentInfo>>, persister: K) -> Self {
+		let touched_keys = HashSet::new();
+		Self { inner, touched_keys, persister }
+	}
+
+	pub fn entry(
+		&mut self, payment_hash: PaymentHash,
+	) -> hash_map::Entry<PaymentHash, PaymentInfo> {
+		self.touched_keys.insert(payment_hash);
+		self.inner.entry(payment_hash)
+	}
+}
+
+impl<'a, K: Deref> Drop for PaymentInfoGuard<'a, K>
+where
+	K::Target: KVStorePersister + KVStoreUnpersister,
+{
+	fn drop(&mut self) {
+		for key in self.touched_keys.iter() {
+			let store_key =
+				format!("{}/{}", PAYMENT_INFO_PERSISTENCE_PREFIX, hex_utils::to_string(&key.0));
+
+			match self.inner.entry(*key) {
+				hash_map::Entry::Vacant(_) => {
+					self.persister.unpersist(&store_key).expect("Persistence failed");
+				}
+				hash_map::Entry::Occupied(e) => {
+					self.persister.persist(&store_key, e.get()).expect("Persistence failed");
+				}
+			};
+		}
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use crate::tests::test_utils::TestPersister;
+	use std::sync::Arc;
+
+	#[test]
+	fn persistence_guard_persists_on_drop() {
+		let persister = Arc::new(TestPersister::new());
+		let payment_info_store = PaymentInfoStorage::new(Arc::clone(&persister));
+
+		let payment_hash = PaymentHash([42u8; 32]);
+		assert!(!payment_info_store.contains(&payment_hash));
+
+		let payment_info = PaymentInfo {
+			payment_hash,
+			preimage: None,
+			secret: None,
+			amount_msat: None,
+			direction: PaymentDirection::Inbound,
+			status: PaymentStatus::Pending,
+		};
+
+		assert!(!persister.get_and_clear_did_persist());
+		payment_info_store.lock().unwrap().entry(payment_hash).or_insert(payment_info);
+		assert!(persister.get_and_clear_did_persist());
+	}
+}

src/tests/test_utils.rs

@@ -1,3 +1,4 @@
+use crate::io_utils::KVStoreUnpersister;
 use lightning::util::persist::KVStorePersister;
 use lightning::util::ser::Writeable;
 
@@ -53,3 +54,11 @@ impl KVStorePersister for TestPersister {
 		Ok(())
 	}
 }
+
+impl KVStoreUnpersister for TestPersister {
+	fn unpersist(&self, key: &str) -> std::io::Result<bool> {
+		let mut persisted_bytes_lock = self.persisted_bytes.lock().unwrap();
+		self.did_persist.store(true, Ordering::SeqCst);
+		Ok(persisted_bytes_lock.remove(key).is_some())
+	}
+}