diff --git a/packages/wasm-privacy-coin/README.md b/packages/wasm-privacy-coin/README.md index c335c100081..2b5bd8f894f 100644 --- a/packages/wasm-privacy-coin/README.md +++ b/packages/wasm-privacy-coin/README.md @@ -130,7 +130,9 @@ memory and reads a `Response` proto from the `LAST_RESULT` buffer. **Persistence.** `save()` / `fromState()` use serde JSON internally (the `PersistedShardTreeState` format). This is the on-disk/DB format, not the Java↔WASM -wire format. The Java layer sees it as opaque `TreeState` bytes. +wire format. The Java layer sees it as opaque `TreeState` bytes. `PersistedShardTreeState` +carries `max_checkpoints`, so a restored tree keeps its original checkpoint retention +capacity across `save()`/`fromState()` round-trips. **One instance = one tree.** Each `ShieldedMerkleTree` owns a dedicated Chicory `Instance` with its own WASM linear memory. Two instances never share state. @@ -150,10 +152,11 @@ Implements `AutoCloseable`. Always use in try-with-resources. #### Factory methods -| Method | Description | -| -------------------------------------------------------- | ------------------------------------------------------------------------------ | -| `static fromFrontier(byte[] frontier, long blockHeight)` | Initialize from a CommitmentTree v0 frontier (raw bytes from `z_gettreestate`) | -| `static fromState(TreeState state)` | Restore from a `TreeState` previously returned by `save()` | +| Method | Description | +| -------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------- | +| `static fromFrontier(byte[] frontier, long blockHeight)` | Initialize from a CommitmentTree v0 frontier (raw bytes from `z_gettreestate`); convenience overload, `maxCheckpoints` defaults to 100 | +| `static fromFrontier(byte[] frontier, long blockHeight, Integer maxCheckpoints)` | Initialize from a frontier with a custom checkpoint retention capacity; `null` = default (100) | +| `static fromState(TreeState state)` | Restore from a `TreeState` previously returned by `save()` | #### Instance methods @@ -216,11 +219,11 @@ TreeState restored = TreeState.of(blob); // load from DB Immutable snapshot returned by `getInfo()`. -| Field | Type | Description | -| ----------------- | ------ | ------------------------------------------------------------------------ | -| `tipHeight` | `Long` | Most recently checkpointed block height; `null` if no block appended yet | -| `leafCount` | `long` | Total Orchard commitments appended across all blocks | -| `checkpointCount` | `int` | Number of checkpoints currently retained (max 100) | +| Field | Type | Description | +| ----------------- | ------ | ---------------------------------------------------------------------------------- | +| `tipHeight` | `Long` | Most recently checkpointed block height; `null` if no block appended yet | +| `leafCount` | `long` | Total Orchard commitments appended across all blocks | +| `checkpointCount` | `int` | Number of checkpoints currently retained, capped by `maxCheckpoints` (default 100) | --- @@ -255,12 +258,21 @@ try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromFrontier(frontier, blockHe } ``` +To use a non-default checkpoint retention capacity (e.g. a shallower reorg-depth +tolerance), pass an explicit `maxCheckpoints`: + +```java +try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromFrontier(frontier, blockHeight, 20)) { + // tree retains at most 20 checkpoints before pruning the oldest +} +``` + ### 2. Initialize from an empty state ```java TreeState emptyState = new TreeState( "{\"shards\":[],\"cap\":{\"type\":\"Nil\"},\"checkpoints\":[]," - + "\"tip_height\":null,\"leaf_count\":0}"); + + "\"tip_height\":null,\"leaf_count\":0,\"max_checkpoints\":100}"); try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(emptyState)) { tree.ping(); @@ -330,7 +342,7 @@ The wire format between Java and Rust is declared in `proto/privacy_coin.proto`. | Message | Export | Fields | | -------------------------- | ------------------------ | -------------------------------------------------------------------------------------- | -| `FromFrontierRequest` | `from_frontier` | `frontier: bytes`, `block_height: uint32` | +| `FromFrontierRequest` | `from_frontier` | `frontier: bytes`, `block_height: uint32`, `max_checkpoints: optional uint32` | | `AppendCommitmentsRequest` | `append_commitments` | `block_height: uint32`, `commitments: repeated bytes`, `expected_root: optional bytes` | | `TruncateRequest` | `truncate_to_checkpoint` | `block_height: uint32` | diff --git a/packages/wasm-privacy-coin/proto/privacy_coin.proto b/packages/wasm-privacy-coin/proto/privacy_coin.proto index 691fb10cb02..09cdb38f5a4 100644 --- a/packages/wasm-privacy-coin/proto/privacy_coin.proto +++ b/packages/wasm-privacy-coin/proto/privacy_coin.proto @@ -9,6 +9,7 @@ option java_multiple_files = true; message FromFrontierRequest { bytes frontier = 1; // raw CommitmentTree v0 bytes uint32 block_height = 2; + optional uint32 max_checkpoints = 3; // shardtree checkpoint capacity; absent = default (100) } message AppendCommitmentsRequest { diff --git a/packages/wasm-privacy-coin/src/lib.rs b/packages/wasm-privacy-coin/src/lib.rs index f0a08be359f..3aff44d018a 100644 --- a/packages/wasm-privacy-coin/src/lib.rs +++ b/packages/wasm-privacy-coin/src/lib.rs @@ -122,7 +122,11 @@ pub unsafe extern "C" fn from_frontier(ptr: *const u8, len: u32) -> i32 { let bytes = unsafe { std::slice::from_raw_parts(ptr, len as usize) }; match FromFrontierRequest::decode(bytes) { Err(e) => write_error("DECODE_ERROR", &e.to_string()), - Ok(req) => match zcash::tree::OwnedTree::from_frontier(&req.frontier, req.block_height) { + Ok(req) => match zcash::tree::OwnedTree::from_frontier( + &req.frontier, + req.block_height, + req.max_checkpoints.map(|v| v as usize), + ) { Ok(tree) => { TREE.with(|t| *t.borrow_mut() = Some(tree)); write_ok(); diff --git a/packages/wasm-privacy-coin/src/main/java/com/bitgo/wasm/privacycoin/zcash/ShieldedMerkleTree.java b/packages/wasm-privacy-coin/src/main/java/com/bitgo/wasm/privacycoin/zcash/ShieldedMerkleTree.java index 7b25ae40fa7..16d83e791a0 100644 --- a/packages/wasm-privacy-coin/src/main/java/com/bitgo/wasm/privacycoin/zcash/ShieldedMerkleTree.java +++ b/packages/wasm-privacy-coin/src/main/java/com/bitgo/wasm/privacycoin/zcash/ShieldedMerkleTree.java @@ -59,16 +59,36 @@ private static ShieldedMerkleTree create(Consumer init) { * @throws WasmException if the frontier is invalid */ public static ShieldedMerkleTree fromFrontier(byte[] frontier, long blockHeight) { + return fromFrontier(frontier, blockHeight, null); + } + + /** + * Initializes a new tree from a CommitmentTree v0 frontier + * (the {@code orchardTree} value from {@code z_gettreestate}), with a custom + * checkpoint retention capacity. + * + * @param frontier raw CommitmentTree v0 bytes + * @param blockHeight block height at which the frontier was captured (u32 range) + * @param maxCheckpoints number of checkpoints to retain before pruning the oldest; + * {@code null} to use the default (100) + * @return initialized tree instance + * @throws WasmException if the frontier is invalid + */ + public static ShieldedMerkleTree fromFrontier(byte[] frontier, long blockHeight, Integer maxCheckpoints) { Objects.requireNonNull(frontier, "frontier must not be null"); requireU32(blockHeight, "blockHeight"); + if (maxCheckpoints != null) { + requireU32(maxCheckpoints, "maxCheckpoints"); + } return create(bridge -> { - byte[] reqBytes = FromFrontierRequest.newBuilder() + FromFrontierRequest.Builder builder = FromFrontierRequest.newBuilder() .setFrontier(ByteString.copyFrom(frontier)) // safe: requireU32 guarantees blockHeight is in [0, 0xFFFF_FFFF] - .setBlockHeight((int) blockHeight) - .build() - .toByteArray(); - unwrapVoid(bridge.call("from_frontier", reqBytes)); + .setBlockHeight((int) blockHeight); + if (maxCheckpoints != null) { + builder.setMaxCheckpoints(maxCheckpoints); + } + unwrapVoid(bridge.call("from_frontier", builder.build().toByteArray())); }); } diff --git a/packages/wasm-privacy-coin/src/test/java/com/bitgo/wasm/privacycoin/zcash/ShieldedMerkleTreeTest.java b/packages/wasm-privacy-coin/src/test/java/com/bitgo/wasm/privacycoin/zcash/ShieldedMerkleTreeTest.java index ee24892892b..0ad5a5e4e26 100644 --- a/packages/wasm-privacy-coin/src/test/java/com/bitgo/wasm/privacycoin/zcash/ShieldedMerkleTreeTest.java +++ b/packages/wasm-privacy-coin/src/test/java/com/bitgo/wasm/privacycoin/zcash/ShieldedMerkleTreeTest.java @@ -24,7 +24,7 @@ class ShieldedMerkleTreeTest { */ private static final TreeState EMPTY_STATE = new TreeState( "{\"shards\":[],\"cap\":{\"type\":\"Nil\"},\"checkpoints\":[]," - + "\"tip_height\":null,\"leaf_count\":0}"); + + "\"tip_height\":null,\"leaf_count\":0,\"max_checkpoints\":100}"); /** * CommitmentTree v0 frontier encoding for a single-leaf tree. @@ -102,6 +102,45 @@ void fromFrontier_setsCheckpointCountToOne() { } } + @Test + void fromFrontier_withMaxCheckpoints_capsCheckpointRetention() { + try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromFrontier(FRONTIER, 1L, 2)) { + tree.appendCommitments(2L, Collections.emptyList(), List.of(), null); + tree.appendCommitments(3L, Collections.emptyList(), List.of(), null); + tree.appendCommitments(4L, Collections.emptyList(), List.of(), null); + + MerkleTreeInfo info = tree.getInfo(); + assertEquals(2, info.checkpointCount); + + WasmException ex = assertThrows(WasmException.class, () -> tree.truncateToCheckpoint(1L)); + assertEquals("CHECKPOINT_NOT_FOUND", ex.getErrorCode()); + } + } + + @Test + void fromFrontier_negativeMaxCheckpoints_throwsIllegalArgumentException() { + assertThrows(IllegalArgumentException.class, () -> + ShieldedMerkleTree.fromFrontier(FRONTIER, 1L, -1)); + } + + @Test + void saveAndLoad_preservesMaxCheckpoints() { + TreeState savedState; + try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromFrontier(FRONTIER, 1L, 2)) { + tree.appendCommitments(2L, Collections.emptyList(), List.of(), null); + tree.appendCommitments(3L, Collections.emptyList(), List.of(), null); + savedState = tree.save(); + } + try (ShieldedMerkleTree restored = ShieldedMerkleTree.fromState(savedState)) { + restored.appendCommitments(4L, Collections.emptyList(), List.of(), null); + MerkleTreeInfo info = restored.getInfo(); + assertEquals(2, info.checkpointCount); + + WasmException ex = assertThrows(WasmException.class, () -> restored.truncateToCheckpoint(1L)); + assertEquals("CHECKPOINT_NOT_FOUND", ex.getErrorCode()); + } + } + // ------------------------------------------------------------------------- // fromState // ------------------------------------------------------------------------- diff --git a/packages/wasm-privacy-coin/src/zcash/tree.rs b/packages/wasm-privacy-coin/src/zcash/tree.rs index 59a80d14ca0..eb038a96bfb 100644 --- a/packages/wasm-privacy-coin/src/zcash/tree.rs +++ b/packages/wasm-privacy-coin/src/zcash/tree.rs @@ -62,6 +62,7 @@ pub struct PersistedShardTreeState { pub checkpoints: Vec, pub tip_height: Option, pub leaf_count: u64, + pub max_checkpoints: usize, } // --------------------------------------------------------------------------- @@ -121,6 +122,7 @@ pub fn extract_state( tree: &ShieldedShardTree, tip_height: Option, leaf_count: u64, + max_checkpoints: usize, ) -> Result { let store = tree.store(); @@ -175,6 +177,7 @@ pub fn extract_state( checkpoints, tip_height, leaf_count, + max_checkpoints, }) } @@ -212,7 +215,7 @@ pub fn restore_state(state: &PersistedShardTreeState) -> Result, leaf_count: u64, + max_checkpoints: usize, } impl OwnedTree { @@ -334,11 +338,17 @@ impl OwnedTree { tree, tip_height: persisted.tip_height, leaf_count: persisted.leaf_count, + max_checkpoints: persisted.max_checkpoints, }) } /// Initialize from a CommitmentTree v0 frontier (raw bytes, not hex-encoded). - pub fn from_frontier(frontier: &[u8], block_height: u32) -> Result { + pub fn from_frontier( + frontier: &[u8], + block_height: u32, + max_checkpoints: Option, + ) -> Result { + let max_checkpoints = max_checkpoints.unwrap_or(MAX_CHECKPOINTS); use incrementalmerkletree::frontier::NonEmptyFrontier; let mut offset = 0; @@ -395,7 +405,7 @@ impl OwnedTree { let leaf_count = u64::from(nef.position()) + 1; - let mut tree = ShardTree::new(MemoryShardStore::empty(), MAX_CHECKPOINTS); + let mut tree = ShardTree::new(MemoryShardStore::empty(), max_checkpoints); tree.insert_frontier_nodes( nef, Retention::Checkpoint { @@ -409,12 +419,18 @@ impl OwnedTree { tree, tip_height: Some(block_height), leaf_count, + max_checkpoints, }) } /// Serialize the tree state to bytes (UTF-8 JSON of `PersistedShardTreeState`). pub fn save(&self) -> Result, String> { - let state = extract_state(&self.tree, self.tip_height, self.leaf_count)?; + let state = extract_state( + &self.tree, + self.tip_height, + self.leaf_count, + self.max_checkpoints, + )?; serde_json::to_vec(&state).map_err(|e| format!("JSON serialize error: {}", e)) } @@ -540,7 +556,7 @@ mod tests { const F_CHECKPOINT_MARKED: u8 = 3; fn empty_tree() -> OwnedTree { - let json = r#"{"shards":[],"cap":{"type":"Nil"},"checkpoints":[],"tip_height":null,"leaf_count":0}"#; + let json = r#"{"shards":[],"cap":{"type":"Nil"},"checkpoints":[],"tip_height":null,"leaf_count":0,"max_checkpoints":100}"#; OwnedTree::from_state(json.as_bytes()).expect("empty state") } @@ -582,7 +598,13 @@ mod tests { None, ) .unwrap(); - let state = extract_state(&tree.tree, tree.tip_height, tree.leaf_count).unwrap(); + let state = extract_state( + &tree.tree, + tree.tip_height, + tree.leaf_count, + tree.max_checkpoints, + ) + .unwrap(); assert_eq!( find_f_in_state(&state, cmx1_hex), @@ -603,7 +625,13 @@ mod tests { let mut tree = empty_tree(); tree.append_commitments(1, vec![cmx(1)], vec![false], None) .unwrap(); - let state = extract_state(&tree.tree, tree.tip_height, tree.leaf_count).unwrap(); + let state = extract_state( + &tree.tree, + tree.tip_height, + tree.leaf_count, + tree.max_checkpoints, + ) + .unwrap(); assert_eq!( find_f_in_state(&state, cmx_hex), @@ -619,4 +647,53 @@ mod tests { tree.append_commitments(1, vec![cmx(1), cmx(2)], vec![true, false, true], None); assert!(result.is_err()); } + + // ------------------------------------------------------------------------- + // max_checkpoints + // ------------------------------------------------------------------------- + + fn frontier_bytes() -> Vec { + hex::decode("0101000000000000000000000000000000000000000000000000000000000000000000") + .unwrap() + } + + #[test] + fn from_frontier_defaults_max_checkpoints_when_not_specified() { + let tree = OwnedTree::from_frontier(&frontier_bytes(), 1, None).unwrap(); + assert_eq!(tree.max_checkpoints, MAX_CHECKPOINTS); + } + + #[test] + fn from_frontier_uses_custom_max_checkpoints_when_specified() { + let tree = OwnedTree::from_frontier(&frontier_bytes(), 1, Some(5)).unwrap(); + assert_eq!(tree.max_checkpoints, 5); + } + + #[test] + fn max_checkpoints_enforced_evicts_oldest_checkpoint() { + let mut tree = OwnedTree::from_frontier(&frontier_bytes(), 1, Some(2)).unwrap(); + tree.append_commitments(2, vec![], vec![], None).unwrap(); + tree.append_commitments(3, vec![], vec![], None).unwrap(); + tree.append_commitments(4, vec![], vec![], None).unwrap(); + + let (_, _, checkpoint_count) = tree.get_info().unwrap(); + assert_eq!( + checkpoint_count, 2, + "checkpoint_count should be capped at 2" + ); + + let result = tree.truncate_to_checkpoint(1); + assert!( + result.is_err(), + "oldest checkpoint should have been evicted" + ); + } + + #[test] + fn save_round_trip_preserves_max_checkpoints() { + let tree = OwnedTree::from_frontier(&frontier_bytes(), 1, Some(7)).unwrap(); + let bytes = tree.save().unwrap(); + let restored = OwnedTree::from_state(&bytes).unwrap(); + assert_eq!(restored.max_checkpoints, 7); + } }