From c97baf086f6b4a28c0646fa75b9aa5fce1a4d8b7 Mon Sep 17 00:00:00 2001
From: Steven Roose <steven@stevenroose.org>
Date: Tue, 6 Feb 2024 15:50:21 +0000
Subject: [PATCH] Don't accidentally rediscover an already forfeited vtxo

---
 noah/src/database.rs | 11 +++++++++++
 noah/src/lib.rs      | 10 ++++++++++
 2 files changed, 21 insertions(+)

diff --git a/noah/src/database.rs b/noah/src/database.rs
index c987363..e53f065 100644
--- a/noah/src/database.rs
+++ b/noah/src/database.rs
@@ -12,6 +12,7 @@ use crate::exit;
 // Trees
 
 const VTXO_TREE: &str = "noah_vtxos";
+const FORFEIT_VTXO_TREE: &str = "noah_forfeited_vtxos";
 
 // Top-level entries
 
@@ -93,4 +94,14 @@ impl Db {
 		self.db.insert(LAST_ARK_SYNC_HEIGHT, height.to_be_bytes().to_vec())?;
 		Ok(())
 	}
+
+	pub fn store_forfeited_vtxo(&self, id: VtxoId, height: u32) -> anyhow::Result<()> {
+		self.db.open_tree(VTXO_TREE)?.insert(id, height.to_be_bytes().to_vec())?;
+		Ok(())
+	}
+
+	pub fn has_forfeited_vtxo(&self, id: VtxoId) -> anyhow::Result<bool> {
+		Ok(self.db.open_tree(VTXO_TREE)?.get(id)?.is_some())
+	}
+	//TODO(stevenroose) regularly prune forfeit vtxos based on height
 }
diff --git a/noah/src/lib.rs b/noah/src/lib.rs
index 632910c..356800b 100644
--- a/noah/src/lib.rs
+++ b/noah/src/lib.rs
@@ -226,6 +226,12 @@ impl Wallet {
 			leaf_idx: leaf_idx,
 			exit_branch: exit_branch,
 		};
+
+		if self.db.has_forfeited_vtxo(vtxo.id())? {
+			debug!("Not adding vtxo {} because we previously forfeited it", vtxo.id());
+			return Ok(());
+		}
+
 		if self.db.get_vtxo(vtxo.id()).ok().flatten().is_none() {
 			debug!("Storing new vtxo {} with value {}", vtxo.id(), vtxo.spec().amount);
 			self.db.store_vtxo(vtxo).context("failed to store vtxo")?;
@@ -290,6 +296,7 @@ impl Wallet {
 
 	pub async fn send_payment(&mut self, destination: Destination) -> anyhow::Result<()> {
 		self.sync_ark().await.context("ark sync error")?;
+		let current_height = self.onchain.tip()?.0;
 
 		//TODO(stevenroose) impl key derivation
 		let vtxo_key = self.vtxo_seed.to_keypair(&SECP);
@@ -297,6 +304,7 @@ impl Wallet {
 		// Prepare the payment.
 		let input_vtxos = self.db.get_all_vtxos()?;
 		let vtxo_ids = input_vtxos.iter().map(|v| v.id()).collect::<HashSet<_>>();
+		debug!("Spending vtxos: {:?}", vtxo_ids);
 		let change = {
 			let sum = input_vtxos.iter().map(|v| v.amount()).sum::<Amount>();
 			if sum < destination.amount {
@@ -515,6 +523,8 @@ impl Wallet {
 			// And remove the input vtxos.
 			for v in input_vtxos {
 				self.db.remove_vtxo(v.id()).context("failed to drop input vtxo")?;
+				self.db.store_forfeited_vtxo(v.id(), current_height)
+					.context("failed to store forfeited vtxo")?;
 			}
 
 			info!("Finished payment");