Implement MPT

pkg/core/blockchain.go

@@ -17,6 +17,7 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/core/state"
 	"github.com/nspcc-dev/neo-go/pkg/core/storage"
 	"github.com/nspcc-dev/neo-go/pkg/core/transaction"
+	"github.com/nspcc-dev/neo-go/pkg/crypto/hash"
 	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
 	"github.com/nspcc-dev/neo-go/pkg/encoding/bigint"
 	"github.com/nspcc-dev/neo-go/pkg/io"
@@ -221,6 +222,9 @@ func (bc *Blockchain) init() error {
 	}
 	bc.blockHeight = bHeight
 	bc.persistedHeight = bHeight
+	if err = bc.dao.InitMPT(bHeight); err != nil {
+		return errors.Wrapf(err, "can't init MPT at height %d", bHeight)
+	}
 
 	hashes, err := bc.dao.GetHeaderHashes()
 	if err != nil {
@@ -550,6 +554,11 @@ func (bc *Blockchain) processHeader(h *block.Header, batch storage.Batch, header
 	return nil
 }
 
+// GetStateRoot returns state root for a given height.
+func (bc *Blockchain) GetStateRoot(height uint32) (*state.MPTRootState, error) {
+	return bc.dao.GetStateRoot(height)
+}
+
 // storeBlock performs chain update using the block given, it executes all
 // transactions with all appropriate side-effects and updates Blockchain state.
 // This is the only way to change Blockchain state.
@@ -633,17 +642,44 @@ func (bc *Blockchain) storeBlock(block *block.Block) error {
 		}
 	}
 
+	root := bc.dao.MPT.StateRoot()
+	var prevHash util.Uint256
+	if block.Index > 0 {
+		prev, err := bc.dao.GetStateRoot(block.Index - 1)
+		if err != nil {
+			return errors.WithMessagef(err, "can't get previous state root")
+		}
+		prevHash = hash.DoubleSha256(prev.GetSignedPart())
+	}
+	err := bc.AddStateRoot(&state.MPTRoot{
+		MPTRootBase: state.MPTRootBase{
+			Index:    block.Index,
+			PrevHash: prevHash,
+			Root:     root,
+		},
+	})
+	if err != nil {
+		return err
+	}
+
 	if bc.config.SaveStorageBatch {
 		bc.lastBatch = cache.DAO.GetBatch()
 	}
 
 	bc.lock.Lock()
-	_, err := cache.Persist()
+	_, err = cache.Persist()
 	if err != nil {
 		bc.lock.Unlock()
 		return err
 	}
 	bc.contracts.Policy.OnPersistEnd(bc.dao)
+	bc.dao.MPT.Flush()
+	// Every persist cycle we also compact our in-memory MPT.
+	persistedHeight := atomic.LoadUint32(&bc.persistedHeight)
+	if persistedHeight == block.Index-1 {
+		// 10 is good and roughly estimated to fit remaining trie into 1M of memory.
+		bc.dao.MPT.Collapse(10)
+	}
 	bc.topBlock.Store(block)
 	atomic.StoreUint32(&bc.blockHeight, block.Index)
 	bc.memPool.RemoveStale(bc.isTxStillRelevant, bc)
@@ -1194,6 +1230,82 @@ func (bc *Blockchain) isTxStillRelevant(t *transaction.Transaction) bool {
 
 }
 
+// AddStateRoot add new (possibly unverified) state root to the blockchain.
+func (bc *Blockchain) AddStateRoot(r *state.MPTRoot) error {
+	our, err := bc.GetStateRoot(r.Index)
+	if err == nil {
+		if our.Flag == state.Verified {
+			return bc.updateStateHeight(r.Index)
+		} else if r.Witness == nil && our.Witness != nil {
+			r.Witness = our.Witness
+		}
+	}
+	if err := bc.verifyStateRoot(r); err != nil {
+		return errors.WithMessage(err, "invalid state root")
+	}
+	if r.Index > bc.BlockHeight() { // just put it into the store for future checks
+		return bc.dao.PutStateRoot(&state.MPTRootState{
+			MPTRoot: *r,
+			Flag:    state.Unverified,
+		})
+	}
+
+	flag := state.Unverified
+	if r.Witness != nil {
+		if err := bc.verifyStateRootWitness(r); err != nil {
+			return errors.WithMessage(err, "can't verify signature")
+		}
+		flag = state.Verified
+	}
+	err = bc.dao.PutStateRoot(&state.MPTRootState{
+		MPTRoot: *r,
+		Flag:    flag,
+	})
+	if err != nil {
+		return err
+	}
+	return bc.updateStateHeight(r.Index)
+}
+
+func (bc *Blockchain) updateStateHeight(newHeight uint32) error {
+	h, err := bc.dao.GetCurrentStateRootHeight()
+	if err != nil {
+		return errors.WithMessage(err, "can't get current state root height")
+	} else if newHeight == h+1 {
+		updateStateHeightMetric(newHeight)
+		return bc.dao.PutCurrentStateRootHeight(h + 1)
+	}
+	return nil
+}
+
+// verifyStateRoot checks if state root is valid.
+func (bc *Blockchain) verifyStateRoot(r *state.MPTRoot) error {
+	if r.Index == 0 {
+		return nil
+	}
+	prev, err := bc.GetStateRoot(r.Index - 1)
+	if err != nil {
+		return errors.New("can't get previous state root")
+	} else if !r.PrevHash.Equals(hash.DoubleSha256(prev.GetSignedPart())) {
+		return errors.New("previous hash mismatch")
+	} else if prev.Version != r.Version {
+		return errors.New("version mismatch")
+	}
+	return nil
+}
+
+// verifyStateRootWitness verifies that state root signature is correct.
+func (bc *Blockchain) verifyStateRootWitness(r *state.MPTRoot) error {
+	b, err := bc.GetBlock(bc.GetHeaderHash(int(r.Index)))
+	if err != nil {
+		return err
+	}
+	interopCtx := bc.newInteropContext(trigger.Verification, bc.dao, nil, nil)
+	interopCtx.Container = r
+	return bc.verifyHashAgainstScript(b.NextConsensus, r.Witness, interopCtx, true,
+		bc.contracts.Policy.GetMaxVerificationGas(interopCtx.DAO))
+}
+
 // VerifyTx verifies whether a transaction is bonafide or not. Block parameter
 // is used for easy interop access and can be omitted for transactions that are
 // not yet added into any block.

pkg/core/blockchainer/blockchainer.go

@@ -20,6 +20,7 @@ type Blockchainer interface {
 	GetConfig() config.ProtocolConfiguration
 	AddHeaders(...*block.Header) error
 	AddBlock(*block.Block) error
+	AddStateRoot(r *state.MPTRoot) error
 	BlockHeight() uint32
 	CalculateClaimable(value *big.Int, startHeight, endHeight uint32) *big.Int
 	Close()
@@ -42,6 +43,7 @@ type Blockchainer interface {
 	GetValidators() ([]*keys.PublicKey, error)
 	GetStandByValidators() keys.PublicKeys
 	GetScriptHashesForVerifying(*transaction.Transaction) ([]util.Uint160, error)
+	GetStateRoot(height uint32) (*state.MPTRootState, error)
 	GetStorageItem(id int32, key []byte) *state.StorageItem
 	GetStorageItems(id int32) (map[string]*state.StorageItem, error)
 	GetTestVM(tx *transaction.Transaction) *vm.VM

pkg/core/dao/dao.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/nspcc-dev/neo-go/pkg/config/netmode"
 	"github.com/nspcc-dev/neo-go/pkg/core/block"
+	"github.com/nspcc-dev/neo-go/pkg/core/mpt"
 	"github.com/nspcc-dev/neo-go/pkg/core/state"
 	"github.com/nspcc-dev/neo-go/pkg/core/storage"
 	"github.com/nspcc-dev/neo-go/pkg/core/transaction"
@@ -29,10 +30,13 @@ type DAO interface {
 	GetContractState(hash util.Uint160) (*state.Contract, error)
 	GetCurrentBlockHeight() (uint32, error)
 	GetCurrentHeaderHeight() (i uint32, h util.Uint256, err error)
+	GetCurrentStateRootHeight() (uint32, error)
 	GetHeaderHashes() ([]util.Uint256, error)
 	GetNEP5Balances(acc util.Uint160) (*state.NEP5Balances, error)
 	GetNEP5TransferLog(acc util.Uint160, index uint32) (*state.NEP5TransferLog, error)
 	GetAndUpdateNextContractID() (int32, error)
+	GetStateRoot(height uint32) (*state.MPTRootState, error)
+	PutStateRoot(root *state.MPTRootState) error
 	GetStorageItem(id int32, key []byte) *state.StorageItem
 	GetStorageItems(id int32) (map[string]*state.StorageItem, error)
 	GetStorageItemsWithPrefix(id int32, prefix []byte) (map[string]*state.StorageItem, error)
@@ -58,13 +62,15 @@ type DAO interface {
 
 // Simple is memCached wrapper around DB, simple DAO implementation.
 type Simple struct {
+	MPT     *mpt.Trie
 	Store   *storage.MemCachedStore
 	network netmode.Magic
 }
 
 // NewSimple creates new simple dao using provided backend store.
 func NewSimple(backend storage.Store, network netmode.Magic) *Simple {
-	return &Simple{Store: storage.NewMemCachedStore(backend), network: network}
+	st := storage.NewMemCachedStore(backend)
+	return &Simple{Store: st, network: network, MPT: mpt.NewTrie(nil, st)}
 }
 
 // GetBatch returns currently accumulated DB changeset.
@@ -75,7 +81,9 @@ func (dao *Simple) GetBatch() *storage.MemBatch {
 // GetWrapped returns new DAO instance with another layer of wrapped
 // MemCachedStore around the current DAO Store.
 func (dao *Simple) GetWrapped() DAO {
-	return NewSimple(dao.Store, dao.network)
+	d := NewSimple(dao.Store, dao.network)
+	d.MPT = dao.MPT
+	return d
 }
 
 // GetAndDecode performs get operation and decoding with serializable structures.
@@ -288,6 +296,63 @@ func (dao *Simple) PutAppExecResult(aer *state.AppExecResult) error {
 
 // -- start storage item.
 
+func makeStateRootKey(height uint32) []byte {
+	key := make([]byte, 5)
+	key[0] = byte(storage.DataMPT)
+	binary.LittleEndian.PutUint32(key[1:], height)
+	return key
+}
+
+// InitMPT initializes MPT at the given height.
+func (dao *Simple) InitMPT(height uint32) error {
+	if height == 0 {
+		dao.MPT = mpt.NewTrie(nil, dao.Store)
+		return nil
+	}
+	r, err := dao.GetStateRoot(height)
+	if err != nil {
+		return err
+	}
+	dao.MPT = mpt.NewTrie(mpt.NewHashNode(r.Root), dao.Store)
+	return nil
+}
+
+// GetCurrentStateRootHeight returns current state root height.
+func (dao *Simple) GetCurrentStateRootHeight() (uint32, error) {
+	key := []byte{byte(storage.DataMPT)}
+	val, err := dao.Store.Get(key)
+	if err != nil {
+		if err == storage.ErrKeyNotFound {
+			err = nil
+		}
+		return 0, err
+	}
+	return binary.LittleEndian.Uint32(val), nil
+}
+
+// PutCurrentStateRootHeight updates current state root height.
+func (dao *Simple) PutCurrentStateRootHeight(height uint32) error {
+	key := []byte{byte(storage.DataMPT)}
+	val := make([]byte, 4)
+	binary.LittleEndian.PutUint32(val, height)
+	return dao.Store.Put(key, val)
+}
+
+// GetStateRoot returns state root of a given height.
+func (dao *Simple) GetStateRoot(height uint32) (*state.MPTRootState, error) {
+	r := new(state.MPTRootState)
+	err := dao.GetAndDecode(r, makeStateRootKey(height))
+	if err != nil {
+		return nil, err
+	}
+	return r, nil
+}
+
+// PutStateRoot puts state root of a given height into the store.
+func (dao *Simple) PutStateRoot(r *state.MPTRootState) error {
+	return dao.Put(r, makeStateRootKey(r.Index))
+}
+
 // GetStorageItem returns StorageItem if it exists in the given store.
 func (dao *Simple) GetStorageItem(id int32, key []byte) *state.StorageItem {
 	b, err := dao.Store.Get(makeStorageItemKey(id, key))
@@ -308,13 +373,27 @@ func (dao *Simple) GetStorageItem(id int32, key []byte) *state.StorageItem {
 // PutStorageItem puts given StorageItem for given id with given
 // key into the given store.
 func (dao *Simple) PutStorageItem(id int32, key []byte, si *state.StorageItem) error {
-	return dao.Put(si, makeStorageItemKey(id, key))
+	stKey := makeStorageItemKey(id, key)
+	buf := io.NewBufBinWriter()
+	si.EncodeBinary(buf.BinWriter)
+	if buf.Err != nil {
+		return buf.Err
+	}
+	v := buf.Bytes()
+	if err := dao.MPT.Put(stKey[1:], v); err != nil && err != mpt.ErrNotFound {
+		return err
+	}
+	return dao.Store.Put(stKey, v)
 }
 
 // DeleteStorageItem drops storage item for the given id with the
 // given key from the store.
 func (dao *Simple) DeleteStorageItem(id int32, key []byte) error {
-	return dao.Store.Delete(makeStorageItemKey(id, key))
+	stKey := makeStorageItemKey(id, key)
+	if err := dao.MPT.Delete(stKey[1:]); err != nil && err != mpt.ErrNotFound {
+		return err
+	}
+	return dao.Store.Delete(stKey)
 }
 
 // GetStorageItems returns all storage items for a given id.