fix/remove-malleated-tx-cache (#834)

* feat: add method in cache to remove txns by key(hash)

* fix: remove original tx from cache if current tx is malleated

* chore: remove redundant type conversion

* chore: typo

* test: add cache test for removing tx by key

* test: fix check for cache map and list length

Author: rahulghangas

mempool/cache.go

@@ -23,6 +23,9 @@ type TxCache interface {
 	// Remove removes the given raw transaction from the cache.
 	Remove(tx types.Tx)
 
+	// RemoveTxByKey removes a given transaction hash from the cache.
+	RemoveTxByKey(key types.TxKey)
+
 	// Has reports whether tx is present in the cache. Checking for presence is
 	// not treated as an access of the value.
 	Has(tx types.Tx) bool
@@ -89,10 +92,14 @@ func (c *LRUTxCache) Push(tx types.Tx) bool {
 }
 
 func (c *LRUTxCache) Remove(tx types.Tx) {
+	key := tx.Key()
+	c.RemoveTxByKey(key)
+}
+
+func (c *LRUTxCache) RemoveTxByKey(key types.TxKey) {
 	c.mtx.Lock()
 	defer c.mtx.Unlock()
 
-	key := tx.Key()
 	e := c.cacheMap[key]
 	delete(c.cacheMap, key)
 
@@ -114,7 +121,8 @@ type NopTxCache struct{}
 
 var _ TxCache = (*NopTxCache)(nil)
 
-func (NopTxCache) Reset()             {}
-func (NopTxCache) Push(types.Tx) bool { return true }
-func (NopTxCache) Remove(types.Tx)    {}
-func (NopTxCache) Has(types.Tx) bool  { return false }
+func (NopTxCache) Reset()                        {}
+func (NopTxCache) Push(types.Tx) bool            { return true }
+func (NopTxCache) Remove(types.Tx)               {}
+func (NopTxCache) RemoveTxByKey(key types.TxKey) {}
+func (NopTxCache) Has(types.Tx) bool             { return false }

mempool/cache_test.go

@@ -5,31 +5,58 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"github.com/tendermint/tendermint/types"
 )
 
-func TestCacheRemove(t *testing.T) {
-	cache := NewLRUTxCache(100)
-	numTxs := 10
+func populate(cache TxCache, numTxs int) ([][]byte, error) {
 
 	txs := make([][]byte, numTxs)
 	for i := 0; i < numTxs; i++ {
 		// probability of collision is 2**-256
 		txBytes := make([]byte, 32)
 		_, err := rand.Read(txBytes)
-		require.NoError(t, err)
+
+		if err != nil {
+			return nil, err
+		}
 
 		txs[i] = txBytes
 		cache.Push(txBytes)
-
-		// make sure its added to both the linked list and the map
-		require.Equal(t, i+1, len(cache.cacheMap))
-		require.Equal(t, i+1, cache.list.Len())
 	}
 
+	return txs, nil
+}
+
+func TestCacheRemove(t *testing.T) {
+	cache := NewLRUTxCache(100)
+	numTxs := 10
+
+	txs, err := populate(cache, numTxs)
+	require.NoError(t, err)
+	require.Equal(t, numTxs, len(cache.cacheMap))
+	require.Equal(t, numTxs, cache.list.Len())
+
 	for i := 0; i < numTxs; i++ {
 		cache.Remove(txs[i])
 		// make sure its removed from both the map and the linked list
 		require.Equal(t, numTxs-(i+1), len(cache.cacheMap))
 		require.Equal(t, numTxs-(i+1), cache.list.Len())
 	}
 }
+
+func TestCacheRemoveByKey(t *testing.T) {
+	cache := NewLRUTxCache(100)
+	numTxs := 10
+
+	txs, err := populate(cache, numTxs)
+	require.NoError(t, err)
+	require.Equal(t, numTxs, len(cache.cacheMap))
+	require.Equal(t, numTxs, cache.list.Len())
+
+	for i := 0; i < numTxs; i++ {
+		cache.RemoveTxByKey(types.Tx(txs[i]).Key())
+		// make sure its removed from both the map and the linked list
+		require.Equal(t, numTxs-(i+1), len(cache.cacheMap))
+		require.Equal(t, numTxs-(i+1), cache.list.Len())
+	}
+}

mempool/v1/mempool.go

@@ -1,7 +1,6 @@
 package v1
 
 import (
-	"crypto/sha256"
 	"fmt"
 	"runtime"
 	"sort"
@@ -401,22 +400,25 @@ func (txmp *TxMempool) Update(
 	}
 
 	for i, tx := range blockTxs {
+		// The tx hash corresponding to the originating transaction. Set to the
+		// current tx hash initially, but can change in case of a malleated tx.
+		originalKey := tx.Key()
+
+		// Regardless of outcome, remove the transaction from the mempool.
+		if err := txmp.removeTxByKey(originalKey); err != nil {
+			if originalHash, _, isMalleated := types.UnwrapMalleatedTx(tx); isMalleated {
+				copy(originalKey[:], originalHash)
+				_ = txmp.removeTxByKey(originalKey)
+			}
+		}
+
 		// Add successful committed transactions to the cache (if they are not
 		// already present).  Transactions that failed to commit are removed from
 		// the cache unless the operator has explicitly requested we keep them.
 		if deliverTxResponses[i].Code == abci.CodeTypeOK {
 			_ = txmp.cache.Push(tx)
 		} else if !txmp.config.KeepInvalidTxsInCache {
-			txmp.cache.Remove(tx)
-		}
-
-		// Regardless of success, remove the transaction from the mempool.
-		if err := txmp.removeTxByKey(tx.Key()); err != nil {
-			if originalHash, _, isMalleated := types.UnwrapMalleatedTx(tx); isMalleated {
-				var originalKey [sha256.Size]byte
-				copy(originalKey[:], originalHash)
-				_ = txmp.removeTxByKey(types.TxKey(originalKey))
-			}
+			txmp.cache.RemoveTxByKey(originalKey)
 		}
 	}