fix: wip

zama-ai · Dec 20, 2024 · 57ece12 · 57ece12
1 parent 3fbc515
commit 57ece12
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 22 deletions.
diff --git a/contracts/token/ERC20/ConfidentialERC20Wrapped.sol b/contracts/token/ERC20/ConfidentialERC20Wrapped.sol
@@ -3,6 +3,7 @@ pragma solidity ^0.8.24;
 
 import { IERC20Metadata } from "@openzeppelin/contracts/token/ERC20/extensions/IERC20Metadata.sol";
 import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import { ReentrancyGuardTransient } from "@openzeppelin/contracts/utils/ReentrancyGuardTransient.sol";
 
 import "fhevm/lib/TFHE.sol";
 import "fhevm/gateway/GatewayCaller.sol";
@@ -18,7 +19,12 @@ import { ConfidentialERC20 } from "./ConfidentialERC20.sol";
  *                    tokens with a fee on transfer. All ERC20 tokens must have decimals
  *                    inferior or equal to 18 decimals but superior or equal to 6 decimals.
  */
-abstract contract ConfidentialERC20Wrapped is ConfidentialERC20, IConfidentialERC20Wrapped, GatewayCaller {
+abstract contract ConfidentialERC20Wrapped is
+    ConfidentialERC20,
+    IConfidentialERC20Wrapped,
+    ReentrancyGuardTransient,
+    GatewayCaller
+{
     using SafeERC20 for IERC20Metadata;
 
     /// @notice Returned if the maximum decryption delay is higher than 1 day.
@@ -114,7 +120,7 @@ abstract contract ConfidentialERC20Wrapped is ConfidentialERC20, IConfidentialER
      * @param requestId   Request id.
      * @param canUnwrap   Whether it can be unwrapped.
      */
-    function callbackUnwrap(uint256 requestId, bool canUnwrap) public virtual onlyGateway {
+    function callbackUnwrap(uint256 requestId, bool canUnwrap) public virtual nonReentrant onlyGateway {
         UnwrapRequest memory unwrapRequest = unwrapRequests[requestId];
         delete unwrapRequests[requestId];
 
@@ -129,7 +135,7 @@ abstract contract ConfidentialERC20Wrapped is ConfidentialERC20, IConfidentialER
 
             emit Unwrap(unwrapRequest.account, unwrapRequest.amount);
         } else {
-            emit UnwrapFail(unwrapRequest.account, unwrapRequest.amount);
+            emit UnwrapFailNotEnoughBalance(unwrapRequest.account, unwrapRequest.amount);
         }
 
         delete isAccountRestricted[unwrapRequest.account];

diff --git a/contracts/token/ERC20/ConfidentialWETH.sol b/contracts/token/ERC20/ConfidentialWETH.sol
@@ -2,6 +2,7 @@
 pragma solidity ^0.8.24;
 
 import { ConfidentialERC20 } from "./ConfidentialERC20.sol";
+import { ReentrancyGuardTransient } from "@openzeppelin/contracts/utils/ReentrancyGuardTransient.sol";
 
 import "fhevm/lib/TFHE.sol";
 import "fhevm/gateway/GatewayCaller.sol";
@@ -13,7 +14,12 @@ import { IConfidentialERC20Wrapped } from "./IConfidentialERC20Wrapped.sol";
  * @notice            This contract allows users to wrap/unwrap trustlessly
  *                    ETH (or other native tokens) to ConfidentialERC20 tokens.
  */
-abstract contract ConfidentialWETH is ConfidentialERC20, IConfidentialERC20Wrapped, GatewayCaller {
+abstract contract ConfidentialWETH is
+    ConfidentialERC20,
+    IConfidentialERC20Wrapped,
+    ReentrancyGuardTransient,
+    GatewayCaller
+{
     /// @notice Returned if ETH transfer fails.
     error ETHTransferFail();
 
@@ -44,13 +50,6 @@ abstract contract ConfidentialWETH is ConfidentialERC20, IConfidentialERC20Wrapp
         }
     }
 
-    /**
-     * @notice         Fallback function calls wrap().
-     */
-    fallback() external payable {
-        wrap();
-    }
-
     /**
      * @notice         Receive function calls wrap().
      */
@@ -106,28 +105,27 @@ abstract contract ConfidentialWETH is ConfidentialERC20, IConfidentialERC20Wrapp
      * @param requestId   Request id.
      * @param canUnwrap   Whether it can be unwrapped.
      */
-    function callbackUnwrap(uint256 requestId, bool canUnwrap) public virtual onlyGateway {
+    function callbackUnwrap(uint256 requestId, bool canUnwrap) public virtual nonReentrant onlyGateway {
         UnwrapRequest memory unwrapRequest = unwrapRequests[requestId];
         delete unwrapRequests[requestId];
 
         if (canUnwrap) {
-            _unsafeBurn(unwrapRequest.account, unwrapRequest.amount);
-            _totalSupply -= unwrapRequest.amount;
-
             /// @dev It does a supply adjustment.
             uint256 amountUint256 = unwrapRequest.amount * (10 ** (18 - decimals()));
 
             /* solhint-disable avoid-call-value*/
             /* solhint-disable avoid-low-level-calls*/
             (bool callSuccess, ) = unwrapRequest.account.call{ value: amountUint256 }("");
 
-            if (!callSuccess) {
-                revert ETHTransferFail();
+            if (callSuccess) {
+                _unsafeBurn(unwrapRequest.account, unwrapRequest.amount);
+                _totalSupply -= unwrapRequest.amount;
+                emit Unwrap(unwrapRequest.account, unwrapRequest.amount);
+            } else {
+                emit UnwrapFailTransferFail(unwrapRequest.account, unwrapRequest.amount);
             }
-
-            emit Unwrap(unwrapRequest.account, unwrapRequest.amount);
         } else {
-            emit UnwrapFail(unwrapRequest.account, unwrapRequest.amount);
+            emit UnwrapFailNotEnoughBalance(unwrapRequest.account, unwrapRequest.amount);
         }
 
         delete isAccountRestricted[unwrapRequest.account];

diff --git a/contracts/token/ERC20/IConfidentialERC20Wrapped.sol b/contracts/token/ERC20/IConfidentialERC20Wrapped.sol
@@ -16,8 +16,11 @@ interface IConfidentialERC20Wrapped {
     /// @notice Emitted when token is unwrapped.
     event Unwrap(address indexed to, uint64 amount);
 
-    /// @notice Emitted if unwrap fails.
-    event UnwrapFail(address account, uint64 amount);
+    /// @notice Emitted if unwrap fails due to lack of funds.
+    event UnwrapFailNotEnoughBalance(address account, uint64 amount);
+
+    /// @notice Emitted if unwrap fails due to fail transfer.
+    event UnwrapFailTransferFail(address account, uint64 amount);
 
     /// @notice Emitted when token is wrapped.
     event Wrap(address indexed to, uint64 amount);