From c6824b26cbd37d3d11c417d3072585a425c9abd6 Mon Sep 17 00:00:00 2001
From: Petar Ivanov <29689712+dartdart26@users.noreply.github.com>
Date: Mon, 18 Sep 2023 10:02:09 +0300
Subject: [PATCH] Check reserved slots in protected storage on SLOAD

First check if the handle to load in SLOAD is in the reserved protected
storage slots list. If yes, do not try to verify it as a ciphertext.

Add unit tests for reserved protected storage slots for both SSTORE and
SLOAD.

Resolves #166.
---
 core/vm/instructions.go      | 20 ++++++--
 core/vm/instructions_test.go | 94 +++++++++++++++++++++++++++++++++++-
 2 files changed, 109 insertions(+), 5 deletions(-)

diff --git a/core/vm/instructions.go b/core/vm/instructions.go
index 6639083fa..e82e8377b 100644
--- a/core/vm/instructions.go
+++ b/core/vm/instructions.go
@@ -573,6 +573,11 @@ func verifyIfCiphertextHandle(val common.Hash, interpreter *EVMInterpreter, cont
 		return nil
 	}
 
+	// If a reserved slot, do not try treat it as ciphertext metadata.
+	if isReservedSlot(val) {
+		return nil
+	}
+
 	protectedStorage := crypto.CreateProtectedStorageContractAddress(contractAddress)
 	metadataInt := newInt(interpreter.evm.StateDB.GetState(protectedStorage, val).Bytes())
 	if !metadataInt.IsZero() {
@@ -680,13 +685,20 @@ func persistIfVerifiedCiphertext(val common.Hash, protectedStorage common.Addres
 // TODO: This list will be removed when we change the way we handle ciphertext handles and refcounts.
 var reservedProtectedStorageSlots []common.Hash = make([]common.Hash, 0)
 
+func isReservedSlot(key common.Hash) bool {
+	for _, slot := range reservedProtectedStorageSlots {
+		if bytes.Equal(key.Bytes(), slot.Bytes()) {
+			return true
+		}
+	}
+	return false
+}
+
 // If references are still left, reduce refCount by 1. Otherwise, zero out the metadata and the ciphertext slots.
 func garbageCollectProtectedStorage(metadataKey common.Hash, protectedStorage common.Address, interpreter *EVMInterpreter) {
 	// If a reserved slot, do not try to garbage collect it.
-	for _, slot := range reservedProtectedStorageSlots {
-		if bytes.Equal(metadataKey.Bytes(), slot.Bytes()) {
-			return
-		}
+	if isReservedSlot(metadataKey) {
+		return
 	}
 	existingMetadataHash := interpreter.evm.StateDB.GetState(protectedStorage, metadataKey)
 	existingMetadataInt := newInt(existingMetadataHash.Bytes())
diff --git a/core/vm/instructions_test.go b/core/vm/instructions_test.go
index 56ac480bc..e2efaa36d 100644
--- a/core/vm/instructions_test.go
+++ b/core/vm/instructions_test.go
@@ -837,6 +837,98 @@ func TestProtectedStorageGarbageCollection(t *testing.T) {
 	}
 }
 
+func TestProtectedStorageGarbageCollectionOnReservedSlot(t *testing.T) {
+	scope := newTestScopeConext()
+	protectedStorage := crypto.CreateProtectedStorageContractAddress(scope.Contract.Address())
+	interpreter := newTestInterpreter()
+	pc := uint64(0)
+	depth := 1
+	interpreter.evm.depth = depth
+
+	// Simulate metadata for a ciphertext at a reserved protected storage slot.
+	metadata := ciphertextMetadata{}
+	metadata.fheUintType = FheUint8
+	metadata.refCount = 1
+	metadata.length = 3
+	metadataSer := metadata.serialize()
+	metadataSlot := reservedProtectedStorageSlots[0]
+	interpreter.evm.StateDB.SetState(protectedStorage, metadataSlot, common.BytesToHash(metadataSer[:]))
+
+	// Simulate a ciphertext in protected storage.
+	slot := uint256FromBig(metadataSlot.Big())
+	nonZero := common.Hash{}
+	nonZero[0] = 1
+	for i := uint64(1); i < metadata.length+1; i++ {
+		slot = slot.AddUint64(slot, i)
+		interpreter.evm.StateDB.SetState(protectedStorage, common.BytesToHash(slot.Bytes()), nonZero)
+	}
+
+	// Simulate SSTORE with a new value that is different the reserved slot.
+	valueHash := metadataSlot
+	valueHash[0]++
+	loc := uint256.NewInt(10)
+	value := uint256FromBig(valueHash.Big())
+
+	// Call SSTORE.
+	scope.Stack.push(value)
+	scope.Stack.push(loc)
+	_, err := opSstore(&pc, interpreter, scope)
+	if err != nil {
+		t.Fatalf(err.Error())
+	}
+
+	// Verify that garbage collection hasn't happened for a reserved protected storage slot.
+	slot = uint256FromBig(metadataSlot.Big())
+	for i := uint64(0); i < metadata.length+1; i++ {
+		slot = slot.AddUint64(slot, i)
+		res := interpreter.evm.StateDB.GetState(protectedStorage, common.BytesToHash(slot.Bytes()))
+		if bytes.Equal(res.Bytes(), common.Hash{}.Bytes()) {
+			t.Fatalf("garbage collection must not have happened")
+		}
+	}
+}
+
+func TestProtectedStorageSloadOnReservedSlot(t *testing.T) {
+	scope := newTestScopeConext()
+	interpreter := newTestInterpreter()
+	pc := uint64(0)
+	depth := 1
+	interpreter.evm.depth = depth
+
+	handle := verifyCiphertextInTestMemory(interpreter, 2, depth, FheUint8).getHash()
+	loc := uint256.NewInt(10)
+	value := uint256FromBig(handle.Big())
+
+	// Consider the returned handle as a reserved slot.
+	reservedProtectedStorageSlots = append(reservedProtectedStorageSlots, handle)
+
+	// Persist the ciphertext in protected storage.
+	scope.Stack.push(value)
+	scope.Stack.push(loc)
+	_, err := opSstore(&pc, interpreter, scope)
+	if err != nil {
+		t.Fatalf(err.Error())
+	}
+
+	// Clear verified ciphertexts.
+	interpreter.verifiedCiphertexts = make(map[common.Hash]*verifiedCiphertext)
+
+	// Call SLOAD.
+	scope.Stack.push(loc)
+	_, err = opSload(&pc, interpreter, scope)
+	if err != nil {
+		t.Fatalf(err.Error())
+	}
+
+	// Remove the handle from reserved slots.
+	reservedProtectedStorageSlots = reservedProtectedStorageSlots[:len(reservedProtectedStorageSlots)-1]
+
+	// Expect no verified ciphertexts.
+	if len(interpreter.verifiedCiphertexts) != 0 {
+		t.Fatalf("expected no verified ciphetexts")
+	}
+}
+
 func TestProtectedStorageSloadDoesNotVerifyNonHandle(t *testing.T) {
 	pc := uint64(0)
 	interpreter := newTestInterpreter()
@@ -857,7 +949,7 @@ func TestProtectedStorageSloadDoesNotVerifyNonHandle(t *testing.T) {
 		t.Fatalf(err.Error())
 	}
 
-	// Expect no verified ciphertexts
+	// Expect no verified ciphertexts.
 	if len(interpreter.verifiedCiphertexts) != 0 {
 		t.Fatalf("expected no verified ciphetexts")
 	}