Merge pull request #658 from zama-ai/custom-errors-tfhe

refactor: add custom errors for TFHE
zama-ai · Dec 19, 2024 · 269931d · 269931d
2 parents d4523fc + 655dc24
commit 269931d
Show file tree

Hide file tree

Showing 4 changed files with 137 additions and 18 deletions.
diff --git a/codegen/templates.ts b/codegen/templates.ts
@@ -265,6 +265,16 @@ ${commonSolLib()}
  *          that interact with TFHE.
  */
 library TFHE {
+
+/// @notice Returned if the input's length is greater than 64 bytes.
+error InputLengthAbove64Bytes(uint256 inputLength);
+
+/// @notice Returned if the input's length is greater than 128 bytes.
+error InputLengthAbove128Bytes(uint256 inputLength);
+
+/// @notice Returned if the input's length is greater than 256 bytes.
+error InputLengthAbove256Bytes(uint256 inputLength);
+
   function setFHEVM(FHEVMConfigStruct memory fhevmConfig) internal {
       Impl.setFHEVM(fhevmConfig);
   }
@@ -1011,13 +1021,19 @@ function tfheCustomMethods(): string {
 
     // Left-pad a bytes array with zeros such that it becomes of length 64.
     function padToBytes64(bytes memory input) internal pure returns (bytes memory) {
-      require(input.length <= 64, "Input exceeds 64 bytes");
+      uint256 inputLength = input.length;
+
+      if (inputLength > 64) {
+          revert InputLengthAbove64Bytes(inputLength);
+      }
+
       bytes memory result = new bytes(64);
-      uint256 paddingLength = 64 - input.length;
+      uint256 paddingLength = 64 - inputLength;
+
       for (uint256 i = 0; i < paddingLength; i++) {
           result[i] = 0;
       }
-      for (uint256 i = 0; i < input.length; i++) {
+      for (uint256 i = 0; i < inputLength; i++) {
           result[paddingLength + i] = input[i];
       }
       return result;
@@ -1035,13 +1051,18 @@ function tfheCustomMethods(): string {
 
     // Left-pad a bytes array with zeros such that it becomes of length 128.
     function padToBytes128(bytes memory input) internal pure returns (bytes memory) {
-      require(input.length <= 128, "Input exceeds 128 bytes");
+      uint256 inputLength = input.length;
+
+      if (inputLength > 128) {
+          revert InputLengthAbove128Bytes(inputLength);    
+      }
+
       bytes memory result = new bytes(128);
-      uint256 paddingLength = 128 - input.length;
+      uint256 paddingLength = 128 - inputLength;
       for (uint256 i = 0; i < paddingLength; i++) {
           result[i] = 0;
       }
-      for (uint256 i = 0; i < input.length; i++) {
+      for (uint256 i = 0; i < inputLength; i++) {
           result[paddingLength + i] = input[i];
       }
       return result;
@@ -1059,13 +1080,18 @@ function tfheCustomMethods(): string {
 
     // Left-pad a bytes array with zeros such that it becomes of length 256.
     function padToBytes256(bytes memory input) internal pure returns (bytes memory) {
-      require(input.length <= 256, "Input exceeds 256 bytes");
+      uint256 inputLength = input.length;
+
+      if (inputLength > 256) {
+          revert InputLengthAbove256Bytes(inputLength);    
+      }
+
       bytes memory result = new bytes(256);
-      uint256 paddingLength = 256 - input.length;
+      uint256 paddingLength = 256 - inputLength;
       for (uint256 i = 0; i < paddingLength; i++) {
           result[i] = 0;
       }
-      for (uint256 i = 0; i < input.length; i++) {
+      for (uint256 i = 0; i < inputLength; i++) {
           result[paddingLength + i] = input[i];
       }
       return result;

diff --git a/examples/tests/TFHERevertTest.sol b/examples/tests/TFHERevertTest.sol
@@ -0,0 +1,23 @@
+// SPDX-License-Identifier: BSD-3-Clause-Clear
+pragma solidity ^0.8.24;
+
+import "../../lib/TFHE.sol";
+import "../FHEVMConfig.sol";
+
+contract TFHERevertTest {
+    constructor() {
+        TFHE.setFHEVM(FHEVMConfig.defaultConfig());
+    }
+
+    function padToBytes64(bytes memory input) public pure {
+        TFHE.padToBytes64(input);
+    }
+
+    function padToBytes128(bytes memory input) public pure {
+        TFHE.padToBytes128(input);
+    }
+
+    function padToBytes256(bytes memory input) public pure {
+        TFHE.padToBytes256(input);
+    }
+}
diff --git a/lib/TFHE.sol b/lib/TFHE.sol
@@ -42,6 +42,15 @@ library Common {
  *          that interact with TFHE.
  */
 library TFHE {
+    /// @notice Returned if the input's length is greater than 64 bytes.
+    error InputLengthAbove64Bytes(uint256 inputLength);
+
+    /// @notice Returned if the input's length is greater than 128 bytes.
+    error InputLengthAbove128Bytes(uint256 inputLength);
+
+    /// @notice Returned if the input's length is greater than 256 bytes.
+    error InputLengthAbove256Bytes(uint256 inputLength);
+
     function setFHEVM(FHEVMConfigStruct memory fhevmConfig) internal {
         Impl.setFHEVM(fhevmConfig);
     }
@@ -10430,13 +10439,19 @@ library TFHE {
 
     // Left-pad a bytes array with zeros such that it becomes of length 64.
     function padToBytes64(bytes memory input) internal pure returns (bytes memory) {
-        require(input.length <= 64, "Input exceeds 64 bytes");
+        uint256 inputLength = input.length;
+
+        if (inputLength > 64) {
+            revert InputLengthAbove64Bytes(inputLength);
+        }
+
         bytes memory result = new bytes(64);
-        uint256 paddingLength = 64 - input.length;
+        uint256 paddingLength = 64 - inputLength;
+
         for (uint256 i = 0; i < paddingLength; i++) {
             result[i] = 0;
         }
-        for (uint256 i = 0; i < input.length; i++) {
+        for (uint256 i = 0; i < inputLength; i++) {
             result[paddingLength + i] = input[i];
         }
         return result;
@@ -10454,13 +10469,18 @@ library TFHE {
 
     // Left-pad a bytes array with zeros such that it becomes of length 128.
     function padToBytes128(bytes memory input) internal pure returns (bytes memory) {
-        require(input.length <= 128, "Input exceeds 128 bytes");
+        uint256 inputLength = input.length;
+
+        if (inputLength > 128) {
+            revert InputLengthAbove128Bytes(inputLength);
+        }
+
         bytes memory result = new bytes(128);
-        uint256 paddingLength = 128 - input.length;
+        uint256 paddingLength = 128 - inputLength;
         for (uint256 i = 0; i < paddingLength; i++) {
             result[i] = 0;
         }
-        for (uint256 i = 0; i < input.length; i++) {
+        for (uint256 i = 0; i < inputLength; i++) {
             result[paddingLength + i] = input[i];
         }
         return result;
@@ -10478,13 +10498,18 @@ library TFHE {
 
     // Left-pad a bytes array with zeros such that it becomes of length 256.
     function padToBytes256(bytes memory input) internal pure returns (bytes memory) {
-        require(input.length <= 256, "Input exceeds 256 bytes");
+        uint256 inputLength = input.length;
+
+        if (inputLength > 256) {
+            revert InputLengthAbove256Bytes(inputLength);
+        }
+
         bytes memory result = new bytes(256);
-        uint256 paddingLength = 256 - input.length;
+        uint256 paddingLength = 256 - inputLength;
         for (uint256 i = 0; i < paddingLength; i++) {
             result[i] = 0;
         }
-        for (uint256 i = 0; i < input.length; i++) {
+        for (uint256 i = 0; i < inputLength; i++) {
             result[paddingLength + i] = input[i];
         }
         return result;

diff --git a/test/tfheOperations/tfheRevertPaths.ts b/test/tfheOperations/tfheRevertPaths.ts
@@ -0,0 +1,45 @@
+import { expect } from 'chai';
+import { randomBytes } from 'crypto';
+import { ethers } from 'hardhat';
+
+import { getSigners, initSigners } from '../signers';
+
+describe('TFHE revert paths', function () {
+  before(async function () {
+    await initSigners(1);
+    this.signers = await getSigners();
+
+    const contractFactory = await ethers.getContractFactory('TFHERevertTest');
+    const contract = await contractFactory.connect(this.signers.alice).deploy();
+    await contract.waitForDeployment();
+
+    this.contract = contract;
+  });
+
+  it('padToBytes64 reverts if input > 64 bytes', async function () {
+    const numberBytes = 65;
+    const input = randomBytes(numberBytes);
+
+    await expect(this.contract.padToBytes64(input))
+      .to.be.revertedWithCustomError(this.contract, 'InputLengthAbove64Bytes')
+      .withArgs(numberBytes);
+  });
+
+  it('padToBytes128 reverts if input > 128 bytes', async function () {
+    const numberBytes = 129;
+    const input = randomBytes(numberBytes);
+
+    await expect(this.contract.padToBytes128(input))
+      .to.be.revertedWithCustomError(this.contract, 'InputLengthAbove128Bytes')
+      .withArgs(numberBytes);
+  });
+
+  it('padToBytes256 reverts if input > 256 bytes', async function () {
+    const numberBytes = 257;
+    const input = randomBytes(numberBytes);
+
+    await expect(this.contract.padToBytes256(input))
+      .to.be.revertedWithCustomError(this.contract, 'InputLengthAbove256Bytes')
+      .withArgs(numberBytes);
+  });
+});