BUG - StackingERC721 unstake amount

[pldespaigne]

The Bug

When unstaking user is asked to fill an amount parameter, but the code does not take it into account:

function unstake(uint256 amount, bool isTokenA) external { // ask user for amount to unstake
  stakingContractEth.unstakeFrom(msg.sender, amount);

  uint256[] storage tokens;
  if(isTokenA)
    tokens = tokenOwnershipA[msg.sender];
  else
    tokens = tokenOwnershipB[msg.sender];

  for(uint i = tokens.length - 1; i > 0; i--) { // unstake every tokens no matter the amount specified by the user
    uint256 id = tokens[i];
    tokens.pop();
    if(isTokenA)
      tokenA.transferFrom(address(this), msg.sender, id);
    else
      tokenB.transferFrom(address(this), msg.sender, id);
  }

Proposed Fix

We want to remove amount number of elements from the array. Because we use .pop(), we remove elements starting from the end of the array.

  for(uint i = 0; i < amount; i++) { // we run the loop 'amount' times (if amount > length, the function revert at the beginning)
    uint256 id = tokens[tokens.length - 1]; // retrieve the last elements before poping it
    tokens.pop();
    if(isTokenA)
      tokenA.transferFrom(address(this), msg.sender, id);
    else
      tokenB.transferFrom(address(this), msg.sender, id);
  }

[stalker474]

Please add a require for amount to be less or equal to tokens length also

[pldespaigne]

If amount is greater than tokens.length, then the first line (stakingContractEth.unstakeFrom(msg.sender, amount); ) will already revert with "Distribute: Dont have enough staked".

Do you still want me to add a require() ?

[stalker474]

No you're right. Maybe just add a comment then, saying that at this point amount cant be greater than length (for auditors)