refactor: support for custom maxDecryptionDelay

zama-ai · Dec 5, 2024 · e3d76e4 · e3d76e4
1 parent 9e5ceef
commit e3d76e4
Show file tree

Hide file tree

Showing 6 changed files with 41 additions and 9 deletions.
diff --git a/contracts/test/token/ERC20/TestConfidentialERC20Wrapped.sol b/contracts/test/token/ERC20/TestConfidentialERC20Wrapped.sol
@@ -6,5 +6,5 @@ import { MockZamaFHEVMConfig } from "fhevm/config/ZamaFHEVMConfig.sol";
 import { MockZamaGatewayConfig } from "fhevm/config/ZamaGatewayConfig.sol";
 
 contract TestConfidentialERC20Wrapped is MockZamaFHEVMConfig, MockZamaGatewayConfig, ConfidentialERC20Wrapped {
-    constructor(address erc20_) ConfidentialERC20Wrapped(erc20_) {}
+    constructor(address erc20_, uint256 maxDecryptionDelay_) ConfidentialERC20Wrapped(erc20_, maxDecryptionDelay_) {}
 }
diff --git a/contracts/test/token/ERC20/TestConfidentialWETH.sol b/contracts/test/token/ERC20/TestConfidentialWETH.sol
@@ -5,5 +5,6 @@ import { ConfidentialWETH } from "../../../token/ERC20/ConfidentialWETH.sol";
 import { MockZamaFHEVMConfig } from "fhevm/config/ZamaFHEVMConfig.sol";
 import { MockZamaGatewayConfig } from "fhevm/config/ZamaGatewayConfig.sol";
 
-/* solhint-disable no-empty-blocks*/
-contract TestConfidentialWETH is MockZamaFHEVMConfig, MockZamaGatewayConfig, ConfidentialWETH {}
+contract TestConfidentialWETH is MockZamaFHEVMConfig, MockZamaGatewayConfig, ConfidentialWETH {
+    constructor(uint256 maxDecryptionDelay_) ConfidentialWETH(maxDecryptionDelay_) {}
+}
diff --git a/contracts/token/ERC20/ConfidentialERC20Wrapped.sol b/contracts/token/ERC20/ConfidentialERC20Wrapped.sol
@@ -21,6 +21,9 @@ import { ConfidentialERC20 } from "./ConfidentialERC20.sol";
 abstract contract ConfidentialERC20Wrapped is ConfidentialERC20, IConfidentialERC20Wrapped, GatewayCaller {
     using SafeERC20 for IERC20Metadata;
 
+    /// @notice Returned if the maximum decryption delay is higher than 1 day.
+    error MaxDecryptionDelayTooHigh();
+
     /// @notice ERC20 token that is wrapped.
     IERC20Metadata public immutable ERC20_TOKEN;
 
@@ -37,16 +40,27 @@ abstract contract ConfidentialERC20Wrapped is ConfidentialERC20, IConfidentialER
      *                 For instance,
      *                 "Wrapped Ether" --> "Encrypted Wrapped Ether"
      *                 "WETH" --> "eWETH".
+     * @param maxDecryptionDelay_       Maximum delay for the Gateway to decrypt.
+     * @dev                             Do not use a small value in production to avoid security issues if the response
+     *                                  cannot be processed because the block time is higher than the delay.
+     *                                  The current implementation expects the Gateway to always return a decrypted
+     *                                  value within the delay specified, as long as it is sufficient enough.
      */
     constructor(
-        address erc20_
+        address erc20_,
+        uint256 maxDecryptionDelay_
     )
         ConfidentialERC20(
             string(abi.encodePacked("Encrypted ", IERC20Metadata(erc20_).name())),
             string(abi.encodePacked("e", IERC20Metadata(erc20_).symbol()))
         )
     {
         ERC20_TOKEN = IERC20Metadata(erc20_);
+
+        /// @dev The maximum delay is set to 1 day.
+        if (maxDecryptionDelay_ > 1 days) {
+            revert MaxDecryptionDelayTooHigh();
+        }
     }
 
     /**

diff --git a/contracts/token/ERC20/ConfidentialWETH.sol b/contracts/token/ERC20/ConfidentialWETH.sol
@@ -17,6 +17,9 @@ abstract contract ConfidentialWETH is ConfidentialERC20, IConfidentialERC20Wrapp
     /// @notice Returned if ETH transfer fails.
     error ETHTransferFail();
 
+    /// @notice Returned if the maximum decryption delay is higher than 1 day.
+    error MaxDecryptionDelayTooHigh();
+
     /// @notice Tracks whether the account can move funds.
     mapping(address account => bool isRestricted) public isAccountRestricted;
 
@@ -26,10 +29,20 @@ abstract contract ConfidentialWETH is ConfidentialERC20, IConfidentialERC20Wrapp
     /**
      * @notice         Deposit/withdraw ethers (or native tokens).
      * @dev            The name/symbol are autogenerated.
+     * @param maxDecryptionDelay_       Maximum delay for the Gateway to decrypt.
+     * @dev                             Do not use a small value in production to avoid security issues if the response
+     *                                  cannot be processed because the block time is higher than the delay.
+     *                                  The current implementation expects the Gateway to always return a decrypted
+     *                                  value within the delay specified, as long as it is sufficient enough.
      */
-    constructor()
-        ConfidentialERC20(string(abi.encodePacked("Encrypted Wrapped Ether")), string(abi.encodePacked("eWETH")))
-    {}
+    constructor(
+        uint256 maxDecryptionDelay_
+    ) ConfidentialERC20(string(abi.encodePacked("Encrypted Wrapped Ether")), string(abi.encodePacked("eWETH"))) {
+        /// @dev The maximum delay is set to 1 day.
+        if (maxDecryptionDelay_ > 1 days) {
+            revert MaxDecryptionDelayTooHigh();
+        }
+    }
 
     /**
      * @notice         Fallback function calls wrap().

diff --git a/test/confidentialERC20/ConfidentialERC20Wrapped.fixture.ts b/test/confidentialERC20/ConfidentialERC20Wrapped.fixture.ts
@@ -9,6 +9,8 @@ export async function deployERC20AndConfidentialERC20WrappedFixture(
   symbol: string,
   decimals: number,
 ): Promise<[ERC20Mintable, TestConfidentialERC20Wrapped]> {
+  // @dev We use 5 minutes for the maximum decryption delay (from the Gateway).
+  const maxDecryptionDelay = 60 * 5;
   const contractFactoryERC20Mintable = await ethers.getContractFactory("ERC20Mintable");
   const contractERC20 = await contractFactoryERC20Mintable
     .connect(signers.alice)
@@ -18,7 +20,7 @@ export async function deployERC20AndConfidentialERC20WrappedFixture(
   const contractFactory = await ethers.getContractFactory("TestConfidentialERC20Wrapped");
   const contractConfidentialERC20Wrapped = await contractFactory
     .connect(signers.alice)
-    .deploy(contractERC20.getAddress());
+    .deploy(contractERC20.getAddress(), maxDecryptionDelay);
   await contractConfidentialERC20Wrapped.waitForDeployment();
 
   return [contractERC20, contractConfidentialERC20Wrapped];

diff --git a/test/confidentialERC20/ConfidentialWETH.fixture.ts b/test/confidentialERC20/ConfidentialWETH.fixture.ts
@@ -4,8 +4,10 @@ import type { TestConfidentialWETH } from "../../types";
 import { Signers } from "../signers";
 
 export async function deployConfidentialWETHFixture(signers: Signers): Promise<TestConfidentialWETH> {
+  // @dev We use 5 minutes for the maximum decryption delay (from the Gateway).
+  const maxDecryptionDelay = 60 * 5;
   const contractFactory = await ethers.getContractFactory("TestConfidentialWETH");
-  const confidentialWETH = await contractFactory.connect(signers.alice).deploy();
+  const confidentialWETH = await contractFactory.connect(signers.alice).deploy(maxDecryptionDelay);
   await confidentialWETH.waitForDeployment();
 
   return confidentialWETH;