darkpool_client/
persistence.rs

1//! Save/load wallet state (UTXOs, Merkle leaves, sync checkpoint) to JSON.
2
3use ethers::types::U256;
4use serde::{Deserialize, Serialize};
5use std::path::Path;
6use tracing::{debug, info};
7
8use crate::merkle_tree::LocalMerkleTree;
9use crate::utxo_store::{OwnedNote, UtxoStore};
10
11#[derive(Debug, Serialize, Deserialize)]
12struct UtxoSnapshot {
13    notes: Vec<(String, OwnedNote)>,
14    spent_nullifiers: Vec<String>,
15    nullifier_to_commitment: Vec<(String, String)>,
16}
17
18#[derive(Debug, Serialize, Deserialize)]
19pub struct WalletState {
20    version: u32,
21    pub last_synced_block: u64,
22    utxo_snapshot: UtxoSnapshot,
23    merkle_leaves: Vec<String>,
24}
25
26const CURRENT_VERSION: u32 = 1;
27
28fn u256_to_hex(v: &U256) -> String {
29    format!("{v:#066x}")
30}
31
32fn hex_to_u256(s: &str) -> Result<U256, PersistenceError> {
33    U256::from_str_radix(s.strip_prefix("0x").unwrap_or(s), 16)
34        .map_err(|e| PersistenceError::Deserialization(format!("invalid U256 hex '{s}': {e}")))
35}
36
37#[derive(Debug, thiserror::Error)]
38pub enum PersistenceError {
39    #[error("IO error: {0}")]
40    Io(#[from] std::io::Error),
41    #[error("JSON serialization error: {0}")]
42    Serialization(String),
43    #[error("JSON deserialization error: {0}")]
44    Deserialization(String),
45    #[error("Version mismatch: file version {file} > supported {supported}")]
46    VersionMismatch { file: u32, supported: u32 },
47}
48
49impl WalletState {
50    #[must_use]
51    pub fn capture(utxos: &UtxoStore, tree: &LocalMerkleTree, last_synced_block: u64) -> Self {
52        let notes: Vec<_> = utxos
53            .get_unspent()
54            .into_iter()
55            .map(|n| (u256_to_hex(&n.commitment), n.clone()))
56            .collect();
57
58        let spent_nullifiers: Vec<_> = utxos.spent_nullifiers_iter().map(u256_to_hex).collect();
59
60        let nullifier_to_commitment: Vec<_> = utxos
61            .nullifier_map_iter()
62            .map(|(k, v)| (u256_to_hex(k), u256_to_hex(v)))
63            .collect();
64
65        let merkle_leaves: Vec<_> = tree.leaves().iter().map(u256_to_hex).collect();
66
67        Self {
68            version: CURRENT_VERSION,
69            last_synced_block,
70            utxo_snapshot: UtxoSnapshot {
71                notes,
72                spent_nullifiers,
73                nullifier_to_commitment,
74            },
75            merkle_leaves,
76        }
77    }
78
79    pub fn restore(self) -> Result<(UtxoStore, LocalMerkleTree, u64), PersistenceError> {
80        if self.version > CURRENT_VERSION {
81            return Err(PersistenceError::VersionMismatch {
82                file: self.version,
83                supported: CURRENT_VERSION,
84            });
85        }
86
87        let mut utxos = UtxoStore::new();
88        for (commitment_hex, note) in &self.utxo_snapshot.notes {
89            let _commitment = hex_to_u256(commitment_hex)?;
90            // Find the nullifier_hash for this commitment from the mapping
91            let nullifier_hash = self
92                .utxo_snapshot
93                .nullifier_to_commitment
94                .iter()
95                .find(|(_, c)| c == commitment_hex)
96                .map(|(n, _)| hex_to_u256(n))
97                .transpose()?
98                .unwrap_or(U256::zero());
99            utxos.add_note(note.clone(), nullifier_hash);
100        }
101        for nullifier_hex in &self.utxo_snapshot.spent_nullifiers {
102            let nullifier_hash = hex_to_u256(nullifier_hex)?;
103            utxos.mark_spent(nullifier_hash);
104        }
105
106        let mut tree = LocalMerkleTree::new();
107        let leaves: Vec<U256> = self
108            .merkle_leaves
109            .iter()
110            .map(|h| hex_to_u256(h))
111            .collect::<Result<Vec<_>, _>>()?;
112        tree.load_from_leaves(&leaves);
113
114        info!(
115            notes = utxos.count(),
116            leaves = tree.size(),
117            block = self.last_synced_block,
118            "Wallet state restored from snapshot"
119        );
120
121        Ok((utxos, tree, self.last_synced_block))
122    }
123
124    pub fn save_to_file(&self, path: &Path) -> Result<(), PersistenceError> {
125        let json = serde_json::to_string_pretty(self)
126            .map_err(|e| PersistenceError::Serialization(e.to_string()))?;
127
128        let tmp_path = path.with_extension("tmp");
129        std::fs::write(&tmp_path, json)?;
130        std::fs::rename(&tmp_path, path)?;
131
132        debug!(
133            path = %path.display(),
134            notes = self.utxo_snapshot.notes.len(),
135            leaves = self.merkle_leaves.len(),
136            block = self.last_synced_block,
137            "Wallet state saved"
138        );
139        Ok(())
140    }
141
142    pub fn load_from_file(path: &Path) -> Result<Self, PersistenceError> {
143        let json = std::fs::read_to_string(path)?;
144        let state: Self = serde_json::from_str(&json)
145            .map_err(|e| PersistenceError::Deserialization(e.to_string()))?;
146
147        if state.version > CURRENT_VERSION {
148            return Err(PersistenceError::VersionMismatch {
149                file: state.version,
150                supported: CURRENT_VERSION,
151            });
152        }
153
154        info!(
155            path = %path.display(),
156            version = state.version,
157            block = state.last_synced_block,
158            "Wallet state loaded"
159        );
160        Ok(state)
161    }
162}
163
164pub fn save_wallet_state(
165    path: &Path,
166    utxos: &UtxoStore,
167    tree: &LocalMerkleTree,
168    last_synced_block: u64,
169) -> Result<(), PersistenceError> {
170    let state = WalletState::capture(utxos, tree, last_synced_block);
171    state.save_to_file(path)
172}
173
174/// Returns `None` if file doesn't exist.
175pub fn load_wallet_state(
176    path: &Path,
177) -> Result<Option<(UtxoStore, LocalMerkleTree, u64)>, PersistenceError> {
178    if !path.exists() {
179        debug!(path = %path.display(), "No wallet state file found, starting fresh");
180        return Ok(None);
181    }
182
183    let state = WalletState::load_from_file(path)?;
184    let (utxos, tree, block) = state.restore()?;
185    Ok(Some((utxos, tree, block)))
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191    use crate::proof_inputs::NotePlaintext;
192
193    fn test_note(value: u64, commitment: U256, leaf_index: u64) -> OwnedNote {
194        OwnedNote {
195            plaintext: NotePlaintext {
196                value: U256::from(value),
197                asset_id: U256::from(1),
198                secret: U256::from(42),
199                nullifier: U256::from(leaf_index + 1000),
200                timelock: U256::zero(),
201                hashlock: U256::zero(),
202            },
203            commitment,
204            leaf_index,
205            spending_secret: U256::from(leaf_index + 2000),
206            is_transfer: false,
207            received_block: 100,
208        }
209    }
210
211    #[test]
212    fn test_roundtrip_empty_state() {
213        let utxos = UtxoStore::new();
214        let tree = LocalMerkleTree::new();
215
216        let state = WalletState::capture(&utxos, &tree, 0);
217        let (restored_utxos, restored_tree, block) = state.restore().unwrap();
218
219        assert_eq!(restored_utxos.count(), 0);
220        assert_eq!(restored_tree.size(), 0);
221        assert_eq!(block, 0);
222    }
223
224    #[test]
225    fn test_roundtrip_with_notes_and_leaves() {
226        let mut utxos = UtxoStore::new();
227        let mut tree = LocalMerkleTree::new();
228
229        // Add some notes
230        let note1 = test_note(100, U256::from(1), 0);
231        let note2 = test_note(200, U256::from(2), 1);
232        utxos.add_note(note1, U256::from(5001));
233        utxos.add_note(note2, U256::from(5002));
234
235        // Add some leaves
236        tree.insert(U256::from(1));
237        tree.insert(U256::from(2));
238        tree.insert(U256::from(3));
239
240        let original_root = tree.root();
241        let original_balance = utxos.count();
242
243        // Capture and restore
244        let state = WalletState::capture(&utxos, &tree, 42);
245        let (restored_utxos, restored_tree, block) = state.restore().unwrap();
246
247        assert_eq!(restored_utxos.count(), original_balance);
248        assert_eq!(restored_tree.root(), original_root);
249        assert_eq!(restored_tree.size(), 3);
250        assert_eq!(block, 42);
251    }
252
253    #[test]
254    fn test_roundtrip_with_spent_nullifiers() {
255        let mut utxos = UtxoStore::new();
256        let tree = LocalMerkleTree::new();
257
258        let note = test_note(500, U256::from(10), 0);
259        let nullifier_hash = U256::from(9999);
260        utxos.add_note(note, nullifier_hash);
261        utxos.mark_spent(nullifier_hash);
262
263        assert_eq!(utxos.count(), 0);
264        assert!(utxos.is_spent(&nullifier_hash));
265
266        let state = WalletState::capture(&utxos, &tree, 100);
267        let (restored_utxos, _, _) = state.restore().unwrap();
268
269        assert_eq!(restored_utxos.count(), 0);
270        assert!(restored_utxos.is_spent(&nullifier_hash));
271    }
272
273    #[test]
274    fn test_file_roundtrip() {
275        let mut utxos = UtxoStore::new();
276        let mut tree = LocalMerkleTree::new();
277
278        utxos.add_note(test_note(100, U256::from(1), 0), U256::from(5001));
279        tree.insert(U256::from(1));
280
281        let dir = std::env::temp_dir().join("nox_test_persistence");
282        std::fs::create_dir_all(&dir).unwrap();
283        let path = dir.join("test_state.json");
284
285        // Save
286        save_wallet_state(&path, &utxos, &tree, 50).unwrap();
287        assert!(path.exists());
288
289        // Load
290        let loaded = load_wallet_state(&path).unwrap();
291        assert!(loaded.is_some());
292        let (r_utxos, r_tree, r_block) = loaded.unwrap();
293        assert_eq!(r_utxos.count(), 1);
294        assert_eq!(r_tree.root(), tree.root());
295        assert_eq!(r_block, 50);
296
297        // Cleanup
298        std::fs::remove_dir_all(&dir).ok();
299    }
300
301    #[test]
302    fn test_load_nonexistent_returns_none() {
303        let path = Path::new("/tmp/nox_test_nonexistent_state.json");
304        let result = load_wallet_state(path).unwrap();
305        assert!(result.is_none());
306    }
307
308    #[test]
309    fn test_version_mismatch() {
310        let json = r#"{"version": 999, "last_synced_block": 0, "utxo_snapshot": {"notes": [], "spent_nullifiers": [], "nullifier_to_commitment": []}, "merkle_leaves": []}"#;
311        let state: WalletState = serde_json::from_str(json).unwrap();
312        let result = state.restore();
313        assert!(result.is_err());
314        assert!(matches!(
315            result.unwrap_err(),
316            PersistenceError::VersionMismatch { .. }
317        ));
318    }
319
320    #[test]
321    fn test_pending_spends_not_persisted() {
322        let mut utxos = UtxoStore::new();
323        let tree = LocalMerkleTree::new();
324
325        let note = test_note(100, U256::from(1), 0);
326        utxos.add_note(note, U256::from(5001));
327        utxos.mark_pending_spend(&U256::from(1));
328        assert!(utxos.is_pending_spend(&U256::from(1)));
329
330        let state = WalletState::capture(&utxos, &tree, 10);
331        let (restored_utxos, _, _) = state.restore().unwrap();
332
333        assert!(!restored_utxos.is_pending_spend(&U256::from(1)));
334        assert_eq!(restored_utxos.count(), 1);
335    }
336}
darkpool_client/persistence.rs

darkpool_client/
persistence.rs
