From 3d8be699594b2fc87dd13dfe34ad6a090601660b Mon Sep 17 00:00:00 2001
From: Thiago Coimbra Lemos <tclemos@users.noreply.github.com>
Date: Wed, 19 Apr 2023 13:51:04 -0300
Subject: [PATCH] Improve claim checks (#2027)

---
 pool/pgpoolstorage/pgpoolstorage.go | 12 ++++++++----
 pool/pool.go                        | 26 +++++++++++++++++++++++++-
 pool/pool_test.go                   | 20 ++++++++++++++++----
 3 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/pool/pgpoolstorage/pgpoolstorage.go b/pool/pgpoolstorage/pgpoolstorage.go
index 883b597de2..1ef337f0bb 100644
--- a/pool/pgpoolstorage/pgpoolstorage.go
+++ b/pool/pgpoolstorage/pgpoolstorage.go
@@ -683,12 +683,16 @@ func (p *PostgresPoolStorage) GetAllAddressesBlocked(ctx context.Context) ([]com
 	return addrs, nil
 }
 
-// DepositCountExists checks if already exists a transaction in the pool with the
-// provided deposit count
+// DepositCountExists checks if already exists a `pending` or `selected` transaction
+// in the pool with the provided deposit count
 func (p *PostgresPoolStorage) DepositCountExists(ctx context.Context, depositCount uint64) (bool, error) {
 	var exists bool
-	req := "SELECT exists (SELECT 1 FROM pool.transaction WHERE deposit_count = $1)"
-	err := p.db.QueryRow(ctx, req, depositCount).Scan(&exists)
+	req := `
+        SELECT EXISTS (SELECT 1
+                         FROM pool.transaction
+                        WHERE deposit_count = $1
+                          AND status IN ($2, $3))`
+	err := p.db.QueryRow(ctx, req, depositCount, pool.TxStatusPending, pool.TxStatusSelected).Scan(&exists)
 	if err != nil && err != sql.ErrNoRows {
 		return false, err
 	}
diff --git a/pool/pool.go b/pool/pool.go
index 28c678bcfe..95a0c98f3e 100644
--- a/pool/pool.go
+++ b/pool/pool.go
@@ -179,11 +179,15 @@ func (p *Pool) StoreTx(ctx context.Context, tx types.Transaction, ip string, isW
 		}
 	}
 
+	// CLAIM CHECK
 	if poolTx.IsClaims {
 		isFreeTx := poolTx.GasPrice().Cmp(big.NewInt(0)) <= 0
+		// if the tx is free and it was reverted in the pre execution
+		// the transaction gets rejected
 		if isFreeTx && preExecutionResponse.isReverted {
 			return fmt.Errorf("free claim reverted")
-		} else {
+		} else { // otherwise
+			// DEPOSIT COUNT CHECK
 			depositCount, err := p.extractDepositCountFromClaimTx(poolTx)
 			if err != nil {
 				return err
@@ -192,10 +196,30 @@ func (p *Pool) StoreTx(ctx context.Context, tx types.Transaction, ip string, isW
 			if err != nil && !errors.Is(err, ErrNotFound) {
 				return err
 			}
+			// if the claim deposit count already exist in the pool,
+			// the transaction gets rejected
 			if exists {
 				return fmt.Errorf("deposit count already exists")
 			}
 
+			// CURRENT NONCE CHECK
+			from, err := state.GetSender(poolTx.Transaction)
+			if err != nil {
+				return ErrInvalidSender
+			}
+			lastL2Block, err := p.state.GetLastL2Block(ctx, nil)
+			if err != nil {
+				return err
+			}
+			nonce, err := p.state.GetNonce(ctx, from, lastL2Block.Root())
+			if err != nil {
+				return err
+			}
+			// if the nonce is different from the current nonce for the
+			// account sending the claim, the transaction gets rejected
+			if poolTx.Nonce() != nonce {
+				return fmt.Errorf("invalid nonce")
+			}
 			poolTx.DepositCount = depositCount
 		}
 	}
diff --git a/pool/pool_test.go b/pool/pool_test.go
index a1e7bd4870..f23b3de568 100644
--- a/pool/pool_test.go
+++ b/pool/pool_test.go
@@ -1112,7 +1112,7 @@ func Test_AddTx_GasPriceErr(t *testing.T) {
 		},
 		{
 			name:          "NoGasPriceTooLowErr_ForClaims",
-			nonce:         1,
+			nonce:         0,
 			to:            &l2BridgeAddr,
 			gasLimit:      cfg.FreeClaimGasLimit,
 			gasPrice:      big.NewInt(0),
@@ -1430,16 +1430,28 @@ func Test_AvoidDuplicatedClaims(t *testing.T) {
 	auth.GasPrice = big.NewInt(0)
 	auth.Nonce = big.NewInt(0)
 
-	signedTx, err := bridgeSC.ClaimAsset(auth, [32][32]byte{}, uint32(123456789), [32]byte{}, [32]byte{}, 69, common.Address{}, uint32(20), common.Address{}, big.NewInt(0), []byte{})
+	depositCount := uint32(123456789)
+	signedTx, err := bridgeSC.ClaimAsset(auth, [32][32]byte{}, depositCount, [32]byte{}, [32]byte{}, 69, common.Address{}, uint32(20), common.Address{}, big.NewInt(0), []byte{})
 	require.NoError(t, err)
 
+	// add claim
 	err = p.AddTx(ctx, *signedTx, "")
 	require.NoError(t, err)
 
-	auth.Nonce = big.NewInt(1)
-	signedTx, err = bridgeSC.ClaimAsset(auth, [32][32]byte{}, uint32(123456789), [32]byte{}, [32]byte{}, 69, common.Address{}, uint32(20), common.Address{}, big.NewInt(0), []byte{})
+	auth.GasLimit++
+	signedTx, err = bridgeSC.ClaimAsset(auth, [32][32]byte{}, depositCount, [32]byte{}, [32]byte{}, 69, common.Address{}, uint32(20), common.Address{}, big.NewInt(0), []byte{})
 	require.NoError(t, err)
 
+	// add claim with same deposit count
 	err = p.AddTx(ctx, *signedTx, "")
 	require.Equal(t, err.Error(), "deposit count already exists")
+
+	auth.Nonce = big.NewInt(1)
+	depositCount = uint32(12345678)
+	signedTx, err = bridgeSC.ClaimAsset(auth, [32][32]byte{}, depositCount, [32]byte{}, [32]byte{}, 69, common.Address{}, uint32(20), common.Address{}, big.NewInt(0), []byte{})
+	require.NoError(t, err)
+
+	// add claim with wrong nonce
+	err = p.AddTx(ctx, *signedTx, "")
+	require.Equal(t, err.Error(), "invalid nonce")
 }