From ff4b754db95ecf7cc5f923836eaa4d23b773f815 Mon Sep 17 00:00:00 2001
From: Ganesha Upadhyaya <ganeshrvce@gmail.com>
Date: Mon, 30 Oct 2023 07:38:13 -0500
Subject: [PATCH] loop through all txs before returning (#1287)

Fixes #1283

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

Refactor:
- Improvement: Enhanced the transaction handling process in the
`TxMempool` module. Transactions are now more accurately filtered based
on their gas and data size limits, ensuring that only valid transactions
are processed. This update improves the overall efficiency and
reliability of transaction processing.

Tests:
- Test: Added a new test, `TestTxMempoolTxLargerThanMaxBytes`, to verify
the correct handling of large transactions. This ensures the robustness
of the system when dealing with transactions of varying sizes.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Ganesha Upadhyaya <gupadhyaya@Ganeshas-MacBook-Pro-2.local>
---
 mempool/v1/mempool.go      | 14 ++++++++------
 mempool/v1/mempool_test.go | 22 ++++++++++++++++++++++
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/mempool/v1/mempool.go b/mempool/v1/mempool.go
index 9a09dfe0bdf..c95caa6c6b0 100644
--- a/mempool/v1/mempool.go
+++ b/mempool/v1/mempool.go
@@ -336,15 +336,17 @@ func (txmp *TxMempool) allEntriesSorted() []*WrappedTx {
 func (txmp *TxMempool) ReapMaxBytesMaxGas(maxBytes, maxGas int64) types.Txs {
 	var totalGas, totalBytes int64
 
-	var keep []types.Tx //nolint:prealloc
+	keep := make([]types.Tx, 0, len(txmp.allEntriesSorted()))
 	for _, w := range txmp.allEntriesSorted() {
 		// N.B. When computing byte size, we need to include the overhead for
-		// encoding as protobuf to send to the application.
-		totalGas += w.gasWanted
-		totalBytes += types.ComputeProtoSizeForTxs([]types.Tx{w.tx})
-		if (maxGas >= 0 && totalGas > maxGas) || (maxBytes >= 0 && totalBytes > maxBytes) {
-			break
+		// encoding as protobuf to send to the application. This actually overestimates it
+		// as we add the proto overhead to each transaction
+		txBytes := types.ComputeProtoSizeForTxs([]types.Tx{w.tx})
+		if (maxGas >= 0 && totalGas+w.gasWanted > maxGas) || (maxBytes >= 0 && totalBytes+txBytes > maxBytes) {
+			continue
 		}
+		totalBytes += txBytes
+		totalGas += w.gasWanted
 		keep = append(keep, w.tx)
 	}
 	return keep
diff --git a/mempool/v1/mempool_test.go b/mempool/v1/mempool_test.go
index 8ebb94b326c..4c894013b45 100644
--- a/mempool/v1/mempool_test.go
+++ b/mempool/v1/mempool_test.go
@@ -241,6 +241,28 @@ func TestTxMempool_ReapMaxBytesMaxGas(t *testing.T) {
 	require.Len(t, reapedTxs, 25)
 }
 
+func TestTxMempoolTxLargerThanMaxBytes(t *testing.T) {
+	rng := rand.New(rand.NewSource(time.Now().UnixNano()))
+	txmp := setup(t, 0)
+	bigPrefix := make([]byte, 100)
+	_, err := rng.Read(bigPrefix)
+	require.NoError(t, err)
+	// large high priority tx
+	bigTx := []byte(fmt.Sprintf("sender-1-1=%X=2", bigPrefix))
+	smallPrefix := make([]byte, 20)
+	_, err = rng.Read(smallPrefix)
+	require.NoError(t, err)
+	// smaller low priority tx with different sender
+	smallTx := []byte(fmt.Sprintf("sender-2-1=%X=1", smallPrefix))
+	require.NoError(t, txmp.CheckTx(bigTx, nil, mempool.TxInfo{SenderID: 1}))
+	require.NoError(t, txmp.CheckTx(smallTx, nil, mempool.TxInfo{SenderID: 2}))
+
+	// reap by max bytes less than the large tx
+	reapedTxs := txmp.ReapMaxBytesMaxGas(100, -1)
+	require.Len(t, reapedTxs, 1)
+	require.Equal(t, types.Tx(smallTx), reapedTxs[0])
+}
+
 func TestTxMempool_ReapMaxTxs(t *testing.T) {
 	txmp := setup(t, 0)
 	tTxs := checkTxs(t, txmp, 100, 0)