From 07fee591c3c2d08e4db706076277b89fb045cc0f Mon Sep 17 00:00:00 2001
From: Steven Landers <steven@seinetwork.io>
Date: Mon, 29 Jan 2024 18:07:03 -0500
Subject: [PATCH] [EVM] Avoid returning nil transactions on ForEach (#197)

* remove heapIndex to avoid nil scenario

* avoid returning nil in loop (mimic Peek)
---
 internal/mempool/priority_queue.go      | 15 +++++++++------
 internal/mempool/priority_queue_test.go | 14 ++++++++++++++
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/internal/mempool/priority_queue.go b/internal/mempool/priority_queue.go
index 3de4e810f..abcbfe6a7 100644
--- a/internal/mempool/priority_queue.go
+++ b/internal/mempool/priority_queue.go
@@ -293,6 +293,11 @@ func (pq *TxPriorityQueue) popTxUnsafe() *WrappedTx {
 	}
 	tx := x.(*WrappedTx)
 
+	// this situation is primarily for a test case that inserts nils
+	if tx == nil {
+		return nil
+	}
+
 	// non-evm transactions do not have txs waiting on a nonce
 	if !tx.isEVM {
 		return tx
@@ -338,6 +343,9 @@ func (pq *TxPriorityQueue) ForEachTx(handler func(wtx *WrappedTx) bool) {
 
 	for i := 0; i < numTxs; i++ {
 		popped := pq.popTxUnsafe()
+		if popped == nil {
+			break
+		}
 		txs = append(txs, popped)
 		if !handler(popped) {
 			return
@@ -377,9 +385,7 @@ func (pq *TxPriorityQueue) PeekTxs(max int) []*WrappedTx {
 //
 // NOTE: A caller should never call Push. Use PushTx instead.
 func (pq *TxPriorityQueue) Push(x interface{}) {
-	n := len(pq.txs)
 	item := x.(*WrappedTx)
-	item.heapIndex = n
 	pq.txs = append(pq.txs, item)
 }
 
@@ -390,8 +396,7 @@ func (pq *TxPriorityQueue) Pop() interface{} {
 	old := pq.txs
 	n := len(old)
 	item := old[n-1]
-	old[n-1] = nil      // avoid memory leak
-	item.heapIndex = -1 // for safety
+	old[n-1] = nil // avoid memory leak
 	pq.txs = old[0 : n-1]
 	return item
 }
@@ -420,6 +425,4 @@ func (pq *TxPriorityQueue) Less(i, j int) bool {
 // Swap implements the Heap interface. It swaps two transactions in the queue.
 func (pq *TxPriorityQueue) Swap(i, j int) {
 	pq.txs[i], pq.txs[j] = pq.txs[j], pq.txs[i]
-	pq.txs[i].heapIndex = i
-	pq.txs[j].heapIndex = j
 }
diff --git a/internal/mempool/priority_queue_test.go b/internal/mempool/priority_queue_test.go
index 8cb6d2e1a..eb340b170 100644
--- a/internal/mempool/priority_queue_test.go
+++ b/internal/mempool/priority_queue_test.go
@@ -47,6 +47,20 @@ func TestTxPriorityQueue_ReapHalf(t *testing.T) {
 	}
 }
 
+func TestAvoidPanicIfTransactionIsNil(t *testing.T) {
+	pq := NewTxPriorityQueue()
+	pq.Push(&WrappedTx{sender: "1", isEVM: true, evmAddress: "0xabc", evmNonce: 1, priority: 10})
+	pq.txs = append(pq.txs, nil)
+
+	var count int
+	pq.ForEachTx(func(tx *WrappedTx) bool {
+		count++
+		return true
+	})
+
+	require.Equal(t, 1, count)
+}
+
 func TestTxPriorityQueue_PriorityAndNonceOrdering(t *testing.T) {
 	testCases := []TxTestCase{
 		{