Skip to content

Commit

Permalink
feat: add support for euints in sstore/sload
Browse files Browse the repository at this point in the history 
Add an implementation of SLOAD that supports euints. Also, update the
implementation of SSTORE with latest one from fhEVM.
  • Loading branch information
@dartdart26
dartdart26 committed Oct 20, 2023
1 parent b863ebb commit 762deb0
Show file tree
Hide file tree
Showing 2 changed files with 134 additions and 44 deletions.
177 changes: 133 additions & 44 deletions fhevm/instructions.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,17 +3,14 @@ package fhevm
import (
"bytes"
"encoding/hex"
"errors"

"github.com/ethereum/go-ethereum/common"
crypto "github.com/ethereum/go-ethereum/crypto"
"github.com/holiman/uint256"
"github.com/zama-ai/fhevm-go/crypto"
ps "github.com/zama-ai/fhevm-go/crypto"
)

// A list of slots that we consider reserved in protected storage.
// Namely, we won't treat them as ciphertext metadata and we won't garbage collect them.
// 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)

var zero = uint256.NewInt(0).Bytes32()

func newInt(buf []byte) *uint256.Int {
Expand Down Expand Up @@ -51,21 +48,41 @@ func newCiphertextMetadata(buf [32]byte) *ciphertextMetadata {
return m.deserialize(buf)
}

// 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, environment EVMEnvironment) {
// If a reserved slot, do not try to garbage collect it.
for _, slot := range reservedProtectedStorageSlots {
if bytes.Equal(metadataKey.Bytes(), slot.Bytes()) {
return
}
func minUint64(a, b uint64) uint64 {
if a < b {
return a
}
existingMetadataHash := environment.GetState(protectedStorage, metadataKey)
return b
}

// If references are still left, reduce refCount by 1. Otherwise, zero out the metadata and the ciphertext slots.
func garbageCollectProtectedStorage(flagHandleLocation common.Hash, handle common.Hash, protectedStorage common.Address, env EVMEnvironment) {
// The location of ciphertext metadata is at Keccak256(handle). Doing so avoids attacks from users trying to garbage
// collect arbitrary locations in protected storage. Hashing the handle makes it hard to find a preimage such that
// it ends up in arbitrary non-zero places in protected stroage.
metadataKey := crypto.Keccak256Hash(handle.Bytes())

existingMetadataHash := env.GetState(protectedStorage, metadataKey)
existingMetadataInt := newInt(existingMetadataHash.Bytes())
if !existingMetadataInt.IsZero() {
logger := environment.GetLogger()
logger := env.GetLogger()

// If no flag in protected storage for the location, ignore garbage collection.
// Else, set the value at the location to zero.
foundFlag := env.GetState(protectedStorage, flagHandleLocation)
if !bytes.Equal(foundFlag.Bytes(), flag.Bytes()) {
logger.Error("opSstore location flag not found for a ciphertext handle, ignoring garbage collection",
"expectedFlag", hex.EncodeToString(flag[:]),
"foundFlag", hex.EncodeToString(foundFlag[:]),
"flagHandleLocation", hex.EncodeToString(flagHandleLocation[:]))
return
} else {
env.SetState(protectedStorage, flagHandleLocation, zero)
}

metadata := newCiphertextMetadata(existingMetadataInt.Bytes32())
if metadata.refCount == 1 {
if environment.IsCommitting() {
if env.IsCommitting() {
logger.Info("opSstore garbage collecting ciphertext",
"protectedStorage", hex.EncodeToString(protectedStorage[:]),
"metadataKey", hex.EncodeToString(metadataKey[:]),
Expand All @@ -74,7 +91,7 @@ func garbageCollectProtectedStorage(metadataKey common.Hash, protectedStorage co
}

// Zero the metadata key-value.
environment.SetState(protectedStorage, metadataKey, zero)
env.SetState(protectedStorage, metadataKey, zero)

// Set the slot to the one after the metadata one.
slot := newInt(metadataKey.Bytes())
Expand All @@ -86,19 +103,19 @@ func garbageCollectProtectedStorage(metadataKey common.Hash, protectedStorage co
slotsToZero++
}
for i := uint64(0); i < slotsToZero; i++ {
environment.SetState(protectedStorage, slot.Bytes32(), zero)
env.SetState(protectedStorage, slot.Bytes32(), zero)
slot.AddUint64(slot, 1)
}
} else if metadata.refCount > 1 {
if environment.IsCommitting() {
if env.IsCommitting() {
logger.Info("opSstore decrementing ciphertext refCount",
"protectedStorage", hex.EncodeToString(protectedStorage[:]),
"metadataKey", hex.EncodeToString(metadataKey[:]),
"type", metadata.fheUintType,
"len", metadata.length)
}
metadata.refCount--
environment.SetState(protectedStorage, existingMetadataHash, metadata.serialize())
env.SetState(protectedStorage, existingMetadataHash, metadata.serialize())
}
}
}
Expand All @@ -117,29 +134,93 @@ func getVerifiedCiphertextFromEVM(environment EVMEnvironment, ciphertextHash com
return nil
}

func verifyIfCiphertextHandle(handle common.Hash, env EVMEnvironment, contractAddress common.Address) error {
ct, ok := env.GetFhevmData().verifiedCiphertexts[handle]
if ok {
// If already existing in memory, skip storage and import the same ciphertext at the current depth.
//
// Also works for gas estimation - we don't persist anything to protected storage during gas estimation.
// However, ciphertexts remain in memory for the duration of the call, allowing for this lookup to find it.
// Note that even if a ciphertext has an empty verification depth set, it still remains in memory.
importCiphertextToEVM(env, ct.ciphertext)
return nil
}

metadataKey := crypto.Keccak256Hash(handle.Bytes())
protectedStorage := ps.CreateProtectedStorageContractAddress(contractAddress)
metadataInt := newInt(env.GetState(protectedStorage, metadataKey).Bytes())
if !metadataInt.IsZero() {
metadata := newCiphertextMetadata(metadataInt.Bytes32())
ctBytes := make([]byte, 0)
left := metadata.length
protectedSlotIdx := newInt(metadataKey.Bytes())
protectedSlotIdx.AddUint64(protectedSlotIdx, 1)
for {
if left == 0 {
break
}
bytes := env.GetState(protectedStorage, protectedSlotIdx.Bytes32())
toAppend := minUint64(uint64(len(bytes)), left)
left -= toAppend
ctBytes = append(ctBytes, bytes[0:toAppend]...)
protectedSlotIdx.AddUint64(protectedSlotIdx, 1)
}

ct := new(tfheCiphertext)
err := ct.deserialize(ctBytes, metadata.fheUintType)
if err != nil {
msg := "opSload failed to deserialize a ciphertext"
env.GetLogger().Error(msg, "err", err)
return errors.New(msg)
}
importCiphertextToEVM(env, ct)
}
return nil
}

func OpSload(pc *uint64, env EVMEnvironment, scope ScopeContext) ([]byte, error) {
loc := scope.GetStack().Peek()
hash := common.Hash(loc.Bytes32())
val := env.GetState(scope.GetContract().Address(), hash)
if err := verifyIfCiphertextHandle(val, env, scope.GetContract().Address()); err != nil {
return nil, err
}
loc.SetBytes(val.Bytes())
return nil, nil
}

// An arbitrary constant value to flag locations in protected storage.
var flag = common.HexToHash("0xa145ffde0100a145ffde0100a145ffde0100a145ffde0100a145ffde0100fab3")

// If a verified ciphertext:
// * if the ciphertext does not exist in protected storage, persist it with a refCount = 1
// * if the ciphertexts exists in protected, bump its refCount by 1
func persistIfVerifiedCiphertext(val common.Hash, protectedStorage common.Address, environment EVMEnvironment) {
verifiedCiphertext := getVerifiedCiphertextFromEVM(environment, val)
func persistIfVerifiedCiphertext(flagHandleLocation common.Hash, handle common.Hash, protectedStorage common.Address, env EVMEnvironment) {
verifiedCiphertext := getVerifiedCiphertextFromEVM(env, handle)
if verifiedCiphertext == nil {
return
}
logger := environment.GetLogger()
logger := env.GetLogger()

// Try to read ciphertext metadata from protected storage.
metadataInt := newInt(environment.GetState(protectedStorage, val).Bytes())
metadataKey := crypto.Keccak256Hash(handle.Bytes())
metadataInt := newInt(env.GetState(protectedStorage, metadataKey).Bytes())
metadata := ciphertextMetadata{}

// Set flag in protected storage to mark the location as containing a handle.
env.SetState(protectedStorage, flagHandleLocation, flag)

if metadataInt.IsZero() {
// If no metadata, it means this ciphertext itself hasn't been persisted to protected storage yet. We do that as part of SSTORE.
metadata.refCount = 1
metadata.length = uint64(expandedFheCiphertextSize[verifiedCiphertext.ciphertext.fheUintType])
metadata.fheUintType = verifiedCiphertext.ciphertext.fheUintType
ciphertextSlot := newInt(val.Bytes())
ciphertextSlot := newInt(metadataKey.Bytes())
ciphertextSlot.AddUint64(ciphertextSlot, 1)
if environment.IsCommitting() {
if env.IsCommitting() {
logger.Info("opSstore persisting new ciphertext",
"protectedStorage", hex.EncodeToString(protectedStorage[:]),
"handle", hex.EncodeToString(val.Bytes()),
"handle", hex.EncodeToString(handle.Bytes()),
"type", metadata.fheUintType,
"len", metadata.length,
"ciphertextSlot", hex.EncodeToString(ciphertextSlot.Bytes()))
Expand All @@ -149,7 +230,7 @@ func persistIfVerifiedCiphertext(val common.Hash, protectedStorage common.Addres
ctBytes := verifiedCiphertext.ciphertext.serialize()
for i, b := range ctBytes {
if i%32 == 0 && i != 0 {
environment.SetState(protectedStorage, ciphertextSlot.Bytes32(), common.BytesToHash(ctPart32))
env.SetState(protectedStorage, ciphertextSlot.Bytes32(), common.BytesToHash(ctPart32))
ciphertextSlot.AddUint64(ciphertextSlot, 1)
ctPart32 = make([]byte, 32)
partIdx = 0
Expand All @@ -158,44 +239,52 @@ func persistIfVerifiedCiphertext(val common.Hash, protectedStorage common.Addres
partIdx++
}
if len(ctPart32) != 0 {
environment.SetState(protectedStorage, ciphertextSlot.Bytes32(), common.BytesToHash(ctPart32))
env.SetState(protectedStorage, ciphertextSlot.Bytes32(), common.BytesToHash(ctPart32))
}
} else {
// If metadata exists, bump the refcount by 1.
metadata = *newCiphertextMetadata(environment.GetState(protectedStorage, val))
metadata = *newCiphertextMetadata(env.GetState(protectedStorage, metadataKey))
metadata.refCount++
if environment.IsCommitting() {
if env.IsCommitting() {
logger.Info("opSstore bumping refcount of existing ciphertext",
"protectedStorage", hex.EncodeToString(protectedStorage[:]),
"handle", hex.EncodeToString(val.Bytes()),
"handle", hex.EncodeToString(handle.Bytes()),
"type", metadata.fheUintType,
"len", metadata.length,
"refCount", metadata.refCount)
}
}
// Save the metadata in protected storage.
environment.SetState(protectedStorage, val, metadata.serialize())
env.SetState(protectedStorage, metadataKey, metadata.serialize())
}

func OpSstore(pc *uint64, environment EVMEnvironment, scope ScopeContext) ([]byte, error) {
if environment.IsReadOnly() {
func OpSstore(pc *uint64, env EVMEnvironment, scope ScopeContext) ([]byte, error) {
if env.IsReadOnly() {
return nil, ErrWriteProtection
}
loc := scope.GetStack().Pop()
locHash := common.BytesToHash(loc.Bytes())
newVal := scope.GetStack().Pop()
newValBytes := newVal.Bytes()
newValHash := common.BytesToHash(newValBytes)
oldValHash := environment.GetState(scope.GetContract().Address(), common.Hash(loc.Bytes32()))
protectedStorage := crypto.CreateProtectedStorageContractAddress(scope.GetContract().Address())
newValHash := common.BytesToHash(newVal.Bytes())
oldValHash := env.GetState(scope.GetContract().Address(), common.Hash(loc.Bytes32()))
// If the value is the same or if we are not going to commit, don't do anything to protected storage.
if newValHash != oldValHash && environment.IsCommitting() {
if newValHash != oldValHash && env.IsCommitting() {
protectedStorage := ps.CreateProtectedStorageContractAddress(scope.GetContract().Address())

// Define flag location as keccak256(keccak256(loc)) in protected storage. Used to mark the location as containing a handle.
// Note: We apply the hash function twice to make sure a flag location in protected storage cannot clash with a ciphertext
// metadata location that is keccak256(keccak256(ciphertext)). Since a location is 32 bytes, it cannot clash with a well-formed
// ciphertext. Therefore, there needs to be a hash collistion for a clash to happen. If hash is applied only once, there could
// be a collision, since malicous users could store at loc = keccak256(ciphertext), making the flag clash with metadata.
flagHandleLocation := crypto.Keccak256Hash(crypto.Keccak256Hash(locHash[:]).Bytes())

// Since the old value is no longer stored in actual contract storage, run garbage collection on protected storage.
garbageCollectProtectedStorage(oldValHash, protectedStorage, environment)
garbageCollectProtectedStorage(flagHandleLocation, oldValHash, protectedStorage, env)

// If a verified ciphertext, persist to protected storage.
persistIfVerifiedCiphertext(newValHash, protectedStorage, environment)
persistIfVerifiedCiphertext(flagHandleLocation, newValHash, protectedStorage, env)
}
// Set the SSTORE's value in the actual contract.
environment.SetState(scope.GetContract().Address(),
loc.Bytes32(), newValHash)
env.SetState(scope.GetContract().Address(), loc.Bytes32(), newValHash)
return nil, nil
}
1 change: 1 addition & 0 deletions fhevm/stack.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,4 +4,5 @@ import "github.com/holiman/uint256"

type Stack interface {
Pop() uint256.Int
Peek() *uint256.Int
}

0 comments on commit 762deb0

Please sign in to comment.