Merge pull request #225 from zama-ai/features/remove-decrypt

Features/remove decrypt

examples/BlindAuction.sol

@@ -26,7 +26,7 @@ contract BlindAuction is EIP712WithModifier {
     EncryptedERC20 public tokenContract;
 
     // Whether the auction object has been claimed.
-    bool public objectClaimed;
+    ebool public objectClaimed;
 
     // If the token has been transferred to the beneficiary
     bool public tokenTransferred;
@@ -56,7 +56,7 @@ contract BlindAuction is EIP712WithModifier {
         beneficiary = _beneficiary;
         tokenContract = _tokenContract;
         endTime = block.timestamp + biddingTime;
-        objectClaimed = false;
+        objectClaimed = TFHE.asEbool(false);
         tokenTransferred = false;
         bidCounter = 0;
         stoppable = isStoppable;
@@ -117,12 +117,11 @@ contract BlindAuction is EIP712WithModifier {
 
     // Claim the object. Succeeds only if the caller has the highest bid.
     function claim() public onlyAfterEnd {
-        require(!objectClaimed);
-        require(TFHE.decrypt(TFHE.le(highestBid, bids[msg.sender])));
+        ebool canClaim = TFHE.and(TFHE.le(highestBid, bids[msg.sender]), TFHE.not(objectClaimed));
 
-        objectClaimed = true;
-        bids[msg.sender] = TFHE.NIL32;
-        emit Winner(msg.sender);
+        objectClaimed = canClaim;
+        bids[msg.sender] = TFHE.cmux(canClaim, TFHE.asEuint32(0), bids[msg.sender]);
+        // emit Winner(msg.sender);
     }
 
     // Transfer token to beneficiary
@@ -136,11 +135,10 @@ contract BlindAuction is EIP712WithModifier {
     // Withdraw a bid from the auction to the caller once the auction has stopped.
     function withdraw() public onlyAfterEnd {
         euint32 bidValue = bids[msg.sender];
-        if (!objectClaimed) {
-            require(TFHE.decrypt(TFHE.lt(bidValue, highestBid)));
-        }
-        tokenContract.transfer(msg.sender, bidValue);
-        bids[msg.sender] = TFHE.NIL32;
+        ebool isHighestBid = TFHE.eq(bidValue, highestBid);
+        ebool canWithdraw = TFHE.not(TFHE.and(isHighestBid, TFHE.not(objectClaimed)));
+        tokenContract.transfer(msg.sender, TFHE.cmux(canWithdraw, bidValue, TFHE.asEuint32(0)));
+        bids[msg.sender] = TFHE.cmux(canWithdraw, TFHE.asEuint32(0), bids[msg.sender]);
     }
 
     modifier onlyBeforeEnd() {

examples/EncryptedERC20.sol

@@ -24,8 +24,22 @@ contract EncryptedERC20 is EIP712WithModifier {
     // The owner of the contract.
     address public contractOwner;
 
+    mapping(address => euint8) internal lastError;
+
+    euint8 internal NO_ERROR;
+    euint8 internal NOT_ENOUGH_FUND;
+
     constructor() EIP712WithModifier("Authorization token", "1") {
         contractOwner = msg.sender;
+        NO_ERROR = TFHE.asEuint8(0);
+        NOT_ENOUGH_FUND = TFHE.asEuint8(1);
+    }
+
+    function getLastError(
+        bytes32 publicKey,
+        bytes calldata signature
+    ) public view onlySignedPublicKey(publicKey, signature) returns (bytes memory) {
+        return TFHE.reencrypt(lastError[msg.sender], publicKey, 0);
     }
 
     // Sets the balance of the owner to the given encrypted balance.
@@ -53,11 +67,16 @@ contract EncryptedERC20 is EIP712WithModifier {
     }
 
     // Returns the balance of the caller encrypted under the provided public key.
+
     function balanceOf(
+        address wallet,
         bytes32 publicKey,
         bytes calldata signature
     ) public view onlySignedPublicKey(publicKey, signature) returns (bytes memory) {
-        return TFHE.reencrypt(balances[msg.sender], publicKey, 0);
+        if (wallet == msg.sender) {
+            return TFHE.reencrypt(balances[wallet], publicKey, 0);
+        }
+        return TFHE.reencrypt(TFHE.asEuint32(0), publicKey, 0);
     }
 
     // Sets the `encryptedAmount` as the allowance of `spender` over the caller's tokens.
@@ -104,18 +123,19 @@ contract EncryptedERC20 is EIP712WithModifier {
 
     function _updateAllowance(address owner, address spender, euint32 amount) internal {
         euint32 currentAllowance = _allowance(owner, spender);
-        require(TFHE.decrypt(TFHE.le(amount, currentAllowance)));
-        _approve(owner, spender, currentAllowance - amount);
+        ebool canTransfer = TFHE.le(amount, currentAllowance);
+        _approve(owner, spender, TFHE.cmux(canTransfer, currentAllowance - amount, TFHE.asEuint32(0)));
     }
 
     // Transfers an encrypted amount.
     function _transfer(address from, address to, euint32 amount) internal {
         // Make sure the sender has enough tokens.
-        require(TFHE.decrypt(TFHE.le(amount, balances[from])));
+        ebool canTransfer = TFHE.le(amount, balances[from]);
+        lastError[msg.sender] = TFHE.cmux(canTransfer, NO_ERROR, NOT_ENOUGH_FUND);
 
         // Add to the balance of `to` and subract from the balance of `from`.
-        balances[to] = balances[to] + amount;
-        balances[from] = balances[from] - amount;
+        balances[to] = balances[to] + TFHE.cmux(canTransfer, amount, TFHE.asEuint32(0));
+        balances[from] = balances[from] - TFHE.cmux(canTransfer, amount, TFHE.asEuint32(0));
     }
 
     modifier onlyContractOwner() {

examples/Governor/Comp.sol

@@ -183,18 +183,18 @@ contract Comp is EIP712WithModifier {
 
     function _updateAllowance(address owner, address spender, euint32 amount) internal {
         euint32 currentAllowance = _allowance(owner, spender);
-        require(TFHE.decrypt(TFHE.le(amount, currentAllowance)));
-        _approve(owner, spender, currentAllowance - amount);
+        ebool canApprove = TFHE.le(amount, currentAllowance);
+        _approve(owner, spender, TFHE.cmux(canApprove, currentAllowance - amount, TFHE.asEuint32(0)));
     }
 
     // Transfers an encrypted amount.
     function _transfer(address from, address to, euint32 amount) internal {
         // Make sure the sender has enough tokens.
-        require(TFHE.decrypt(TFHE.le(amount, balances[from])));
+        ebool canTransfer = TFHE.le(amount, balances[from]);
 
         // Add to the balance of `to` and subract from the balance of `from`.
-        balances[to] = balances[to] + amount;
-        balances[from] = balances[from] - amount;
+        balances[to] = balances[to] + TFHE.cmux(canTransfer, amount, TFHE.asEuint32(0));
+        balances[from] = balances[from] - TFHE.cmux(canTransfer, amount, TFHE.asEuint32(0));
         emit Transfer(from, to, amount);
 
         _moveDelegates(delegates[from], delegates[to], amount);

examples/Identity/CompliantERC20.sol

@@ -89,8 +89,8 @@ contract CompliantERC20 is AbstractCompliantERC20 {
 
     function _updateAllowance(address owner, address spender, euint32 amount) internal {
         euint32 currentAllowance = _allowance(owner, spender);
-        require(TFHE.decrypt(TFHE.le(amount, currentAllowance)));
-        _approve(owner, spender, currentAllowance - amount);
+        ebool canTransfer = TFHE.le(amount, currentAllowance);
+        _approve(owner, spender, TFHE.cmux(canTransfer, currentAllowance - amount, TFHE.asEuint32(0)));
     }
 
     modifier onlyContractOwner() {

test/blindAuction/BlindAuction.ts

@@ -99,7 +99,11 @@ describe('BlindAuction', function () {
 
     const instance = await createInstances(this.contractERC20Address, ethers, this.signers);
     const tokenAlice = instance.alice.getTokenSignature(this.contractERC20Address)!;
-    const encryptedBalanceAlice = await this.erc20.balanceOf(tokenAlice.publicKey, tokenAlice.signature);
+    const encryptedBalanceAlice = await this.erc20.balanceOf(
+      this.signers.alice,
+      tokenAlice.publicKey,
+      tokenAlice.signature,
+    );
 
     const balanceAlice = instance.alice.decrypt(this.contractERC20Address, encryptedBalanceAlice);
     expect(balanceAlice).to.equal(1000 - 100 - 100 + 20);

test/encryptedERC20/EncryptedERC20.ts

@@ -27,7 +27,7 @@ describe('EncryptedERC20', function () {
       signature: '',
       publicKey: '',
     };
-    const encryptedBalance = await this.erc20.balanceOf(token.publicKey, token.signature);
+    const encryptedBalance = await this.erc20.balanceOf(this.signers.alice, token.publicKey, token.signature);
     // Decrypt the balance
     const balance = this.instances.alice.decrypt(this.contractAddress, encryptedBalance);
     expect(balance).to.equal(1000);
@@ -49,7 +49,11 @@ describe('EncryptedERC20', function () {
 
     const tokenAlice = this.instances.alice.getTokenSignature(this.contractAddress)!;
 
-    const encryptedBalanceAlice = await this.erc20.balanceOf(tokenAlice.publicKey, tokenAlice.signature);
+    const encryptedBalanceAlice = await this.erc20.balanceOf(
+      this.signers.alice,
+      tokenAlice.publicKey,
+      tokenAlice.signature,
+    );
 
     // Decrypt the balance
     const balanceAlice = this.instances.alice.decrypt(this.contractAddress, encryptedBalanceAlice);
@@ -60,11 +64,49 @@ describe('EncryptedERC20', function () {
 
     const tokenBob = this.instances.bob.getTokenSignature(this.contractAddress)!;
 
-    const encryptedBalanceBob = await bobErc20.balanceOf(tokenBob.publicKey, tokenBob.signature);
+    const encryptedBalanceBob = await bobErc20.balanceOf(this.signers.bob, tokenBob.publicKey, tokenBob.signature);
 
     // Decrypt the balance
     const balanceBob = this.instances.bob.decrypt(this.contractAddress, encryptedBalanceBob);
 
     expect(balanceBob).to.equal(1337);
   });
+
+  it('should not transfer tokens between two users', async function () {
+    const encryptedAmount = this.instances.alice.encrypt32(1000);
+    const transaction = await this.erc20.mint(encryptedAmount);
+    await transaction.wait();
+
+    const encryptedTransferAmount = this.instances.alice.encrypt32(1337);
+    const tx = await this.erc20['transfer(address,bytes)'](this.signers.bob.address, encryptedTransferAmount);
+    await tx.wait();
+
+    const tokenAlice = this.instances.alice.getTokenSignature(this.contractAddress)!;
+
+    const encryptedBalanceAlice = await this.erc20.balanceOf(
+      this.signers.alice,
+      tokenAlice.publicKey,
+      tokenAlice.signature,
+    );
+
+    // Decrypt the balance
+    const balanceAlice = this.instances.alice.decrypt(this.contractAddress, encryptedBalanceAlice);
+
+    expect(balanceAlice).to.equal(1000);
+
+    const bobErc20 = this.erc20.connect(this.signers.bob);
+
+    const tokenBob = this.instances.bob.getTokenSignature(this.contractAddress)!;
+
+    const encryptedBalanceBob = await bobErc20.balanceOf(this.signers.bob, tokenBob.publicKey, tokenBob.signature);
+
+    // Decrypt the balance
+    const balanceBob = this.instances.bob.decrypt(this.contractAddress, encryptedBalanceBob);
+
+    expect(balanceBob).to.equal(0);
+
+    const encryptedErrorAlice = await this.erc20.getLastError(tokenAlice.publicKey, tokenAlice.signature);
+    const errorAlice = this.instances.alice.decrypt(this.contractAddress, encryptedErrorAlice);
+    expect(errorAlice).to.equal(1);
+  });
 });