fix: add missing gas costs for loading ciphertexts

zama-ai · Jun 11, 2024 · cd54c7d · cd54c7d
1 parent 09e1ffe
commit cd54c7d
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 22 deletions.
diff --git a/fhevm/operators_arithmetic_gas.go b/fhevm/operators_arithmetic_gas.go
@@ -59,7 +59,7 @@ func fheMulRequiredGas(environment EVMEnvironment, input []byte) uint64 {
 			logger.Error("fheMul RequiredGas() operand type mismatch", "lhs", lhs.Type(), "rhs", rhs.Type())
 			return 0
 		}
-		return environment.FhevmParams().GasCosts.FheMul[lhs.Type()]
+		return environment.FhevmParams().GasCosts.FheMul[lhs.Type()] + loadGas
 	} else {
 		lhs, _, loadGas, err = getScalarOperands(environment, input)
 		if err != nil {

diff --git a/fhevm/operators_bit_gas.go b/fhevm/operators_bit_gas.go
@@ -28,7 +28,7 @@ func fheShlRequiredGas(environment EVMEnvironment, input []byte) uint64 {
 			logger.Error("fheShift RequiredGas() operand type mismatch", "lhs", lhs.Type(), "rhs", rhs.Type())
 			return 0
 		}
-		return environment.FhevmParams().GasCosts.FheShift[lhs.Type()]
+		return environment.FhevmParams().GasCosts.FheShift[lhs.Type()] + loadGas
 	} else {
 		lhs, _, loadGas, err = getScalarOperands(environment, input)
 		if err != nil {

diff --git a/fhevm/operators_comparison.go b/fhevm/operators_comparison.go
@@ -617,20 +617,22 @@ func init() {
 	}
 }
 
-func getVerifiedCiphertexts(environment EVMEnvironment, unpacked interface{}) ([]*tfhe.TfheCiphertext, error) {
+func getVerifiedCiphertexts(environment EVMEnvironment, unpacked interface{}) ([]*tfhe.TfheCiphertext, uint64, error) {
+	totalLoadGas := uint64(0)
 	big, ok := unpacked.([]*big.Int)
 	if !ok {
-		return nil, fmt.Errorf("fheArrayEq failed to cast to []*big.Int")
+		return nil, 0, fmt.Errorf("fheArrayEq failed to cast to []*big.Int")
 	}
 	ret := make([]*tfhe.TfheCiphertext, 0, len(big))
 	for _, b := range big {
-		ct, _ := loadCiphertext(environment, common.BigToHash(b))
+		ct, loadGas := loadCiphertext(environment, common.BigToHash(b))
 		if ct == nil {
-			return nil, fmt.Errorf("fheArrayEq unverified ciphertext")
+			return nil, 0, fmt.Errorf("fheArrayEq unverified ciphertext")
 		}
+		totalLoadGas += loadGas
 		ret = append(ret, ct)
 	}
-	return ret, nil
+	return ret, totalLoadGas, nil
 }
 
 func fheArrayEqRun(environment EVMEnvironment, caller common.Address, addr common.Address, input []byte, readOnly bool, runSpan trace.Span) ([]byte, error) {
@@ -649,14 +651,14 @@ func fheArrayEqRun(environment EVMEnvironment, caller common.Address, addr commo
 		return nil, err
 	}
 
-	lhs, err := getVerifiedCiphertexts(environment, unpacked[0])
+	lhs, _, err := getVerifiedCiphertexts(environment, unpacked[0])
 	if err != nil {
 		msg := "fheArrayEqRun failed to get lhs to verified ciphertexts"
 		logger.Error(msg, "err", err)
 		return nil, err
 	}
 
-	rhs, err := getVerifiedCiphertexts(environment, unpacked[1])
+	rhs, _, err := getVerifiedCiphertexts(environment, unpacked[1])
 	if err != nil {
 		msg := "fheArrayEqRun failed to get rhs to verified ciphertexts"
 		logger.Error(msg, "err", err)

diff --git a/fhevm/operators_comparison_gas.go b/fhevm/operators_comparison_gas.go
@@ -110,7 +110,7 @@ func fheMinRequiredGas(environment EVMEnvironment, input []byte) uint64 {
 			logger.Error("fheMin/Max RequiredGas() operand type mismatch", "lhs", lhs.Type(), "rhs", rhs.Type())
 			return 0
 		}
-		return environment.FhevmParams().GasCosts.FheMinMax[lhs.Type()]
+		return environment.FhevmParams().GasCosts.FheMinMax[lhs.Type()] + loadGas
 	} else {
 		lhs, _, loadGas, err = getScalarOperands(environment, input)
 		if err != nil {
@@ -162,14 +162,14 @@ func fheArrayEqRequiredGas(environment EVMEnvironment, input []byte) uint64 {
 		return 0
 	}
 
-	lhs, err := getVerifiedCiphertexts(environment, unpacked[0])
+	lhs, lhsLoadGas, err := getVerifiedCiphertexts(environment, unpacked[0])
 	if err != nil {
 		msg := "fheArrayEqRun RequiredGas() failed to get lhs to verified ciphertexts"
 		logger.Error(msg, "err", err)
 		return 0
 	}
 
-	rhs, err := getVerifiedCiphertexts(environment, unpacked[1])
+	rhs, rhsLoadGas, err := getVerifiedCiphertexts(environment, unpacked[1])
 	if err != nil {
 		msg := "fheArrayEqRun RequiredGas() failed to get rhs to verified ciphertexts"
 		logger.Error(msg, "err", err)
@@ -180,6 +180,8 @@ func fheArrayEqRequiredGas(environment EVMEnvironment, input []byte) uint64 {
 		return environment.FhevmParams().GasCosts.FheTrivialEncrypt[tfhe.FheBool]
 	}
 
+	totalLoadGas := lhsLoadGas + rhsLoadGas
+
 	numElements := len(lhs)
 	elementType := lhs[0].Type()
 	// TODO: tie to supported types in tfhe.TfheCiphertext.EqArray()
@@ -188,24 +190,24 @@ func fheArrayEqRequiredGas(environment EVMEnvironment, input []byte) uint64 {
 	}
 	for i := range lhs {
 		if lhs[i].Type() != elementType || rhs[i].Type() != elementType {
-			return 0
+			return totalLoadGas
 		}
 	}
 
 	numBits := elementType.NumBits() * uint(numElements)
 	if numBits <= 4 {
-		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint4]
+		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint4] + totalLoadGas
 	} else if numBits <= 8 {
-		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint8]
+		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint8] + totalLoadGas
 	} else if numBits <= 16 {
-		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint16]
+		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint16] + totalLoadGas
 	} else if numBits <= 32 {
-		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint32]
+		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint32] + totalLoadGas
 	} else if numBits <= 64 {
-		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint64]
+		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint64] + totalLoadGas
 	} else if numBits <= 160 {
-		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint160]
+		return environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint160] + totalLoadGas
 	} else {
-		return (environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint160] + environment.FhevmParams().GasCosts.FheArrayEqBigArrayFactor) * (uint64(numBits) / 160)
+		return ((environment.FhevmParams().GasCosts.FheEq[tfhe.FheUint160] + environment.FhevmParams().GasCosts.FheArrayEqBigArrayFactor) * (uint64(numBits) / 160)) + totalLoadGas
 	}
 }
diff --git a/fhevm/operators_crypto_gas.go b/fhevm/operators_crypto_gas.go
@@ -55,13 +55,20 @@ func getCiphertextRequiredGas(environment EVMEnvironment, input []byte) uint64 {
 func castRequiredGas(environment EVMEnvironment, input []byte) uint64 {
 	input = input[:minInt(33, len(input))]
 
+	logger := environment.GetLogger()
 	if len(input) != 33 {
-		environment.GetLogger().Error(
+		logger.Error(
 			"cast RequiredGas() input needs to contain a ciphertext and one byte for its type",
 			"len", len(input))
 		return 0
 	}
-	return environment.FhevmParams().GasCosts.FheCast
+
+	ct, loadGas := loadCiphertext(environment, common.BytesToHash(input[0:32]))
+	if ct == nil {
+		logger.Error("cast RequiredGas() input doesn't point to verified ciphertext", "input", hex.EncodeToString(input))
+		return 0
+	}
+	return environment.FhevmParams().GasCosts.FheCast + loadGas
 }
 
 func decryptRequiredGas(environment EVMEnvironment, input []byte) uint64 {