Account Compression: Add CMT checks for out of bounds leaf indices and initialization (#3724)

ngundotra · web-flow · commit 5504af4996b7 · 2022-10-21T14:06:47.000-04:00
* cmt: add check for leaf index OOB

* ac: add checks for leaf index OOB

* ac: add tests for leaf index OOB

* nit: fix some poor logging &amp; unused imports

* cmt: add initialization checks before public methods

* cmt: update LeafContentsModified error message

* cmt: make tests easier to read, add PartialEq, Eq to CMTError

* ac: make LeafOOB error the last error in the struct to prevent breaking changes

* ac: fmt fix
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/account-compression/Cargo.lock b/account-compression/Cargo.lock
diff --git a/account-compression/programs/account-compression/Cargo.toml b/account-compression/programs/account-compression/Cargo.toml
@@ -20,7 +20,7 @@ default = []
 [dependencies]
 anchor-lang = "0.25.0"
 bytemuck = "1.8.0"
-spl-concurrent-merkle-tree = { version = "0.1.1", path="../../../libraries/concurrent-merkle-tree", features = [ "sol-log" ]}
+spl-concurrent-merkle-tree = { version="0.1.2", path="../../../libraries/concurrent-merkle-tree", features = [ "sol-log" ]}
 spl-noop = { version = "0.1.3", path="../noop", features = [ "no-entrypoint" ]}
 
 [profile.release]
diff --git a/account-compression/programs/account-compression/src/error.rs b/account-compression/programs/account-compression/src/error.rs
@@ -42,6 +42,11 @@ pub enum AccountCompressionError {
     /// Incorrect account type
     #[msg("Account provided has incorrect account type")]
     IncorrectAccountType,
+
+    /// Tree information cannot be processed because the provided leaf_index
+    /// is out of bounds of tree's maximum leaf capacity
+    #[msg("Leaf index of concurrent merkle tree is out of bounds")]
+    LeafIndexOutOfBounds,
 }
 
 impl From<&ConcurrentMerkleTreeError> for AccountCompressionError {
diff --git a/account-compression/programs/account-compression/src/lib.rs b/account-compression/programs/account-compression/src/lib.rs
@@ -352,6 +352,7 @@ pub mod spl_account_compression {
 
         let header = ConcurrentMerkleTreeHeader::try_from_slice(header_bytes)?;
         header.assert_valid_authority(&ctx.accounts.authority.key())?;
+        header.assert_valid_leaf_index(index)?;
 
         let merkle_tree_size = merkle_tree_get_size(&header)?;
         let (tree_bytes, canopy_bytes) = rest.split_at_mut(merkle_tree_size);
@@ -428,6 +429,7 @@ pub mod spl_account_compression {
 
         let header = ConcurrentMerkleTreeHeader::try_from_slice(header_bytes)?;
         header.assert_valid()?;
+        header.assert_valid_leaf_index(index)?;
 
         let merkle_tree_size = merkle_tree_get_size(&header)?;
         let (tree_bytes, canopy_bytes) = rest.split_at_mut(merkle_tree_size);
@@ -498,6 +500,7 @@ pub mod spl_account_compression {
 
         let header = ConcurrentMerkleTreeHeader::try_from_slice(header_bytes)?;
         header.assert_valid_authority(&ctx.accounts.authority.key())?;
+        header.assert_valid_leaf_index(index)?;
 
         let merkle_tree_size = merkle_tree_get_size(&header)?;
         let (tree_bytes, canopy_bytes) = rest.split_at_mut(merkle_tree_size);
diff --git a/account-compression/programs/account-compression/src/state/concurrent_merkle_tree_header.rs b/account-compression/programs/account-compression/src/state/concurrent_merkle_tree_header.rs
@@ -139,4 +139,11 @@ impl ConcurrentMerkleTreeHeader {
         }
         Ok(())
     }
+
+    pub fn assert_valid_leaf_index(&self, leaf_index: u32) -> Result<()> {
+        if leaf_index >= (1 << self.get_max_depth()) {
+            return Err(AccountCompressionError::LeafIndexOutOfBounds.into());
+        }
+        Ok(())
+    }
 }
diff --git a/account-compression/sdk/idl/spl_account_compression.json b/account-compression/sdk/idl/spl_account_compression.json
@@ -622,6 +622,11 @@
       "code": 6007,
       "name": "IncorrectAccountType",
       "msg": "Account provided has incorrect account type"
+    },
+    {
+      "code": 6008,
+      "name": "LeafIndexOutOfBounds",
+      "msg": "Leaf index of concurrent merkle tree is out of bounds"
     }
   ],
   "metadata": {
diff --git a/account-compression/sdk/package.json b/account-compression/sdk/package.json
@@ -38,8 +38,10 @@
     "run-tests": "jest tests --detectOpenHandles",
     "run-tests:events": "jest tests/events --detectOpenHandles",
     "run-tests:accounts": "jest tests/accounts --detectOpenHandles",
+    "run-tests:e2e": "jest accountCompression.test.ts --detectOpenHandles",
     "test:events": "start-server-and-test start-validator http://localhost:8899/health run-tests:events",
     "test:accounts": "start-server-and-test start-validator http://localhost:8899/health run-tests:accounts",
+    "test:e2e": "start-server-and-test start-validator http://localhost:8899/health run-tests:e2e",
     "test": "start-server-and-test start-validator http://localhost:8899/health run-tests"
   },
   "dependencies": {
diff --git a/account-compression/sdk/src/generated/errors/index.ts b/account-compression/sdk/src/generated/errors/index.ts
@@ -200,6 +200,29 @@ createErrorFromNameLookup.set(
   () => new IncorrectAccountTypeError()
 )
 
+/**
+ * LeafIndexOutOfBounds: 'Leaf index of concurrent merkle tree is out of bounds'
+ *
+ * @category Errors
+ * @category generated
+ */
+export class LeafIndexOutOfBoundsError extends Error {
+  readonly code: number = 0x1778
+  readonly name: string = 'LeafIndexOutOfBounds'
+  constructor() {
+    super('Leaf index of concurrent merkle tree is out of bounds')
+    if (typeof Error.captureStackTrace === 'function') {
+      Error.captureStackTrace(this, LeafIndexOutOfBoundsError)
+    }
+  }
+}
+
+createErrorFromCodeLookup.set(0x1778, () => new LeafIndexOutOfBoundsError())
+createErrorFromNameLookup.set(
+  'LeafIndexOutOfBounds',
+  () => new LeafIndexOutOfBoundsError()
+)
+
 /**
  * Attempts to resolve a custom program error from the provided error code.
  * @category Errors
diff --git a/account-compression/sdk/tests/accountCompression.test.ts b/account-compression/sdk/tests/accountCompression.test.ts
@@ -650,4 +650,45 @@ describe("Account Compression", () => {
       }
     });
   });
+  describe(`Having created a tree with 8 leaves`, () => {
+    beforeEach(async () => {
+      [cmtKeypair, offChainTree] = await createTreeOnChain(
+        provider,
+        payer,
+        1 << 3,
+        3,
+        8,
+      );
+    });
+    it(`Attempt to replace a leaf beyond the tree's capacity`, async () => {
+      // Ensure that this fails
+      let outOfBoundsIndex = 8;
+      const index = outOfBoundsIndex;
+      const newLeaf = hash(
+        payer.publicKey.toBuffer(),
+        Buffer.from(new BN(outOfBoundsIndex).toArray())
+      );
+      let proof;
+      let node;
+      node = offChainTree.leaves[outOfBoundsIndex - 1].node
+      proof = getProofOfLeaf(offChainTree, index - 1);
+
+      const replaceIx = createReplaceIx(
+        payer,
+        cmtKeypair.publicKey,
+        offChainTree.root,
+        node,
+        newLeaf,
+        index,
+        proof.map((treeNode) => {
+          return treeNode.node;
+        })
+      );
+
+      try {
+        await execute(provider, [replaceIx], [payer]);
+        throw Error("This replace instruction should have failed because the leaf index is OOB");
+      } catch (_e) { }
+    });
+  });
 });
diff --git a/account-compression/sdk/tests/utils.ts b/account-compression/sdk/tests/utils.ts
@@ -17,6 +17,7 @@ import {
     createInitEmptyMerkleTreeIx,
     createAllocTreeIx,
     createAppendIx,
+    ALL_DEPTH_SIZE_PAIRS,
 } from '../src'
 import {
     buildTree,
diff --git a/libraries/concurrent-merkle-tree/Cargo.toml b/libraries/concurrent-merkle-tree/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "spl-concurrent-merkle-tree"
-version = "0.1.1"
+version = "0.1.2"
 description = "Solana Program Library Concurrent Merkle Tree"
 authors = ["Solana Maintainers <maintainers@solana.foundation>"]
 repository = "https://github.com/solana-labs/solana-program-library"
diff --git a/libraries/concurrent-merkle-tree/src/concurrent_merkle_tree.rs b/libraries/concurrent-merkle-tree/src/concurrent_merkle_tree.rs
@@ -19,6 +19,13 @@ fn check_bounds(max_depth: usize, max_buffer_size: usize) {
     assert!(max_buffer_size & (max_buffer_size - 1) == 0);
 }
 
+fn check_leaf_index(leaf_index: u32, max_depth: usize) -> Result<(), ConcurrentMerkleTreeError> {
+    if leaf_index >= (1 << max_depth) {
+        return Err(ConcurrentMerkleTreeError::LeafIndexOutOfBounds);
+    }
+    Ok(())
+}
+
 /// Conurrent Merkle Tree is a Merkle Tree that allows
 /// multiple tree operations targeted for the same tree root to succeed.
 ///
@@ -123,6 +130,8 @@ impl<const MAX_DEPTH: usize, const MAX_BUFFER_SIZE: usize>
         index: u32,
     ) -> Result<Node, ConcurrentMerkleTreeError> {
         check_bounds(MAX_DEPTH, MAX_BUFFER_SIZE);
+        check_leaf_index(index, MAX_DEPTH)?;
+
         if self.is_initialized() {
             return Err(ConcurrentMerkleTreeError::TreeAlreadyInitialized);
         }
@@ -148,6 +157,9 @@ impl<const MAX_DEPTH: usize, const MAX_BUFFER_SIZE: usize>
 
     /// Errors if one of the leaves of the current merkle tree is non-EMPTY
     pub fn prove_tree_is_empty(&self) -> Result<(), ConcurrentMerkleTreeError> {
+        if !self.is_initialized() {
+            return Err(ConcurrentMerkleTreeError::TreeNotInitialized);
+        }
         let mut empty_node_cache = Box::new([EMPTY; MAX_DEPTH]);
         if self.get_root()
             != empty_node_cached::<MAX_DEPTH>(MAX_DEPTH as u32, &mut empty_node_cache)
@@ -164,6 +176,10 @@ impl<const MAX_DEPTH: usize, const MAX_BUFFER_SIZE: usize>
 
     /// Returns the most recent changelog
     pub fn get_change_log(&self) -> Box<ChangeLog<MAX_DEPTH>> {
+        if !self.is_initialized() {
+            solana_logging!("Tree is not initialized, returning default change log");
+            return Box::new(ChangeLog::default());
+        }
         Box::new(self.change_logs[self.active_index as usize])
     }
 
@@ -186,6 +202,11 @@ impl<const MAX_DEPTH: usize, const MAX_BUFFER_SIZE: usize>
         leaf_index: u32,
     ) -> Result<(), ConcurrentMerkleTreeError> {
         check_bounds(MAX_DEPTH, MAX_BUFFER_SIZE);
+        check_leaf_index(leaf_index, MAX_DEPTH)?;
+        if !self.is_initialized() {
+            return Err(ConcurrentMerkleTreeError::TreeNotInitialized);
+        }
+
         if leaf_index > self.rightmost_proof.index {
             solana_logging!(
                 "Received an index larger than the rightmost index {} > {}",
@@ -224,6 +245,9 @@ impl<const MAX_DEPTH: usize, const MAX_BUFFER_SIZE: usize>
     /// Appending a non-empty Node will always succeed .
     pub fn append(&mut self, mut node: Node) -> Result<Node, ConcurrentMerkleTreeError> {
         check_bounds(MAX_DEPTH, MAX_BUFFER_SIZE);
+        if !self.is_initialized() {
+            return Err(ConcurrentMerkleTreeError::TreeNotInitialized);
+        }
         if node == EMPTY {
             return Err(ConcurrentMerkleTreeError::CannotAppendEmptyNode);
         }
@@ -289,6 +313,11 @@ impl<const MAX_DEPTH: usize, const MAX_BUFFER_SIZE: usize>
         index: u32,
     ) -> Result<Node, ConcurrentMerkleTreeError> {
         check_bounds(MAX_DEPTH, MAX_BUFFER_SIZE);
+        check_leaf_index(index, MAX_DEPTH)?;
+        if !self.is_initialized() {
+            return Err(ConcurrentMerkleTreeError::TreeNotInitialized);
+        }
+
         let mut proof: [Node; MAX_DEPTH] = [Node::default(); MAX_DEPTH];
         fill_in_proof::<MAX_DEPTH>(proof_vec, &mut proof);
 
@@ -314,6 +343,11 @@ impl<const MAX_DEPTH: usize, const MAX_BUFFER_SIZE: usize>
         index: u32,
     ) -> Result<Node, ConcurrentMerkleTreeError> {
         check_bounds(MAX_DEPTH, MAX_BUFFER_SIZE);
+        check_leaf_index(index, MAX_DEPTH)?;
+        if !self.is_initialized() {
+            return Err(ConcurrentMerkleTreeError::TreeNotInitialized);
+        }
+
         if index > self.rightmost_proof.index {
             Err(ConcurrentMerkleTreeError::LeafIndexOutOfBounds)
         } else {
@@ -433,6 +467,14 @@ impl<const MAX_DEPTH: usize, const MAX_BUFFER_SIZE: usize>
         proof: &[Node; MAX_DEPTH],
         leaf_index: u32,
     ) -> bool {
+        if !self.is_initialized() {
+            solana_logging!("Tree is not initialized, returning false");
+            return false;
+        }
+        if check_leaf_index(leaf_index, MAX_DEPTH).is_err() {
+            solana_logging!("Leaf index out of bounds for max_depth");
+            return false;
+        }
         recompute(leaf, proof, leaf_index) == self.get_root()
     }
 
diff --git a/libraries/concurrent-merkle-tree/src/error.rs b/libraries/concurrent-merkle-tree/src/error.rs
@@ -1,10 +1,10 @@
 use thiserror::Error;
 
 /// Concurrent merkle tree operation errors
-#[derive(Error, Debug)]
+#[derive(Error, Debug, PartialEq, Eq)]
 pub enum ConcurrentMerkleTreeError {
     /// Received an index larger than the rightmost index
-    #[error("Received an index larger than the rightmost index")]
+    #[error("Received an index larger than the rightmost index, or greater than (1 << max_depth)")]
     LeafIndexOutOfBounds,
 
     /// Invalid root recomputed from proof
@@ -23,14 +23,16 @@ pub enum ConcurrentMerkleTreeError {
     #[error("Tree already initialized")]
     TreeAlreadyInitialized,
 
+    /// This tree has not yet been initialized
+    #[error("Tree needs to be initialized before using")]
+    TreeNotInitialized,
+
     /// Root passed as argument cannot be found in stored changelog buffer
     #[error("Root not found in changelog buffer")]
     RootNotFound,
 
-    /// Valid proof was passed to a leaf, but its value has changed since the proof was issued
-    #[error(
-        "Valid proof was passed to a leaf, but its value has changed since the proof was issued"
-    )]
+    /// The tree's current leaf value does not match the supplied proof's leaf value
+    #[error("This tree's current leaf value does not match the supplied proof's leaf value")]
     LeafContentsModified,
 
     /// Tree has at least 1 non-EMTPY leaf
diff --git a/libraries/concurrent-merkle-tree/tests/tests.rs b/libraries/concurrent-merkle-tree/tests/tests.rs
@@ -39,6 +39,56 @@ async fn test_initialize() {
     }
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn test_bypass_initialize() {
+    let (mut cmt, off_chain_tree) = setup();
+    let mut rng = thread_rng();
+    let leaf = rng.gen::<[u8; 32]>();
+
+    assert_eq!(
+        ConcurrentMerkleTreeError::TreeNotInitialized,
+        cmt.append(leaf).unwrap_err(),
+        "Expected TreeNotInitialized error when appending to uninitialized tree"
+    );
+
+    assert_eq!(
+        ConcurrentMerkleTreeError::TreeNotInitialized,
+        cmt.set_leaf(
+            off_chain_tree.get_root(),
+            [0; 32],
+            leaf,
+            &off_chain_tree.get_proof_of_leaf(0),
+            0
+        )
+        .unwrap_err(),
+        "Expected TreeNotInitialized error when setting a leaf on an uninitialized tree"
+    );
+
+    assert_eq!(
+        ConcurrentMerkleTreeError::TreeNotInitialized,
+        cmt.prove_leaf(
+            off_chain_tree.get_root(),
+            leaf,
+            &off_chain_tree.get_proof_of_leaf(0),
+            0,
+        )
+        .unwrap_err(),
+        "Expected TreeNotInitialized error when proving a leaf exists on an uninitialized tree"
+    );
+
+    assert_eq!(
+        ConcurrentMerkleTreeError::TreeNotInitialized,
+        cmt.fill_empty_or_append(
+            off_chain_tree.get_root(),
+            leaf,
+            &off_chain_tree.get_proof_of_leaf(0),
+            0,
+        )
+        .unwrap_err(),
+        "Expected TreeNotInitialized error when filling an empty leaf or appending to uninitialized tree"
+    );
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn test_append() {
     let (mut cmt, mut off_chain_tree) = setup();

Original file line number	Diff line number	Diff line change
`@@ -139,4 +139,11 @@ impl ConcurrentMerkleTreeHeader {`
`139`	`139`	`}`
`140`	`140`	`Ok(())`
`141`	`141`	`}`
	`142`	`+`
	`143`	`+ pub fn assert_valid_leaf_index(&self, leaf_index: u32) -> Result<()> {`
	`144`	`+ if leaf_index >= (1 << self.get_max_depth()) {`
	`145`	`+ return Err(AccountCompressionError::LeafIndexOutOfBounds.into());`
	`146`	`+ }`
	`147`	`+ Ok(())`
	`148`	`+ }`
`142`	`149`	`}`