Merge pull request #160 from zama-ai/feature/fhe-lib

Add fheLib precompile with first fheAdd method
zama-ai · Sep 12, 2023 · 8800c42 · 8800c42
2 parents d9b1715 + dfc0040
commit 8800c42
Show file tree

Hide file tree

Showing 2 changed files with 118 additions and 0 deletions.
diff --git a/core/vm/contracts.go b/core/vm/contracts.go
@@ -94,6 +94,7 @@ var PrecompiledContractsHomestead = map[common.Address]PrecompiledContract{
 	common.BytesToAddress([]byte{90}): &fheNot{},
 	common.BytesToAddress([]byte{91}): &decrypt{},
 	common.BytesToAddress([]byte{92}): &fheDiv{},
+	common.BytesToAddress([]byte{93}): &fheLib{},
 	common.BytesToAddress([]byte{99}): &faucet{},
 }
 
@@ -137,6 +138,7 @@ var PrecompiledContractsByzantium = map[common.Address]PrecompiledContract{
 	common.BytesToAddress([]byte{90}): &fheNot{},
 	common.BytesToAddress([]byte{91}): &decrypt{},
 	common.BytesToAddress([]byte{92}): &fheDiv{},
+	common.BytesToAddress([]byte{93}): &fheLib{},
 	common.BytesToAddress([]byte{99}): &faucet{},
 }
 
@@ -181,6 +183,7 @@ var PrecompiledContractsIstanbul = map[common.Address]PrecompiledContract{
 	common.BytesToAddress([]byte{90}): &fheNot{},
 	common.BytesToAddress([]byte{91}): &decrypt{},
 	common.BytesToAddress([]byte{92}): &fheDiv{},
+	common.BytesToAddress([]byte{93}): &fheLib{},
 	common.BytesToAddress([]byte{99}): &faucet{},
 }
 
@@ -225,6 +228,7 @@ var PrecompiledContractsBerlin = map[common.Address]PrecompiledContract{
 	common.BytesToAddress([]byte{90}): &fheNot{},
 	common.BytesToAddress([]byte{91}): &decrypt{},
 	common.BytesToAddress([]byte{92}): &fheDiv{},
+	common.BytesToAddress([]byte{93}): &fheLib{},
 	common.BytesToAddress([]byte{99}): &faucet{},
 }
 
@@ -269,6 +273,7 @@ var PrecompiledContractsBLS = map[common.Address]PrecompiledContract{
 	common.BytesToAddress([]byte{90}): &fheNot{},
 	common.BytesToAddress([]byte{91}): &decrypt{},
 	common.BytesToAddress([]byte{92}): &fheDiv{},
+	common.BytesToAddress([]byte{93}): &fheLib{},
 	common.BytesToAddress([]byte{99}): &faucet{},
 }
 
@@ -1471,6 +1476,49 @@ func writeResult(ct *tfheCiphertext, fileName string, logger Logger) {
 	os.WriteFile("/tmp/"+fileName, ct.serialize(), 0644)
 }
 
+type fheLib struct{}
+
+func (e *fheLib) RequiredGas(accessibleState PrecompileAccessibleState, input []byte) uint64 {
+	logger := accessibleState.Interpreter().evm.Logger
+	if len(input) < 4 {
+		err := errors.New("input must contain at least 4 bytes for method signature")
+		logger.Error("fheLib precompile error", "err", err, "input", hex.EncodeToString(input))
+		return 0
+	}
+	signature := binary.BigEndian.Uint32(input[0:4])
+	switch signature {
+	// first 4 bytes of keccak256('fheAdd(uint256,uint256,bytes1)')
+	case 0xf953e427:
+		bwCompatBytes := input[4:minInt(69, len(input))]
+		return (*fheAdd)(nil).RequiredGas(accessibleState, bwCompatBytes)
+	default:
+		err := errors.New("precompile method not found")
+		logger.Error("fheLib precompile error", "err", err, "input", hex.EncodeToString(input))
+		return 0
+	}
+}
+
+func (e *fheLib) Run(accessibleState PrecompileAccessibleState, caller common.Address, addr common.Address, input []byte, readOnly bool) ([]byte, error) {
+	logger := accessibleState.Interpreter().evm.Logger
+	if len(input) < 4 {
+		err := errors.New("input must contain at least 4 bytes for method signature")
+		logger.Error("fheLib precompile error", "err", err, "input", hex.EncodeToString(input))
+		return nil, err
+	}
+	signature := binary.BigEndian.Uint32(input[0:4])
+	switch signature {
+	// first 4 bytes of keccak256('fheAdd(uint256,uint256,bytes1)')
+	case 0xf953e427:
+		bwCompatBytes := input[4:minInt(69, len(input))]
+		// state of fheAdd struct is never needed or accessed so we use nil
+		return (*fheAdd)(nil).Run(accessibleState, caller, addr, bwCompatBytes, readOnly)
+	default:
+		err := errors.New("precompile method not found")
+		logger.Error("fheLib precompile error", "err", err, "input", hex.EncodeToString(input))
+		return nil, err
+	}
+}
+
 type fheAdd struct{}
 
 func (e *fheAdd) RequiredGas(accessibleState PrecompileAccessibleState, input []byte) uint64 {

diff --git a/core/vm/contracts_test.go b/core/vm/contracts_test.go
@@ -28,6 +28,7 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/rawdb"
 	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/crypto"
 )
 
 // precompiledTest defines the input/output pairs for precompiled contract tests.
@@ -468,6 +469,28 @@ func toPrecompileInput(isScalar bool, hashes ...common.Hash) []byte {
 	return ret
 }
 
+var scalarBytePadding = make([]byte, 31)
+
+func toLibPrecompileInput(method string, isScalar bool, hashes ...common.Hash) []byte {
+	ret := make([]byte, 0)
+	state := crypto.NewKeccakState()
+	hashRes := crypto.HashData(state, []byte(method))
+	signature := hashRes.Bytes()[0:4]
+	ret = append(ret, signature...)
+	for _, hash := range hashes {
+		ret = append(ret, hash.Bytes()...)
+	}
+	var isScalarByte byte
+	if isScalar {
+		isScalarByte = 1
+	} else {
+		isScalarByte = 0
+	}
+	ret = append(ret, isScalarByte)
+	ret = append(ret, scalarBytePadding...)
+	return ret
+}
+
 func VerifyCiphertext(t *testing.T, fheUintType fheUintType) {
 	var value uint32
 	switch fheUintType {
@@ -562,6 +585,49 @@ func TrivialEncrypt(t *testing.T, fheUintType fheUintType) {
 	}
 }
 
+func FheLibAdd(t *testing.T, fheUintType fheUintType, scalar bool) {
+	var lhs, rhs uint64
+	switch fheUintType {
+	case FheUint8:
+		lhs = 2
+		rhs = 1
+	case FheUint16:
+		lhs = 4283
+		rhs = 1337
+	case FheUint32:
+		lhs = 1333337
+		rhs = 133337
+	}
+	expected := lhs + rhs
+	c := &fheLib{}
+	signature := "fheAdd(uint256,uint256,bytes1)"
+	depth := 1
+	state := newTestState()
+	state.interpreter.evm.depth = depth
+	addr := common.Address{}
+	readOnly := false
+	lhsHash := verifyCiphertextInTestMemory(state.interpreter, lhs, depth, fheUintType).getHash()
+	var rhsHash common.Hash
+	if scalar {
+		rhsHash = common.BytesToHash(big.NewInt(int64(rhs)).Bytes())
+	} else {
+		rhsHash = verifyCiphertextInTestMemory(state.interpreter, rhs, depth, fheUintType).getHash()
+	}
+	input := toLibPrecompileInput(signature, scalar, lhsHash, rhsHash)
+	out, err := c.Run(state, addr, addr, input, readOnly)
+	if err != nil {
+		t.Fatalf(err.Error())
+	}
+	res := getVerifiedCiphertextFromEVM(state.interpreter, common.BytesToHash(out))
+	if res == nil {
+		t.Fatalf("output ciphertext is not found in verifiedCiphertexts")
+	}
+	decrypted, err := res.ciphertext.decrypt()
+	if err != nil || decrypted.Uint64() != expected {
+		t.Fatalf("invalid decrypted result, decrypted %v != expected %v", decrypted.Uint64(), expected)
+	}
+}
+
 func FheAdd(t *testing.T, fheUintType fheUintType, scalar bool) {
 	var lhs, rhs uint64
 	switch fheUintType {
@@ -1838,6 +1904,10 @@ func TestVerifyCiphertextBadCiphertext(t *testing.T) {
 	}
 }
 
+func TestFheLibAdd8(t *testing.T) {
+	FheLibAdd(t, FheUint8, false)
+}
+
 func TestFheAdd8(t *testing.T) {
 	FheAdd(t, FheUint8, false)
 }