diff --git a/README.md b/README.md index 848af9b568..c0e7849bb0 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,15 @@ https://camo.githubusercontent.com/915b7be44ada53c290eb157634330494ebe3e30a/6874 ZK-Rollup adapts the Go Ethereum to run as Layer 2 Sequencer. The codebase is based on v1.10.13. +### ZKTrie Storage + +Another implement for storage trie, base on patricia merkle tree, has been induced. It is feasible to zk proving in the storage part. It is specified as a flag +in gensis, set `config.zktrie` to true for enabling it. Using `genesis_zktrie.json` as an example to create a L2 chain with zktrie sotrage: + +> geth init \/genesis_zktrie.json + +Notice current the snapshot would be disabled by the zktrie implement. + ## Building the source For prerequisites and detailed build instructions please read the [Installation Instructions](https://geth.ethereum.org/docs/install-and-build/installing-geth). diff --git a/cmd/geth/genesis_test.go b/cmd/geth/genesis_test.go index 0bd8468aee..3b077bae8b 100644 --- a/cmd/geth/genesis_test.go +++ b/cmd/geth/genesis_test.go @@ -121,7 +121,7 @@ func TestCustomGenesis(t *testing.T) { runGeth(t, "--datadir", datadir, "init", json).WaitExit() // Query the custom genesis block - geth := runGeth(t, "--networkid", "1337", "--syncmode=full", "--cache", "16", + geth := runGeth(t, "--networkid", "1337", "--syncmode=full", "--snapshot=false", "--cache", "16", "--datadir", datadir, "--maxpeers", "0", "--port", "0", "--nodiscover", "--nat", "none", "--ipcdisable", "--exec", tt.query, "console") diff --git a/common/types.go b/common/types.go index a7f6c3928d..b1eccc651c 100644 --- a/common/types.go +++ b/common/types.go @@ -229,6 +229,12 @@ func IsHexAddress(s string) bool { // Bytes gets the string representation of the underlying address. func (a Address) Bytes() []byte { return a[:] } +func (a Address) Bytes32() []byte { + ret := make([]byte, 32) + copy(ret, a.Bytes()) + return ret +} + // Hash converts an address to a hash by left-padding it with zeros. func (a Address) Hash() Hash { return BytesToHash(a[:]) } diff --git a/core/blockchain.go b/core/blockchain.go index 39d5fcc801..b20219657c 100644 --- a/core/blockchain.go +++ b/core/blockchain.go @@ -236,6 +236,11 @@ func NewBlockChain(db ethdb.Database, cacheConfig *CacheConfig, chainConfig *par if cacheConfig.TraceCacheLimit != 0 { blockResultCache, _ = lru.New(cacheConfig.TraceCacheLimit) } + // override snapshot setting + if chainConfig.Zktrie && cacheConfig.SnapshotLimit > 0 { + log.Warn("snapshot has been disabled by zktrie") + cacheConfig.SnapshotLimit = 0 + } bc := &BlockChain{ chainConfig: chainConfig, @@ -246,6 +251,7 @@ func NewBlockChain(db ethdb.Database, cacheConfig *CacheConfig, chainConfig *par Cache: cacheConfig.TrieCleanLimit, Journal: cacheConfig.TrieCleanJournal, Preimages: cacheConfig.Preimages, + Zktrie: chainConfig.Zktrie, }), quit: make(chan struct{}), chainmu: syncx.NewClosableMutex(), diff --git a/core/genesis.go b/core/genesis.go index ffc4d6c33f..d7e83a7d7a 100644 --- a/core/genesis.go +++ b/core/genesis.go @@ -180,7 +180,21 @@ func SetupGenesisBlockWithOverride(db ethdb.Database, genesis *Genesis, override // We have the genesis block in database(perhaps in ancient database) // but the corresponding state is missing. header := rawdb.ReadHeader(db, stored, 0) - if _, err := state.New(header.Root, state.NewDatabaseWithConfig(db, nil), nil); err != nil { + + var trieCfg *trie.Config + + if genesis == nil { + storedcfg := rawdb.ReadChainConfig(db, stored) + if storedcfg == nil { + log.Warn("Found genesis block without chain config") + } else { + trieCfg = &trie.Config{Zktrie: storedcfg.Zktrie} + } + } else { + trieCfg = &trie.Config{Zktrie: genesis.Config.Zktrie} + } + + if _, err := state.New(header.Root, state.NewDatabaseWithConfig(db, trieCfg), nil); err != nil { if genesis == nil { genesis = DefaultGenesisBlock() } @@ -261,7 +275,11 @@ func (g *Genesis) ToBlock(db ethdb.Database) *types.Block { if db == nil { db = rawdb.NewMemoryDatabase() } - statedb, err := state.New(common.Hash{}, state.NewDatabase(db), nil) + var trieCfg *trie.Config + if g.Config != nil { + trieCfg = &trie.Config{Zktrie: g.Config.Zktrie} + } + statedb, err := state.New(common.Hash{}, state.NewDatabaseWithConfig(db, trieCfg), nil) if err != nil { panic(err) } @@ -441,7 +459,8 @@ func DeveloperGenesisBlock(period uint64, gasLimit uint64, faucet common.Address common.BytesToAddress([]byte{7}): {Balance: big.NewInt(1)}, // ECScalarMul common.BytesToAddress([]byte{8}): {Balance: big.NewInt(1)}, // ECPairing common.BytesToAddress([]byte{9}): {Balance: big.NewInt(1)}, // BLAKE2b - faucet: {Balance: new(big.Int).Sub(new(big.Int).Lsh(big.NewInt(1), 256), big.NewInt(9))}, + // LSH 250 due to finite field limitation + faucet: {Balance: new(big.Int).Sub(new(big.Int).Lsh(big.NewInt(1), 250), big.NewInt(9))}, }, } } diff --git a/core/state/database.go b/core/state/database.go index b42c921215..bb73fcecd2 100644 --- a/core/state/database.go +++ b/core/state/database.go @@ -120,6 +120,7 @@ func NewDatabase(db ethdb.Database) Database { func NewDatabaseWithConfig(db ethdb.Database, config *trie.Config) Database { csc, _ := lru.New(codeSizeCacheSize) return &cachingDB{ + zktrie: config != nil && config.Zktrie, db: trie.NewDatabaseWithConfig(db, config), codeSizeCache: csc, codeCache: fastcache.New(codeCacheSize), @@ -130,10 +131,18 @@ type cachingDB struct { db *trie.Database codeSizeCache *lru.Cache codeCache *fastcache.Cache + zktrie bool } // OpenTrie opens the main account trie at a specific root hash. func (db *cachingDB) OpenTrie(root common.Hash) (Trie, error) { + if db.zktrie { + tr, err := trie.NewZkTrie(root, trie.NewZktrieDatabaseFromTriedb(db.db)) + if err != nil { + return nil, err + } + return tr, nil + } tr, err := trie.NewSecure(root, db.db) if err != nil { return nil, err @@ -143,6 +152,13 @@ func (db *cachingDB) OpenTrie(root common.Hash) (Trie, error) { // OpenStorageTrie opens the storage trie of an account. func (db *cachingDB) OpenStorageTrie(addrHash, root common.Hash) (Trie, error) { + if db.zktrie { + tr, err := trie.NewZkTrie(root, trie.NewZktrieDatabaseFromTriedb(db.db)) + if err != nil { + return nil, err + } + return tr, nil + } tr, err := trie.NewSecure(root, db.db) if err != nil { return nil, err @@ -155,6 +171,8 @@ func (db *cachingDB) CopyTrie(t Trie) Trie { switch t := t.(type) { case *trie.SecureTrie: return t.Copy() + case *trie.ZkTrie: + return t.Copy() default: panic(fmt.Errorf("unknown trie type %T", t)) } diff --git a/core/state/snapshot/snapshot.go b/core/state/snapshot/snapshot.go index e6b004a133..417e742d95 100644 --- a/core/state/snapshot/snapshot.go +++ b/core/state/snapshot/snapshot.go @@ -185,6 +185,9 @@ type Tree struct { // a background thread. func New(diskdb ethdb.KeyValueStore, triedb *trie.Database, cache int, root common.Hash, async bool, rebuild bool, recovery bool) (*Tree, error) { // Create a new, empty snapshot tree + if triedb.Zktrie { + panic("zktrie does not support snapshot yet") + } snap := &Tree{ diskdb: diskdb, triedb: triedb, diff --git a/core/state/state_object.go b/core/state/state_object.go index ad4e197b79..a5294ef602 100644 --- a/core/state/state_object.go +++ b/core/state/state_object.go @@ -107,7 +107,7 @@ func newObject(db *StateDB, address common.Address, data types.StateAccount) *st data.CodeHash = emptyCodeHash } if data.Root == (common.Hash{}) { - data.Root = emptyRoot + data.Root = db.db.TrieDB().EmptyRoot() } return &stateObject{ db: db, @@ -151,7 +151,7 @@ func (s *stateObject) getTrie(db Database) Trie { if s.trie == nil { // Try fetching from prefetcher first // We don't prefetch empty tries - if s.data.Root != emptyRoot && s.db.prefetcher != nil { + if s.data.Root != s.db.db.TrieDB().EmptyRoot() && s.db.prefetcher != nil { // When the miner is creating the pending state, there is no // prefetcher s.trie = s.db.prefetcher.trie(s.data.Root) @@ -245,12 +245,16 @@ func (s *stateObject) GetCommittedState(db Database, key common.Hash) common.Has } } var value common.Hash - if len(enc) > 0 { - _, content, _, err := rlp.Split(enc) - if err != nil { - s.setError(err) + if db.TrieDB().Zktrie { + value = common.BytesToHash(enc) + } else { + if len(enc) > 0 { + _, content, _, err := rlp.Split(enc) + if err != nil { + s.setError(err) + } + value.SetBytes(content) } - value.SetBytes(content) } s.originStorage[key] = value return value @@ -309,7 +313,7 @@ func (s *stateObject) finalise(prefetch bool) { slotsToPrefetch = append(slotsToPrefetch, common.CopyBytes(key[:])) // Copy needed for closure } } - if s.db.prefetcher != nil && prefetch && len(slotsToPrefetch) > 0 && s.data.Root != emptyRoot { + if s.db.prefetcher != nil && prefetch && len(slotsToPrefetch) > 0 && s.data.Root != s.db.db.TrieDB().EmptyRoot() { s.db.prefetcher.prefetch(s.data.Root, slotsToPrefetch) } if len(s.dirtyStorage) > 0 { @@ -348,8 +352,12 @@ func (s *stateObject) updateTrie(db Database) Trie { s.setError(tr.TryDelete(key[:])) s.db.StorageDeleted += 1 } else { - // Encoding []byte cannot fail, ok to ignore the error. - v, _ = rlp.EncodeToBytes(common.TrimLeftZeroes(value[:])) + if db.TrieDB().Zktrie { + v = common.CopyBytes(value[:]) + } else { + // Encoding []byte cannot fail, ok to ignore the error. + v, _ = rlp.EncodeToBytes(common.TrimLeftZeroes(value[:])) + } s.setError(tr.TryUpdate(key[:], v)) s.db.StorageUpdated += 1 } diff --git a/core/state/statedb.go b/core/state/statedb.go index e0d06c3cc5..5ac2e102bb 100644 --- a/core/state/statedb.go +++ b/core/state/statedb.go @@ -40,11 +40,6 @@ type revision struct { journalIndex int } -var ( - // emptyRoot is the known root hash of an empty trie. - emptyRoot = common.HexToHash("56e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421") -) - type proofList [][]byte func (n *proofList) Put(key []byte, value []byte) error { @@ -187,6 +182,10 @@ func (s *StateDB) Error() error { return s.dbErr } +func (s *StateDB) IsZktrie() bool { + return s.db.TrieDB().Zktrie +} + func (s *StateDB) AddLog(log *types.Log) { s.journal.append(addLogChange{txhash: s.thash}) @@ -315,11 +314,19 @@ func (s *StateDB) GetState(addr common.Address, hash common.Hash) common.Hash { // GetProof returns the Merkle proof for a given account. func (s *StateDB) GetProof(addr common.Address) ([][]byte, error) { + if s.IsZktrie() { + var proof proofList + err := s.trie.Prove(addr.Bytes32(), 0, &proof) + return proof, err + } return s.GetProofByHash(crypto.Keccak256Hash(addr.Bytes())) } // GetProofByHash returns the Merkle proof for a given account. func (s *StateDB) GetProofByHash(addrHash common.Hash) ([][]byte, error) { + if s.IsZktrie() { + panic("unimplemented") + } var proof proofList err := s.trie.Prove(addrHash[:], 0, &proof) return proof, err @@ -339,6 +346,7 @@ func (s *StateDB) GetRootHash() common.Hash { // StorageTrieProof is not in Db interface and used explictily for reading proof in storage trie (not the dirty value) func (s *StateDB) GetStorageTrieProof(a common.Address, key common.Hash) ([][]byte, error) { + // try the trie in stateObject first, else we would create one stateObject := s.getStateObject(a) if stateObject == nil { @@ -356,7 +364,11 @@ func (s *StateDB) GetStorageTrieProof(a common.Address, key common.Hash) ([][]by } var proof proofList - err = trie.Prove(crypto.Keccak256(key.Bytes()), 0, &proof) + if s.IsZktrie() { + err = trie.Prove(key.Bytes(), 0, &proof) + } else { + err = trie.Prove(crypto.Keccak256(key.Bytes()), 0, &proof) + } return proof, err } @@ -367,7 +379,12 @@ func (s *StateDB) GetStorageProof(a common.Address, key common.Hash) ([][]byte, if trie == nil { return proof, errors.New("storage trie for requested address does not exist") } - err := trie.Prove(crypto.Keccak256(key.Bytes()), 0, &proof) + var err error + if s.IsZktrie() { + err = trie.Prove(key.Bytes(), 0, &proof) + } else { + err = trie.Prove(crypto.Keccak256(key.Bytes()), 0, &proof) + } return proof, err } @@ -564,7 +581,7 @@ func (s *StateDB) getDeletedStateObject(addr common.Address) *stateObject { data.CodeHash = emptyCodeHash } if data.Root == (common.Hash{}) { - data.Root = emptyRoot + data.Root = s.db.TrieDB().EmptyRoot() } } } @@ -582,7 +599,12 @@ func (s *StateDB) getDeletedStateObject(addr common.Address) *stateObject { return nil } data = new(types.StateAccount) - if err := rlp.DecodeBytes(enc, data); err != nil { + if s.IsZktrie() { + data, err = types.UnmarshalStateAccount(enc) + } else { + err = rlp.DecodeBytes(enc, data) + } + if err != nil { log.Error("Failed to decode state object", "addr", addr, "err", err) return nil } @@ -977,7 +999,7 @@ func (s *StateDB) Commit(deleteEmptyObjects bool) (common.Hash, error) { if err := rlp.DecodeBytes(leaf, &account); err != nil { return nil } - if account.Root != emptyRoot { + if account.Root != s.db.TrieDB().EmptyRoot() { s.db.TrieDB().Reference(account.Root, parent) } return nil diff --git a/core/types/state_account_marshalling.go b/core/types/state_account_marshalling.go new file mode 100644 index 0000000000..8d4bf0b727 --- /dev/null +++ b/core/types/state_account_marshalling.go @@ -0,0 +1,107 @@ +// Copyright 2021 The go-ethereum Authors +// This file is part of the go-ethereum library. +// +// The go-ethereum library is free software: you can redistribute it and/or modify +// it under the terms of the GNU Lesser General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// The go-ethereum library is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Lesser General Public License for more details. +// +// You should have received a copy of the GNU Lesser General Public License +// along with the go-ethereum library. If not, see . + +package types + +import ( + "encoding/binary" + "errors" + "math/big" + + "github.com/iden3/go-iden3-crypto/poseidon" + "github.com/iden3/go-iden3-crypto/utils" + + "github.com/scroll-tech/go-ethereum/common" + zkt "github.com/scroll-tech/go-ethereum/core/types/zktrie" +) + +var ( + ErrInvalidLength = errors.New("StateAccount: invalid input length") +) + +// Hash of StateAccount +// AccountHash = Hash( +// Hash(nonce, balance), +// Hash( +// Root, +// Hash(codeHashFirst16, codeHashLast16) +// )) +func (s *StateAccount) Hash() (*big.Int, error) { + nonce := new(big.Int).SetUint64(s.Nonce) + if s.Balance == nil { + s.Balance = new(big.Int) + } + hash1, err := poseidon.Hash([]*big.Int{nonce, s.Balance}) + if err != nil { + return nil, err + } + + codeHashFirst16 := new(big.Int).SetBytes(s.CodeHash[0:16]) + codeHashLast16 := new(big.Int).SetBytes(s.CodeHash[16:32]) + hash2, err := poseidon.Hash([]*big.Int{codeHashFirst16, codeHashLast16}) + if err != nil { + return nil, err + } + + rootHash, err := zkt.NewHashFromBytes(s.Root.Bytes()) + if err != nil { + return nil, err + } + hash3, err := poseidon.Hash([]*big.Int{hash2, rootHash.BigInt()}) + if err != nil { + return nil, err + } + + hash4, err := poseidon.Hash([]*big.Int{hash1, hash3}) + if err != nil { + return nil, err + } + return hash4, nil +} + +// MarshalFields, the bytes scheme is: +// [0:32] Nonce uint64 big-endian in 32 byte +// [32:64] Balance +// [64:96] Root +// [96:128] CodeHash +func (s *StateAccount) MarshalFields() ([]zkt.Byte32, uint32) { + fields := make([]zkt.Byte32, 4) + binary.BigEndian.PutUint64(fields[0][24:], s.Nonce) + + if !utils.CheckBigIntInField(s.Balance) { + panic("balance overflow") + } + s.Balance.FillBytes(fields[1][:]) + + copy(fields[2][:], s.CodeHash) + copy(fields[3][:], s.Root.Bytes()) + return fields, 4 +} + +func UnmarshalStateAccount(bytes []byte) (*StateAccount, error) { + if len(bytes) != 128 { + return nil, ErrInvalidLength + } + acc := new(StateAccount) + acc.Nonce = binary.BigEndian.Uint64(bytes[24:]) + acc.Balance = new(big.Int).SetBytes(bytes[32:64]) + acc.CodeHash = make([]byte, 32) + copy(acc.CodeHash, bytes[64:96]) + acc.Root = common.Hash{} + acc.Root.SetBytes(bytes[96:128]) + + return acc, nil +} diff --git a/core/types/state_account_marshalling_test.go b/core/types/state_account_marshalling_test.go new file mode 100644 index 0000000000..02110fb789 --- /dev/null +++ b/core/types/state_account_marshalling_test.go @@ -0,0 +1,44 @@ +package types + +import ( + "math/big" + "testing" + + "github.com/scroll-tech/go-ethereum/common" + "github.com/scroll-tech/go-ethereum/core/types/zktrie" +) + +func TestAccountMarshalling(t *testing.T) { + //ensure the hash scheme consistent with designation + example1 := &StateAccount{ + Nonce: 5, + Balance: big.NewInt(0).SetBytes(common.Hex2Bytes("01fffffffffffffffffffffffffffffffffffffffffffffffff9c8672c6bc7b3")), + CodeHash: common.Hex2Bytes("c5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470"), + } + + example2 := &StateAccount{ + Nonce: 2, + Balance: big.NewInt(0), + CodeHash: common.Hex2Bytes("cc0a77f6e063b4b62eb7d9ed6f427cf687d8d0071d751850cfe5d136bc60d3ab"), + Root: common.HexToHash("22fb59aa5410ed465267023713ab42554c250f394901455a3366e223d5f7d147"), + } + + for i, example := range []*StateAccount{example1, example2} { + fields, flag := example.MarshalFields() + + h1, err := zktrie.PreHandlingElems(flag, fields) + if err != nil { + t.Fatal(err) + } + + h2, err := example.Hash() + if err != nil { + t.Fatal(err) + } + + if h1.BigInt().Cmp(h2) != 0 { + t.Errorf("hash <%d> unmatched, expected [%x], get [%x]", i, h2.Bytes(), h1.Bytes()) + } + } + +} diff --git a/core/types/zktrie/byte32.go b/core/types/zktrie/byte32.go new file mode 100644 index 0000000000..a622b0496d --- /dev/null +++ b/core/types/zktrie/byte32.go @@ -0,0 +1,28 @@ +package zktrie + +import ( + "fmt" + "math/big" + + "github.com/iden3/go-iden3-crypto/poseidon" +) + +type Byte32 [32]byte + +func (b *Byte32) Hash() (*big.Int, error) { + first16 := new(big.Int).SetBytes(b[0:16]) + last16 := new(big.Int).SetBytes(b[16:32]) + hash, err := poseidon.Hash([]*big.Int{first16, last16}) + if err != nil { + return nil, err + } + return hash, nil +} +func NewByte32FromBytesPaddingZero(b []byte) *Byte32 { + if len(b) != 32 && len(b) != 20 { + panic(fmt.Errorf("do not support length except for 120bit and 256bit now. data: %v len: %v", b, len(b))) + } + byte32 := new(Byte32) + copy(byte32[:], b) + return byte32 +} diff --git a/core/types/zktrie/hash.go b/core/types/zktrie/hash.go new file mode 100644 index 0000000000..3432242a2a --- /dev/null +++ b/core/types/zktrie/hash.go @@ -0,0 +1,120 @@ +package zktrie + +import ( + "encoding/hex" + "fmt" + "math/big" + + "github.com/iden3/go-iden3-crypto/utils" + + "github.com/scroll-tech/go-ethereum/common" +) + +const numCharPrint = 8 + +// ElemBytesLen is the length of the Hash byte array +const ElemBytesLen = 32 + +var HashZero = Hash{} + +// Hash is the generic type stored in the MerkleTree +type Hash [32]byte + +// MarshalText implements the marshaler for the Hash type +func (h Hash) MarshalText() ([]byte, error) { + return []byte(h.BigInt().String()), nil +} + +// UnmarshalText implements the unmarshaler for the Hash type +func (h *Hash) UnmarshalText(b []byte) error { + ha, err := NewHashFromString(string(b)) + copy(h[:], ha[:]) + return err +} + +// String returns decimal representation in string format of the Hash +func (h Hash) String() string { + s := h.BigInt().String() + if len(s) < numCharPrint { + return s + } + return s[0:numCharPrint] + "..." +} + +// Hex returns the hexadecimal representation of the Hash +func (h Hash) Hex() string { + return hex.EncodeToString(h[:]) +} + +// BigInt returns the *big.Int representation of the *Hash +func (h *Hash) BigInt() *big.Int { + if new(big.Int).SetBytes(ReverseByteOrder(h[:])) == nil { + return big.NewInt(0) + } + return new(big.Int).SetBytes(ReverseByteOrder(h[:])) +} + +// Bytes returns the little endian []byte representation of the *Hash, which always is 32 +// bytes length. +func (h *Hash) Bytes() []byte { + b := [32]byte{} + copy(b[:], h[:]) + return ReverseByteOrder(b[:]) +} + +// NewBigIntFromHashBytes returns a *big.Int from a byte array, swapping the +// endianness in the process. This is the intended method to get a *big.Int +// from a byte array that previously has ben generated by the Hash.Bytes() +// method. +func NewBigIntFromHashBytes(b []byte) (*big.Int, error) { + if len(b) != ElemBytesLen { + return nil, fmt.Errorf("expected 32 bytes, found %d bytes", len(b)) + } + bi := new(big.Int).SetBytes(b[:ElemBytesLen]) + if !utils.CheckBigIntInField(bi) { + return nil, fmt.Errorf("NewBigIntFromHashBytes: Value not inside the Finite Field") + } + return bi, nil +} + +// NewHashFromBigInt returns a *Hash representation of the given *big.Int +func NewHashFromBigInt(b *big.Int) *Hash { + r := &Hash{} + copy(r[:], ReverseByteOrder(b.Bytes())) + return r +} + +// NewHashFromBytes returns a *Hash from a byte array, swapping the endianness +// in the process. This is the intended method to get a *Hash from a byte array +// that previously has ben generated by the Hash.Bytes() method. +func NewHashFromBytes(b []byte) (*Hash, error) { + if len(b) != ElemBytesLen { + return nil, fmt.Errorf("expected 32 bytes, found %d bytes", len(b)) + } + var h Hash + copy(h[:], ReverseByteOrder(b)) + return &h, nil +} + +// NewHashFromHex returns a *Hash representation of the given hex string +func NewHashFromHex(h string) (*Hash, error) { + return NewHashFromBytes(ReverseByteOrder(common.FromHex(h))) +} + +// NewHashFromString returns a *Hash representation of the given decimal string +func NewHashFromString(s string) (*Hash, error) { + bi, ok := new(big.Int).SetString(s, 10) + if !ok { + return nil, fmt.Errorf("can not parse string to Hash") + } + return NewHashFromBigInt(bi), nil +} + +// ReverseByteOrder swaps the order of the bytes in the slice. +func ReverseByteOrder(b []byte) []byte { + o := make([]byte, len(b)) + for i := range b { + o[len(b)-1-i] = b[i] + } + return o +} diff --git a/core/types/zktrie/util.go b/core/types/zktrie/util.go new file mode 100644 index 0000000000..657b207bc8 --- /dev/null +++ b/core/types/zktrie/util.go @@ -0,0 +1,93 @@ +package zktrie + +import ( + "math/big" + + "github.com/iden3/go-iden3-crypto/poseidon" +) + +// HashElems performs a recursive poseidon hash over the array of ElemBytes, each hash +// reduce 2 fieds into one +func HashElems(fst, snd *big.Int, elems ...*big.Int) (*Hash, error) { + + l := len(elems) + baseH, err := poseidon.Hash([]*big.Int{fst, snd}) + if err != nil { + return nil, err + } + if l == 0 { + return NewHashFromBigInt(baseH), nil + } else if l == 1 { + return HashElems(baseH, elems[0]) + } + + tmp := make([]*big.Int, l/2) + for i := range tmp { + if (i+1)*2 > l { + tmp[i] = elems[i*2+1] + } else { + h, err := poseidon.Hash(elems[i*2 : (i+1)*2]) + if err != nil { + return nil, err + } + tmp[i] = h + } + } + + return HashElems(baseH, tmp[0], tmp[1:]...) +} + +// PreHandlingElems turn persisted byte32 elements into field arrays for our hashElem +// it also has the compressed byte32 +func PreHandlingElems(flagArray uint32, elems []Byte32) (*Hash, error) { + + ret := make([]*big.Int, len(elems)) + var err error + + for i, elem := range elems { + if flagArray&(1<. + +package trie + +import ( + "fmt" + "math/big" + + "github.com/scroll-tech/go-ethereum/common" + "github.com/scroll-tech/go-ethereum/core/types" + zkt "github.com/scroll-tech/go-ethereum/core/types/zktrie" + "github.com/scroll-tech/go-ethereum/ethdb" + "github.com/scroll-tech/go-ethereum/log" +) + +// ZkTrie wraps a trie with key hashing. In a secure trie, all +// access operations hash the key using keccak256. This prevents +// calling code from creating long chains of nodes that +// increase the access time. +// +// Contrary to a regular trie, a ZkTrie can only be created with +// New and must have an attached database. The database also stores +// the preimage of each key. +// +// ZkTrie is not safe for concurrent use. +type ZkTrie struct { + tree *ZkTrieImpl +} + +// NewSecure creates a trie +// SecureBinaryTrie bypasses all the buffer mechanism in *Database, it directly uses the +// underlying diskdb +func NewZkTrie(root common.Hash, db *ZktrieDatabase) (*ZkTrie, error) { + rootHash, err := zkt.NewHashFromBytes(root.Bytes()) + if err != nil { + return nil, err + } + tree, err := NewZkTrieImplWithRoot((db), rootHash, 256) + if err != nil { + return nil, err + } + return &ZkTrie{ + tree: tree, + }, nil +} + +// Get returns the value for key stored in the trie. +// The value bytes must not be modified by the caller. +func (t *ZkTrie) Get(key []byte) []byte { + res, err := t.TryGet(key) + if err != nil { + log.Error(fmt.Sprintf("Unhandled trie error: %v", err)) + } + return res +} + +// TryGet returns the value for key stored in the trie. +// The value bytes must not be modified by the caller. +// If a node was not found in the database, a MissingNodeError is returned. +func (t *ZkTrie) TryGet(key []byte) ([]byte, error) { + word := zkt.NewByte32FromBytesPaddingZero(key) + k, err := word.Hash() + if err != nil { + return nil, err + } + + return t.tree.TryGet(k.Bytes()) +} + +// TryGetNode attempts to retrieve a trie node by compact-encoded path. It is not +// possible to use keybyte-encoding as the path might contain odd nibbles. +func (t *ZkTrie) TryGetNode(path []byte) ([]byte, int, error) { + panic("unimplemented") +} + +func (t *ZkTrie) updatePreimage(preimage []byte, hashField *big.Int) { + db := t.tree.db.db + if db.preimages != nil { // Ugly direct check but avoids the below write lock + db.lock.Lock() + // we must copy the input key + db.insertPreimage(common.BytesToHash(hashField.Bytes()), common.CopyBytes(preimage)) + db.lock.Unlock() + } +} + +// TryUpdateAccount will abstract the write of an account to the +// secure trie. +func (t *ZkTrie) TryUpdateAccount(key []byte, acc *types.StateAccount) error { + keyPreimage := zkt.NewByte32FromBytesPaddingZero(key) + k, err := keyPreimage.Hash() + if err != nil { + return err + } + t.updatePreimage(key, k) + return t.tree.TryUpdateAccount(k.Bytes(), acc) +} + +// Update associates key with value in the trie. Subsequent calls to +// Get will return value. If value has length zero, any existing value +// is deleted from the trie and calls to Get will return nil. +// +// The value bytes must not be modified by the caller while they are +// stored in the trie. +func (t *ZkTrie) Update(key, value []byte) { + if err := t.TryUpdate(key, value); err != nil { + log.Error(fmt.Sprintf("Unhandled trie error: %v", err)) + } +} + +// TryUpdate associates key with value in the trie. Subsequent calls to +// Get will return value. If value has length zero, any existing value +// is deleted from the trie and calls to Get will return nil. +// +// The value bytes must not be modified by the caller while they are +// stored in the trie. +// +// If a node was not found in the database, a MissingNodeError is returned. +// +// NOTE: value is restricted to length of bytes32. +func (t *ZkTrie) TryUpdate(key, value []byte) error { + keyPreimage := zkt.NewByte32FromBytesPaddingZero(key) + k, err := keyPreimage.Hash() + if err != nil { + return err + } + t.updatePreimage(key, k) + return t.tree.TryUpdate(k.Bytes(), value) +} + +// Delete removes any existing value for key from the trie. +func (t *ZkTrie) Delete(key []byte) { + if err := t.TryDelete(key); err != nil { + log.Error(fmt.Sprintf("Unhandled trie error: %v", err)) + } +} + +// TryDelete removes any existing value for key from the trie. +// If a node was not found in the database, a MissingNodeError is returned. +func (t *ZkTrie) TryDelete(key []byte) error { + keyPreimage := zkt.NewByte32FromBytesPaddingZero(key) + k, err := keyPreimage.Hash() + if err != nil { + return err + } + + //mitigate the create-delete issue: do not delete unexisted key + if r := t.tree.Get(k.Bytes()); r == nil { + return nil + } + + zeroBt := common.Hash{} + // FIXME: delete should not be implemented as Update(0) + return t.tree.TryUpdate(k.Bytes(), zeroBt[:]) + //kPreimage := smt.NewByte32FromBytesPadding(key) + //return t.tree.DeleteWord(kPreimage) +} + +// GetKey returns the preimage of a hashed key that was +// previously used to store a value. +func (t *ZkTrie) GetKey(kHashBytes []byte) []byte { + // TODO: use a kv cache in memory + k, err := zkt.NewBigIntFromHashBytes(kHashBytes) + if err != nil { + log.Error(fmt.Sprintf("Unhandled trie error: %v", err)) + } + + return t.tree.db.db.preimage(common.BytesToHash(k.Bytes())) + +} + +// Commit writes all nodes and the secure hash pre-images to the trie's database. +// Nodes are stored with their sha3 hash as the key. +// +// Committing flushes nodes from memory. Subsequent Get calls will load nodes +// from the database. +func (t *ZkTrie) Commit(LeafCallback) (common.Hash, int, error) { + // in current implmentation, every update of trie already writes into database + // so Commmit does nothing + return t.Hash(), 0, nil +} + +// Hash returns the root hash of SecureBinaryTrie. It does not write to the +// database and can be used even if the trie doesn't have one. +func (t *ZkTrie) Hash() common.Hash { + var hash common.Hash + hash.SetBytes(t.tree.rootKey.Bytes()) + return hash +} + +// Copy returns a copy of SecureBinaryTrie. +func (t *ZkTrie) Copy() *ZkTrie { + cpy, err := NewZkTrieImplWithRoot(t.tree.db, t.tree.rootKey, t.tree.maxLevels) + if err != nil { + panic("clone trie failed") + } + return &ZkTrie{ + tree: cpy, + } +} + +// NodeIterator returns an iterator that returns nodes of the underlying trie. Iteration +// starts at the key after the given start key. +func (t *ZkTrie) NodeIterator(start []byte) NodeIterator { + /// FIXME + panic("not implemented") +} + +// hashKey returns the hash of key as an ephemeral buffer. +// The caller must not hold onto the return value because it will become +// invalid on the next call to hashKey or secKey. +/*func (t *ZkTrie) hashKey(key []byte) []byte { + if len(key) != 32 { + panic("non byte32 input to hashKey") + } + low16 := new(big.Int).SetBytes(key[:16]) + high16 := new(big.Int).SetBytes(key[16:]) + hash, err := poseidon.Hash([]*big.Int{low16, high16}) + if err != nil { + panic(err) + } + return hash.Bytes() +} +*/ + +// Prove constructs a merkle proof for key. The result contains all encoded nodes +// on the path to the value at key. The value itself is also included in the last +// node and can be retrieved by verifying the proof. +// +// If the trie does not contain a value for key, the returned proof contains all +// nodes of the longest existing prefix of the key (at least the root node), ending +// with the node that proves the absence of the key. +func (t *ZkTrie) Prove(key []byte, fromLevel uint, proofDb ethdb.KeyValueWriter) error { + word := zkt.NewByte32FromBytesPaddingZero(key) + k, err := word.Hash() + if err != nil { + return err + } + err = t.tree.prove(zkt.NewHashFromBigInt(k), fromLevel, func(n *Node) error { + key, err := n.Key() + if err != nil { + return err + } + + if n.Type == NodeTypeLeaf { + preImage := t.GetKey(n.NodeKey.Bytes()) + if len(preImage) > 0 { + n.KeyPreimage = &zkt.Byte32{} + copy(n.KeyPreimage[:], preImage) + //return fmt.Errorf("key preimage not found for [%x] ref %x", n.NodeKey.Bytes(), k.Bytes()) + } + } + return proofDb.Put(key.Bytes(), n.Value()) + }) + if err != nil { + return err + } + + // we put this special kv pair in db so we can distinguish the type and + // make suitable Proof + return proofDb.Put(magicHash, magicSMTBytes) +} diff --git a/trie/zk_trie_database.go b/trie/zk_trie_database.go new file mode 100644 index 0000000000..b8a4fe9292 --- /dev/null +++ b/trie/zk_trie_database.go @@ -0,0 +1,86 @@ +package trie + +import ( + "github.com/syndtr/goleveldb/leveldb" + + "github.com/scroll-tech/go-ethereum/ethdb" +) + +// TODO: we should refactor codes, so ZktrieDatabase and Database become two implementation of a +// interface later, making codes less surprising.. +// ZktrieDatabase Database adaptor +type ZktrieDatabase struct { + db *Database + prefix []byte +} + +func NewZktrieDatabase(diskdb ethdb.KeyValueStore) *ZktrieDatabase { + return &ZktrieDatabase{db: NewDatabase(diskdb), prefix: []byte{}} +} + +// adhoc wrapper... +func NewZktrieDatabaseFromTriedb(db *Database) *ZktrieDatabase { + db.Zktrie = true + return &ZktrieDatabase{db: db, prefix: []byte{}} +} + +// Put saves a key:value into the Storage +func (l *ZktrieDatabase) Put(k, v []byte) error { + l.db.lock.Lock() + l.db.rawDirties.Put(Concat(l.prefix, k[:]), v) + l.db.lock.Unlock() + return nil +} + +// Get retrieves a value from a key in the Storage +func (l *ZktrieDatabase) Get(key []byte) ([]byte, error) { + concatKey := Concat(l.prefix, key[:]) + l.db.lock.RLock() + value, ok := l.db.rawDirties.Get(concatKey) + l.db.lock.RUnlock() + if ok { + return value, nil + } + v, err := l.db.diskdb.Get(concatKey) + if err == leveldb.ErrNotFound { + return nil, ErrNotFound + } + return v, err +} + +// Iterate implements the method Iterate of the interface Storage +func (l *ZktrieDatabase) Iterate(f func([]byte, []byte) (bool, error)) error { + iter := l.db.diskdb.NewIterator(l.prefix, nil) + defer iter.Release() + for iter.Next() { + localKey := iter.Key()[len(l.prefix):] + if cont, err := f(localKey, iter.Value()); err != nil { + return err + } else if !cont { + break + } + } + iter.Release() + return iter.Error() +} + +// Close implements the method Close of the interface Storage +func (l *ZktrieDatabase) Close() { + // FIXME: is this correct? + if err := l.db.diskdb.Close(); err != nil { + panic(err) + } +} + +// List implements the method List of the interface Storage +func (l *ZktrieDatabase) List(limit int) ([]KV, error) { + ret := []KV{} + err := l.Iterate(func(key []byte, value []byte) (bool, error) { + ret = append(ret, KV{K: Clone(key), V: Clone(value)}) + if len(ret) == limit { + return false, nil + } + return true, nil + }) + return ret, err +} diff --git a/trie/zk_trie_impl.go b/trie/zk_trie_impl.go new file mode 100644 index 0000000000..2e8104a78e --- /dev/null +++ b/trie/zk_trie_impl.go @@ -0,0 +1,972 @@ +package trie + +import ( + "bytes" + "errors" + "fmt" + "io" + + cryptoUtils "github.com/iden3/go-iden3-crypto/utils" + + "github.com/scroll-tech/go-ethereum/common" + "github.com/scroll-tech/go-ethereum/core/types" + zkt "github.com/scroll-tech/go-ethereum/core/types/zktrie" + "github.com/scroll-tech/go-ethereum/log" +) + +const ( + // proofFlagsLen is the byte length of the flags in the proof header + // (first 32 bytes). + proofFlagsLen = 2 +) + +var ( + // ErrNodeKeyAlreadyExists is used when a node key already exists. + ErrNodeKeyAlreadyExists = errors.New("key already exists") + // ErrKeyNotFound is used when a key is not found in the ZkTrieImpl. + ErrKeyNotFound = errors.New("key not found in ZkTrieImpl") + // ErrNodeBytesBadSize is used when the data of a node has an incorrect + // size and can't be parsed. + ErrNodeBytesBadSize = errors.New("node data has incorrect size in the DB") + // ErrReachedMaxLevel is used when a traversal of the MT reaches the + // maximum level. + ErrReachedMaxLevel = errors.New("reached maximum level of the merkle tree") + // ErrInvalidNodeFound is used when an invalid node is found and can't + // be parsed. + ErrInvalidNodeFound = errors.New("found an invalid node in the DB") + // ErrInvalidProofBytes is used when a serialized proof is invalid. + ErrInvalidProofBytes = errors.New("the serialized proof is invalid") + // ErrEntryIndexAlreadyExists is used when the entry index already + // exists in the tree. + ErrEntryIndexAlreadyExists = errors.New("the entry index already exists in the tree") + // ErrNotWritable is used when the ZkTrieImpl is not writable and a + // write function is called + ErrNotWritable = errors.New("merkle Tree not writable") + + dbKeyRootNode = []byte("currentroot") +) + +// ZkTrieImpl is the struct with the main elements of the ZkTrieImpl +type ZkTrieImpl struct { + db *ZktrieDatabase + rootKey *zkt.Hash + writable bool + maxLevels int + Debug bool +} + +func NewZkTrieImpl(storage *ZktrieDatabase, maxLevels int) (*ZkTrieImpl, error) { + return NewZkTrieImplWithRoot(storage, &zkt.HashZero, maxLevels) +} + +// NewZkTrieImplWithRoot loads a new ZkTrieImpl. If in the storage already exists one +// will open that one, if not, will create a new one. +func NewZkTrieImplWithRoot(storage *ZktrieDatabase, root *zkt.Hash, maxLevels int) (*ZkTrieImpl, error) { + mt := ZkTrieImpl{db: storage, maxLevels: maxLevels, writable: true} + mt.rootKey = root + if *root != zkt.HashZero { + _, err := mt.GetNode(mt.rootKey) + if err != nil { + return nil, err + } + } + return &mt, nil +} + +// DB returns the ZkTrieImpl.DB() +func (mt *ZkTrieImpl) DB() *ZktrieDatabase { + return mt.db +} + +// Root returns the MerkleRoot +func (mt *ZkTrieImpl) Root() *zkt.Hash { + if mt.Debug { + _, err := mt.GetNode(mt.rootKey) + if err != nil { + var hash common.Hash + hash.SetBytes(mt.rootKey.Bytes()) + panic(fmt.Errorf("load trie root failed hash %v", hash)) + } + } + return mt.rootKey +} + +// MaxLevels returns the MT maximum level +func (mt *ZkTrieImpl) MaxLevels() int { + return mt.maxLevels +} + +// tryUpdate update a Key & Value into the ZkTrieImpl. Where the `k` determines the +// path from the Root to the Leaf. This also return the updated leaf node +func (mt *ZkTrieImpl) tryUpdate(kHash *zkt.Hash, vFlag uint32, vPreimage []zkt.Byte32) error { + // verify that the ZkTrieImpl is writable + if !mt.writable { + return ErrNotWritable + } + + // verify that k are valid and fit inside the Finite Field. + if !cryptoUtils.CheckBigIntInField(kHash.BigInt()) { + return errors.New("Key not inside the Finite Field") + } + + newNodeLeaf := NewNodeLeaf(kHash, vFlag, vPreimage) + path := getPath(mt.maxLevels, kHash[:]) + + // precalc Key of new leaf here + if _, err := newNodeLeaf.Key(); err != nil { + return err + } + + newRootKey, err := mt.addLeaf(newNodeLeaf, mt.rootKey, 0, path, true) + // sanity check + if err == ErrEntryIndexAlreadyExists { + panic("Encounter unexpected errortype: ErrEntryIndexAlreadyExists") + } else if err != nil { + return err + } + mt.rootKey = newRootKey + err = mt.dbInsert(dbKeyRootNode, DBEntryTypeRoot, mt.rootKey[:]) + if err != nil { + return err + } + + return nil +} + +// AddWord +// Deprecated: Add a Bytes32 kv to ZkTrieImpl, only for testing +func (mt *ZkTrieImpl) AddWord(kPreimage, vPreimage *zkt.Byte32) error { + k, err := kPreimage.Hash() + if err != nil { + return err + } + + kHash := zkt.NewHashFromBigInt(k) + newNodeLeaf := NewNodeLeaf(kHash, 1, []zkt.Byte32{*vPreimage}) + path := getPath(mt.maxLevels, kHash[:]) + + // precalc Key of new leaf here + if _, err := newNodeLeaf.Key(); err != nil { + return err + } + + newRootKey, err := mt.addLeaf(newNodeLeaf, mt.rootKey, 0, path, false) + if err != nil { + return err + } + mt.rootKey = newRootKey + err = mt.dbInsert(dbKeyRootNode, DBEntryTypeRoot, mt.rootKey[:]) + if err != nil { + return err + } + + return nil +} + +// pushLeaf recursively pushes an existing oldLeaf down until its path diverges +// from newLeaf, at which point both leafs are stored, all while updating the +// path. +func (mt *ZkTrieImpl) pushLeaf(newLeaf *Node, oldLeaf *Node, lvl int, + pathNewLeaf []bool, pathOldLeaf []bool) (*zkt.Hash, error) { + if lvl > mt.maxLevels-2 { + return nil, ErrReachedMaxLevel + } + var newNodeMiddle *Node + if pathNewLeaf[lvl] == pathOldLeaf[lvl] { // We need to go deeper! + nextKey, err := mt.pushLeaf(newLeaf, oldLeaf, lvl+1, pathNewLeaf, pathOldLeaf) + if err != nil { + return nil, err + } + if pathNewLeaf[lvl] { // go right + newNodeMiddle = NewNodeMiddle(&zkt.HashZero, nextKey) + } else { // go left + newNodeMiddle = NewNodeMiddle(nextKey, &zkt.HashZero) + } + return mt.addNode(newNodeMiddle) + } + oldLeafKey, err := oldLeaf.Key() + if err != nil { + return nil, err + } + newLeafKey, err := newLeaf.Key() + if err != nil { + return nil, err + } + + if pathNewLeaf[lvl] { + newNodeMiddle = NewNodeMiddle(oldLeafKey, newLeafKey) + } else { + newNodeMiddle = NewNodeMiddle(newLeafKey, oldLeafKey) + } + // We can add newLeaf now. We don't need to add oldLeaf because it's + // already in the tree. + _, err = mt.addNode(newLeaf) + if err != nil { + return nil, err + } + return mt.addNode(newNodeMiddle) +} + +// addLeaf recursively adds a newLeaf in the MT while updating the path. +func (mt *ZkTrieImpl) addLeaf(newLeaf *Node, key *zkt.Hash, + lvl int, path []bool, forceUpdate bool) (*zkt.Hash, error) { + var err error + var nextKey *zkt.Hash + if lvl > mt.maxLevels-1 { + return nil, ErrReachedMaxLevel + } + n, err := mt.GetNode(key) + if err != nil { + fmt.Printf("addLeaf:GetNode err %v key %v root %v level %v\n", err, key, mt.rootKey, lvl) + fmt.Printf("root %v\n", mt.Root()) + return nil, err + } + switch n.Type { + case NodeTypeEmpty: + // We can add newLeaf now + { + r, e := mt.addNode(newLeaf) + if e != nil { + fmt.Println("err on NodeTypeEmpty mt.addNode ", e) + } + return r, e + } + case NodeTypeLeaf: + // Check if leaf node found contains the leaf node we are + // trying to add + if bytes.Equal(n.NodeKey[:], newLeaf.NodeKey[:]) { + k, err := n.Key() + if err != nil { + fmt.Println("err on obtain key of duplicated entry", err) + return nil, err + } + if bytes.Equal(k[:], newLeaf.key[:]) { + // do nothing, duplicate entry + // FIXME more optimization may needed here + return k, nil + } else if forceUpdate { + return mt.updateNode(newLeaf) + } + + fmt.Printf("ErrEntryIndexAlreadyExists nodeKey %v n.Key() %v newLeaf.Key() %v\n", + n.NodeKey, k, newLeaf.key) + return nil, ErrEntryIndexAlreadyExists + + } + pathOldLeaf := getPath(mt.maxLevels, n.NodeKey[:]) + // We need to push newLeaf down until its path diverges from + // n's path + return mt.pushLeaf(newLeaf, n, lvl, path, pathOldLeaf) + case NodeTypeMiddle: + // We need to go deeper, continue traversing the tree, left or + // right depending on path + var newNodeMiddle *Node + if path[lvl] { // go right + nextKey, err = mt.addLeaf(newLeaf, n.ChildR, lvl+1, path, forceUpdate) + newNodeMiddle = NewNodeMiddle(n.ChildL, nextKey) + } else { // go left + nextKey, err = mt.addLeaf(newLeaf, n.ChildL, lvl+1, path, forceUpdate) + newNodeMiddle = NewNodeMiddle(nextKey, n.ChildR) + } + if err != nil { + fmt.Printf("addLeaf:GetNode err %v level %v\n", err, lvl) + return nil, err + } + // Update the node to reflect the modified child + return mt.addNode(newNodeMiddle) + default: + return nil, ErrInvalidNodeFound + } +} + +// addNode adds a node into the MT. Empty nodes are not stored in the tree; +// they are all the same and assumed to always exist. +func (mt *ZkTrieImpl) addNode(n *Node) (*zkt.Hash, error) { + // verify that the ZkTrieImpl is writable + if !mt.writable { + return nil, ErrNotWritable + } + if n.Type == NodeTypeEmpty { + return n.Key() + } + k, err := n.Key() + if err != nil { + return nil, err + } + v := n.Value() + // Check that the node key doesn't already exist + oldV, err := mt.db.Get(k[:]) + if err == nil { + if !bytes.Equal(oldV, v) { + return nil, ErrNodeKeyAlreadyExists + } else { + // duplicated + return k, nil + } + } + err = mt.db.Put(k[:], v) + return k, err +} + +// updateNode updates an existing node in the MT. Empty nodes are not stored +// in the tree; they are all the same and assumed to always exist. +func (mt *ZkTrieImpl) updateNode(n *Node) (*zkt.Hash, error) { + // verify that the ZkTrieImpl is writable + if !mt.writable { + return nil, ErrNotWritable + } + if n.Type == NodeTypeEmpty { + return n.Key() + } + k, err := n.Key() + if err != nil { + return nil, err + } + v := n.Value() + err = mt.db.Put(k[:], v) + return k, err +} + +/* +// Get returns the value of the leaf for the given key +func (mt *ZkTrieImpl) Get(k *big.Int) (*big.Int, *big.Int, []*zkt.Hash, error) { + // verify that k is valid and fit inside the Finite Field. + if !cryptoUtils.CheckBigIntInField(k) { + return nil, nil, nil, errors.New("Key not inside the Finite Field") + } + + kHash := zkt.NewHashFromBigInt(k) + path := getPath(mt.maxLevels, kHash[:]) + + nextKey := mt.rootKey + siblings := []*zkt.Hash{} + for i := 0; i < mt.maxLevels; i++ { + n, err := mt.GetNode(nextKey) + if err != nil { + return nil, nil, nil, err + } + switch n.Type { + case NodeTypeEmpty: + return big.NewInt(0), big.NewInt(0), siblings, ErrKeyNotFound + case NodeTypeLeaf: + if bytes.Equal(kHash[:], n.NodeKey[:]) { + return n.Entry[0].BigInt(), n.Entry[1].BigInt(), siblings, nil + } + return n.Entry[0].BigInt(), n.Entry[1].BigInt(), siblings, ErrKeyNotFound + case NodeTypeMiddle: + if path[i] { + nextKey = n.ChildR + siblings = append(siblings, n.ChildL) + } else { + nextKey = n.ChildL + siblings = append(siblings, n.ChildR) + } + default: + return nil, nil, nil, ErrInvalidNodeFound + } + } + + return nil, nil, nil, ErrReachedMaxLevel +} +*/ + +func (mt *ZkTrieImpl) tryGet(kHash *zkt.Hash) (*Node, []*zkt.Hash, error) { + + path := getPath(mt.maxLevels, kHash[:]) + nextKey := mt.rootKey + var siblings []*zkt.Hash + for i := 0; i < mt.maxLevels; i++ { + n, err := mt.GetNode(nextKey) + if err != nil { + return nil, nil, err + } + switch n.Type { + case NodeTypeEmpty: + return NewNodeEmpty(), siblings, ErrKeyNotFound + case NodeTypeLeaf: + if bytes.Equal(kHash[:], n.NodeKey[:]) { + return n, siblings, nil + } + return n, siblings, ErrKeyNotFound + case NodeTypeMiddle: + if path[i] { + nextKey = n.ChildR + siblings = append(siblings, n.ChildL) + } else { + nextKey = n.ChildL + siblings = append(siblings, n.ChildR) + } + default: + return nil, nil, ErrInvalidNodeFound + } + } + + return nil, siblings, ErrReachedMaxLevel +} + +// Get returns the value for key stored in the trie. +// The value bytes must not be modified by the caller. +func (mt *ZkTrieImpl) Get(key []byte) []byte { + res, err := mt.TryGet(key) + if err != nil { + log.Error(fmt.Sprintf("Unhandled trie error: %v", err)) + } + return res +} + +// TryGet returns the value for key stored in the trie. +// The value bytes must not be modified by the caller. +// If a node was not found in the database, a MissingNodeError is returned. +func (mt *ZkTrieImpl) TryGet(key []byte) ([]byte, error) { + kHash, err := zkt.NewHashFromBytes(common.BytesToHash(key).Bytes()) + if err != nil { + return nil, err + } + + node, _, err := mt.tryGet(kHash) + if err == ErrKeyNotFound { + // according to https://github.com/ethereum/go-ethereum/blob/37f9d25ba027356457953eab5f181c98b46e9988/trie/trie.go#L135 + return nil, nil + } else if err != nil { + return nil, err + } + return node.Data(), nil +} + +// GetLeafNodeByWord +// Deprecated: Get a Bytes32 kv to ZkTrieImpl, only for testing +func (mt *ZkTrieImpl) GetLeafNodeByWord(kPreimage *zkt.Byte32) (*Node, error) { + k, err := kPreimage.Hash() + if err != nil { + return nil, err + } + n, _, err := mt.tryGet(zkt.NewHashFromBigInt(k)) + return n, err +} + +/* +// Update function updates the value of a specified key in the ZkTrieImpl, and updates +// the path from the leaf to the Root with the new values,and returns the +// CircomProcessorProof. +func (mt *ZkTrieImpl) Update(k, v *big.Int, kPreimage *zkt.Byte32, vPreimage []byte) error { + // verify that the ZkTrieImpl is writable + if !mt.writable { + return ErrNotWritable + } + + // verify that k & are valid and fit inside the Finite Field. + if !cryptoUtils.CheckBigIntInField(k) { + return errors.New("Key not inside the Finite Field") + } + if !cryptoUtils.CheckBigIntInField(v) { + return errors.New("Key not inside the Finite Field") + } + + kHash := zkt.NewHashFromBigInt(k) + vHash := zkt.NewHashFromBigInt(v) + path := getPath(mt.maxLevels, kHash[:]) + + nextKey := mt.rootKey + siblings := []*zkt.Hash{} + for i := 0; i < mt.maxLevels; i++ { + n, err := mt.GetNode(nextKey) + if err != nil { + return err + } + switch n.Type { + case NodeTypeEmpty: + return ErrKeyNotFound + case NodeTypeLeaf: + if bytes.Equal(kHash[:], n.NodeKey[:]) { + // update leaf and upload to the root + newNodeLeaf := NewNodeLeaf(kHash, vHash, kPreimage, vPreimage) + _, err := mt.updateNode(newNodeLeaf) + if err != nil { + return err + } + newRootKey, err := + mt.recalculatePathUntilRoot(path, newNodeLeaf, siblings) + if err != nil { + return err + } + mt.rootKey = newRootKey + err = mt.dbInsert(dbKeyRootNode, DBEntryTypeRoot, mt.rootKey[:]) + if err != nil { + return err + } + return nil + } + return ErrKeyNotFound + case NodeTypeMiddle: + if path[i] { + nextKey = n.ChildR + siblings = append(siblings, n.ChildL) + } else { + nextKey = n.ChildL + siblings = append(siblings, n.ChildR) + } + default: + return ErrInvalidNodeFound + } + } + + return ErrKeyNotFound +} +*/ + +// Deprecated: only for testing +func (mt *ZkTrieImpl) UpdateWord(kPreimage, vPreimage *zkt.Byte32) error { + k, err := kPreimage.Hash() + if err != nil { + return err + } + + err = mt.tryUpdate(zkt.NewHashFromBigInt(k), 1, []zkt.Byte32{*vPreimage}) + return err +} + +// TryUpdateAccount will abstract the write of an account to the trie +func (mt *ZkTrieImpl) TryUpdateAccount(key []byte, acc *types.StateAccount) error { + kHash, err := zkt.NewHashFromBytes(common.BytesToHash(key).Bytes()) + if err != nil { + return err + } + + value, flag := acc.MarshalFields() + return mt.tryUpdate(kHash, flag, value) +} + +// Update associates key with value in the trie. Subsequent calls to +// Get will return value. +// TODO: If value has length zero, any existing value +// is deleted from the trie and calls to Get will return nil. +// +// The value bytes must not be modified by the caller while they are +// stored in the trie. +func (mt *ZkTrieImpl) Update(key, value []byte) { + if err := mt.TryUpdate(key, value); err != nil { + log.Error(fmt.Sprintf("Unhandled trie error: %v", err)) + } +} + +// TryUpdate associates key with value in the trie. Subsequent calls to +// Get will return value. +// TODO: If value has length zero, any existing value +// is deleted from the trie and calls to Get will return nil. +// +// The value bytes must not be modified by the caller while they are +// stored in the trie. +// +// Currently the value must shorter than 32 bytes +// +// NOTE: value is restricted to length of bytes32. +func (mt *ZkTrieImpl) TryUpdate(key, value []byte) error { + kHash, err := zkt.NewHashFromBytes(common.BytesToHash(key).Bytes()) + if err != nil { + return err + } + vPreimage := zkt.NewByte32FromBytesPaddingZero(value) + return mt.tryUpdate(kHash, 1, []zkt.Byte32{*vPreimage}) +} + +/* +func (mt *ZkTrieImpl) UpdateVarWord(kPreimage *zkt.Byte32, vHash *big.Int, vPreimage []byte) error { + k, err := kPreimage.Hash() + if err != nil { + return err + } + err = mt.Update(k, vHash, kPreimage, vPreimage[:]) + if err == ErrKeyNotFound { + err = mt.Add(k, vHash, kPreimage, vPreimage[:]) + if err != nil { + log.Error("UpdateVarWord, inset still failed %v root %v", err, mt.rootKey) + } + return err + } else if err != nil { + log.Error("UpdateVarWord err %v %v", err, reflect.TypeOf(err)) + } + return err +} +*/ + +// Delete removes the specified Key from the ZkTrieImpl and updates the path +// from the deleted key to the Root with the new values. This method removes +// the key from the ZkTrieImpl, but does not remove the old nodes from the +// key-value database; this means that if the tree is accessed by an old Root +// where the key was not deleted yet, the key will still exist. If is desired +// to remove the key-values from the database that are not under the current +// Root, an option could be to dump all the leafs (using mt.DumpLeafs) and +// import them in a new ZkTrieImpl in a new database (using +// mt.ImportDumpedLeafs), but this will lose all the Root history of the +// ZkTrieImpl +func (mt *ZkTrieImpl) tryDelete(kHash *zkt.Hash) error { + // verify that the ZkTrieImpl is writable + if !mt.writable { + return ErrNotWritable + } + + // verify that k is valid and fit inside the Finite Field. + if !cryptoUtils.CheckBigIntInField(kHash.BigInt()) { + return errors.New("Key not inside the Finite Field") + } + + path := getPath(mt.maxLevels, kHash[:]) + + nextKey := mt.rootKey + siblings := []*zkt.Hash{} + for i := 0; i < mt.maxLevels; i++ { + n, err := mt.GetNode(nextKey) + if err != nil { + return err + } + switch n.Type { + case NodeTypeEmpty: + return ErrKeyNotFound + case NodeTypeLeaf: + if bytes.Equal(kHash[:], n.NodeKey[:]) { + // remove and go up with the sibling + err = mt.rmAndUpload(path, kHash, siblings) + return err + } + return ErrKeyNotFound + case NodeTypeMiddle: + if path[i] { + nextKey = n.ChildR + siblings = append(siblings, n.ChildL) + } else { + nextKey = n.ChildL + siblings = append(siblings, n.ChildR) + } + default: + return ErrInvalidNodeFound + } + } + + return ErrKeyNotFound +} + +// TryDelete removes any existing value for key from the trie. +// If a node was not found in the database, a MissingNodeError is returned. +func (mt *ZkTrieImpl) TryDelete(key []byte) error { + kHash, err := zkt.NewHashFromBytes(common.BytesToHash(key).Bytes()) + if err != nil { + return err + } + + return mt.tryDelete(kHash) +} + +// Deprecated: only for testing +func (mt *ZkTrieImpl) DeleteWord(kPreimage *zkt.Byte32) error { + k, err := kPreimage.Hash() + if err != nil { + return err + } + return mt.tryDelete(zkt.NewHashFromBigInt(k)) +} + +// rmAndUpload removes the key, and goes up until the root updating all the +// nodes with the new values. +func (mt *ZkTrieImpl) rmAndUpload(path []bool, kHash *zkt.Hash, siblings []*zkt.Hash) error { + if len(siblings) == 0 { + mt.rootKey = &zkt.HashZero + err := mt.dbInsert(dbKeyRootNode, DBEntryTypeRoot, mt.rootKey[:]) + if err != nil { + return err + } + return nil + } + + toUpload := siblings[len(siblings)-1] + if len(siblings) < 2 { //nolint:gomnd + mt.rootKey = siblings[0] + err := mt.dbInsert(dbKeyRootNode, DBEntryTypeRoot, mt.rootKey[:]) + if err != nil { + return err + } + return nil + } + for i := len(siblings) - 2; i >= 0; i-- { //nolint:gomnd + if !bytes.Equal(siblings[i][:], zkt.HashZero[:]) { + var newNode *Node + if path[i] { + newNode = NewNodeMiddle(siblings[i], toUpload) + } else { + newNode = NewNodeMiddle(toUpload, siblings[i]) + } + _, err := mt.addNode(newNode) + if err != ErrNodeKeyAlreadyExists && err != nil { + return err + } + // go up until the root + newRootKey, err := mt.recalculatePathUntilRoot(path, newNode, + siblings[:i]) + if err != nil { + return err + } + mt.rootKey = newRootKey + err = mt.dbInsert(dbKeyRootNode, DBEntryTypeRoot, mt.rootKey[:]) + if err != nil { + return err + } + break + } + // if i==0 (root position), stop and store the sibling of the + // deleted leaf as root + if i == 0 { + mt.rootKey = toUpload + err := mt.dbInsert(dbKeyRootNode, DBEntryTypeRoot, mt.rootKey[:]) + if err != nil { + return err + } + break + } + } + return nil +} + +// recalculatePathUntilRoot recalculates the nodes until the Root +func (mt *ZkTrieImpl) recalculatePathUntilRoot(path []bool, node *Node, + siblings []*zkt.Hash) (*zkt.Hash, error) { + for i := len(siblings) - 1; i >= 0; i-- { + nodeKey, err := node.Key() + if err != nil { + return nil, err + } + if path[i] { + node = NewNodeMiddle(siblings[i], nodeKey) + } else { + node = NewNodeMiddle(nodeKey, siblings[i]) + } + _, err = mt.addNode(node) + if err != ErrNodeKeyAlreadyExists && err != nil { + return nil, err + } + } + + // return last node added, which is the root + nodeKey, err := node.Key() + return nodeKey, err +} + +// dbInsert is a helper function to insert a node into a key in an open db +// transaction. +func (mt *ZkTrieImpl) dbInsert(k []byte, t NodeType, data []byte) error { + v := append([]byte{byte(t)}, data...) + return mt.db.Put(k, v) +} + +// GetNode gets a node by key from the MT. Empty nodes are not stored in the +// tree; they are all the same and assumed to always exist. +// for non exist key, return (NewNodeEmpty(), nil) +func (mt *ZkTrieImpl) GetNode(key *zkt.Hash) (*Node, error) { + if bytes.Equal(key[:], zkt.HashZero[:]) { + return NewNodeEmpty(), nil + } + nBytes, err := mt.db.Get(key[:]) + if err == ErrNotFound { + //return NewNodeEmpty(), nil + return nil, ErrKeyNotFound + } else if err != nil { + return nil, err + } + return NewNodeFromBytes(nBytes) +} + +// getPath returns the binary path, from the root to the leaf. +func getPath(numLevels int, k []byte) []bool { + path := make([]bool, numLevels) + for n := 0; n < numLevels; n++ { + path[n] = zkt.TestBit(k[:], uint(n)) + } + return path +} + +// NodeAux contains the auxiliary node used in a non-existence proof. +type NodeAux struct { + Key *zkt.Hash + Value *zkt.Hash +} + +// Proof defines the required elements for a MT proof of existence or +// non-existence. +type Proof struct { + // existence indicates wether this is a proof of existence or + // non-existence. + Existence bool + // depth indicates how deep in the tree the proof goes. + depth uint + // notempties is a bitmap of non-empty Siblings found in Siblings. + notempties [zkt.ElemBytesLen - proofFlagsLen]byte + // Siblings is a list of non-empty sibling keys. + Siblings []*zkt.Hash + // Key is the key of leaf in existence case + Key *zkt.Hash + NodeAux *NodeAux +} + +// VerifyProof verifies the Merkle Proof for the entry and root. +func VerifyProofZkTrie(rootKey *zkt.Hash, proof *Proof, node *Node) bool { + key, err := node.Key() + if err != nil { + return false + } + + rootFromProof, err := proof.Verify(key, node.NodeKey) + if err != nil { + return false + } + return bytes.Equal(rootKey[:], rootFromProof[:]) +} + +// Verify the proof and calculate the root, key can be nil when try to verify +// an nonexistent proof +func (proof *Proof) Verify(key, kHash *zkt.Hash) (*zkt.Hash, error) { + + if proof.Existence { + if key == nil { + return nil, ErrKeyNotFound + } + return proof.rootFromProof(key, kHash) + } else { + + if proof.NodeAux == nil { + return proof.rootFromProof(&zkt.HashZero, kHash) + } else { + if bytes.Equal(kHash[:], proof.NodeAux.Key[:]) { + return nil, fmt.Errorf("non-existence proof being checked against hIndex equal to nodeAux") + } + midKey, err := LeafKey(proof.NodeAux.Key, proof.NodeAux.Value) + if err != nil { + return nil, err + } + return proof.rootFromProof(midKey, kHash) + } + } + +} + +func (proof *Proof) rootFromProof(key, kHash *zkt.Hash) (*zkt.Hash, error) { + midKey := key + var err error + + sibIdx := len(proof.Siblings) - 1 + path := getPath(int(proof.depth), kHash[:]) + var siblingKey *zkt.Hash + for lvl := int(proof.depth) - 1; lvl >= 0; lvl-- { + if zkt.TestBitBigEndian(proof.notempties[:], uint(lvl)) { + siblingKey = proof.Siblings[sibIdx] + sibIdx-- + } else { + siblingKey = &zkt.HashZero + } + if path[lvl] { + midKey, err = NewNodeMiddle(siblingKey, midKey).Key() + if err != nil { + return nil, err + } + } else { + midKey, err = NewNodeMiddle(midKey, siblingKey).Key() + if err != nil { + return nil, err + } + } + } + return midKey, nil +} + +// walk is a helper recursive function to iterate over all tree branches +func (mt *ZkTrieImpl) walk(key *zkt.Hash, f func(*Node)) error { + n, err := mt.GetNode(key) + if err != nil { + return err + } + switch n.Type { + case NodeTypeEmpty: + f(n) + case NodeTypeLeaf: + f(n) + case NodeTypeMiddle: + f(n) + if err := mt.walk(n.ChildL, f); err != nil { + return err + } + if err := mt.walk(n.ChildR, f); err != nil { + return err + } + default: + return ErrInvalidNodeFound + } + return nil +} + +// Walk iterates over all the branches of a ZkTrieImpl with the given rootKey +// if rootKey is nil, it will get the current RootKey of the current state of +// the ZkTrieImpl. For each node, it calls the f function given in the +// parameters. See some examples of the Walk function usage in the +// ZkTrieImpl.go and merkletree_test.go +func (mt *ZkTrieImpl) Walk(rootKey *zkt.Hash, f func(*Node)) error { + if rootKey == nil { + rootKey = mt.Root() + } + err := mt.walk(rootKey, f) + return err +} + +// GraphViz uses Walk function to generate a string GraphViz representation of +// the tree and writes it to w +func (mt *ZkTrieImpl) GraphViz(w io.Writer, rootKey *zkt.Hash) error { + fmt.Fprintf(w, `digraph hierarchy { +node [fontname=Monospace,fontsize=10,shape=box] +`) + cnt := 0 + var errIn error + err := mt.Walk(rootKey, func(n *Node) { + k, err := n.Key() + if err != nil { + errIn = err + } + switch n.Type { + case NodeTypeEmpty: + case NodeTypeLeaf: + fmt.Fprintf(w, "\"%v\" [style=filled];\n", k.String()) + case NodeTypeMiddle: + lr := [2]string{n.ChildL.String(), n.ChildR.String()} + emptyNodes := "" + for i := range lr { + if lr[i] == "0" { + lr[i] = fmt.Sprintf("empty%v", cnt) + emptyNodes += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n", lr[i]) + cnt++ + } + } + fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", k.String(), lr[0], lr[1]) + fmt.Fprint(w, emptyNodes) + default: + } + }) + fmt.Fprintf(w, "}\n") + if errIn != nil { + return errIn + } + return err +} + +// PrintGraphViz prints directly the GraphViz() output +func (mt *ZkTrieImpl) PrintGraphViz(rootKey *zkt.Hash) error { + if rootKey == nil { + rootKey = mt.Root() + } + w := bytes.NewBufferString("") + fmt.Fprintf(w, + "--------\nGraphViz of the ZkTrieImpl with RootKey "+rootKey.BigInt().String()+"\n") + err := mt.GraphViz(w, nil) + if err != nil { + return err + } + fmt.Fprintf(w, + "End of GraphViz of the ZkTrieImpl with RootKey "+rootKey.BigInt().String()+"\n--------\n") + + fmt.Println(w) + return nil +} diff --git a/trie/zk_trie_impl_test.go b/trie/zk_trie_impl_test.go new file mode 100644 index 0000000000..789c22a7cf --- /dev/null +++ b/trie/zk_trie_impl_test.go @@ -0,0 +1,184 @@ +package trie + +import ( + "math/big" + "testing" + + "github.com/iden3/go-iden3-crypto/constants" + cryptoUtils "github.com/iden3/go-iden3-crypto/utils" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/scroll-tech/go-ethereum/common" + "github.com/scroll-tech/go-ethereum/core/types" + zkt "github.com/scroll-tech/go-ethereum/core/types/zktrie" + "github.com/scroll-tech/go-ethereum/ethdb/memorydb" +) + +type Fatalable interface { + Fatal(args ...interface{}) +} + +func newTestingMerkle(f Fatalable, numLevels int) *ZkTrieImpl { + mt, err := NewZkTrieImpl(NewZktrieDatabase((memorydb.New())), numLevels) + if err != nil { + f.Fatal(err) + return nil + } + return mt +} + +func TestHashParsers(t *testing.T) { + h0 := zkt.NewHashFromBigInt(big.NewInt(0)) + assert.Equal(t, "0", h0.String()) + h1 := zkt.NewHashFromBigInt(big.NewInt(1)) + assert.Equal(t, "1", h1.String()) + h10 := zkt.NewHashFromBigInt(big.NewInt(10)) + assert.Equal(t, "10", h10.String()) + + h7l := zkt.NewHashFromBigInt(big.NewInt(1234567)) + assert.Equal(t, "1234567", h7l.String()) + h8l := zkt.NewHashFromBigInt(big.NewInt(12345678)) + assert.Equal(t, "12345678...", h8l.String()) + + b, ok := new(big.Int).SetString("4932297968297298434239270129193057052722409868268166443802652458940273154854", 10) //nolint:lll + assert.True(t, ok) + h := zkt.NewHashFromBigInt(b) + assert.Equal(t, "4932297968297298434239270129193057052722409868268166443802652458940273154854", h.BigInt().String()) //nolint:lll + assert.Equal(t, "49322979...", h.String()) + assert.Equal(t, "265baaf161e875c372d08e50f52abddc01d32efc93e90290bb8b3d9ceb94e70a", h.Hex()) + + b1, err := zkt.NewBigIntFromHashBytes(b.Bytes()) + assert.Nil(t, err) + assert.Equal(t, new(big.Int).SetBytes(b.Bytes()).String(), b1.String()) + + b2, err := zkt.NewHashFromBytes(b.Bytes()) + assert.Nil(t, err) + assert.Equal(t, b.String(), b2.BigInt().String()) + + h2, err := zkt.NewHashFromHex(h.Hex()) + assert.Nil(t, err) + assert.Equal(t, h, h2) + _, err = zkt.NewHashFromHex("0x12") + assert.NotNil(t, err) + + // check limits + a := new(big.Int).Sub(constants.Q, big.NewInt(1)) + testHashParsers(t, a) + a = big.NewInt(int64(1)) + testHashParsers(t, a) +} + +func testHashParsers(t *testing.T, a *big.Int) { + require.True(t, cryptoUtils.CheckBigIntInField(a)) + h := zkt.NewHashFromBigInt(a) + assert.Equal(t, a, h.BigInt()) + hFromBytes, err := zkt.NewHashFromBytes(h.Bytes()) + assert.Nil(t, err) + assert.Equal(t, h, hFromBytes) + assert.Equal(t, a, hFromBytes.BigInt()) + assert.Equal(t, a.String(), hFromBytes.BigInt().String()) + hFromHex, err := zkt.NewHashFromHex(h.Hex()) + assert.Nil(t, err) + assert.Equal(t, h, hFromHex) + + aBIFromHBytes, err := zkt.NewBigIntFromHashBytes(h.Bytes()) + assert.Nil(t, err) + assert.Equal(t, a, aBIFromHBytes) + assert.Equal(t, new(big.Int).SetBytes(a.Bytes()).String(), aBIFromHBytes.String()) +} + +func TestMerkleTree_AddUpdateGetWord(t *testing.T) { + mt := newTestingMerkle(t, 10) + err := mt.AddWord(&zkt.Byte32{1}, &zkt.Byte32{2}) + assert.Nil(t, err) + err = mt.AddWord(&zkt.Byte32{3}, &zkt.Byte32{4}) + assert.Nil(t, err) + err = mt.AddWord(&zkt.Byte32{5}, &zkt.Byte32{6}) + assert.Nil(t, err) + err = mt.AddWord(&zkt.Byte32{5}, &zkt.Byte32{7}) + assert.Equal(t, ErrEntryIndexAlreadyExists, err) + + node, err := mt.GetLeafNodeByWord(&zkt.Byte32{1}) + assert.Nil(t, err) + assert.Equal(t, len(node.ValuePreimage), 1) + assert.Equal(t, (&zkt.Byte32{2})[:], node.ValuePreimage[0][:]) + node, err = mt.GetLeafNodeByWord(&zkt.Byte32{3}) + assert.Nil(t, err) + assert.Equal(t, len(node.ValuePreimage), 1) + assert.Equal(t, (&zkt.Byte32{4})[:], node.ValuePreimage[0][:]) + node, err = mt.GetLeafNodeByWord(&zkt.Byte32{5}) + assert.Nil(t, err) + assert.Equal(t, len(node.ValuePreimage), 1) + assert.Equal(t, (&zkt.Byte32{6})[:], node.ValuePreimage[0][:]) + + err = mt.UpdateWord(&zkt.Byte32{1}, &zkt.Byte32{7}) + assert.Nil(t, err) + err = mt.UpdateWord(&zkt.Byte32{3}, &zkt.Byte32{8}) + assert.Nil(t, err) + err = mt.UpdateWord(&zkt.Byte32{5}, &zkt.Byte32{9}) + assert.Nil(t, err) + + node, err = mt.GetLeafNodeByWord(&zkt.Byte32{1}) + assert.Nil(t, err) + assert.Equal(t, len(node.ValuePreimage), 1) + assert.Equal(t, (&zkt.Byte32{7})[:], node.ValuePreimage[0][:]) + node, err = mt.GetLeafNodeByWord(&zkt.Byte32{3}) + assert.Nil(t, err) + assert.Equal(t, len(node.ValuePreimage), 1) + assert.Equal(t, (&zkt.Byte32{8})[:], node.ValuePreimage[0][:]) + node, err = mt.GetLeafNodeByWord(&zkt.Byte32{5}) + assert.Nil(t, err) + assert.Equal(t, len(node.ValuePreimage), 1) + assert.Equal(t, (&zkt.Byte32{9})[:], node.ValuePreimage[0][:]) + _, err = mt.GetLeafNodeByWord(&zkt.Byte32{100}) + assert.Equal(t, ErrKeyNotFound, err) +} + +func TestMerkleTree_UpdateAccount(t *testing.T) { + + mt := newTestingMerkle(t, 10) + + acc1 := &types.StateAccount{ + Nonce: 1, + Balance: big.NewInt(10000000), + Root: common.HexToHash("22fb59aa5410ed465267023713ab42554c250f394901455a3366e223d5f7d147"), + CodeHash: common.HexToHash("cc0a77f6e063b4b62eb7d9ed6f427cf687d8d0071d751850cfe5d136bc60d3ab").Bytes(), + } + + err := mt.TryUpdateAccount(common.HexToAddress("0x05fDbDfaE180345C6Cff5316c286727CF1a43327").Bytes(), acc1) + assert.Nil(t, err) + + acc2 := &types.StateAccount{ + Nonce: 5, + Balance: big.NewInt(50000000), + Root: common.HexToHash("0"), + CodeHash: common.HexToHash("c5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470").Bytes(), + } + err = mt.TryUpdateAccount(common.HexToAddress("0x4cb1aB63aF5D8931Ce09673EbD8ae2ce16fD6571").Bytes(), acc2) + assert.Nil(t, err) + + bt, err := mt.TryGet(common.HexToAddress("0x05fDbDfaE180345C6Cff5316c286727CF1a43327").Bytes()) + assert.Nil(t, err) + + acc, err := types.UnmarshalStateAccount(bt) + assert.Nil(t, err) + assert.Equal(t, acc1.Nonce, acc.Nonce) + assert.Equal(t, acc1.Balance.Uint64(), acc.Balance.Uint64()) + assert.Equal(t, acc1.Root.Bytes(), acc.Root.Bytes()) + assert.Equal(t, acc1.CodeHash, acc.CodeHash) + + bt, err = mt.TryGet(common.HexToAddress("0x4cb1aB63aF5D8931Ce09673EbD8ae2ce16fD6571").Bytes()) + assert.Nil(t, err) + + acc, err = types.UnmarshalStateAccount(bt) + assert.Nil(t, err) + assert.Equal(t, acc2.Nonce, acc.Nonce) + assert.Equal(t, acc2.Balance.Uint64(), acc.Balance.Uint64()) + assert.Equal(t, acc2.Root.Bytes(), acc.Root.Bytes()) + assert.Equal(t, acc2.CodeHash, acc.CodeHash) + + bt, err = mt.TryGet(common.HexToAddress("0x8dE13967F19410A7991D63c2c0179feBFDA0c261").Bytes()) + assert.Nil(t, err) + assert.Nil(t, bt) +} diff --git a/trie/zk_trie_node.go b/trie/zk_trie_node.go new file mode 100644 index 0000000000..487f922378 --- /dev/null +++ b/trie/zk_trie_node.go @@ -0,0 +1,224 @@ +package trie + +import ( + "encoding/binary" + "fmt" + "math/big" + "reflect" + "unsafe" + + zkt "github.com/scroll-tech/go-ethereum/core/types/zktrie" +) + +// NodeType defines the type of node in the MT. +type NodeType byte + +const ( + // NodeTypeMiddle indicates the type of middle Node that has children. + NodeTypeMiddle NodeType = 0 + // NodeTypeLeaf indicates the type of a leaf Node that contains a key & + // value. + NodeTypeLeaf NodeType = 1 + // NodeTypeEmpty indicates the type of an empty Node. + NodeTypeEmpty NodeType = 2 + + // DBEntryTypeRoot indicates the type of a DB entry that indicates the + // current Root of a MerkleTree + DBEntryTypeRoot NodeType = 3 +) + +// Node is the struct that represents a node in the MT. The node should not be +// modified after creation because the cached key won't be updated. +type Node struct { + // Type is the type of node in the tree. + Type NodeType + // ChildL is the left child of a middle node. + ChildL *zkt.Hash + // ChildR is the right child of a middle node. + ChildR *zkt.Hash + // NodeKey is the node's key stored in a leaf node. + NodeKey *zkt.Hash + // ValuePreimage can store at most 256 byte32 as fields (represnted by BIG-ENDIAN integer) + // and the first 24 can be compressed (each bytes32 consider as 2 fields), in hashing the compressed + // elemments would be calculated first + ValuePreimage []zkt.Byte32 + // CompressedFlags use each bit for indicating the compressed flag for the first 24 fields + CompressedFlags uint32 + // key is the cache of entry key used to avoid recalculating + key *zkt.Hash + // valueHash is the cache of hashes of valuePreimage, used to avoid recalculating + valueHash *zkt.Hash + // KeyPreimage is kept here only for proof + KeyPreimage *zkt.Byte32 +} + +// NewNodeLeaf creates a new leaf node. +func NewNodeLeaf(k *zkt.Hash, valueFlags uint32, valuePreimage []zkt.Byte32) *Node { + return &Node{Type: NodeTypeLeaf, NodeKey: k, CompressedFlags: valueFlags, ValuePreimage: valuePreimage} +} + +// NewNodeMiddle creates a new middle node. +func NewNodeMiddle(childL *zkt.Hash, childR *zkt.Hash) *Node { + return &Node{Type: NodeTypeMiddle, ChildL: childL, ChildR: childR} +} + +// NewNodeEmpty creates a new empty node. +func NewNodeEmpty() *Node { + return &Node{Type: NodeTypeEmpty} +} + +// NewNodeFromBytes creates a new node by parsing the input []byte. +func NewNodeFromBytes(b []byte) (*Node, error) { + if len(b) < 1 { + return nil, ErrNodeBytesBadSize + } + n := Node{Type: NodeType(b[0])} + b = b[1:] + switch n.Type { + case NodeTypeMiddle: + if len(b) != 2*zkt.ElemBytesLen { + return nil, ErrNodeBytesBadSize + } + n.ChildL, n.ChildR = &zkt.Hash{}, &zkt.Hash{} + copy(n.ChildL[:], b[:zkt.ElemBytesLen]) + copy(n.ChildR[:], b[zkt.ElemBytesLen:zkt.ElemBytesLen*2]) + case NodeTypeLeaf: + if len(b) < zkt.ElemBytesLen+4 { + return nil, ErrNodeBytesBadSize + } + n.NodeKey = &zkt.Hash{} + copy(n.NodeKey[:], b[0:32]) + mark := binary.LittleEndian.Uint32(b[32:36]) + preimageLen := int(mark & 255) + n.CompressedFlags = mark >> 8 + n.ValuePreimage = make([]zkt.Byte32, preimageLen) + curPos := 36 + for i := 0; i < preimageLen; i++ { + copy(n.ValuePreimage[i][:], b[i*32+curPos:(i+1)*32+curPos]) + } + curPos = 36 + preimageLen*32 + preImageSize := int(b[curPos]) + curPos += 1 + if preImageSize != 0 { + n.KeyPreimage = new(zkt.Byte32) + copy(n.KeyPreimage[:], b[curPos:curPos+preImageSize]) + } + case NodeTypeEmpty: + break + default: + return nil, ErrInvalidNodeFound + } + return &n, nil +} + +// LeafKey computes the key of a leaf node given the hIndex and hValue of the +// entry of the leaf. +func LeafKey(k, v *zkt.Hash) (*zkt.Hash, error) { + return zkt.HashElems(big.NewInt(1), k.BigInt(), v.BigInt()) +} + +// Key computes the key of the node by hashing the content in a specific way +// for each type of node. This key is used as the hash of the merklee tree for +// each node. +func (n *Node) Key() (*zkt.Hash, error) { + if n.key == nil { // Cache the key to avoid repeated hash computations. + // NOTE: We are not using the type to calculate the hash! + switch n.Type { + case NodeTypeMiddle: // H(ChildL || ChildR) + var err error + n.key, err = zkt.HashElems(n.ChildL.BigInt(), n.ChildR.BigInt()) + if err != nil { + return nil, err + } + case NodeTypeLeaf: + var err error + n.valueHash, err = zkt.PreHandlingElems(n.CompressedFlags, n.ValuePreimage) + if err != nil { + return nil, err + } + + n.key, err = LeafKey(n.NodeKey, n.valueHash) + if err != nil { + return nil, err + } + + case NodeTypeEmpty: // Zero + n.key = &zkt.HashZero + default: + n.key = &zkt.HashZero + } + } + return n.key, nil +} + +func (n *Node) ValueKey() (*zkt.Hash, error) { + if _, err := n.Key(); err != nil { + return nil, err + } + return n.valueHash, nil +} + +// Data returns the wrapped data inside LeafNode and cast them into bytes +// for other node type it just return nil +func (n *Node) Data() []byte { + switch n.Type { + case NodeTypeLeaf: + var data []byte + hdata := (*reflect.SliceHeader)(unsafe.Pointer(&data)) + //TODO: uintptr(reflect.ValueOf(n.ValuePreimage).UnsafePointer()) should be more elegant but only available until go 1.18 + hdata.Data = uintptr(unsafe.Pointer(&n.ValuePreimage[0])) + hdata.Len = 32 * len(n.ValuePreimage) + hdata.Cap = hdata.Len + return data + default: + return nil + } +} + +// Value returns the value of the node. This is the content that is stored in +// the backend database. +func (n *Node) Value() []byte { + switch n.Type { + case NodeTypeMiddle: // {Type || ChildL || ChildR} + bytes := []byte{byte(n.Type)} + bytes = append(bytes, n.ChildL[:]...) + bytes = append(bytes, n.ChildR[:]...) + return bytes + case NodeTypeLeaf: // {Type || Data...} + bytes := []byte{byte(n.Type)} + bytes = append(bytes, n.NodeKey[:]...) + tmp := make([]byte, 4) + compressedFlag := (n.CompressedFlags << 8) + uint32(len(n.ValuePreimage)) + binary.LittleEndian.PutUint32(tmp, compressedFlag) + bytes = append(bytes, tmp...) + for _, elm := range n.ValuePreimage { + bytes = append(bytes, elm[:]...) + } + if n.KeyPreimage != nil { + bytes = append(bytes, byte(len(n.KeyPreimage))) + bytes = append(bytes, n.KeyPreimage[:]...) + } else { + bytes = append(bytes, 0) + } + + return bytes + case NodeTypeEmpty: // { Type } + return []byte{byte(n.Type)} + default: + return []byte{} + } +} + +// String outputs a string representation of a node (different for each type). +func (n *Node) String() string { + switch n.Type { + case NodeTypeMiddle: // {Type || ChildL || ChildR} + return fmt.Sprintf("Middle L:%s R:%s", n.ChildL, n.ChildR) + case NodeTypeLeaf: // {Type || Data...} + return fmt.Sprintf("Leaf I:%v Items: %d, First:%v", n.NodeKey, len(n.ValuePreimage), n.ValuePreimage[0]) + case NodeTypeEmpty: // {} + return "Empty" + default: + return "Invalid Node" + } +} diff --git a/trie/zk_trie_proof.go b/trie/zk_trie_proof.go new file mode 100644 index 0000000000..c5d84e2c92 --- /dev/null +++ b/trie/zk_trie_proof.go @@ -0,0 +1,195 @@ +package trie + +import ( + "bytes" + "fmt" + "math/big" + + "github.com/scroll-tech/go-ethereum/common" + zkt "github.com/scroll-tech/go-ethereum/core/types/zktrie" + "github.com/scroll-tech/go-ethereum/ethdb" +) + +// TODO: remove this hack +var magicHash []byte +var magicSMTBytes []byte + +func init() { + magicSMTBytes = []byte("THIS IS SOME MAGIC BYTES FOR SMT m1rRXgP2xpDI") + hasher := newHasher(false) + defer returnHasherToPool(hasher) + magicHash = hasher.hashData(magicSMTBytes) +} + +// Prove constructs a merkle proof for SMT, it respect the protocol used by the ethereum-trie +// but save the node data with a compact form +func (mt *ZkTrieImpl) prove(kHash *zkt.Hash, fromLevel uint, writeNode func(*Node) error) error { + + path := getPath(mt.maxLevels, kHash[:]) + var nodes []*Node + tn := mt.rootKey + for i := 0; i < mt.maxLevels; i++ { + n, err := mt.GetNode(tn) + if err != nil { + return err + } + + finished := true + switch n.Type { + case NodeTypeEmpty: + case NodeTypeLeaf: + // notice even we found a leaf whose entry didn't match the expected k, + // we still include it as the proof of absence + case NodeTypeMiddle: + finished = false + if path[i] { + tn = n.ChildR + } else { + tn = n.ChildL + } + default: + return ErrInvalidNodeFound + } + + nodes = append(nodes, n) + if finished { + break + } + } + + for _, n := range nodes { + if fromLevel > 0 { + fromLevel-- + continue + } + + // TODO: notice here we may have broken some implicit on the proofDb: + // the key is not kecca(value) and it even can not be derived from + // the value by any means without a actually decoding + if err := writeNode(n); err != nil { + return err + } + } + + return nil +} + +func (mt *ZkTrieImpl) Prove(key []byte, fromLevel uint, proofDb ethdb.KeyValueWriter) error { + + kHash, err := zkt.NewHashFromBytes(common.BytesToHash(key).Bytes()) + if err != nil { + return err + } + + err = mt.prove(kHash, fromLevel, func(n *Node) error { + key, err := n.Key() + if err != nil { + return err + } + return proofDb.Put(key[:], n.Value()) + }) + + if err != nil { + return err + } + + // we put this special kv pair in db so we can distinguish the type and + // make suitable Proof + return proofDb.Put(magicHash, magicSMTBytes) +} + +func buildZkTrieProof(rootKey *zkt.Hash, k *big.Int, lvl int, getNode func(key *zkt.Hash) (*Node, error)) (*Proof, + *Node, error) { + + p := &Proof{} + var siblingKey *zkt.Hash + + kHash := zkt.NewHashFromBigInt(k) + path := getPath(lvl, kHash[:]) + + nextKey := rootKey + for p.depth = 0; p.depth < uint(lvl); p.depth++ { + n, err := getNode(nextKey) + if err != nil { + return nil, nil, err + } + switch n.Type { + case NodeTypeEmpty: + return p, n, nil + case NodeTypeLeaf: + if bytes.Equal(kHash[:], n.NodeKey[:]) { + p.Existence = true + return p, n, nil + } + // We found a leaf whose entry didn't match hIndex + p.NodeAux = &NodeAux{Key: n.NodeKey, Value: n.valueHash} + return p, n, nil + case NodeTypeMiddle: + if path[p.depth] { + nextKey = n.ChildR + siblingKey = n.ChildL + } else { + nextKey = n.ChildL + siblingKey = n.ChildR + } + default: + return nil, nil, ErrInvalidNodeFound + } + if !bytes.Equal(siblingKey[:], zkt.HashZero[:]) { + zkt.SetBitBigEndian(p.notempties[:], p.depth) + p.Siblings = append(p.Siblings, siblingKey) + } + } + return nil, nil, ErrKeyNotFound + +} + +// DecodeProof try to decode a node bytes, return can be nil for any non-node data (magic code) +func DecodeSMTProof(data []byte) (*Node, error) { + + if bytes.Equal(magicSMTBytes, data) { + //skip magic bytes node + return nil, nil + } + + return NewNodeFromBytes(data) +} + +// VerifyProof checks merkle proofs. The given proof must contain the value for +// key in a trie with the given root hash. VerifyProof returns an error if the +// proof contains invalid trie nodes or the wrong value. +func VerifyProofSMT(rootHash common.Hash, key []byte, proofDb ethdb.KeyValueReader) (value []byte, err error) { + + h, err := zkt.NewHashFromBytes(rootHash.Bytes()) + if err != nil { + return nil, err + } + + word := zkt.NewByte32FromBytesPaddingZero(key) + k, err := word.Hash() + if err != nil { + return nil, err + } + + proof, n, err := buildZkTrieProof(h, k, len(key)*8, func(key *zkt.Hash) (*Node, error) { + buf, _ := proofDb.Get(key[:]) + if buf == nil { + return nil, ErrKeyNotFound + } + n, err := NewNodeFromBytes(buf) + return n, err + }) + + if err != nil { + // do not contain the key + return nil, err + } else if !proof.Existence { + return nil, nil + } + + if VerifyProofZkTrie(h, proof, n) { + return n.Data(), nil + } else { + return nil, fmt.Errorf("bad proof node %v", proof) + } +} diff --git a/trie/zk_trie_proof_test.go b/trie/zk_trie_proof_test.go new file mode 100644 index 0000000000..58d93b725c --- /dev/null +++ b/trie/zk_trie_proof_test.go @@ -0,0 +1,189 @@ +// Copyright 2015 The go-ethereum Authors +// This file is part of the go-ethereum library. +// +// The go-ethereum library is free software: you can redistribute it and/or modify +// it under the terms of the GNU Lesser General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// The go-ethereum library is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Lesser General Public License for more details. +// +// You should have received a copy of the GNU Lesser General Public License +// along with the go-ethereum library. If not, see . + +package trie + +import ( + "bytes" + mrand "math/rand" + "testing" + "time" + + "github.com/stretchr/testify/assert" + + "github.com/scroll-tech/go-ethereum/common" + zkt "github.com/scroll-tech/go-ethereum/core/types/zktrie" + "github.com/scroll-tech/go-ethereum/crypto" + "github.com/scroll-tech/go-ethereum/ethdb/memorydb" +) + +func init() { + mrand.Seed(time.Now().Unix()) +} + +// makeProvers creates Merkle trie provers based on different implementations to +// test all variations. +func makeSMTProvers(mt *ZkTrieImpl) []func(key []byte) *memorydb.Database { + var provers []func(key []byte) *memorydb.Database + + // Create a direct trie based Merkle prover + provers = append(provers, func(key []byte) *memorydb.Database { + word := zkt.NewByte32FromBytesPaddingZero(key) + k, err := word.Hash() + if err != nil { + panic(err) + } + proof := memorydb.New() + mt.Prove(k.Bytes(), 0, proof) + return proof + }) + return provers +} + +func verifyValue(proveVal []byte, vPreimage []byte) bool { + return bytes.Equal(proveVal, vPreimage) +} + +func TestSMTOneElementProof(t *testing.T) { + mt, _ := NewZkTrieImpl(NewZktrieDatabase((memorydb.New())), 64) + err := mt.UpdateWord( + zkt.NewByte32FromBytesPaddingZero(bytes.Repeat([]byte("k"), 32)), + zkt.NewByte32FromBytesPaddingZero(bytes.Repeat([]byte("v"), 32)), + ) + assert.Nil(t, err) + for i, prover := range makeSMTProvers(mt) { + keyBytes := bytes.Repeat([]byte("k"), 32) + proof := prover(keyBytes) + if proof == nil { + t.Fatalf("prover %d: nil proof", i) + } + if proof.Len() != 2 { + t.Errorf("prover %d: proof should have 1+1 element (including the magic kv)", i) + } + val, err := VerifyProof(common.BytesToHash(mt.Root().Bytes()), keyBytes, proof) + if err != nil { + t.Fatalf("prover %d: failed to verify proof: %v\nraw proof: %x", i, err, proof) + } + if !verifyValue(val, bytes.Repeat([]byte("v"), 32)) { + t.Fatalf("prover %d: verified value mismatch: want 'v' get %x", i, val) + } + } +} + +func TestSMTProof(t *testing.T) { + mt, vals := randomZktrie(t, 500) + root := mt.Root() + for i, prover := range makeSMTProvers(mt) { + for _, kv := range vals { + proof := prover(kv.k) + if proof == nil { + t.Fatalf("prover %d: missing key %x while constructing proof", i, kv.k) + } + val, err := VerifyProof(common.BytesToHash(root.Bytes()), kv.k, proof) + if err != nil { + t.Fatalf("prover %d: failed to verify proof for key %x: %v\nraw proof: %x\n", i, kv.k, err, proof) + } + if !verifyValue(val, zkt.NewByte32FromBytesPaddingZero(kv.v)[:]) { + t.Fatalf("prover %d: verified value mismatch for key %x, want %x, get %x", i, kv.k, kv.v, val) + } + } + } +} + +func TestSMTBadProof(t *testing.T) { + mt, vals := randomZktrie(t, 500) + root := mt.Root() + for i, prover := range makeSMTProvers(mt) { + for _, kv := range vals { + proof := prover(kv.k) + if proof == nil { + t.Fatalf("prover %d: nil proof", i) + } + it := proof.NewIterator(nil, nil) + for i, d := 0, mrand.Intn(proof.Len()); i <= d; i++ { + it.Next() + } + key := it.Key() + val, _ := proof.Get(key) + proof.Delete(key) + it.Release() + + mutateByte(val) + proof.Put(crypto.Keccak256(val), val) + + if _, err := VerifyProof(common.BytesToHash(root.Bytes()), kv.k, proof); err == nil { + t.Fatalf("prover %d: expected proof to fail for key %x", i, kv.k) + } + } + } +} + +// Tests that missing keys can also be proven. The test explicitly uses a single +// entry trie and checks for missing keys both before and after the single entry. +func TestSMTMissingKeyProof(t *testing.T) { + mt, _ := NewZkTrieImpl(NewZktrieDatabase((memorydb.New())), 64) + err := mt.UpdateWord( + zkt.NewByte32FromBytesPaddingZero(bytes.Repeat([]byte("k"), 20)), + zkt.NewByte32FromBytesPaddingZero(bytes.Repeat([]byte("v"), 20)), + ) + assert.Nil(t, err) + + prover := makeSMTProvers(mt)[0] + + for i, key := range []string{"a", "j", "l", "z"} { + keyBytes := bytes.Repeat([]byte(key), 32) + proof := prover(keyBytes) + + if proof.Len() != 2 { + t.Errorf("test %d: proof should have 2 element (with magic kv)", i) + } + val, err := VerifyProof(common.BytesToHash(mt.Root().Bytes()), keyBytes, proof) + if err != nil { + t.Fatalf("test %d: failed to verify proof: %v\nraw proof: %x", i, err, proof) + } + if val != nil { + t.Fatalf("test %d: verified value mismatch: have %x, want nil", i, val) + } + } +} + +func randomZktrie(t *testing.T, n int) (*ZkTrieImpl, map[string]*kv) { + mt, err := NewZkTrieImpl(NewZktrieDatabase((memorydb.New())), 64) + if err != nil { + panic(err) + } + vals := make(map[string]*kv) + for i := byte(0); i < 100; i++ { + + value := &kv{common.LeftPadBytes([]byte{i}, 32), bytes.Repeat([]byte{i}, 32), false} + value2 := &kv{common.LeftPadBytes([]byte{i + 10}, 32), bytes.Repeat([]byte{i}, 32), false} + + err = mt.UpdateWord(zkt.NewByte32FromBytesPaddingZero(value.k), zkt.NewByte32FromBytesPaddingZero(value.v)) + assert.Nil(t, err) + err = mt.UpdateWord(zkt.NewByte32FromBytesPaddingZero(value2.k), zkt.NewByte32FromBytesPaddingZero(value2.v)) + assert.Nil(t, err) + vals[string(value.k)] = value + vals[string(value2.k)] = value2 + } + for i := 0; i < n; i++ { + value := &kv{randBytes(32), randBytes(20), false} + err = mt.UpdateWord(zkt.NewByte32FromBytesPaddingZero(value.k), zkt.NewByte32FromBytesPaddingZero(value.v)) + assert.Nil(t, err) + vals[string(value.k)] = value + } + + return mt, vals +} diff --git a/trie/zk_trie_test.go b/trie/zk_trie_test.go new file mode 100644 index 0000000000..5ebb020ab7 --- /dev/null +++ b/trie/zk_trie_test.go @@ -0,0 +1,149 @@ +// Copyright 2015 The go-ethereum Authors +// This file is part of the go-ethereum library. +// +// The go-ethereum library is free software: you can redistribute it and/or modify +// it under the terms of the GNU Lesser General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// The go-ethereum library is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Lesser General Public License for more details. +// +// You should have received a copy of the GNU Lesser General Public License +// along with the go-ethereum library. If not, see . + +package trie + +import ( + "bytes" + "runtime" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/scroll-tech/go-ethereum/common" + zkt "github.com/scroll-tech/go-ethereum/core/types/zktrie" + "github.com/scroll-tech/go-ethereum/ethdb/memorydb" +) + +func newEmptyZkTrie() *ZkTrie { + trie, _ := NewZkTrie(common.Hash{}, NewZktrieDatabase(memorydb.New())) + return trie +} + +// makeTestSecureTrie creates a large enough secure trie for testing. +func makeTestZkTrie() (*ZktrieDatabase, *ZkTrie, map[string][]byte) { + // Create an empty trie + triedb := NewZktrieDatabase(memorydb.New()) + trie, _ := NewZkTrie(common.Hash{}, triedb) + + // Fill it with some arbitrary data + content := make(map[string][]byte) + for i := byte(0); i < 255; i++ { + // Map the same data under multiple keys + key, val := common.LeftPadBytes([]byte{1, i}, 32), bytes.Repeat([]byte{i}, 32) + content[string(key)] = val + trie.Update(key, val) + + key, val = common.LeftPadBytes([]byte{2, i}, 32), bytes.Repeat([]byte{i}, 32) + content[string(key)] = val + trie.Update(key, val) + + // Add some other data to inflate the trie + for j := byte(3); j < 13; j++ { + key, val = common.LeftPadBytes([]byte{j, i}, 32), bytes.Repeat([]byte{j, i}, 16) + content[string(key)] = val + trie.Update(key, val) + } + } + trie.Commit(nil) + + // Return the generated trie + return triedb, trie, content +} + +func TestZktrieDelete(t *testing.T) { + t.Skip("var-len kv not supported") + trie := newEmptyZkTrie() + vals := []struct{ k, v string }{ + {"do", "verb"}, + {"ether", "wookiedoo"}, + {"horse", "stallion"}, + {"shaman", "horse"}, + {"doge", "coin"}, + {"ether", ""}, + {"dog", "puppy"}, + {"shaman", ""}, + } + for _, val := range vals { + if val.v != "" { + trie.Update([]byte(val.k), []byte(val.v)) + } else { + trie.Delete([]byte(val.k)) + } + } + hash := trie.Hash() + exp := common.HexToHash("29b235a58c3c25ab83010c327d5932bcf05324b7d6b1185e650798034783ca9d") + if hash != exp { + t.Errorf("expected %x got %x", exp, hash) + } +} + +func TestZktrieGetKey(t *testing.T) { + trie := newEmptyZkTrie() + key := []byte("0a1b2c3d4e5f6g7h8i9j0a1b2c3d4e5f") + value := []byte("9j8i7h6g5f4e3d2c1b0a9j8i7h6g5f4e") + trie.Update(key, value) + + kPreimage := zkt.NewByte32FromBytesPaddingZero(key) + kHash, err := kPreimage.Hash() + assert.Nil(t, err) + + if !bytes.Equal(trie.Get(key), value) { + t.Errorf("Get did not return bar") + } + if k := trie.GetKey(kHash.Bytes()); !bytes.Equal(k, key) { + t.Errorf("GetKey returned %q, want %q", k, key) + } +} + +func TestZkTrieConcurrency(t *testing.T) { + // Create an initial trie and copy if for concurrent access + _, trie, _ := makeTestZkTrie() + + threads := runtime.NumCPU() + tries := make([]*ZkTrie, threads) + for i := 0; i < threads; i++ { + cpy := *trie + tries[i] = &cpy + } + // Start a batch of goroutines interactng with the trie + pend := new(sync.WaitGroup) + pend.Add(threads) + for i := 0; i < threads; i++ { + go func(index int) { + defer pend.Done() + + for j := byte(0); j < 255; j++ { + // Map the same data under multiple keys + key, val := common.LeftPadBytes([]byte{byte(index), 1, j}, 32), bytes.Repeat([]byte{j}, 32) + tries[index].Update(key, val) + + key, val = common.LeftPadBytes([]byte{byte(index), 2, j}, 32), bytes.Repeat([]byte{j}, 32) + tries[index].Update(key, val) + + // Add some other data to inflate the trie + for k := byte(3); k < 13; k++ { + key, val = common.LeftPadBytes([]byte{byte(index), k, j}, 32), bytes.Repeat([]byte{k, j}, 16) + tries[index].Update(key, val) + } + } + tries[index].Commit(nil) + }(i) + } + // Wait for all threads to finish + pend.Wait() +}