Skip to content

Commit

Permalink
feat: add upper bound for encrypted random integers
Browse files Browse the repository at this point in the history
Add overloads of TFHE.randEuint() methods that accept an upper bound,
e.g. TFHE.randEuint32(uint32 upperBound). The returned integer will be
in the [0, upperBound) range. Moreover, upperBound must be a power of
2 - that is a design choice for performance reasons.

Note that, as of now, the random numbers are generated in the plain and
are completely public and predictable. An FHE version is coming soon.

Add tests to cover above functionality.
  • Loading branch information
dartdart26 committed Jan 9, 2024
1 parent 96de5da commit fe1e970
Show file tree
Hide file tree
Showing 7 changed files with 227 additions and 5 deletions.
50 changes: 48 additions & 2 deletions codegen/templates.ts
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,8 @@ export function implSol(ctx: CodegenContext, operators: Operator[]): string {
pragma solidity ^0.8.20;
import "./TFHE.sol";
${fheLibInterface}
address constant EXT_TFHE_LIBRARY = address(${ctx.libFheAddress});
Expand Down Expand Up @@ -132,8 +134,9 @@ function fheLibCustomInterfaceFunctions(): string {
function cast(uint256 ct, bytes1 toType) external pure returns (uint256 result);
function trivialEncrypt(uint256 ct, bytes1 toType) external pure returns (uint256 result);
function decrypt(uint256 ct) external view returns (uint256 result);
function fheRand(bytes1 inp) external view returns (uint256 result);
function fheIfThenElse(uint256 control, uint256 ifTrue, uint256 ifFalse) external pure returns (uint256 result);
function fheRand(bytes1 randType) external view returns (uint256 result);
function fheRandBounded(uint256 upperBound, bytes1 randType) external view returns (uint256 result);
`;
}

Expand Down Expand Up @@ -706,17 +709,38 @@ function tfheCustomMethods(ctx: CodegenContext): string {
return euint8.wrap(Impl.rand(Common.euint8_t));
}
// Generates a random encrypted 8-bit unsigned integer in the [0, upperBound) range.
// The upperBound must be a power of 2.
// Important: The random integer is generated in the plain! An FHE-based version is coming soon.
function randEuint8(uint8 upperBound) internal view returns (euint8) {
return euint8.wrap(Impl.randBounded(upperBound, Common.euint8_t));
}
// Generates a random encrypted 16-bit unsigned integer.
// Important: The random integer is generated in the plain! An FHE-based version is coming soon.
function randEuint16() internal view returns (euint16) {
return euint16.wrap(Impl.rand(Common.euint16_t));
}
// Generates a random encrypted 16-bit unsigned integer in the [0, upperBound) range.
// The upperBound must be a power of 2.
// Important: The random integer is generated in the plain! An FHE-based version is coming soon.
function randEuint16(uint16 upperBound) internal view returns (euint16) {
return euint16.wrap(Impl.randBounded(upperBound, Common.euint16_t));
}
// Generates a random encrypted 32-bit unsigned integer.
// Important: The random integer is generated in the plain! An FHE-based version is coming soon.
function randEuint32() internal view returns (euint32) {
return euint32.wrap(Impl.rand(Common.euint32_t));
}
// Generates a random encrypted 32-bit unsigned integer in the [0, upperBound) range.
// The upperBound must be a power of 2.
// Important: The random integer is generated in the plain! An FHE-based version is coming soon.
function randEuint32(uint32 upperBound) internal view returns (euint32) {
return euint32.wrap(Impl.randBounded(upperBound, Common.euint32_t));
}
`;
}

Expand Down Expand Up @@ -781,6 +805,10 @@ function implCustomMethods(ctx: CodegenContext): string {
function rand(uint8 randType) internal view returns(uint256 result) {
result = FhevmLib(address(EXT_TFHE_LIBRARY)).fheRand(bytes1(randType));
}
function randBounded(uint256 upperBound, uint8 randType) internal view returns(uint256 result) {
result = FhevmLib(address(EXT_TFHE_LIBRARY)).fheRandBounded(upperBound, bytes1(randType));
}
`;
}

Expand All @@ -792,6 +820,8 @@ export function implSolMock(ctx: CodegenContext, operators: Operator[]): string
pragma solidity ^0.8.20;
import "./TFHE.sol";
library Impl {
function add(uint256 lhs, uint256 rhs, bool scalar) internal pure returns (uint256 result) {
unchecked {
Expand Down Expand Up @@ -947,7 +977,23 @@ library Impl {
}
function rand(uint8 randType) internal view returns (uint256 result) {
result = uint256(keccak256(abi.encodePacked(block.number, gasleft(), msg.sender))); // assuming no duplicated tx by same sender in a single block
uint256 randomness = uint256(keccak256(abi.encodePacked(block.number, gasleft(), msg.sender))); // assuming no duplicated tx by same sender in a single block
if (randType == Common.euint8_t) {
result = uint8(randomness);
} else if (randType == Common.euint16_t) {
result = uint16(randomness);
} else if (randType == Common.euint32_t) {
result = uint32(randomness);
} else {
revert("rand() mock invalid type");
}
}
function randBounded(uint256 upperBound, uint8 randType) internal view returns (uint256 result) {
// Here, we assume upperBound is a power of 2. Therefore, using modulo is secure.
// If not a power of 2, we might have to do something else (though might not matter
// much as this is a mock).
result = rand(randType) % upperBound;
}
`);

Expand Down
27 changes: 27 additions & 0 deletions examples/Rand.sol
Original file line number Diff line number Diff line change
Expand Up @@ -16,14 +16,26 @@ contract Rand {
value8 = TFHE.randEuint8();
}

function generate8UpperBound(uint8 upperBound) public {
value8 = TFHE.randEuint8(upperBound);
}

function generate16() public {
value16 = TFHE.randEuint16();
}

function generate16UpperBound(uint16 upperBound) public {
value16 = TFHE.randEuint16(upperBound);
}

function generate32() public {
value32 = TFHE.randEuint32();
}

function generate32UpperBound(uint32 upperBound) public {
value32 = TFHE.randEuint32(upperBound);
}

function decrypt8() public view returns (uint8) {
return TFHE.decrypt(value8);
}
Expand Down Expand Up @@ -53,13 +65,28 @@ contract Rand {
TFHE.randEuint8();
}

// Must fail.
function generate8UpperBoundInView(uint8 upperBound) public view {
TFHE.randEuint8(upperBound);
}

// Must fail.
function generate16InView() public view {
TFHE.randEuint16();
}

// Must fail.
function generate16UpperBoundInView(uint16 upperBound) public view {
TFHE.randEuint16(upperBound);
}

// Must fail.
function generate32InView() public view {
TFHE.randEuint32();
}

// Must fail.
function generate32UpperBoundInView(uint32 upperBound) public view {
TFHE.randEuint32(upperBound);
}
}
12 changes: 10 additions & 2 deletions lib/Impl.sol
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@

pragma solidity ^0.8.20;

import "./TFHE.sol";

interface FhevmLib {
function fheAdd(uint256 lhs, uint256 rhs, bytes1 scalarByte) external pure returns (uint256 result);

Expand Down Expand Up @@ -55,9 +57,11 @@ interface FhevmLib {

function decrypt(uint256 ct) external view returns (uint256 result);

function fheRand(bytes1 inp) external view returns (uint256 result);

function fheIfThenElse(uint256 control, uint256 ifTrue, uint256 ifFalse) external pure returns (uint256 result);

function fheRand(bytes1 randType) external view returns (uint256 result);

function fheRandBounded(uint256 upperBound, bytes1 randType) external view returns (uint256 result);
}

address constant EXT_TFHE_LIBRARY = address(0x000000000000000000000000000000000000005d);
Expand Down Expand Up @@ -283,4 +287,8 @@ library Impl {
function rand(uint8 randType) internal view returns (uint256 result) {
result = FhevmLib(address(EXT_TFHE_LIBRARY)).fheRand(bytes1(randType));
}

function randBounded(uint256 upperBound, uint8 randType) internal view returns (uint256 result) {
result = FhevmLib(address(EXT_TFHE_LIBRARY)).fheRandBounded(upperBound, bytes1(randType));
}
}
21 changes: 21 additions & 0 deletions lib/TFHE.sol
Original file line number Diff line number Diff line change
Expand Up @@ -2426,17 +2426,38 @@ library TFHE {
return euint8.wrap(Impl.rand(Common.euint8_t));
}

// Generates a random encrypted 8-bit unsigned integer in the [0, upperBound) range.
// The upperBound must be a power of 2.
// Important: The random integer is generated in the plain! An FHE-based version is coming soon.
function randEuint8(uint8 upperBound) internal view returns (euint8) {
return euint8.wrap(Impl.randBounded(upperBound, Common.euint8_t));
}

// Generates a random encrypted 16-bit unsigned integer.
// Important: The random integer is generated in the plain! An FHE-based version is coming soon.
function randEuint16() internal view returns (euint16) {
return euint16.wrap(Impl.rand(Common.euint16_t));
}

// Generates a random encrypted 16-bit unsigned integer in the [0, upperBound) range.
// The upperBound must be a power of 2.
// Important: The random integer is generated in the plain! An FHE-based version is coming soon.
function randEuint16(uint16 upperBound) internal view returns (euint16) {
return euint16.wrap(Impl.randBounded(upperBound, Common.euint16_t));
}

// Generates a random encrypted 32-bit unsigned integer.
// Important: The random integer is generated in the plain! An FHE-based version is coming soon.
function randEuint32() internal view returns (euint32) {
return euint32.wrap(Impl.rand(Common.euint32_t));
}

// Generates a random encrypted 32-bit unsigned integer in the [0, upperBound) range.
// The upperBound must be a power of 2.
// Important: The random integer is generated in the plain! An FHE-based version is coming soon.
function randEuint32(uint32 upperBound) internal view returns (euint32) {
return euint32.wrap(Impl.randBounded(upperBound, Common.euint32_t));
}
}

using {tfheBinaryOperatorAdd8 as +} for euint8 global;
Expand Down
20 changes: 19 additions & 1 deletion mocks/Impl.sol
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@

pragma solidity ^0.8.20;

import "./TFHE.sol";

library Impl {
function add(uint256 lhs, uint256 rhs, bool scalar) internal pure returns (uint256 result) {
unchecked {
Expand Down Expand Up @@ -157,6 +159,22 @@ library Impl {
}

function rand(uint8 randType) internal view returns (uint256 result) {
result = uint256(keccak256(abi.encodePacked(block.number, gasleft(), msg.sender))); // assuming no duplicated tx by same sender in a single block
uint256 randomness = uint256(keccak256(abi.encodePacked(block.number, gasleft(), msg.sender))); // assuming no duplicated tx by same sender in a single block
if (randType == Common.euint8_t) {
result = uint8(randomness);
} else if (randType == Common.euint16_t) {
result = uint16(randomness);
} else if (randType == Common.euint32_t) {
result = uint32(randomness);
} else {
revert("rand() mock invalid type");
}
}

function randBounded(uint256 upperBound, uint8 randType) internal view returns (uint256 result) {
// Here, we assume upperBound is a power of 2. Therefore, using modulo is secure.
// If not a power of 2, we might have to do something else (though might not matter
// much as this is a mock).
result = rand(randType) % upperBound;
}
}
21 changes: 21 additions & 0 deletions mocks/TFHE.sol
Original file line number Diff line number Diff line change
Expand Up @@ -2426,17 +2426,38 @@ library TFHE {
return euint8.wrap(Impl.rand(Common.euint8_t));
}

// Generates a random encrypted 8-bit unsigned integer in the [0, upperBound) range.
// The upperBound must be a power of 2.
// Important: The random integer is generated in the plain! An FHE-based version is coming soon.
function randEuint8(uint8 upperBound) internal view returns (euint8) {
return euint8.wrap(Impl.randBounded(upperBound, Common.euint8_t));
}

// Generates a random encrypted 16-bit unsigned integer.
// Important: The random integer is generated in the plain! An FHE-based version is coming soon.
function randEuint16() internal view returns (euint16) {
return euint16.wrap(Impl.rand(Common.euint16_t));
}

// Generates a random encrypted 16-bit unsigned integer in the [0, upperBound) range.
// The upperBound must be a power of 2.
// Important: The random integer is generated in the plain! An FHE-based version is coming soon.
function randEuint16(uint16 upperBound) internal view returns (euint16) {
return euint16.wrap(Impl.randBounded(upperBound, Common.euint16_t));
}

// Generates a random encrypted 32-bit unsigned integer.
// Important: The random integer is generated in the plain! An FHE-based version is coming soon.
function randEuint32() internal view returns (euint32) {
return euint32.wrap(Impl.rand(Common.euint32_t));
}

// Generates a random encrypted 32-bit unsigned integer in the [0, upperBound) range.
// The upperBound must be a power of 2.
// Important: The random integer is generated in the plain! An FHE-based version is coming soon.
function randEuint32(uint32 upperBound) internal view returns (euint32) {
return euint32.wrap(Impl.randBounded(upperBound, Common.euint32_t));
}
}

using {tfheBinaryOperatorAdd8 as +} for euint8 global;
Expand Down
Loading

0 comments on commit fe1e970

Please sign in to comment.