From 750a842c8806ac37776ec3029c851a93ed47fe4e Mon Sep 17 00:00:00 2001
From: p4u <pau@dabax.net>
Date: Tue, 16 Jan 2024 18:36:58 +0100
Subject: [PATCH] vochain: remove invalid transactions from mempool

On process proposal, if the proposal is rejected, remove the
invalid transactions. This way we do not contribute on propagate
invalid transactions on the network.

Fix prepare proposal invalid transaction deletion by correctly
increasing the failedCount counter.

Signed-off-by: p4u <pau@dabax.net>
---
 vochain/app.go      |  6 ++++--
 vochain/cometbft.go | 21 +++++++++++----------
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/vochain/app.go b/vochain/app.go
index 9c2899cd9..035a9c785 100644
--- a/vochain/app.go
+++ b/vochain/app.go
@@ -113,6 +113,7 @@ type DeliverTxResponse struct {
 	Log  string
 	Info string
 	Data []byte
+	TxID [32]byte
 }
 
 // ExecuteBlockResponse is the response returned by ExecuteBlock after executing the block.
@@ -174,7 +175,7 @@ func (app *BaseApplication) ExecuteBlock(txs [][]byte, height uint32, blockTime
 				"data", string(resp.Data),
 				"info", resp.Info,
 				"log", resp.Log)
-			invalidTxs = append(invalidTxs, [32]byte{})
+			invalidTxs = append(invalidTxs, resp.TxID)
 		}
 		result = append(result, resp)
 	}
@@ -236,7 +237,7 @@ func (app *BaseApplication) deliverTx(rawTx []byte) *DeliverTxResponse {
 	response, err := app.TransactionHandler.CheckTx(tx, true)
 	if err != nil {
 		log.Errorw(err, "rejected tx")
-		return &DeliverTxResponse{Code: 1, Data: []byte(err.Error())}
+		return &DeliverTxResponse{Code: 1, TxID: tx.TxID, Data: []byte(err.Error())}
 	}
 	app.txReferences.Delete(tx.TxID)
 	// call event listeners
@@ -248,6 +249,7 @@ func (app *BaseApplication) deliverTx(rawTx []byte) *DeliverTxResponse {
 		Data: response.Data,
 		Info: fmt.Sprintf("%x", response.TxHash),
 		Log:  response.Log,
+		TxID: tx.TxID,
 	}
 }
 
diff --git a/vochain/cometbft.go b/vochain/cometbft.go
index abde9fe5d..2a2cdcc60 100644
--- a/vochain/cometbft.go
+++ b/vochain/cometbft.go
@@ -414,8 +414,13 @@ func (app *BaseApplication) PrepareProposal(ctx context.Context,
 					log.Debugf("transaction %x has reached max attempts, remove from mempool", txInfo.DecodedTx.TxID)
 					app.MempoolDeleteTx(txInfo.DecodedTx.TxID)
 					app.txReferences.Delete(txInfo.DecodedTx.TxID)
+				} else {
+					app.txReferences.Store(txInfo.DecodedTx.TxID, val)
 				}
 			}
+			log.Warnw("transaction reference not found on prepare proposal!",
+				"hash", fmt.Sprintf("%x", txInfo.DecodedTx.TxID),
+			)
 			continue
 		}
 		validTxs = append(validTxs, txInfo.Data)
@@ -459,23 +464,19 @@ func (app *BaseApplication) ProcessProposal(_ context.Context,
 	}
 
 	startTime := time.Now()
-
-	// TODO: check if we can enable this check without breaking consensus
-	//
-	// Check if the time difference is within the allowed threshold
-	//	timeDiff := startTime.Sub(req.Time)
-	//	if timeDiff > allowedConsensusTimeDiff {
-	//		return nil, fmt.Errorf("the time difference for the proposal exceeds the threshold")
-	//	}
-
 	resp, err := app.ExecuteBlock(req.Txs, uint32(req.GetHeight()), req.GetTime())
 	if err != nil {
 		return nil, fmt.Errorf("cannot execute block on process proposal: %w", err)
 	}
-	// invalid txx on a proposed block, should actually never happened if proposer acts honestly
+	// invalid tx on a proposed block, should actually never happened if proposer acts honestly
 	if len(resp.InvalidTransactions) > 0 {
 		log.Warnw("invalid transactions on process proposal", "height", app.Height(), "count", len(resp.InvalidTransactions),
 			"proposer", hex.EncodeToString(req.ProposerAddress), "action", "reject")
+		for _, tx := range resp.InvalidTransactions {
+			// remove transaction from mempool, just in case we have it
+			app.MempoolDeleteTx(tx)
+			log.Debugw("remove invalid tx from mempool", "hash", fmt.Sprintf("%x", tx))
+		}
 		return &cometabcitypes.ProcessProposalResponse{
 			Status: cometabcitypes.PROCESS_PROPOSAL_STATUS_REJECT,
 		}, nil