From 4905aac22f01af50ebdf144db0227682d3bf7438 Mon Sep 17 00:00:00 2001 From: Ranjna Ganesh Ram Date: Thu, 16 Jul 2026 17:52:06 +0530 Subject: [PATCH] feat(wasm-privacy-coin): thread owned flag into shardtree retention Per-commitment ownership booleans are now threaded from the proto wire format into the Rust shardtree append path. Owned leaves are assigned Retention::Marked so the shardtree preserves their witness paths; non-owned leaves remain Retention::Ephemeral (or Checkpoint for the last commitment in a block). Ticket: CSHLD-1208 --- packages/wasm-privacy-coin/README.md | 29 +++-- packages/wasm-privacy-coin/pom.xml | 6 + .../proto/privacy_coin.proto | 1 + packages/wasm-privacy-coin/src/lib.rs | 3 +- .../privacycoin/zcash/ShieldedMerkleTree.java | 12 +- .../zcash/ShieldedMerkleTreeTest.java | 71 ++++++------ packages/wasm-privacy-coin/src/zcash/tree.rs | 108 +++++++++++++++++- 7 files changed, 172 insertions(+), 58 deletions(-) diff --git a/packages/wasm-privacy-coin/README.md b/packages/wasm-privacy-coin/README.md index 13c94c59b4b..c335c100081 100644 --- a/packages/wasm-privacy-coin/README.md +++ b/packages/wasm-privacy-coin/README.md @@ -157,15 +157,14 @@ Implements `AutoCloseable`. Always use in try-with-resources. #### Instance methods -| Method | Returns | Description | -| ------------------------------------------------------------------------------------------------------ | ---------------- | --------------------------------------------------------------- | -| `ping()` | `void` | Verifies the WASM module is alive | -| `appendCommitments(long blockHeight, List commitments)` | `ShieldedRoot` | Append cmx values, checkpoint the tree, return the new root | -| `appendCommitments(long blockHeight, List commitments, ShieldedRoot expectedRoot)` | `ShieldedRoot` | Same, with optional root verification | -| `truncateToCheckpoint(long blockHeight)` | `ShieldedRoot` | Roll back to a prior checkpoint, return the root at that height | -| `save()` | `TreeState` | Serialize tree state for persistence | -| `getInfo()` | `MerkleTreeInfo` | Return tip height, leaf count, checkpoint count | -| `close()` | `void` | Drop the in-WASM tree and release the Chicory instance | +| Method | Returns | Description | +| --------------------------------------------------------------------------------------------------------------------------- | ---------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `ping()` | `void` | Verifies the WASM module is alive | +| `appendCommitments(long blockHeight, List commitments, List owned, ShieldedRoot expectedRoot)` | `ShieldedRoot` | Append cmx values, checkpoint the tree; `owned` marks which commitments belong to this wallet (pass `List.of()` if unused); `expectedRoot` is optional root verification (pass `null` to skip) | +| `truncateToCheckpoint(long blockHeight)` | `ShieldedRoot` | Roll back to a prior checkpoint, return the root at that height | +| `save()` | `TreeState` | Serialize tree state for persistence | +| `getInfo()` | `MerkleTreeInfo` | Return tip height, leaf count, checkpoint count | +| `close()` | `void` | Drop the in-WASM tree and release the Chicory instance | **`blockHeight`** must be in the range `[0, 4_294_967_295]` (Rust `u32`). Passing a negative value or a value above `0xFFFFFFFFL` throws `IllegalArgumentException` @@ -275,9 +274,9 @@ ShieldedCommitment cmx = ShieldedCommitment.of(HexFormat.of().parseHex( "0100000000000000000000000000000000000000000000000000000000000000")); try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(savedState)) { - ShieldedRoot root = tree.appendCommitments(2_500_001L, List.of(cmx)); + ShieldedRoot root = tree.appendCommitments(2_500_001L, List.of(cmx), List.of(), null); // Empty block — still creates a checkpoint - tree.appendCommitments(2_500_002L, Collections.emptyList()); + tree.appendCommitments(2_500_002L, Collections.emptyList(), List.of(), null); } ``` @@ -287,7 +286,7 @@ try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(savedState)) { ShieldedRoot expected = ShieldedRoot.of(expectedRootBytes); try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(savedState)) { - ShieldedRoot root = tree.appendCommitments(2_500_001L, cmxList, expected); + ShieldedRoot root = tree.appendCommitments(2_500_001L, cmxList, List.of(), expected); // root.equals(expected) is guaranteed } ``` @@ -299,7 +298,7 @@ try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(savedState)) { ```java byte[] snapshot; try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(emptyState)) { - tree.appendCommitments(100L, List.of(cmx)); + tree.appendCommitments(100L, List.of(cmx), List.of(), null); snapshot = tree.save().bytes(); } @@ -313,8 +312,8 @@ try (ShieldedMerkleTree restored = ShieldedMerkleTree.fromState(TreeState.of(sna ```java try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(savedState)) { - ShieldedRoot root100 = tree.appendCommitments(100L, List.of(cmx)); - tree.appendCommitments(101L, List.of(cmx)); + ShieldedRoot root100 = tree.appendCommitments(100L, List.of(cmx), List.of(), null); + tree.appendCommitments(101L, List.of(cmx), List.of(), null); ShieldedRoot restoredRoot = tree.truncateToCheckpoint(100L); // restoredRoot.equals(root100) diff --git a/packages/wasm-privacy-coin/pom.xml b/packages/wasm-privacy-coin/pom.xml index 0d4f7cd3550..4e8c71fbaa2 100644 --- a/packages/wasm-privacy-coin/pom.xml +++ b/packages/wasm-privacy-coin/pom.xml @@ -41,6 +41,12 @@ 5.10.3 test + + com.fasterxml.jackson.core + jackson-databind + 2.17.0 + test + diff --git a/packages/wasm-privacy-coin/proto/privacy_coin.proto b/packages/wasm-privacy-coin/proto/privacy_coin.proto index 8268bfd1ab9..691fb10cb02 100644 --- a/packages/wasm-privacy-coin/proto/privacy_coin.proto +++ b/packages/wasm-privacy-coin/proto/privacy_coin.proto @@ -15,6 +15,7 @@ message AppendCommitmentsRequest { uint32 block_height = 1; repeated bytes commitments = 2; // each exactly 32 bytes (cmx) optional bytes expected_root = 3; // 32 bytes; absent = skip verification + repeated bool owned = 4; // per-commitment ownership flag; absent/empty = all not-owned } message TruncateRequest { diff --git a/packages/wasm-privacy-coin/src/lib.rs b/packages/wasm-privacy-coin/src/lib.rs index 56467db1377..f0a08be359f 100644 --- a/packages/wasm-privacy-coin/src/lib.rs +++ b/packages/wasm-privacy-coin/src/lib.rs @@ -186,8 +186,9 @@ pub unsafe extern "C" fn append_commitments(ptr: *const u8, len: u32) -> i32 { Some(tree) => { let commitments: Vec> = req.commitments.iter().map(|b| b.to_vec()).collect(); + let owned = req.owned; let exp = expected_root.as_deref(); - match tree.append_commitments(req.block_height, commitments, exp) { + match tree.append_commitments(req.block_height, commitments, owned, exp) { Ok(root) => write_ok_bytes(root), Err(e) => { let (code, msg) = split_error(&e); 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 8017c574195..596784b9347 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 @@ -100,12 +100,14 @@ public void ping() { * * @param blockHeight block height (u32 range) * @param commitments shielded note commitment values (cmx) for this block + * @param owned per-commitment ownership flags; {@code null} or shorter list → remaining are false * @param expectedRoot root to verify against; {@code null} to skip * @return computed root after appending * @throws WasmException with code {@code ROOT_MISMATCH} if verification fails */ public ShieldedRoot appendCommitments( - long blockHeight, List commitments, ShieldedRoot expectedRoot) { + long blockHeight, List commitments, List owned, + ShieldedRoot expectedRoot) { requireU32(blockHeight, "blockHeight"); Objects.requireNonNull(commitments, "commitments must not be null"); @@ -115,6 +117,9 @@ public ShieldedRoot appendCommitments( .addAllCommitments(commitments.stream() .map(c -> ByteString.copyFrom(c.bytes())) .toList()); + if (owned != null && !owned.isEmpty()) { + req.addAllOwned(owned); + } if (expectedRoot != null) { req.setExpectedRoot(ByteString.copyFrom(expectedRoot.bytes())); } @@ -123,11 +128,6 @@ public ShieldedRoot appendCommitments( return ShieldedRoot.of(unwrap(r, resp -> resp.getBytesValue().toByteArray())); } - /** Convenience overload — appends without root verification. */ - public ShieldedRoot appendCommitments(long blockHeight, List commitments) { - return appendCommitments(blockHeight, commitments, null); - } - /** * Rolls the tree back to the checkpoint at the given block height. * 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 6257925f15a..ee24892892b 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 @@ -35,11 +35,11 @@ class ShieldedMerkleTreeTest { "0101000000000000000000000000000000000000000000000000000000000000000000"); /** - * A valid 32-byte Orchard commitment (Pallas base field element = 1, LE-encoded). + * Valid 32-byte Orchard commitments (Pallas base field elements 1, 2, 3 LE-encoded). + * All three are below the Pallas base field prime and therefore valid field elements. */ private static final ShieldedCommitment CMX = ShieldedCommitment.of(HexFormat.of().parseHex( "0100000000000000000000000000000000000000000000000000000000000000")); - // ------------------------------------------------------------------------- // ping // ------------------------------------------------------------------------- @@ -130,14 +130,14 @@ void fromState_invalidJson_throwsWasmException() { void appendCommitments_nullCommitmentsList_throwsNullPointerException() { try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { assertThrows(NullPointerException.class, () -> - tree.appendCommitments(1L, null)); + tree.appendCommitments(1L, null, List.of(), null)); } } @Test void appendCommitments_emptyBlock_returnsRoot() { try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - ShieldedRoot root = tree.appendCommitments(1L, Collections.emptyList()); + ShieldedRoot root = tree.appendCommitments(1L, Collections.emptyList(), List.of(), null); assertNotNull(root); assertEquals(ShieldedRoot.SIZE, root.bytes().length); } @@ -146,7 +146,7 @@ void appendCommitments_emptyBlock_returnsRoot() { @Test void appendCommitments_emptyBlock_doesNotChangeLeafCount() { try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - tree.appendCommitments(1L, Collections.emptyList()); + tree.appendCommitments(1L, Collections.emptyList(), List.of(), null); MerkleTreeInfo info = tree.getInfo(); assertEquals(0L, info.leafCount); assertEquals(Long.valueOf(1L), info.tipHeight); @@ -156,9 +156,9 @@ void appendCommitments_emptyBlock_doesNotChangeLeafCount() { @Test void appendCommitments_multipleEmptyBlocks_incrementsCheckpointCount() { try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - tree.appendCommitments(1L, Collections.emptyList()); - tree.appendCommitments(2L, Collections.emptyList()); - tree.appendCommitments(3L, Collections.emptyList()); + tree.appendCommitments(1L, Collections.emptyList(), List.of(), null); + tree.appendCommitments(2L, Collections.emptyList(), List.of(), null); + tree.appendCommitments(3L, Collections.emptyList(), List.of(), null); MerkleTreeInfo info = tree.getInfo(); assertEquals(Long.valueOf(3L), info.tipHeight); assertEquals(3, info.checkpointCount); @@ -173,7 +173,7 @@ void appendCommitments_multipleEmptyBlocks_incrementsCheckpointCount() { @Test void appendCommitments_withCommitment_increasesLeafCount() { try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - tree.appendCommitments(1L, List.of(CMX)); + tree.appendCommitments(1L, List.of(CMX), List.of(), null); MerkleTreeInfo info = tree.getInfo(); assertEquals(1L, info.leafCount); assertEquals(Long.valueOf(1L), info.tipHeight); @@ -183,8 +183,8 @@ void appendCommitments_withCommitment_increasesLeafCount() { @Test void appendCommitments_twoBlocks_leafCountAccumulates() { try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - tree.appendCommitments(1L, List.of(CMX)); - tree.appendCommitments(2L, List.of(CMX)); + tree.appendCommitments(1L, List.of(CMX), List.of(), null); + tree.appendCommitments(2L, List.of(CMX), List.of(), null); MerkleTreeInfo info = tree.getInfo(); assertEquals(2L, info.leafCount); assertEquals(Long.valueOf(2L), info.tipHeight); @@ -195,8 +195,8 @@ void appendCommitments_twoBlocks_leafCountAccumulates() { @Test void appendCommitments_twoBlocks_rootChanges() { try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - ShieldedRoot root1 = tree.appendCommitments(1L, List.of(CMX)); - ShieldedRoot root2 = tree.appendCommitments(2L, List.of(CMX)); + ShieldedRoot root1 = tree.appendCommitments(1L, List.of(CMX), List.of(), null); + ShieldedRoot root2 = tree.appendCommitments(2L, List.of(CMX), List.of(), null); assertNotEquals(root1, root2, "Root must change when new leaves are appended"); } } @@ -206,10 +206,10 @@ void appendCommitments_returnsRoot_isDeterministic() { // Two fresh trees given the same commitment produce the same root. ShieldedRoot root1, root2; try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - root1 = tree.appendCommitments(1L, List.of(CMX)); + root1 = tree.appendCommitments(1L, List.of(CMX), List.of(), null); } try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - root2 = tree.appendCommitments(1L, List.of(CMX)); + root2 = tree.appendCommitments(1L, List.of(CMX), List.of(), null); } assertEquals(root1, root2); } @@ -219,10 +219,10 @@ void appendCommitments_correctExpectedRoot_succeeds() { // Capture actual root first, then confirm passing it as expectedRoot doesn't throw. ShieldedRoot actualRoot; try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - actualRoot = tree.appendCommitments(1L, List.of(CMX)); + actualRoot = tree.appendCommitments(1L, List.of(CMX), List.of(), null); } try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - ShieldedRoot returned = tree.appendCommitments(1L, List.of(CMX), actualRoot); + ShieldedRoot returned = tree.appendCommitments(1L, List.of(CMX), List.of(), actualRoot); assertEquals(actualRoot, returned); } } @@ -232,7 +232,7 @@ void appendCommitments_wrongExpectedRoot_throwsRootMismatch() { ShieldedRoot wrongRoot = ShieldedRoot.of(new byte[ShieldedRoot.SIZE]); WasmException ex = assertThrows(WasmException.class, () -> { try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - tree.appendCommitments(1L, List.of(CMX), wrongRoot); + tree.appendCommitments(1L, List.of(CMX), List.of(), wrongRoot); } }); assertEquals("ROOT_MISMATCH", ex.getErrorCode()); @@ -241,7 +241,7 @@ void appendCommitments_wrongExpectedRoot_throwsRootMismatch() { @Test void appendCommitments_multipleCommitmentsInOneBlock_allLeavesCounted() { try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - tree.appendCommitments(1L, List.of(CMX, CMX, CMX)); + tree.appendCommitments(1L, List.of(CMX, CMX, CMX), List.of(), null); MerkleTreeInfo info = tree.getInfo(); assertEquals(3L, info.leafCount); assertEquals(1, info.checkpointCount); @@ -255,8 +255,8 @@ void appendCommitments_multipleCommitmentsInOneBlock_allLeavesCounted() { @Test void truncateToCheckpoint_rollsBackTipHeightAndLeafCount() { try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - tree.appendCommitments(1L, List.of(CMX)); - tree.appendCommitments(2L, List.of(CMX)); + tree.appendCommitments(1L, List.of(CMX), List.of(), null); + tree.appendCommitments(2L, List.of(CMX), List.of(), null); tree.truncateToCheckpoint(1L); @@ -270,8 +270,8 @@ void truncateToCheckpoint_rollsBackTipHeightAndLeafCount() { void truncateToCheckpoint_returnsRootMatchingOriginalAppend() { ShieldedRoot rootAt1; try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - rootAt1 = tree.appendCommitments(1L, List.of(CMX)); - tree.appendCommitments(2L, List.of(CMX)); + rootAt1 = tree.appendCommitments(1L, List.of(CMX), List.of(), null); + tree.appendCommitments(2L, List.of(CMX), List.of(), null); ShieldedRoot restoredRoot = tree.truncateToCheckpoint(1L); assertEquals(rootAt1, restoredRoot); @@ -282,7 +282,7 @@ void truncateToCheckpoint_returnsRootMatchingOriginalAppend() { void truncateToCheckpoint_unknownHeight_throwsCheckpointNotFound() { WasmException ex = assertThrows(WasmException.class, () -> { try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - tree.appendCommitments(100L, Collections.emptyList()); + tree.appendCommitments(100L, Collections.emptyList(), List.of(), null); tree.truncateToCheckpoint(999L); } }); @@ -297,7 +297,7 @@ void truncateToCheckpoint_unknownHeight_throwsCheckpointNotFound() { void saveAndLoad_roundTrip_infoFieldsSurvive() { TreeState savedState; try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - tree.appendCommitments(100L, Collections.emptyList()); + tree.appendCommitments(100L, Collections.emptyList(), List.of(), null); savedState = tree.save(); } try (ShieldedMerkleTree restored = ShieldedMerkleTree.fromState(savedState)) { @@ -312,7 +312,7 @@ void saveAndLoad_roundTrip_infoFieldsSurvive() { void saveAndLoad_withLeaves_preservesLeafCount() { TreeState savedState; try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - tree.appendCommitments(1L, List.of(CMX, CMX)); + tree.appendCommitments(1L, List.of(CMX, CMX), List.of(), null); savedState = tree.save(); } try (ShieldedMerkleTree restored = ShieldedMerkleTree.fromState(savedState)) { @@ -326,11 +326,11 @@ void saveAndLoad_withLeaves_preservesLeafCount() { void saveAndLoad_restoredTreeCanContinueAppending() { TreeState savedState; try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - tree.appendCommitments(1L, List.of(CMX)); + tree.appendCommitments(1L, List.of(CMX), List.of(), null); savedState = tree.save(); } try (ShieldedMerkleTree restored = ShieldedMerkleTree.fromState(savedState)) { - assertDoesNotThrow(() -> restored.appendCommitments(2L, List.of(CMX))); + assertDoesNotThrow(() -> restored.appendCommitments(2L, List.of(CMX), List.of(), null)); MerkleTreeInfo info = restored.getInfo(); assertEquals(2L, info.leafCount); assertEquals(Long.valueOf(2L), info.tipHeight); @@ -342,7 +342,7 @@ void saveAndLoad_stateBytesRoundTrip() { // Verify TreeState.bytes() / TreeState.of() round-trips without data loss. TreeState saved; try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - tree.appendCommitments(5L, List.of(CMX)); + tree.appendCommitments(5L, List.of(CMX), List.of(), null); saved = tree.save(); } TreeState restored = TreeState.of(saved.bytes()); @@ -362,11 +362,11 @@ void appendCommitments_wrongSizeCommitment_throwsIllegalArgumentException() { @Test void truncateToCheckpoint_treeIsUsableAfterRollback() { try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - tree.appendCommitments(1L, List.of(CMX)); - tree.appendCommitments(2L, List.of(CMX)); + tree.appendCommitments(1L, List.of(CMX), List.of(), null); + tree.appendCommitments(2L, List.of(CMX), List.of(), null); tree.truncateToCheckpoint(1L); - ShieldedRoot root = tree.appendCommitments(3L, List.of(CMX)); + ShieldedRoot root = tree.appendCommitments(3L, List.of(CMX), List.of(), null); assertEquals(ShieldedRoot.SIZE, root.bytes().length); MerkleTreeInfo info = tree.getInfo(); assertEquals(Long.valueOf(3L), info.tipHeight); @@ -382,7 +382,7 @@ void truncateToCheckpoint_treeIsUsableAfterRollback() { void appendCommitments_negativeBlockHeight_throwsIllegalArgumentException() { try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { assertThrows(IllegalArgumentException.class, () -> - tree.appendCommitments(-1L, Collections.emptyList())); + tree.appendCommitments(-1L, Collections.emptyList(), List.of(), null)); } } @@ -390,7 +390,7 @@ void appendCommitments_negativeBlockHeight_throwsIllegalArgumentException() { void appendCommitments_blockHeightAboveU32Max_throwsIllegalArgumentException() { try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(EMPTY_STATE)) { assertThrows(IllegalArgumentException.class, () -> - tree.appendCommitments(0x1_0000_0000L, Collections.emptyList())); + tree.appendCommitments(0x1_0000_0000L, Collections.emptyList(), List.of(), null)); } } @@ -411,11 +411,12 @@ void twoInstances_shareNoState() { try (ShieldedMerkleTree tree1 = ShieldedMerkleTree.fromState(EMPTY_STATE); ShieldedMerkleTree tree2 = ShieldedMerkleTree.fromState(EMPTY_STATE)) { - tree1.appendCommitments(1L, List.of(CMX)); + tree1.appendCommitments(1L, List.of(CMX), List.of(), null); MerkleTreeInfo info2 = tree2.getInfo(); assertNull(info2.tipHeight); assertEquals(0L, info2.leafCount); } } + } diff --git a/packages/wasm-privacy-coin/src/zcash/tree.rs b/packages/wasm-privacy-coin/src/zcash/tree.rs index 2e06f66c366..59a80d14ca0 100644 --- a/packages/wasm-privacy-coin/src/zcash/tree.rs +++ b/packages/wasm-privacy-coin/src/zcash/tree.rs @@ -425,6 +425,7 @@ impl OwnedTree { &mut self, block_height: u32, commitments: Vec>, + owned: Vec, expected_root: Option<&[u8]>, ) -> Result, String> { if commitments.is_empty() { @@ -439,6 +440,14 @@ impl OwnedTree { return get_checkpoint_root_bytes(&self.tree); } + if !owned.is_empty() && owned.len() != commitments.len() { + return Err(format!( + "owned length {} does not match commitments length {}", + owned.len(), + commitments.len() + )); + } + // Validate ALL commitments before mutating any tree state. // This prevents a partial-append scenario where some leaves are inserted and then // an invalid commitment causes an error, leaving the tree in an inconsistent state @@ -449,11 +458,18 @@ impl OwnedTree { let last_idx = hashes.len() - 1; for (i, hash) in hashes.into_iter().enumerate() { + let is_owned = owned.get(i).copied().unwrap_or(false); let retention = if i == last_idx { Retention::Checkpoint { id: block_height, - marking: Marking::None, + marking: if is_owned { + Marking::Marked + } else { + Marking::None + }, } + } else if is_owned { + Retention::Marked } else { Retention::Ephemeral }; @@ -514,3 +530,93 @@ impl OwnedTree { Ok((self.tip_height, self.leaf_count, checkpoint_count as u32)) } } + +#[cfg(test)] +mod tests { + use super::*; + + const F_CHECKPOINT: u8 = 1; + const F_MARKED: u8 = 2; + const F_CHECKPOINT_MARKED: u8 = 3; + + fn empty_tree() -> OwnedTree { + let json = r#"{"shards":[],"cap":{"type":"Nil"},"checkpoints":[],"tip_height":null,"leaf_count":0}"#; + OwnedTree::from_state(json.as_bytes()).expect("empty state") + } + + fn cmx(byte: u8) -> Vec { + let mut v = vec![0u8; 32]; + v[0] = byte; + v + } + + fn find_leaf_f(node: &TreeNode, target_h: &str) -> Option { + match node { + TreeNode::Leaf { h, f } if h == target_h => Some(*f), + TreeNode::Parent { l, r, .. } => { + find_leaf_f(l, target_h).or_else(|| find_leaf_f(r, target_h)) + } + _ => None, + } + } + + fn find_f_in_state(state: &PersistedShardTreeState, target_h: &str) -> Option { + for shard in &state.shards { + if let Some(f) = find_leaf_f(&shard.tree, target_h) { + return Some(f); + } + } + find_leaf_f(&state.cap, target_h) + } + + #[test] + fn owned_flags_set_correct_retention() { + let cmx1_hex = "0100000000000000000000000000000000000000000000000000000000000000"; + let cmx3_hex = "0300000000000000000000000000000000000000000000000000000000000000"; + + let mut tree = empty_tree(); + tree.append_commitments( + 1, + vec![cmx(1), cmx(2), cmx(3)], + vec![true, false, true], + None, + ) + .unwrap(); + let state = extract_state(&tree.tree, tree.tip_height, tree.leaf_count).unwrap(); + + assert_eq!( + find_f_in_state(&state, cmx1_hex), + Some(F_MARKED), + "CMX: owned, not last → MARKED" + ); + assert_eq!( + find_f_in_state(&state, cmx3_hex), + Some(F_CHECKPOINT_MARKED), + "CMX_3: owned, last → CHECKPOINT|MARKED" + ); + } + + #[test] + fn not_owned_last_leaf_is_checkpoint_only() { + let cmx_hex = "0100000000000000000000000000000000000000000000000000000000000000"; + + 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(); + + assert_eq!( + find_f_in_state(&state, cmx_hex), + Some(F_CHECKPOINT), + "not owned, last → CHECKPOINT" + ); + } + + #[test] + fn owned_length_mismatch_returns_err() { + let mut tree = empty_tree(); + let result = + tree.append_commitments(1, vec![cmx(1), cmx(2)], vec![true, false, true], None); + assert!(result.is_err()); + } +}