From b1bf48867c0439bfb3610e1789c572631fca59a9 Mon Sep 17 00:00:00 2001
From: testinginprod <98415576+testinginprod@users.noreply.github.com>
Date: Mon, 9 Sep 2024 10:37:59 +0200
Subject: [PATCH] fix(stf/branch/memiter): Fix Iter validity (#21556)

---
 server/v2/stf/branch/changeset.go      |  7 +++----
 server/v2/stf/branch/changeset_test.go | 28 ++++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 4 deletions(-)
 create mode 100644 server/v2/stf/branch/changeset_test.go

diff --git a/server/v2/stf/branch/changeset.go b/server/v2/stf/branch/changeset.go
index 13c016725130..c409b1b7becf 100644
--- a/server/v2/stf/branch/changeset.go
+++ b/server/v2/stf/branch/changeset.go
@@ -5,8 +5,6 @@ import (
 	"errors"
 
 	"github.com/tidwall/btree"
-
-	"cosmossdk.io/core/store"
 )
 
 const (
@@ -55,7 +53,7 @@ func (bt changeSet) delete(key []byte) {
 
 // iterator returns a new iterator over the key-value pairs in the changeSet
 // that have keys greater than or equal to the start key and less than the end key.
-func (bt changeSet) iterator(start, end []byte) (store.Iterator, error) {
+func (bt changeSet) iterator(start, end []byte) (*memIterator, error) {
 	if (start != nil && len(start) == 0) || (end != nil && len(end) == 0) {
 		return nil, errKeyEmpty
 	}
@@ -65,7 +63,7 @@ func (bt changeSet) iterator(start, end []byte) (store.Iterator, error) {
 // reverseIterator returns a new iterator that iterates over the key-value pairs in reverse order
 // within the specified range [start, end) in the changeSet's tree.
 // If start or end is an empty byte slice, it returns an error indicating that the key is empty.
-func (bt changeSet) reverseIterator(start, end []byte) (store.Iterator, error) {
+func (bt changeSet) reverseIterator(start, end []byte) (*memIterator, error) {
 	if (start != nil && len(start) == 0) || (end != nil && len(end) == 0) {
 		return nil, errKeyEmpty
 	}
@@ -158,6 +156,7 @@ func (mi *memIterator) Domain() (start, end []byte) {
 // Close releases any resources held by the iterator.
 func (mi *memIterator) Close() error {
 	mi.iter.Release()
+	mi.valid = false
 	return nil
 }
 
diff --git a/server/v2/stf/branch/changeset_test.go b/server/v2/stf/branch/changeset_test.go
new file mode 100644
index 000000000000..6a820c241c3d
--- /dev/null
+++ b/server/v2/stf/branch/changeset_test.go
@@ -0,0 +1,28 @@
+package branch
+
+import (
+	"testing"
+)
+
+func Test_memIterator(t *testing.T) {
+	t.Run("iter is invalid after close", func(t *testing.T) {
+		cs := newChangeSet()
+		for i := byte(0); i < 32; i++ {
+			cs.set([]byte{0, i}, []byte{i})
+		}
+
+		it, err := cs.iterator(nil, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		err = it.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if it.Valid() {
+			t.Fatal("iterator must be invalid")
+		}
+	})
+}