From a31b8652574d8fdcbd9b09aee8d83bfbfe5f57d3 Mon Sep 17 00:00:00 2001
From: yihuang <huang@crypto.com>
Date: Fri, 1 Sep 2023 15:58:41 +0800
Subject: [PATCH] Problem: memiavl's root hash is not safe to retain in commit
 info (#1150)

* Problem: memiavl's root hash is not safe to retain in commit info

Solution:
- clone the byte slices when nesserary

* Update CHANGELOG.md

Signed-off-by: yihuang <huang@crypto.com>

* add unit test that reproduce the issue

---------

Signed-off-by: yihuang <huang@crypto.com>
---
 CHANGELOG.md              |  1 +
 memiavl/db_test.go        | 20 ++++++++++++++++++--
 memiavl/mem_node.go       |  4 ++++
 memiavl/node.go           |  3 +++
 memiavl/persisted_node.go |  4 ++++
 memiavl/tree.go           |  5 +++--
 6 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9513367c05..8b3def4e61 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -37,6 +37,7 @@
 - [#1123](https://github.com/crypto-org-chain/cronos/pull/1123) Fix memiavl snapshot switching
 - [#1125](https://github.com/crypto-org-chain/cronos/pull/1125) Fix genesis migrate for feeibc, evm, feemarket and gravity.
 - [#1130](https://github.com/crypto-org-chain/cronos/pull/1130) Fix lock issues when state-sync with memiavl.
+- [#1150](https://github.com/crypto-org-chain/cronos/pull/1150) Fix memiavl's unsafe retain of the root hashes.
 
 ### Features
 
diff --git a/memiavl/db_test.go b/memiavl/db_test.go
index 003a03cf65..469b1a9bd8 100644
--- a/memiavl/db_test.go
+++ b/memiavl/db_test.go
@@ -280,16 +280,27 @@ func TestLoadVersion(t *testing.T) {
 }
 
 func TestZeroCopy(t *testing.T) {
-	db, err := Load(t.TempDir(), Options{InitialStores: []string{"test"}, CreateIfMissing: true, ZeroCopy: true})
+	db, err := Load(t.TempDir(), Options{InitialStores: []string{"test", "test2"}, CreateIfMissing: true, ZeroCopy: true})
 	require.NoError(t, err)
-	db.Commit([]*NamedChangeSet{
+	_, _, err = db.Commit([]*NamedChangeSet{
 		{Name: "test", Changeset: ChangeSets[0]},
+		{Name: "test2"},
 	})
+	require.NoError(t, err)
 	require.NoError(t, errors.Join(
 		db.RewriteSnapshot(),
 		db.Reload(),
 	))
 
+	// the test tree's root hash will reference the zero-copy value
+	_, _, err = db.Commit([]*NamedChangeSet{
+		{Name: "test"},
+		{Name: "test2", Changeset: ChangeSets[0]},
+	})
+	require.NoError(t, err)
+
+	commitInfo := *db.LastCommitInfo()
+
 	value := db.TreeByName("test").Get([]byte("hello"))
 	require.Equal(t, []byte("world"), value)
 
@@ -297,6 +308,8 @@ func TestZeroCopy(t *testing.T) {
 	valueCloned := db.TreeByName("test").Get([]byte("hello"))
 	require.Equal(t, []byte("world"), valueCloned)
 
+	_ = commitInfo.CommitID()
+
 	require.NoError(t, db.Close())
 
 	require.Equal(t, []byte("world"), valueCloned)
@@ -307,6 +320,9 @@ func TestZeroCopy(t *testing.T) {
 	require.Panics(t, func() {
 		require.Equal(t, []byte("world"), value)
 	})
+
+	// it's ok to access after db closed
+	_ = commitInfo.CommitID()
 }
 
 func TestWalIndexConversion(t *testing.T) {
diff --git a/memiavl/mem_node.go b/memiavl/mem_node.go
index 6f8eb375bf..f56e345aa6 100644
--- a/memiavl/mem_node.go
+++ b/memiavl/mem_node.go
@@ -70,6 +70,10 @@ func (node *MemNode) Mutate(version, cowVersion uint32) *MemNode {
 	return n
 }
 
+func (node *MemNode) SafeHash() []byte {
+	return node.Hash()
+}
+
 // Computes the hash of the node without computing its descendants. Must be
 // called on nodes which have descendant node hashes already computed.
 func (node *MemNode) Hash() []byte {
diff --git a/memiavl/node.go b/memiavl/node.go
index c756ecdbeb..6f659313c8 100644
--- a/memiavl/node.go
+++ b/memiavl/node.go
@@ -20,6 +20,9 @@ type Node interface {
 	Right() Node
 	Hash() []byte
 
+	// SafeHash returns byte slice that's safe to retain
+	SafeHash() []byte
+
 	// PersistedNode clone a new node, MemNode modify in place
 	Mutate(version, cowVersion uint32) *MemNode
 
diff --git a/memiavl/persisted_node.go b/memiavl/persisted_node.go
index 50579dd3b9..59f8c2ac98 100644
--- a/memiavl/persisted_node.go
+++ b/memiavl/persisted_node.go
@@ -131,6 +131,10 @@ func (node PersistedNode) Right() Node {
 	return PersistedNode{snapshot: node.snapshot, index: node.index - 1}
 }
 
+func (node PersistedNode) SafeHash() []byte {
+	return bytes.Clone(node.Hash())
+}
+
 func (node PersistedNode) Hash() []byte {
 	if node.isLeaf {
 		return node.leafNode().Hash()
diff --git a/memiavl/tree.go b/memiavl/tree.go
index 3d10a01e7f..79067f27ed 100644
--- a/memiavl/tree.go
+++ b/memiavl/tree.go
@@ -162,12 +162,13 @@ func (t *Tree) Version() int64 {
 	return int64(t.version)
 }
 
-// RootHash updates the hashes and return the current root hash
+// RootHash updates the hashes and return the current root hash,
+// it clones the persisted node's bytes, so the returned bytes is safe to retain.
 func (t *Tree) RootHash() []byte {
 	if t.root == nil {
 		return emptyHash
 	}
-	return t.root.Hash()
+	return t.root.SafeHash()
 }
 
 func (t *Tree) GetWithIndex(key []byte) (int64, []byte) {