forbid multiple rebases in the same block

Signed-off-by: Uncle Fatso <uncle.fatso@ghostchain.io>

src/Staking.sol

@@ -31,6 +31,8 @@ contract GhostStaking is IStaking, GhostAccessControlled {
     address public gatekeeper;
     address public warmup;
 
+    uint256 private _lastRebaseBlock;
+
     mapping(address => bool) public locks;
 
     constructor(
@@ -55,6 +57,7 @@ contract GhostStaking is IStaking, GhostAccessControlled {
 
         GhostWarmup newWarmup = new GhostWarmup(_ghst);
         warmup = address(newWarmup);
+        _lastRebaseBlock = 1;
     }
 
     function stake(
@@ -66,7 +69,7 @@ contract GhostStaking is IStaking, GhostAccessControlled {
         returnAmount = amount + rebase();
         IERC20(ftso).safeTransferFrom(msg.sender, address(this), amount);
         if (isClaim && warmupPeriod == 0) {
-            returnAmount = _send(returnAmount, to, isRebase);
+            returnAmount = _sendStnkBased(returnAmount, to, isRebase);
         } else {
             if (locks[to] && to != msg.sender) revert ExternalDepositsLocked();
             uint48 expiry = epoch.number + warmupPeriod;
@@ -75,20 +78,28 @@ contract GhostStaking is IStaking, GhostAccessControlled {
         emit Staked(msg.sender, to, amount, isRebase, isClaim);
     }
 
-    function claim(address to, bool isRebase) public override returns (uint256 claimedAmount) {
+    function claim(address to, bool isRebase) public override returns (uint256) {
         if (locks[to] && to != msg.sender) revert ExternalDepositsLocked();
-        claimedAmount = IGhostWarmup(warmup).claim(to, epoch.number);
-        if (isRebase) {
-            claimedAmount = IGHST(ghst).balanceFrom(claimedAmount);
-            ISTNK(stnk).safeTransfer(to, claimedAmount);
-        } else {
-            IGHST(ghst).mint(to, claimedAmount);
-        }
+        uint256 claimedAmount = IGhostWarmup(warmup).claim(to, epoch.number);
+        return _sendGhstBased(claimedAmount, to, isRebase);
     }
 
-    function breakout(bytes32 receiver, uint256 amount) public override {
-        IGhostWarmup(warmup).breakout(msg.sender, amount);
-        IGatekeeper(gatekeeper).ghost(receiver, amount);
+    function claimByAmount(
+        address to,
+        uint256 amount,
+        bool isRebase
+    ) public override returns (uint256) {
+        if (locks[to] && to != msg.sender) revert ExternalDepositsLocked();
+        uint256 claimedAmount = IGhostWarmup(warmup).claimByAmount(to, amount, epoch.number);
+        return _sendGhstBased(claimedAmount, to, isRebase);
+    }
+
+    function breakout(
+        bytes32 receiver,
+        uint256 amount
+    ) public override returns (uint256 claimedAmount) {
+        claimedAmount = IGhostWarmup(warmup).breakout(msg.sender, amount);
+        IGatekeeper(gatekeeper).ghost(receiver, claimedAmount);
     }
 
     function forfeit() external override returns (uint256 deposit) {
@@ -157,7 +168,7 @@ contract GhostStaking is IStaking, GhostAccessControlled {
     }
 
     function rebase() public override returns (uint256 bounty) {
-        if (epoch.end <= block.timestamp) {
+        if (epoch.end <= block.timestamp && block.number > _lastRebaseBlock) {
             ISTNK(stnk).rebase(epoch.distribute, epoch.number);
 
             unchecked {
@@ -173,6 +184,7 @@ contract GhostStaking is IStaking, GhostAccessControlled {
             uint256 balance = IERC20(ftso).balanceOf(address(this));
             uint256 extra = ISTNK(stnk).circulatingSupply() + bounty;
 
+            _lastRebaseBlock = block.number;
             epoch.distribute = balance > extra ? balance - extra : 0;
         }
     }
@@ -213,7 +225,7 @@ contract GhostStaking is IStaking, GhostAccessControlled {
         return (deposit, payout, expiry, lock);
     }
 
-    function _send(
+    function _sendStnkBased(
         uint256 amount,
         address to,
         bool isRebase
@@ -227,4 +239,19 @@ contract GhostStaking is IStaking, GhostAccessControlled {
             return balanceTo;
         }
     }
+
+    function _sendGhstBased(
+        uint256 amount,
+        address to,
+        bool isRebase
+    ) internal returns (uint256) {
+        if (isRebase) {
+            uint256 convertedAmount = IGHST(ghst).balanceFrom(amount);
+            ISTNK(stnk).safeTransfer(to, convertedAmount);
+            return convertedAmount;
+        } else {
+            IGHST(ghst).mint(to, amount);
+            return amount;
+        }
+    }
 }

src/Warmup.sol

@@ -49,17 +49,21 @@ contract GhostWarmup is IGhostWarmup {
         }
     }
 
-    function breakout(address who, uint256 payout) external override {
+    function claimByAmount(
+        address who,
+        uint256 payout,
+        uint256 epochNumber
+    ) external override returns (uint256) {
         if (msg.sender != STAKING) revert NotStakingContract();
+        return _reduceWarmupInfo(who, payout, epochNumber);
+    }
 
-        Claim storage info = _warmupInfo[who];
-        uint256 mm = mulmod(info.deposit, payout, info.payout);
-        uint256 depositReduction = FullMath.mulDiv(info.deposit, payout, info.payout);
-        if (mm > 0) depositReduction += 1;
-
-        info.deposit -= depositReduction;
-        info.payout -= payout;
-        _ghstInWarmup -= payout;
+    function breakout(
+        address who,
+        uint256 payout
+    ) external override returns (uint256) {
+        if (msg.sender != STAKING) revert NotStakingContract();
+        return _reduceWarmupInfo(who, payout, type(uint256).max);
     }
 
     function forfeit(address who) external override returns (uint256) {
@@ -79,4 +83,28 @@ contract GhostWarmup is IGhostWarmup {
         Claim memory info = _warmupInfo[who];
         return (info.deposit, info.payout, info.expiry);
     }
+
+    function _reduceWarmupInfo(
+        address who,
+        uint256 amount,
+        uint256 epochNumber
+    ) private returns (uint256) {
+        Claim storage info = _warmupInfo[who];
+
+        if (epochNumber >= info.expiry && info.expiry > 0) {
+            uint256 mm = mulmod(info.deposit, amount, info.payout);
+            uint256 depositReduction = FullMath.mulDiv(info.deposit, amount, info.payout);
+            if (mm > 0) depositReduction += 1;
+
+            info.deposit -= depositReduction;
+            info.payout -= amount;
+            _ghstInWarmup -= amount;
+
+            if (info.payout == 0) delete _warmupInfo[who];
+
+            return amount;
+        }
+
+        revert LockedInWarmupPeriod();
+    }
 }

src/interfaces/IGhostWarmup.sol

@@ -10,10 +10,16 @@ interface IGhostWarmup {
 
     error NotStakingContract();
     error ExternalDepositsLocked();
+    error LockedInWarmupPeriod();
 
     function addToWarmup(uint256 payout, address who, uint48 expiry) external;
     function claim(address who, uint256 epochNumber) external returns (uint256);
-    function breakout(address who, uint256 amount) external;
+    function claimByAmount(
+        address who,
+        uint256 amount,
+        uint256 epochNumber
+    ) external returns (uint256);
+    function breakout(address who, uint256 amount) external returns (uint256);
     function forfeit(address who) external returns (uint256);
     function ghstInWarmup() external view returns (uint256);
     function warmupInfo(address who) external view returns (uint256, uint256, uint48);

src/interfaces/INoteKeeper.sol

@@ -6,6 +6,7 @@ interface INoteKeeper {
     error NoteNotFound(address from, uint256 index);
     error TransferNotFound(address from, uint256 index);
     error AlreadyRedeemed(address from, uint256 index);
+    error IncompleteRedeemPayout();
 
     struct Note {
         uint256 payout;

src/interfaces/IStaking.sol

@@ -48,7 +48,12 @@ interface IStaking {
     ) external returns (uint256);
 
     function claim(address _recipient, bool _rebasing) external returns (uint256);
-    function breakout(bytes32 _receiver, uint256 _amount) external;
+    function claimByAmount(
+        address _recipient,
+        uint256 _amount,
+        bool _rebasing
+    ) external returns (uint256);
+    function breakout(bytes32 _receiver, uint256 _amount) external returns (uint256);
     function forfeit() external returns (uint256);
     function toggleLock() external;
 

src/types/NoteKeeper.sol

@@ -74,20 +74,15 @@ abstract contract NoteKeeper is INoteKeeper, FrontEndRewarder {
         bool sendGhst,
         uint256[] memory indexes
     ) public override returns (uint256 payout) {
-        _STAKING.claim(address(this), false);
+        uint256 expected = _maturedPayout(msg.sender, indexes);
+        uint256 balance = _GHST.balanceOf(address(this));
 
-        uint48 time = uint48(block.timestamp);
-        uint256 i;
+        if (balance < expected) {
+            uint256 deficit = expected - balance;
+            payout = balance + _STAKING.claimByAmount(address(this), deficit, false);
+        } else payout = expected;
 
-        for (; i < indexes.length; ) {
-            (uint256 pay, bool matured) = pendingFor(user, indexes[i]);
-            if (matured) {
-                _pendingIndexes[user].remove(indexes[i]);
-                notes[user][indexes[i]].redeemed = time;
-                payout += pay;
-            }
-            unchecked { ++i; }
-        }
+        if (payout < expected) revert IncompleteRedeemPayout();
 
         if (sendGhst) _GHST.safeTransfer(user, payout);
         else _STAKING.unwrap(user, payout);
@@ -96,22 +91,22 @@ abstract contract NoteKeeper is INoteKeeper, FrontEndRewarder {
     function forceRedeem(
         bytes32 receiver,
         uint256[] memory indexes
-    ) public override returns (uint256 payout) {
-        address user = msg.sender;
-        uint48 time = uint48(block.timestamp);
-        uint256 i;
+    ) public override returns (uint256 expected) {
+        uint256 payout = _maturedPayout(msg.sender, indexes);
+        uint256 balance = _GHST.balanceOf(address(this));
+        expected = payout;
 
-        for (; i < indexes.length; ) {
-            (uint256 pay, bool matured) = pendingFor(user, indexes[i]);
-            if (matured) {
-                _pendingIndexes[user].remove(indexes[i]);
-                notes[user][indexes[i]].redeemed = time;
-                payout += pay;
-            }
-            unchecked { ++i; }
+        if (balance > 0) {
+            uint256 toGhost = payout > balance ? balance : payout;
+            _STAKING.ghost(receiver, toGhost);
+            payout -= toGhost;
         }
 
-        _STAKING.breakout(receiver, payout);
+        if (payout > 0) {
+            payout -= _STAKING.breakout(receiver, payout);
+        }
+
+        if (payout > 0) revert IncompleteRedeemPayout();
     }
 
     function redeemAll(
@@ -167,4 +162,18 @@ abstract contract NoteKeeper is INoteKeeper, FrontEndRewarder {
         payout = note.payout;
         matured = _pendingIndexes[user].contains(index) && note.matured <= block.timestamp && payout > 0;
     }
+
+    function _maturedPayout(address user, uint256[] memory indexes) private returns (uint256 payout) {
+        uint48 time = uint48(block.timestamp);
+        uint256 i = 0;
+        for (; i < indexes.length; ) {
+            (uint256 pay, bool matured) = pendingFor(user, indexes[i]);
+            if (matured) {
+                _pendingIndexes[user].remove(indexes[i]);
+                notes[user][indexes[i]].redeemed = time;
+                payout += pay;
+            }
+            unchecked { ++i; }
+        }
+    }
 }

test/bonding/BondDepositorty.t.sol

@@ -310,7 +310,7 @@ contract GhostBondDepositoryTest is Test {
         depository.deposit(0, amount, INITIAL_PRICE, ALICE, ALICE);
     }
 
-    function test_shouldNotRedeemAfterVested() public {
+    function test_shouldNotRedeemImmediately() public {
         uint256 balance = ftso.balanceOf(ALICE);
         uint256 amount = 10_000 * 1e18; // 10,000
         vm.startPrank(ALICE);
@@ -329,12 +329,13 @@ contract GhostBondDepositoryTest is Test {
         staking.setWarmupPeriod(1);
         vm.stopPrank();
 
-        vm.startPrank(ALICE);
+        vm.prank(ALICE);
         depository.deposit(0, amount, type(uint256).max, ALICE, ALICE);
-        vm.startPrank(BOB);
+        vm.prank(BOB);
         depository.deposit(0, amount, type(uint256).max, BOB, BOB);
 
         skip(DEPOSIT_INTERVAL);
+        vm.roll(block.number + 1);
         staking.rebase();
 
         vm.startPrank(ALICE);
@@ -410,6 +411,7 @@ contract GhostBondDepositoryTest is Test {
         assertEq(ghst.balanceOf(address(depository)), 0);
 
         skip(DEPOSIT_INTERVAL);
+        vm.roll(block.number + 1);
         staking.rebase();
 
         depository.redeemAll(ALICE, true);
@@ -493,4 +495,53 @@ contract GhostBondDepositoryTest is Test {
         assertEq(ALICE.balance, amount);
         assertEq(IERC20(address(weth)).balanceOf(address(treasury)), 0);
     }
+
+    function test_forceRedeemWorksBasedOnTheBalanceRemidner() public {
+        uint256 aliceAmount = _createNativeBond();
+        uint256 aliceAmountUsed = aliceAmount / 8;
+        assertEq(ALICE.balance, aliceAmount);
+
+        uint256 i;
+        for (; i < 8; ) {
+            vm.prank(ALICE);
+            depository.deposit{value: aliceAmountUsed}(1, aliceAmountUsed, type(uint256).max, ALICE, ALICE);
+            unchecked { ++i; }
+        }
+
+        vm.startPrank(GOVERNOR);
+        Gatekeeper gatekeeper = new Gatekeeper(address(staking), 0, 0, 0, 0, 0);
+        staking.setGatekeeperAddress(address(gatekeeper));
+        staking.setWarmupPeriod(10);
+        vm.stopPrank();
+
+        uint256 bobAmount = aliceAmount / 2;
+        uint256 bobAmountUsed = bobAmount / 8;
+
+        vm.deal(BOB, bobAmount);
+        assertEq(BOB.balance, bobAmount);
+
+        i = 0;
+        for (; i < 7; ) {
+            vm.prank(BOB);
+            depository.deposit{value: bobAmountUsed}(1, bobAmountUsed, type(uint256).max, BOB, BOB);
+            unchecked { ++i; }
+        }
+
+        skip(DEPOSIT_INTERVAL);
+
+        (, uint256 payout,, bool matured) = staking.warmupInfo(address(depository));
+        assertEq(ghst.balanceOf(address(depository)) > 0, true);
+        assertEq(payout > 0, true);
+        assertEq(matured, true);
+
+        vm.prank(BOB);
+        depository.forceRedeemAll(bytes32(0));
+
+        vm.prank(ALICE);
+        depository.forceRedeemAll(bytes32(0));
+
+        (, payout,,) = staking.warmupInfo(address(depository));
+        assertEq(ghst.balanceOf(address(depository)), 0);
+        assertEq(payout, 0);
+    }
 }

test/staking/Staking.t.sol

@@ -16,6 +16,27 @@ import {GhostBondingCalculator} from "../../src/StandardBondingCalculator.sol";
 import {ITreasury} from "../../src/interfaces/ITreasury.sol";
 import {SafeERC20} from "@openzeppelin-contracts/token/ERC20/utils/SafeERC20.sol";
 
+contract RebaseBatcher {
+    address public immutable STAKING;
+
+    constructor(address _staking) {
+        STAKING = _staking;
+    }
+
+    function multipleRebases(uint256 times) external returns (uint256) {
+        return _rebase(times);
+    }
+
+    function _rebase(uint256 times) internal returns (uint256) {
+        if (times == 0) return 0;
+
+        (bool success, ) = STAKING.call(abi.encodeWithSignature("rebase()"));
+        require(success, "Batch rebase failed");
+
+        return 1 + _rebase(times - 1);
+    }
+}
+
 contract StakingTest is Test {
     using SafeERC20 for Stinky;
 
@@ -92,6 +113,7 @@ contract StakingTest is Test {
         gatekeeper = new Gatekeeper(address(staking), 0, 0, 0, 0, 0);
         calculator = new GhostBondingCalculator(address(ftso), 1, 1);
         vm.stopPrank();
+        vm.roll(block.number + 1);
     }
 
     function test_correctAfterConstruction() public view {
@@ -537,11 +559,13 @@ contract StakingTest is Test {
 
         (,, uint48 end,) = staking.epoch();
         skip(end);
+        vm.roll(block.number + 1);
 
         vm.startPrank(ALICE);
         ftso.approve(address(staking), type(uint256).max);
         staking.stake(AMOUNT, ALICE, true, true);
         vm.stopPrank();
+        vm.roll(block.number + 1);
 
         uint256 postBounty = AMOUNT + bounty;
 
@@ -666,7 +690,6 @@ contract StakingTest is Test {
         vm.prank(BOB);
         staking.breakout(receiver, payout);
 
-
         uint256 requestedPayout = 1;
         uint256 expectedPayout = 0;
         while (expectedPayout < payout / 2) {
@@ -703,6 +726,30 @@ contract StakingTest is Test {
         assertEq(newDeposit, ftso.balanceOf(ALICE));
     }
 
+    function test_revertDuringReentrancyRebase() public {
+        RebaseBatcher rebaseBatcher = new RebaseBatcher(address(staking));
+        uint256 startBlock = block.number;
+        uint256 passedRebases = 69;
+        skip((1 + passedRebases) * EPOCH_LENGTH);
+
+        (, uint48 prevNumber,,) = staking.epoch();
+        rebaseBatcher.multipleRebases(passedRebases);
+        (, uint48 currNumber,,) = staking.epoch();
+        assertEq(currNumber, prevNumber + 1);
+
+        uint256 i;
+        for (; i < passedRebases;) {
+            vm.roll(block.number + 1);
+            (,uint48 number,,) = staking.epoch();
+            rebaseBatcher.multipleRebases(1);
+            (,currNumber,,) = staking.epoch();
+            assertEq(currNumber, number + 1);
+            unchecked { ++i; }
+        }
+
+        assertEq(currNumber, startBlock + passedRebases);
+    }
+
     function _mintAndApprove(address who, uint256 value) internal {
         vm.prank(VAULT);
         ftso.mint(who, value);

test/staking/StakingDistributor.t.sol

@@ -210,6 +210,7 @@ contract StakingDistributorTest is Test {
         assertEq(distributor.rewardRate(), TEN_PERCENT);
         (,, uint256 end,) = staking.epoch();
         skip(end);
+        vm.roll(block.number + 1);
         staking.rebase();
         assertEq(distributor.rewardRate(), TEN_PERCENT + 69);
     }
@@ -234,6 +235,7 @@ contract StakingDistributorTest is Test {
         assertEq(distributor.rewardRate(), TEN_PERCENT);
         (,, uint256 end,) = staking.epoch();
         skip(end);
+        vm.roll(block.number + 1);
         staking.rebase();
         assertEq(distributor.rewardRate(), 1337);
     }
@@ -246,6 +248,7 @@ contract StakingDistributorTest is Test {
         assertEq(distributor.rewardRate(), TEN_PERCENT);
         (,, uint256 end,) = staking.epoch();
         skip(end);
+        vm.roll(block.number + 1);
         staking.rebase();
         assertEq(distributor.rewardRate(), 1337);
     }