Revert when wrapper is full (#268) (#271)

* Revert when wrapper full

* up

* simplify and add tests

* update docs

* cei

* add doc

* up

* add custom error and fix tests

* comment clarification

* Rename `_checkTotalSupply` to `_checkConfidentialTotalSupply`

* add changeset

* Update contracts/token/ERC7984/extensions/ERC7984ERC20Wrapper.sol



* docs

---------

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

Author: arr00

.changeset/nasty-walls-taste.md

@@ -0,0 +1,5 @@
+---
+'openzeppelin-confidential-contracts': patch
+---
+
+`ERC7984ERC20Wrapper`: revert on wrap if there is a chance of total supply overflow.

contracts/token/ERC7984/extensions/ERC7984ERC20Wrapper.sol

@@ -30,6 +30,7 @@ abstract contract ERC7984ERC20Wrapper is ERC7984, IERC1363Receiver {
     event UnwrapFinalized(address indexed receiver, euint64 encryptedAmount, uint64 cleartextAmount);
 
     error InvalidUnwrapRequest(euint64 amount);
+    error ERC7984TotalSupplyOverflow();
 
     constructor(IERC20 underlying_) {
         _underlying = underlying_;
@@ -45,28 +46,10 @@ abstract contract ERC7984ERC20Wrapper is ERC7984, IERC1363Receiver {
         }
     }
 
-    /// @inheritdoc ERC7984
-    function decimals() public view virtual override returns (uint8) {
-        return _decimals;
-    }
-
-    /**
-     * @dev Returns the rate at which the underlying token is converted to the wrapped token.
-     * For example, if the `rate` is 1000, then 1000 units of the underlying token equal 1 unit of the wrapped token.
-     */
-    function rate() public view virtual returns (uint256) {
-        return _rate;
-    }
-
-    /// @dev Returns the address of the underlying ERC-20 token that is being wrapped.
-    function underlying() public view returns (IERC20) {
-        return _underlying;
-    }
-
     /**
      * @dev `ERC1363` callback function which wraps tokens to the address specified in `data` or
      * the address `from` (if no address is specified in `data`). This function refunds any excess tokens
-     * sent beyond the nearest multiple of {rate}. See {wrap} from more details on wrapping tokens.
+     * sent beyond the nearest multiple of {rate} to `from`. See {wrap} from more details on wrapping tokens.
      */
     function onTransferReceived(
         address /*operator*/,
@@ -149,6 +132,61 @@ abstract contract ERC7984ERC20Wrapper is ERC7984, IERC1363Receiver {
         emit UnwrapFinalized(to, burntAmount, burntAmountCleartext);
     }
 
+    /// @inheritdoc ERC7984
+    function decimals() public view virtual override returns (uint8) {
+        return _decimals;
+    }
+
+    /**
+     * @dev Returns the rate at which the underlying token is converted to the wrapped token.
+     * For example, if the `rate` is 1000, then 1000 units of the underlying token equal 1 unit of the wrapped token.
+     */
+    function rate() public view virtual returns (uint256) {
+        return _rate;
+    }
+
+    /// @dev Returns the address of the underlying ERC-20 token that is being wrapped.
+    function underlying() public view returns (IERC20) {
+        return _underlying;
+    }
+
+    /**
+     * @dev Returns the underlying balance divided by the {rate}, a value greater or equal to the actual
+     * {confidentialTotalSupply}.
+     *
+     * NOTE: The return value of this function can be inflated by directly sending underlying tokens to the wrapper contract.
+     * Reductions will lag compared to {confidentialTotalSupply} since it is updated on {unwrap} while this function updates
+     * on {finalizeUnwrap}.
+     */
+    function totalSupply() public view virtual returns (uint256) {
+        return underlying().balanceOf(address(this)) / rate();
+    }
+
+    /// @dev Returns the maximum total supply of wrapped tokens supported by the encrypted datatype.
+    function maxTotalSupply() public view virtual returns (uint256) {
+        return type(uint64).max;
+    }
+
+    /**
+     * @dev This function must revert if the new {confidentialTotalSupply} is invalid (overflow occurred).
+     *
+     * NOTE: Overflow can be detected here since the wrapper holdings are non-confidential. In other cases, it may be impossible
+     * to infer total supply overflow synchronously. This function may revert even if the {confidentialTotalSupply} did
+     * not overflow.
+     */
+    function _checkConfidentialTotalSupply() internal virtual {
+        if (totalSupply() > maxTotalSupply()) {
+            revert ERC7984TotalSupplyOverflow();
+        }
+    }
+
+    function _update(address from, address to, euint64 amount) internal virtual override returns (euint64) {
+        if (from == address(0)) {
+            _checkConfidentialTotalSupply();
+        }
+        return super._update(from, to, amount);
+    }
+
     function _unwrap(address from, address to, euint64 amount) internal virtual {
         require(to != address(0), ERC7984InvalidReceiver(to));
         require(from == msg.sender || isOperator(from, msg.sender), ERC7984UnauthorizedSpender(from, msg.sender));

test/token/ERC7984/extensions/ERC7984Wrapper.test.ts

@@ -82,6 +82,47 @@ describe('ERC7984Wrapper', function () {
           ).to.eventually.equal(10);
         });
 
+        it('max amount works', async function () {
+          await this.token.$_mint(this.holder.address, ethers.MaxUint256 / 2n); // mint a lot of tokens
+
+          const rate = await this.wrapper.rate();
+          const maxConfidentialSupply = await this.wrapper.maxTotalSupply();
+          const maxUnderlyingBalance = maxConfidentialSupply * rate;
+
+          if (viaCallback) {
+            await this.token.connect(this.holder).transferAndCall(this.wrapper, maxUnderlyingBalance);
+          } else {
+            await this.wrapper.connect(this.holder).wrap(this.holder.address, maxUnderlyingBalance);
+          }
+
+          await expect(
+            fhevm.userDecryptEuint(
+              FhevmType.euint64,
+              await this.wrapper.confidentialBalanceOf(this.holder.address),
+              this.wrapper.target,
+              this.holder,
+            ),
+          ).to.eventually.equal(maxConfidentialSupply);
+        });
+
+        it('amount exceeding max fails', async function () {
+          await this.token.$_mint(this.holder.address, ethers.MaxUint256 / 2n); // mint a lot of tokens
+
+          const rate = await this.wrapper.rate();
+          const maxConfidentialSupply = await this.wrapper.maxTotalSupply();
+          const maxUnderlyingBalance = maxConfidentialSupply * rate;
+
+          // first deposit close to the max
+          await this.wrapper.connect(this.holder).wrap(this.holder.address, maxUnderlyingBalance);
+
+          // try to deposit more, causing the total supply to exceed the max supported amount
+          await expect(
+            viaCallback
+              ? this.token.connect(this.holder).transferAndCall(this.wrapper, rate)
+              : this.wrapper.connect(this.holder).wrap(this.holder.address, rate),
+          ).to.be.revertedWithCustomError(this.wrapper, 'ERC7984TotalSupplyOverflow');
+        });
+
         if (viaCallback) {
           it('to another address', async function () {
             const amountToWrap = ethers.parseUnits('100', 18);