From 4806960bbc7125e6889de867311a024e6a6ac936 Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Tue, 26 Nov 2024 12:24:39 +0100
Subject: [PATCH] fix: Add Extra Check for Reformatted Root Node in GetNode
 (backport #1007) (#1010)

Co-authored-by: cool-developer <51834436+cool-develope@users.noreply.github.com>
---
 CHANGELOG.md         |  4 ++++
 mutable_tree_test.go | 24 ++++++++++++++++++++++++
 nodedb.go            | 14 ++++++++++++++
 3 files changed, 42 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 202672fb7..bfca8b858 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## v1.3.0, July 31, 2024
 
+### Bug Fixes
+
+- [#1007](https://github.com/cosmos/iavl/pull/1007) Add the extra check for the reformatted root node in `GetNode`
+
 ### Improvements
 
 - [#952](https://github.com/cosmos/iavl/pull/952) Add `DeleteVersionsFrom(int64)` API.
diff --git a/mutable_tree_test.go b/mutable_tree_test.go
index db12846fa..52bde6418 100644
--- a/mutable_tree_test.go
+++ b/mutable_tree_test.go
@@ -1483,6 +1483,30 @@ func TestMutableTreeClose(t *testing.T) {
 	require.NoError(t, tree.Close())
 }
 
+func TestReferenceRootPruning(t *testing.T) {
+	memDB := dbm.NewMemDB()
+	tree := NewMutableTree(memDB, 0, true, NewNopLogger())
+
+	_, err := tree.Set([]byte("foo"), []byte("bar"))
+	require.NoError(t, err)
+	_, _, err = tree.SaveVersion()
+	require.NoError(t, err)
+
+	_, _, err = tree.SaveVersion()
+	require.NoError(t, err)
+
+	_, err = tree.Set([]byte("foo1"), []byte("bar"))
+	require.NoError(t, err)
+	_, _, err = tree.SaveVersion()
+	require.NoError(t, err)
+
+	err = tree.DeleteVersionsTo(1)
+	require.NoError(t, err)
+
+	_, err = tree.Set([]byte("foo"), []byte("bar*"))
+	require.NoError(t, err)
+}
+
 func TestMutableTree_InitialVersionZero(t *testing.T) {
 	db := dbm.NewMemDB()
 
diff --git a/nodedb.go b/nodedb.go
index 2112c6d54..14d74d043 100644
--- a/nodedb.go
+++ b/nodedb.go
@@ -160,6 +160,20 @@ func (ndb *nodeDB) GetNode(nk []byte) (*Node, error) {
 	if err != nil {
 		return nil, fmt.Errorf("can't get node %v: %v", nk, err)
 	}
+	if buf == nil && !isLegcyNode {
+		// if the node is reformatted by pruning, check against (version, 0)
+		nKey := GetNodeKey(nk)
+		if nKey.nonce == 1 {
+			nodeKey = ndb.nodeKey((&NodeKey{
+				version: nKey.version,
+				nonce:   0,
+			}).GetKey())
+			buf, err = ndb.db.Get(nodeKey)
+			if err != nil {
+				return nil, fmt.Errorf("can't get the reformatted node %v: %v", nk, err)
+			}
+		}
+	}
 	if buf == nil {
 		return nil, fmt.Errorf("Value missing for key %v corresponding to nodeKey %x", nk, nodeKey)
 	}