From c2ca6755a195d0bd743335aa1180fac73d1e64a8 Mon Sep 17 00:00:00 2001
From: David <deividas.kazlauskas@zama.ai>
Date: Mon, 13 Nov 2023 13:45:48 +0800
Subject: [PATCH] Add feature to disable decryptions in transaction

---
 fhevm/contracts_test.go    | 19 +++++++++++++++++++
 fhevm/instructions_test.go | 17 +++++++++--------
 fhevm/params.go            |  6 ++++--
 fhevm/precompiles.go       |  6 ++++++
 4 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/fhevm/contracts_test.go b/fhevm/contracts_test.go
index ca01194..0ea0dbd 100644
--- a/fhevm/contracts_test.go
+++ b/fhevm/contracts_test.go
@@ -3381,3 +3381,22 @@ func TestDecryptWithTrueOptimisticRequire(t *testing.T) {
 		t.Fatalf("expected that there are no optimistic requires after decrypt")
 	}
 }
+
+func TestDecryptInTransactionDisabled(t *testing.T) {
+	depth := 0
+	environment := newTestEVMEnvironment()
+	environment.depth = depth
+	environment.commit = true
+	environment.ethCall = false
+	environment.fhevmParams.DisableDecryptionsInTransaction = true
+	addr := common.Address{}
+	readOnly := false
+	hash := verifyCiphertextInTestMemory(environment, 1, depth, FheUint8).getHash()
+	// Call decrypt and expect it to fail due to disabling of decryptions during commit
+	_, err := decryptRun(environment, addr, addr, hash.Bytes(), readOnly)
+	if err == nil {
+		t.Fatalf("expected to error out in test")
+	} else if err.Error() != "decryptions during transaction are disabled" {
+		t.Fatalf("unexpected error for disabling decryption transactions, got %s", err.Error())
+	}
+}
diff --git a/fhevm/instructions_test.go b/fhevm/instructions_test.go
index efa2f23..ba39bcd 100644
--- a/fhevm/instructions_test.go
+++ b/fhevm/instructions_test.go
@@ -167,12 +167,13 @@ func uint256FromBig(b *big.Int) *uint256.Int {
 }
 
 type MockEVMEnvironment struct {
-	fhevmData *FhevmData
-	depth     int
-	stateDb   *state.StateDB
-	commit    bool
-	ethCall   bool
-	readOnly  bool
+	fhevmData   *FhevmData
+	depth       int
+	stateDb     *state.StateDB
+	commit      bool
+	ethCall     bool
+	readOnly    bool
+	fhevmParams FhevmParams
 }
 
 func (environment *MockEVMEnvironment) GetState(addr common.Address, hash common.Hash) common.Hash {
@@ -232,7 +233,7 @@ func (environment *MockEVMEnvironment) FhevmData() *FhevmData {
 }
 
 func (environment *MockEVMEnvironment) FhevmParams() *FhevmParams {
-	return &FhevmParams{}
+	return &environment.fhevmParams
 }
 
 func (environment *MockEVMEnvironment) EVMEnvironment() EVMEnvironment {
@@ -243,7 +244,7 @@ func newTestEVMEnvironment() *MockEVMEnvironment {
 	fhevmData := NewFhevmData()
 	db := rawdb.NewMemoryDatabase()
 	state, _ := state.New(common.Hash{}, state.NewDatabase(db), nil)
-	return &MockEVMEnvironment{fhevmData: &fhevmData, stateDb: state, commit: true}
+	return &MockEVMEnvironment{fhevmData: &fhevmData, stateDb: state, commit: true, fhevmParams: DefaultFhevmParams()}
 }
 
 func TestProtectedStorageSstoreSload(t *testing.T) {
diff --git a/fhevm/params.go b/fhevm/params.go
index 99ed23d..1147c15 100644
--- a/fhevm/params.go
+++ b/fhevm/params.go
@@ -44,12 +44,14 @@ var (
 
 func DefaultFhevmParams() FhevmParams {
 	return FhevmParams{
-		GasCosts: DefaultGasCosts(),
+		GasCosts:                        DefaultGasCosts(),
+		DisableDecryptionsInTransaction: false,
 	}
 }
 
 type FhevmParams struct {
-	GasCosts GasCosts
+	GasCosts                        GasCosts
+	DisableDecryptionsInTransaction bool
 }
 
 type GasCosts struct {
diff --git a/fhevm/precompiles.go b/fhevm/precompiles.go
index b7a3162..97dcde8 100644
--- a/fhevm/precompiles.go
+++ b/fhevm/precompiles.go
@@ -1956,6 +1956,12 @@ func optimisticRequireRun(environment EVMEnvironment, caller common.Address, add
 
 func decryptRun(environment EVMEnvironment, caller common.Address, addr common.Address, input []byte, readOnly bool) ([]byte, error) {
 	logger := environment.GetLogger()
+	// if not gas estimation and not view function fail if decryptions are disabled in transactions
+	if environment.IsCommitting() && !environment.IsEthCall() && environment.FhevmParams().DisableDecryptionsInTransaction {
+		msg := "decryptions during transaction are disabled"
+		logger.Error(msg, "input", hex.EncodeToString(input))
+		return nil, errors.New(msg)
+	}
 	if len(input) != 32 {
 		msg := "decrypt input len must be 32 bytes"
 		logger.Error(msg, "input", hex.EncodeToString(input), "len", len(input))