From c75fa3fe6ea3885f76bfd61e358fb69089cf643a Mon Sep 17 00:00:00 2001
From: istae <14264581+istae@users.noreply.github.com>
Date: Tue, 22 Oct 2024 18:23:59 +0300
Subject: [PATCH] fix: update height if staked when height decreases (#4873)

---
 pkg/node/node.go                              | 34 ++++-----
 pkg/storageincentives/staking/contract.go     |  4 +
 .../staking/contract_test.go                  | 73 +++++++++++++++++--
 3 files changed, 87 insertions(+), 24 deletions(-)

diff --git a/pkg/node/node.go b/pkg/node/node.go
index 1fad6195e26..37500cb2a37 100644
--- a/pkg/node/node.go
+++ b/pkg/node/node.go
@@ -1006,12 +1006,14 @@ func NewBee(
 
 	if chainEnabled {
 
-		if changedOverlay {
-			stake, err := stakingContract.GetPotentialStake(ctx)
-			if err != nil {
-				return nil, err
-			}
-			if stake.Cmp(big.NewInt(0)) > 0 {
+		stake, err := stakingContract.GetPotentialStake(ctx)
+		if err != nil {
+			return nil, err
+		}
+
+		if stake.Cmp(big.NewInt(0)) > 0 {
+
+			if changedOverlay {
 				logger.Debug("changing overlay address in staking contract")
 				tx, err := stakingContract.ChangeStakeOverlay(ctx, common.BytesToHash(nonce))
 				if err != nil {
@@ -1019,23 +1021,17 @@ func NewBee(
 				}
 				logger.Info("overlay address changed in staking contract", "transaction", tx)
 			}
-		}
-
-		// make sure that the staking contract has the up to date height
-		tx, updated, err := stakingContract.UpdateHeight(ctx)
-		if err != nil {
-			return nil, err
-		}
-		if updated {
-			logger.Info("updated new reserve capacity doubling height in the staking contract", "transaction", tx, "new_height", o.ReserveCapacityDoubling)
-		}
 
-		if o.ReserveCapacityDoubling > 0 {
-			stake, err := stakingContract.GetPotentialStake(ctx)
+			// make sure that the staking contract has the up to date height
+			tx, updated, err := stakingContract.UpdateHeight(ctx)
 			if err != nil {
 				return nil, err
 			}
-			if stake.Cmp(big.NewInt(0)) > 0 {
+			if updated {
+				logger.Info("updated new reserve capacity doubling height in the staking contract", "transaction", tx, "new_height", o.ReserveCapacityDoubling)
+			}
+
+			if o.ReserveCapacityDoubling > 0 {
 				// Check if the staked amount is sufficient to cover the additional neighborhoods.
 				// The staked amount must be at least 2^h * MinimumStake.
 				if stake.Cmp(big.NewInt(0).Mul(big.NewInt(1<<o.ReserveCapacityDoubling), staking.MinimumStakeAmount)) < 0 {
diff --git a/pkg/storageincentives/staking/contract.go b/pkg/storageincentives/staking/contract.go
index cf95baec874..1ff87c6e2a4 100644
--- a/pkg/storageincentives/staking/contract.go
+++ b/pkg/storageincentives/staking/contract.go
@@ -103,6 +103,10 @@ func (c *contract) DepositStake(ctx context.Context, stakedAmount *big.Int) (com
 		}
 	}
 
+	if big.NewInt(0).Add(prevStakedAmount, stakedAmount).Cmp(big.NewInt(0).Mul(big.NewInt(1<<c.height), MinimumStakeAmount)) < 0 {
+		return common.Hash{}, fmt.Errorf("stake amount does not sufficiently cover the additional reserve capacity: %w", ErrInsufficientStakeAmount)
+	}
+
 	balance, err := c.getBalance(ctx)
 	if err != nil {
 		return common.Hash{}, err
diff --git a/pkg/storageincentives/staking/contract_test.go b/pkg/storageincentives/staking/contract_test.go
index 080ed13c324..23cb9ac548d 100644
--- a/pkg/storageincentives/staking/contract_test.go
+++ b/pkg/storageincentives/staking/contract_test.go
@@ -290,10 +290,73 @@ func TestDepositStake(t *testing.T) {
 		}
 	})
 
-	t.Run("insufficient stake amount", func(t *testing.T) {
+	t.Run("sufficient stake amount extra height", func(t *testing.T) {
 		t.Parallel()
 
-		totalAmount := big.NewInt(0)
+		balance := big.NewInt(0).Mul(staking.MinimumStakeAmount, big.NewInt(2))
+		prevStake := staking.MinimumStakeAmount
+
+		expectedCallData, err := stakingContractABI.Pack("manageStake", nonce, staking.MinimumStakeAmount, stakingHeight)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		contract := staking.New(
+			owner,
+			stakingContractAddress,
+			stakingContractABI,
+			bzzTokenAddress,
+			transactionMock.New(
+				transactionMock.WithSendFunc(func(ctx context.Context, request *transaction.TxRequest, boost int) (txHash common.Hash, err error) {
+					if *request.To == bzzTokenAddress {
+						return txHashApprove, nil
+					}
+					if *request.To == stakingContractAddress {
+						if !bytes.Equal(expectedCallData[:80], request.Data[:80]) {
+							return common.Hash{}, fmt.Errorf("got wrong call data. wanted %x, got %x", expectedCallData, request.Data)
+						}
+						return txHashDeposited, nil
+					}
+					return common.Hash{}, errors.New("sent to wrong contract")
+				}),
+				transactionMock.WithWaitForReceiptFunc(func(ctx context.Context, txHash common.Hash) (receipt *types.Receipt, err error) {
+					if txHash == txHashDeposited {
+						return &types.Receipt{
+							Status: 1,
+						}, nil
+					}
+					if txHash == txHashApprove {
+						return &types.Receipt{
+							Status: 1,
+						}, nil
+					}
+					return nil, errors.New("unknown tx hash")
+				}),
+				transactionMock.WithCallFunc(func(ctx context.Context, request *transaction.TxRequest) (result []byte, err error) {
+					if *request.To == bzzTokenAddress {
+						return balance.FillBytes(make([]byte, 32)), nil
+					}
+					if *request.To == stakingContractAddress {
+						return getPotentialStakeResponse(t, prevStake), nil
+					}
+					return nil, errors.New("unexpected call")
+				}),
+			),
+			nonce,
+			false,
+			1,
+		)
+
+		_, err = contract.DepositStake(ctx, stakedAmount)
+		if err != nil {
+			t.Fatal(err)
+		}
+	})
+
+	t.Run("insufficient stake amount extra height", func(t *testing.T) {
+		t.Parallel()
+
+		totalAmount := big.NewInt(0).Mul(staking.MinimumStakeAmount, big.NewInt(10))
 		prevStake := big.NewInt(0)
 
 		contract := staking.New(
@@ -314,11 +377,11 @@ func TestDepositStake(t *testing.T) {
 			),
 			nonce,
 			false,
-			stakingHeight,
+			1,
 		)
 
-		_, err := contract.DepositStake(ctx, big.NewInt(100000000000000000))
-		if !errors.Is(err, staking.ErrInsufficientFunds) {
+		_, err := contract.DepositStake(ctx, stakedAmount)
+		if !errors.Is(err, staking.ErrInsufficientStakeAmount) {
 			t.Fatal(fmt.Errorf("wanted %w, got %w", staking.ErrInsufficientStakeAmount, err))
 		}
 	})