feat: Add checks for address(0)

zama-ai · Oct 31, 2024 · 1be227f · 1be227f
1 parent 1ff5f4d
commit 1be227f
Show file tree

Hide file tree

Showing 5 changed files with 69 additions and 16 deletions.
diff --git a/contracts/token/ERC20/EncryptedERC20.sol b/contracts/token/ERC20/EncryptedERC20.sol
@@ -167,6 +167,19 @@ abstract contract EncryptedERC20 is IEncryptedERC20 {
     }
 
     function _transfer(address from, address to, euint64 amount, ebool isTransferable) internal virtual {
+        _transferNoEvent(from, to, amount, isTransferable);
+        emit Transfer(from, to);
+    }
+
+    function _transferNoEvent(address from, address to, euint64 amount, ebool isTransferable) internal virtual {
+        if (from == address(0)) {
+            revert SenderAddressNull();
+        }
+
+        if (to == address(0)) {
+            revert ReceiverAddressNull();
+        }
+
         // Add to the balance of `to` and subract from the balance of `from`.
         euint64 transferValue = TFHE.select(isTransferable, amount, TFHE.asEuint64(0));
         euint64 newBalanceTo = TFHE.add(_balances[to], transferValue);
@@ -177,7 +190,6 @@ abstract contract EncryptedERC20 is IEncryptedERC20 {
         _balances[from] = newBalanceFrom;
         TFHE.allowThis(newBalanceFrom);
         TFHE.allow(newBalanceFrom, from);
-        emit Transfer(from, to);
     }
 
     function _updateAllowance(address owner, address spender, euint64 amount) internal virtual returns (ebool) {

diff --git a/contracts/token/ERC20/IEncryptedERC20.sol b/contracts/token/ERC20/IEncryptedERC20.sol
@@ -20,6 +20,16 @@ interface IEncryptedERC20 {
      */
     event Transfer(address indexed from, address indexed to);
 
+    /**
+     * @notice Emitted when receiver is address(0).
+     */
+    error ReceiverAddressNull();
+
+    /**
+     * @notice Emitted when sender is address(0).
+     */
+    error SenderAddressNull();
+
     /**
      * @notice Sets the `encryptedAmount` as the allowance of `spender` over the caller's tokens.
      */

diff --git a/contracts/token/ERC20/extensions/EncryptedERC20WithErrors.sol b/contracts/token/ERC20/extensions/EncryptedERC20WithErrors.sol
@@ -6,7 +6,7 @@ import { EncryptedERC20 } from "../EncryptedERC20.sol";
 import { EncryptedErrors } from "../../../utils/EncryptedErrors.sol";
 
 /**
- * @title EncryptedERC20WithErrors
+ * @title       EncryptedERC20WithErrors
  * @notice      This contract implements an encrypted ERC20-like token with confidential balances using
  *              Zama's FHE (Fully Homomorphic Encryption) library.
  * @dev         It supports standard ERC20 functions such as transferring tokens, minting,
@@ -95,20 +95,7 @@ abstract contract EncryptedERC20WithErrors is EncryptedERC20, EncryptedErrors {
         ebool isTransferable,
         euint8 errorCode
     ) internal virtual {
-        // Add to the balance of `to` and subract from the balance of `from`.
-        euint64 transferValue = TFHE.select(isTransferable, amount, TFHE.asEuint64(0));
-        euint64 newBalanceTo = TFHE.add(_balances[to], transferValue);
-        _balances[to] = newBalanceTo;
-
-        TFHE.allow(newBalanceTo, address(this));
-        TFHE.allow(newBalanceTo, to);
-
-        euint64 newBalanceFrom = TFHE.sub(_balances[from], transferValue);
-        _balances[from] = newBalanceFrom;
-
-        TFHE.allow(newBalanceFrom, address(this));
-        TFHE.allow(newBalanceFrom, from);
-
+        _transferNoEvent(from, to, amount, isTransferable);
         emit TransferWithErrorHandling(from, to, _transferIdCounter);
 
         // Set error code in the storage and increment

diff --git a/test/encryptedERC20/EncryptedERC20.test.ts b/test/encryptedERC20/EncryptedERC20.test.ts
@@ -270,6 +270,28 @@ describe("EncryptedERC20", function () {
     }
   });
 
+  it("receiver cannot be null address", async function () {
+    const NULL_ADDRESS = "0x0000000000000000000000000000000000000000";
+    const mintAmount = 100_000;
+    const transferAmount = 50_000;
+    let tx = await this.encryptedERC20.connect(this.signers.alice).mint(mintAmount);
+    await tx.wait();
+
+    const input = this.instances.alice.createEncryptedInput(this.encryptedERC20Address, this.signers.alice.address);
+    input.add64(transferAmount);
+    const encryptedTransferAmount = await input.encrypt();
+
+    await expect(
+      this.encryptedERC20
+        .connect(this.signers.alice)
+        ["transfer(address,bytes32,bytes)"](
+          NULL_ADDRESS,
+          encryptedTransferAmount.handles[0],
+          encryptedTransferAmount.inputProof,
+        ),
+    ).to.be.revertedWithCustomError(this.encryptedERC20, "ReceiverAddressNull");
+  });
+
   it("sender who is not allowed cannot transfer using a handle from another account", async function () {
     const mintAmount = 100_000;
     const transferAmount = 50_000;

diff --git a/test/encryptedERC20/EncryptedERC20WithErrors.test.ts b/test/encryptedERC20/EncryptedERC20WithErrors.test.ts
@@ -337,6 +337,28 @@ describe("EncryptedERC20WithErrors", function () {
     }
   });
 
+  it("receiver cannot be null address", async function () {
+    const NULL_ADDRESS = "0x0000000000000000000000000000000000000000";
+    const mintAmount = 100_000;
+    const transferAmount = 50_000;
+    let tx = await this.encryptedERC20.connect(this.signers.alice).mint(mintAmount);
+    await tx.wait();
+
+    const input = this.instances.alice.createEncryptedInput(this.encryptedERC20Address, this.signers.alice.address);
+    input.add64(transferAmount);
+    const encryptedTransferAmount = await input.encrypt();
+
+    await expect(
+      this.encryptedERC20
+        .connect(this.signers.alice)
+        ["transfer(address,bytes32,bytes)"](
+          NULL_ADDRESS,
+          encryptedTransferAmount.handles[0],
+          encryptedTransferAmount.inputProof,
+        ),
+    ).to.be.revertedWithCustomError(this.encryptedERC20, "ReceiverAddressNull");
+  });
+
   it("sender who is not allowed cannot transfer using a handle from another account", async function () {
     const mintAmount = 100_000;
     const transferAmount = 50_000;