chore: add re-encrypt

fhevm/fhelib.go

@@ -248,6 +248,169 @@ var fhelibMethods = []*FheLibMethod{
 		requiredGasFunction: getCiphertextRequiredGas,
 		runFunction:         getCiphertextRun,
 	},
+	// TEE operations
+	{
+		name:                "teeEncrypt",
+		argTypes:            "(uint256,bytes1)",
+		requiredGasFunction: teeEncryptRequiredGas,
+		runFunction:         teeEncryptRun,
+	},
+	{
+		name:                "teeDecrypt",
+		argTypes:            "(uint256)",
+		requiredGasFunction: teeDecryptRequiredGas,
+		runFunction:         teeDecryptRun,
+	},
+	{
+		name:                "teeAdd",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeAddSubRequiredGas,
+		runFunction:         teeAddRun,
+	},
+	{
+		name:                "teeSub",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeAddSubRequiredGas,
+		runFunction:         teeSubRun,
+	},
+	{
+		name:                "teeMul",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeMulRequiredGas,
+		runFunction:         teeMulRun,
+	},
+	{
+		name:                "teeDiv",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeDivRequiredGas,
+		runFunction:         teeDivRun,
+	},
+	{
+		name:                "teeRem",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeRemRequiredGas,
+		runFunction:         teeRemRun,
+	},
+	{
+		name:                "teeLe",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeComparisonRequiredGas,
+		runFunction:         teeLeRun,
+	},
+	{
+		name:                "teeLt",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeComparisonRequiredGas,
+		runFunction:         teeLtRun,
+	},
+	{
+		name:                "teeEq",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeComparisonRequiredGas,
+		runFunction:         teeEqRun,
+	},
+	{
+		name:                "teeGe",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeComparisonRequiredGas,
+		runFunction:         teeGeRun,
+	},
+	{
+		name:                "teeGt",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeComparisonRequiredGas,
+		runFunction:         teeGtRun,
+	},
+	{
+		name:                "teeNe",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeComparisonRequiredGas,
+		runFunction:         teeNeRun,
+	},
+	{
+		name:                "teeMin",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeComparisonRequiredGas,
+		runFunction:         teeMinRun,
+	},
+	{
+		name:                "teeMax",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeComparisonRequiredGas,
+		runFunction:         teeMaxRun,
+	},
+	{
+		name:                "teeSelect",
+		argTypes:            "(uint256,uint256,uint256)",
+		requiredGasFunction: teeComparisonRequiredGas,
+		runFunction:         teeSelectRun,
+	},
+	{
+		name:                "teeShl",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeShiftRequiredGas,
+		runFunction:         teeShlRun,
+	},
+	{
+		name:                "teeShr",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeShiftRequiredGas,
+		runFunction:         teeShrRun,
+	},
+	{
+		name:                "teeRotl",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeShiftRequiredGas,
+		runFunction:         teeRotlRun,
+	},
+	{
+		name:                "teeRotr",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeShiftRequiredGas,
+		runFunction:         teeRotrRun,
+	},
+	{
+		name:                "teeBitAnd",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeBitwiseOpRequiredGas,
+		runFunction:         teeBitAndRun,
+	},
+	{
+		name:                "teeBitOr",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeBitwiseOpRequiredGas,
+		runFunction:         teeBitOrRun,
+	},
+	{
+		name:                "teeBitXor",
+		argTypes:            "(uint256,uint256,bytes1)",
+		requiredGasFunction: teeBitwiseOpRequiredGas,
+		runFunction:         teeBitXorRun,
+	},
+	{
+		name:                "teeNeg",
+		argTypes:            "(uint256)",
+		requiredGasFunction: teeNegRequiredGas,
+		runFunction:         teeNegRun,
+	},
+	{
+		name:                "teeNot",
+		argTypes:            "(uint256)",
+		requiredGasFunction: teeNotRequiredGas,
+		runFunction:         teeNotRun,
+	},
+	{
+		name:                "teeCast",
+		argTypes:            "(uint256,bytes1)",
+		requiredGasFunction: teeCastRequiredGas,
+		runFunction:         teeCastRun,
+	},
+	{
+		name:                "teeReencrypt",
+		argTypes:            "(uint256,uint256)",
+		requiredGasFunction: teeReencryptRequiredGas,
+		runFunction:         teeReencryptRun,
+	},
 }
 
 func init() {

fhevm/instructions_test.go

@@ -238,6 +238,12 @@ func newTestEVMEnvironment() *MockEVMEnvironment {
 	return &MockEVMEnvironment{fhevmData: &fhevmData, stateDb: state, commit: true, fhevmParams: DefaultFhevmParams()}
 }
 
+func newTestEVMEnvironmentWithEthCall() *MockEVMEnvironment {
+	mockEVM := newTestEVMEnvironment()
+	mockEVM.ethCall = true
+	return mockEVM
+}
+
 func TestProtectedStorageSstoreSload(t *testing.T) {
 	environment := newTestEVMEnvironment()
 	pc := uint64(0)

fhevm/params.go

@@ -78,6 +78,7 @@ type GasCosts struct {
 	TeeNeg        map[tfhe.FheUintType]uint64
 	TeeBitwiseOp  map[tfhe.FheUintType]uint64
 	TeeCast       uint64
+	TeeReencrypt  map[tfhe.FheUintType]uint64
 }
 
 func DefaultGasCosts() GasCosts {
@@ -308,6 +309,14 @@ func DefaultGasCosts() GasCosts {
 			tfhe.FheUint32: 150,
 			tfhe.FheUint64: 189,
 		},
+		// TODO: Costs will depend on the complexity of doing reencryption/decryption by the oracle.
+		TeeReencrypt: map[tfhe.FheUintType]uint64{
+			tfhe.FheBool:   1000,
+			tfhe.FheUint4:  1000,
+			tfhe.FheUint8:  1000,
+			tfhe.FheUint16: 1100,
+			tfhe.FheUint32: 1200,
+		},
 	}
 }
 

fhevm/tee_crypto.go

@@ -88,3 +88,49 @@ func teeDecryptRun(environment EVMEnvironment, caller common.Address, addr commo
 	copy(ret[32-len(plaintext):], plaintext)
 	return ret, nil
 }
+
+func teeReencryptRun(environment EVMEnvironment, caller common.Address, addr common.Address, input []byte, readOnly bool, runSpan trace.Span) ([]byte, error) {
+	input = input[:minInt(64, len(input))]
+	logger := environment.GetLogger()
+	if !environment.IsEthCall() {
+		msg := "reencrypt only supported on EthCall"
+		logger.Error(msg)
+		return nil, errors.New(msg)
+	}
+	if len(input) != 64 {
+		msg := "reencrypt input len must be 64 bytes"
+		logger.Error(msg, "input", hex.EncodeToString(input), "len", len(input))
+		return nil, errors.New(msg)
+	}
+	ct := getVerifiedCiphertext(environment, common.BytesToHash(input[0:32]))
+	if ct != nil {
+		otelDescribeOperandsFheTypes(runSpan, ct.fheUintType())
+
+		plainText, err := tee.Decrypt(ct.ciphertext)
+		if err != nil {
+			logger.Error("reencrypt decryption failed", "err", err)
+			return nil, err
+		}
+		decryptedValue := new(big.Int).SetBytes(plainText.Value)
+
+		pubKey := input[32:64]
+		reencryptedValue, err := encryptToUserKey(decryptedValue, pubKey)
+		if err != nil {
+			logger.Error("reencrypt failed to encrypt to user key", "err", err)
+			return nil, err
+		}
+
+		// TODO: decide if `res.Signature` should be verified here
+
+		logger.Info("reencrypt success", "input", hex.EncodeToString(input), "callerAddr", caller, "reencryptedValue", reencryptedValue, "len", len(reencryptedValue))
+		reencryptedValue = toEVMBytes(reencryptedValue)
+		// pad according to abi specification, first add offset to the dynamic bytes argument
+		outputBytes := make([]byte, 32, len(reencryptedValue)+32)
+		outputBytes[31] = 0x20
+		outputBytes = append(outputBytes, reencryptedValue...)
+		return padArrayTo32Multiple(outputBytes), nil
+	}
+	msg := "reencrypt unverified ciphertext handle"
+	logger.Error(msg, "input", hex.EncodeToString(input))
+	return nil, errors.New(msg)
+}

fhevm/tee_crypto_gas.go

@@ -34,3 +34,19 @@ func teeDecryptRequiredGas(environment EVMEnvironment, input []byte) uint64 {
 	}
 	return environment.FhevmParams().GasCosts.TeeDecrypt[ct.fheUintType()]
 }
+
+func teeReencryptRequiredGas(environment EVMEnvironment, input []byte) uint64 {
+	input = input[:minInt(64, len(input))]
+
+	logger := environment.GetLogger()
+	if len(input) != 64 {
+		logger.Error("reencrypt RequiredGas() input len must be 64 bytes", "input", hex.EncodeToString(input), "len", len(input))
+		return 0
+	}
+	ct := getVerifiedCiphertext(environment, common.BytesToHash(input[0:32]))
+	if ct == nil {
+		logger.Error("reencrypt RequiredGas() input doesn't point to verified ciphertext", "input", hex.EncodeToString(input))
+		return 0
+	}
+	return environment.FhevmParams().GasCosts.TeeReencrypt[ct.fheUintType()]
+}

fhevm/tee_crypto_test.go

@@ -46,3 +46,36 @@ func TestTeeDecryptRun(t *testing.T) {
 		}
 	})
 }
+
+func TestTeeReEncryptRun(t *testing.T) {
+	signature := "teeReencrypt(uint256,uint256)"
+	rapid.Check(t, func(t *rapid.T) {
+		testcases := []struct {
+			typ      tfhe.FheUintType
+			expected uint64
+		}{
+			{tfhe.FheUint4, uint64(rapid.Uint8().Draw(t, "expected"))},
+			{tfhe.FheUint8, uint64(rapid.Uint8().Draw(t, "expected"))},
+			{tfhe.FheUint16, uint64(rapid.Uint16().Draw(t, "expected"))},
+			{tfhe.FheUint32, uint64(rapid.Uint32().Draw(t, "expected"))},
+			{tfhe.FheUint64, uint64(rapid.Uint64().Draw(t, "expected"))},
+		}
+		for _, tc := range testcases {
+			depth := 1
+			environment := newTestEVMEnvironmentWithEthCall()
+			environment.depth = depth
+			addr := common.Address{}
+			readOnly := false
+			ct, err := importTeePlaintextToEVM(environment, depth, tc.expected, tc.typ)
+			if err != nil {
+				t.Fatalf(err.Error())
+			}
+
+			input := toLibPrecompileInput(signature, false, ct.GetHash())
+			_, err = FheLibRun(environment, addr, addr, input, readOnly)
+			if err != nil {
+				t.Fatalf("Reencrypt error: %s", err.Error())
+			}
+		}
+	})
+}

tee/tee_mock.go

@@ -5,6 +5,7 @@ import (
 	"encoding/binary"
 	"encoding/json"
 	"fmt"
+	"math/big"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto"
@@ -118,3 +119,17 @@ func Decrypt(ct *tfhe.TfheCiphertext) (TeePlaintext, error) {
 
 	return plaintext, nil
 }
+
+func DecryptToBigInt(ct *tfhe.TfheCiphertext) (*big.Int, error) {
+	decryptedValue, err := Decrypt(ct)
+	if err != nil {
+		return nil, err
+	}
+
+	plaintext := decryptedValue.Value
+	// Always return a 32-byte big-endian integer.
+	ret := make([]byte, 32)
+	copy(ret[32-len(plaintext):], plaintext)
+
+	return new(big.Int).SetBytes(ret), nil
+}