How to fix the ERC777 re-entrancy attack?

[3for]

The first weakness mentioned in Hexens audit report：

But the codes are still the same:

               // In order to support fee tokens check the amount received, not the transferred
                uint256 balanceBefore = IERC20Upgradeable(token).balanceOf(
                    address(this)
                );
                IERC20Upgradeable(token).safeTransferFrom(
                    msg.sender,
                    address(this),
                    amount
                );
                uint256 balanceAfter = IERC20Upgradeable(token).balanceOf(
                    address(this)
                );

                // Override leafAmount with the received amount
                leafAmount = balanceAfter - balanceBefore;

I want to learn how we fixed this attack?

[invocamanman]

Hey!, This type of attacks usually are solved with a reentrancy lib, and that exactly what we add in the function:
https://github.com/0xPolygonHermez/zkevm-contracts/blob/main/contracts/PolygonZkEVMBridge.sol#L149
You can check more info about this in the nonReentrant modifier here: https://github.com/OpenZeppelin/openzeppelin-contracts/blob/master/contracts/security/ReentrancyGuard.sol#L50

[3for]

@invocamanman Thanks a lot.