Fix ERC20._update (#3921)

Co-authored-by: Francisco <frangio.1@gmail.com>
Co-authored-by: Hadrien Croubois <hadrien.croubois@gmail.com>

contracts/mocks/ERC20Mock.sol

@@ -30,4 +30,9 @@ contract ERC20Mock is ERC20 {
     function approveInternal(address owner, address spender, uint256 value) public {
         _approve(owner, spender, value);
     }
+
+    // Exposed for testing purposes
+    function update(address from, address to, uint256 amount) public {
+        _update(from, to, amount);
+    }
 }

contracts/token/ERC20/ERC20.sol

@@ -229,25 +229,23 @@ contract ERC20 is Context, IERC20, IERC20Metadata {
     function _update(address from, address to, uint256 amount) internal virtual {
         if (from == address(0)) {
             _totalSupply += amount;
-            unchecked {
-                // Overflow not possible: balance + amount is at most totalSupply + amount, which is checked above.
-                _balances[to] += amount;
-            }
-        } else if (to == address(0)) {
-            uint256 fromBalance = _balances[from];
-            require(fromBalance >= amount, "ERC20: burn amount exceeds balance");
-            _totalSupply -= amount;
-            unchecked {
-                // Overflow not possible: amount <= fromBalance <= totalSupply.
-                _balances[from] = fromBalance - amount;
-            }
         } else {
             uint256 fromBalance = _balances[from];
             require(fromBalance >= amount, "ERC20: transfer amount exceeds balance");
             unchecked {
+                // Overflow not possible: amount <= fromBalance <= totalSupply.
                 _balances[from] = fromBalance - amount;
-                // Overflow not possible: the sum of all balances is capped by totalSupply, and the sum is preserved by
-                // decrementing then incrementing.
+            }
+        }
+
+        if (to == address(0)) {
+            unchecked {
+                // Overflow not possible: amount <= totalSupply or amount <= fromBalance <= totalSupply.
+                _totalSupply -= amount;
+            }
+        } else {
+            unchecked {
+                // Overflow not possible: balance + amount is at most totalSupply, which we know fits into a uint256.
                 _balances[to] += amount;
             }
         }

contracts/token/ERC20/extensions/ERC20Capped.sol

@@ -32,7 +32,7 @@ abstract contract ERC20Capped is ERC20 {
      */
     function _update(address from, address to, uint256 amount) internal virtual override {
         if (from == address(0)) {
-            require(ERC20.totalSupply() + amount <= cap(), "ERC20Capped: cap exceeded");
+            require(totalSupply() + amount <= cap(), "ERC20Capped: cap exceeded");
         }
 
         super._update(from, to, amount);

contracts/token/ERC20/extensions/ERC20Snapshot.sol

@@ -122,18 +122,17 @@ abstract contract ERC20Snapshot is ERC20 {
     // for _mint, _burn, and _transfer operations.
     function _update(address from, address to, uint256 amount) internal virtual override {
         if (from == address(0)) {
-            // mint
-            _updateAccountSnapshot(to);
-            _updateTotalSupplySnapshot();
-        } else if (to == address(0)) {
-            // burn
-            _updateAccountSnapshot(from);
             _updateTotalSupplySnapshot();
         } else {
-            // transfer
             _updateAccountSnapshot(from);
+        }
+
+        if (to == address(0)) {
+            _updateTotalSupplySnapshot();
+        } else {
             _updateAccountSnapshot(to);
         }
+
         super._update(from, to, amount);
     }
 

test/token/ERC20/ERC20.test.js

@@ -250,7 +250,7 @@ contract('ERC20', function (accounts) {
     describe('for a non zero account', function () {
       it('rejects burning more than balance', async function () {
         await expectRevert(this.token.burn(
-          initialHolder, initialSupply.addn(1)), 'ERC20: burn amount exceeds balance',
+          initialHolder, initialSupply.addn(1)), 'ERC20: transfer amount exceeds balance',
         );
       });
 
@@ -287,6 +287,49 @@ contract('ERC20', function (accounts) {
     });
   });
 
+  describe('_update', function () {
+    const amount = new BN(1);
+
+    it('from is the zero address', async function () {
+      const balanceBefore = await this.token.balanceOf(initialHolder);
+      const totalSupply = await this.token.totalSupply();
+
+      expectEvent(
+        await this.token.update(ZERO_ADDRESS, initialHolder, amount),
+        'Transfer',
+        { from: ZERO_ADDRESS, to: initialHolder, value: amount },
+      );
+      expect(await this.token.totalSupply()).to.be.bignumber.equal(totalSupply.add(amount));
+      expect(await this.token.balanceOf(initialHolder)).to.be.bignumber.equal(balanceBefore.add(amount));
+    });
+
+    it('to is the zero address', async function () {
+      const balanceBefore = await this.token.balanceOf(initialHolder);
+      const totalSupply = await this.token.totalSupply();
+
+      expectEvent(
+        await this.token.update(initialHolder, ZERO_ADDRESS, amount),
+        'Transfer',
+        { from: initialHolder, to: ZERO_ADDRESS, value: amount },
+      );
+      expect(await this.token.totalSupply()).to.be.bignumber.equal(totalSupply.sub(amount));
+      expect(await this.token.balanceOf(initialHolder)).to.be.bignumber.equal(balanceBefore.sub(amount));
+    });
+
+    it('from and to are the zero address', async function () {
+      const totalSupply = await this.token.totalSupply();
+
+      await this.token.update(ZERO_ADDRESS, ZERO_ADDRESS, amount);
+
+      expect(await this.token.totalSupply()).to.be.bignumber.equal(totalSupply);
+      expectEvent(
+        await this.token.update(ZERO_ADDRESS, ZERO_ADDRESS, amount),
+        'Transfer',
+        { from: ZERO_ADDRESS, to: ZERO_ADDRESS, value: amount },
+      );
+    });
+  });
+
   describe('_transfer', function () {
     shouldBehaveLikeERC20Transfer('ERC20', initialHolder, recipient, initialSupply, function (from, to, amount) {
       return this.token.transferInternal(from, to, amount);

test/token/ERC20/extensions/ERC20Burnable.behavior.js

@@ -38,7 +38,7 @@ function shouldBehaveLikeERC20Burnable (owner, initialBalance, [burner]) {
 
       it('reverts', async function () {
         await expectRevert(this.token.burn(amount, { from: owner }),
-          'ERC20: burn amount exceeds balance',
+          'ERC20: transfer amount exceeds balance',
         );
       });
     });
@@ -86,7 +86,7 @@ function shouldBehaveLikeERC20Burnable (owner, initialBalance, [burner]) {
       it('reverts', async function () {
         await this.token.approve(burner, amount, { from: owner });
         await expectRevert(this.token.burnFrom(owner, amount, { from: burner }),
-          'ERC20: burn amount exceeds balance',
+          'ERC20: transfer amount exceeds balance',
         );
       });
     });

test/token/ERC20/extensions/ERC20FlashMint.test.js

@@ -82,7 +82,7 @@ contract('ERC20FlashMint', function (accounts) {
       const data = this.token.contract.methods.transfer(other, 10).encodeABI();
       await expectRevert(
         this.token.flashLoan(receiver.address, this.token.address, loanAmount, data),
-        'ERC20: burn amount exceeds balance',
+        'ERC20: transfer amount exceeds balance',
       );
     });
 

test/token/ERC20/extensions/ERC20Wrapper.test.js

@@ -105,7 +105,7 @@ contract('ERC20', function (accounts) {
     it('missing balance', async function () {
       await expectRevert(
         this.token.withdrawTo(initialHolder, MAX_UINT256, { from: initialHolder }),
-        'ERC20: burn amount exceeds balance',
+        'ERC20: transfer amount exceeds balance',
       );
     });