From 10189b6ab3fb9f63ead5fbb9088cc7d67ec9ef67 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Fri, 29 May 2020 17:54:41 +0300
Subject: [PATCH] consensus: extend payloads with StateRoot info

Create and verify witness after block processing.
---
 pkg/consensus/commit.go          |  3 ++
 pkg/consensus/consensus.go       | 55 +++++++++++++++++++++++++++++---
 pkg/consensus/prepare_request.go |  4 +++
 3 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/pkg/consensus/commit.go b/pkg/consensus/commit.go
index 492a1a156..372448576 100644
--- a/pkg/consensus/commit.go
+++ b/pkg/consensus/commit.go
@@ -8,6 +8,7 @@ import (
 // commit represents dBFT Commit message.
 type commit struct {
 	signature [signatureSize]byte
+	stateSig  [signatureSize]byte
 }
 
 // signatureSize is an rfc6989 signature size in bytes
@@ -19,11 +20,13 @@ var _ payload.Commit = (*commit)(nil)
 // EncodeBinary implements io.Serializable interface.
 func (c *commit) EncodeBinary(w *io.BinWriter) {
 	w.WriteBytes(c.signature[:])
+	w.WriteBytes(c.stateSig[:])
 }
 
 // DecodeBinary implements io.Serializable interface.
 func (c *commit) DecodeBinary(r *io.BinReader) {
 	r.ReadBytes(c.signature[:])
+	r.ReadBytes(c.stateSig[:])
 }
 
 // Signature implements payload.Commit interface.
diff --git a/pkg/consensus/consensus.go b/pkg/consensus/consensus.go
index e12b9d37e..6d21bdeaa 100644
--- a/pkg/consensus/consensus.go
+++ b/pkg/consensus/consensus.go
@@ -13,6 +13,7 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/core"
 	coreb "github.com/nspcc-dev/neo-go/pkg/core/block"
 	"github.com/nspcc-dev/neo-go/pkg/core/mempool"
+	"github.com/nspcc-dev/neo-go/pkg/core/state"
 	"github.com/nspcc-dev/neo-go/pkg/core/transaction"
 	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract"
@@ -135,10 +136,10 @@ func NewService(cfg Config) (Service, error) {
 		dbft.WithGetConsensusAddress(srv.getConsensusAddress),
 
 		dbft.WithNewConsensusPayload(func() payload.ConsensusPayload { p := new(Payload); p.message = &message{}; return p }),
-		dbft.WithNewPrepareRequest(func() payload.PrepareRequest { return new(prepareRequest) }),
+		dbft.WithNewPrepareRequest(srv.newPrepareRequest),
 		dbft.WithNewPrepareResponse(func() payload.PrepareResponse { return new(prepareResponse) }),
 		dbft.WithNewChangeView(func() payload.ChangeView { return new(changeView) }),
-		dbft.WithNewCommit(func() payload.Commit { return new(commit) }),
+		dbft.WithNewCommit(srv.newCommit),
 		dbft.WithNewRecoveryRequest(func() payload.RecoveryRequest { return new(recoveryRequest) }),
 		dbft.WithNewRecoveryMessage(func() payload.RecoveryMessage { return new(recoveryMessage) }),
 	)
@@ -209,6 +210,33 @@ func (s *service) eventLoop() {
 	}
 }
 
+func (s *service) newPrepareRequest() payload.PrepareRequest {
+	sr, err := s.Chain.GetStateRoot(s.Chain.BlockHeight())
+	if err != nil {
+		return new(prepareRequest)
+	}
+	return &prepareRequest{
+		proposalStateRoot: sr.MPTRootBase,
+	}
+}
+
+func (s *service) newCommit() payload.Commit {
+	for _, p := range s.dbft.Context.PreparationPayloads {
+		if p != nil && p.ViewNumber() == s.dbft.ViewNumber && p.Type() == payload.PrepareRequestType {
+			pr := p.GetPrepareRequest().(*prepareRequest)
+			data := pr.proposalStateRoot.GetSignedPart()
+			sign, err := s.dbft.Priv.Sign(data)
+			if err == nil {
+				var c commit
+				copy(c.stateSig[:], sign)
+				return &c
+			}
+			break
+		}
+	}
+	return new(commit)
+}
+
 func (s *service) validatePayload(p *Payload) bool {
 	validators := s.getValidators()
 	if int(p.validatorIndex) >= len(validators) {
@@ -350,16 +378,35 @@ func (s *service) processBlock(b block.Block) {
 			s.log.Warn("error on add block", zap.Error(err))
 		}
 	}
+
+	var rb *state.MPTRootBase
+	for _, p := range s.dbft.PreparationPayloads {
+		if p != nil && p.Type() == payload.PrepareRequestType {
+			rb = &p.GetPrepareRequest().(*prepareRequest).proposalStateRoot
+		}
+	}
+	w := s.getWitness(func(p payload.Commit) []byte { return p.(*commit).stateSig[:] })
+	r := &state.MPTRoot{
+		MPTRootBase: *rb,
+		Witness:     w,
+	}
+	if err := s.Chain.AddStateRoot(r); err != nil {
+		s.log.Warn("errors while adding state root", zap.Error(err))
+	}
 }
 
-func (s *service) getBlockWitness(b *coreb.Block) *transaction.Witness {
+func (s *service) getBlockWitness(_ *coreb.Block) *transaction.Witness {
+	return s.getWitness(func(p payload.Commit) []byte { return p.Signature() })
+}
+
+func (s *service) getWitness(f func(p payload.Commit) []byte) *transaction.Witness {
 	dctx := s.dbft.Context
 	pubs := convertKeys(dctx.Validators)
 	sigs := make(map[*keys.PublicKey][]byte)
 
 	for i := range pubs {
 		if p := dctx.CommitPayloads[i]; p != nil && p.ViewNumber() == dctx.ViewNumber {
-			sigs[pubs[i]] = p.GetCommit().Signature()
+			sigs[pubs[i]] = f(p.GetCommit())
 		}
 	}
 
diff --git a/pkg/consensus/prepare_request.go b/pkg/consensus/prepare_request.go
index bbc7ce90b..0cb18c760 100644
--- a/pkg/consensus/prepare_request.go
+++ b/pkg/consensus/prepare_request.go
@@ -2,6 +2,7 @@ package consensus
 
 import (
 	"github.com/nspcc-dev/dbft/payload"
+	"github.com/nspcc-dev/neo-go/pkg/core/state"
 	"github.com/nspcc-dev/neo-go/pkg/core/transaction"
 	"github.com/nspcc-dev/neo-go/pkg/io"
 	"github.com/nspcc-dev/neo-go/pkg/util"
@@ -14,6 +15,7 @@ type prepareRequest struct {
 	transactionHashes []util.Uint256
 	minerTx           transaction.Transaction
 	nextConsensus     util.Uint160
+	proposalStateRoot state.MPTRootBase
 }
 
 var _ payload.PrepareRequest = (*prepareRequest)(nil)
@@ -25,6 +27,7 @@ func (p *prepareRequest) EncodeBinary(w *io.BinWriter) {
 	w.WriteBytes(p.nextConsensus[:])
 	w.WriteArray(p.transactionHashes)
 	p.minerTx.EncodeBinary(w)
+	p.proposalStateRoot.EncodeBinary(w)
 }
 
 // DecodeBinary implements io.Serializable interface.
@@ -34,6 +37,7 @@ func (p *prepareRequest) DecodeBinary(r *io.BinReader) {
 	r.ReadBytes(p.nextConsensus[:])
 	r.ReadArray(&p.transactionHashes)
 	p.minerTx.DecodeBinary(r)
+	p.proposalStateRoot.DecodeBinary(r)
 }
 
 // Timestamp implements payload.PrepareRequest interface.