core: validate roundHeight mapping in governance contract

5 files changed, 43 insertions, 4 deletions

diff --git a/core/blockchain.go b/core/blockchain.go
index fe9e137c8..04b154ade 100644
--- a/core/blockchain.go
+++ b/core/blockchain.go
@@ -143,6 +143,8 @@ type BlockChain struct {
 	badBlocks      *lru.Cache              // Bad block cache
 	shouldPreserve func(*types.Block) bool // Function used to determine whether should preserve the given block.
 
+	roundHeightMap sync.Map
+
 	confirmedBlockInitMu sync.Mutex
 	confirmedBlocks      map[uint32]map[coreCommon.Hash]*blockInfo
 	addressNonce         map[uint32]map[common.Address]uint64
@@ -1664,6 +1666,11 @@ func (bc *BlockChain) processPendingBlock(block *types.Block, witness *coreTypes
 
 		cache, _ := bc.stateCache.TrieDB().Size()
 		stats.report([]*types.Block{pendingIns.block}, 0, cache)
+
+		_, ok := bc.roundHeightMap.Load(pendingIns.block.Round())
+		if !ok {
+			bc.roundHeightMap.Store(pendingIns.block.Round(), pendingHeight)
+		}
 	}
 	// Append a single chain head event if we've progressed the chain
 	if lastCanon != nil && bc.CurrentBlock().Hash() == lastCanon.Hash() {
@@ -2068,3 +2075,8 @@ func (bc *BlockChain) SubscribeChainSideEvent(ch chan<- ChainSideEvent) event.Su
 func (bc *BlockChain) SubscribeLogsEvent(ch chan<- []*types.Log) event.Subscription {
 	return bc.scope.Track(bc.logsFeed.Subscribe(ch))
 }
+
+// GetRoundHeightMap returns the mapping between round and height.
+func (bc *BlockChain) GetRoundHeightMap() sync.Map {
+	return bc.roundHeightMap
+}
diff --git a/core/evm.go b/core/evm.go
index 2d58e33f8..7e25e1d52 100644
--- a/core/evm.go
+++ b/core/evm.go
@@ -18,6 +18,7 @@ package core
 
 import (
 	"math/big"
+	"sync"
 
 	"github.com/dexon-foundation/dexon/common"
 	"github.com/dexon-foundation/dexon/consensus"
@@ -33,6 +34,9 @@ type ChainContext interface {
 
 	// GetHeader returns the hash corresponding to their hash.
 	GetHeader(common.Hash, uint64) *types.Header
+
+	// GetRoundHeightMap returns the mapping between round and height.
+	GetRoundHeightMap() sync.Map
 }
 
 // NewEVMContext creates a new context for use in the EVM.
@@ -54,6 +58,7 @@ func NewEVMContext(msg Message, header *types.Header, chain ChainContext, author
 		Time:        new(big.Int).SetUint64(header.Time),
 		Randomness:  header.Randomness,
 		Difficulty:  new(big.Int).Set(header.Difficulty),
+		RoundHeight: chain.GetRoundHeightMap(),
 		GasLimit:    header.GasLimit,
 		GasPrice:    new(big.Int).Set(msg.GasPrice()),
 	}
diff --git a/core/types/block.go b/core/types/block.go
index 6186c66bb..fefcbad83 100644
--- a/core/types/block.go
+++ b/core/types/block.go
@@ -310,6 +310,10 @@ func (b *Block) TxHash() common.Hash      { return b.header.TxHash }
 func (b *Block) ReceiptHash() common.Hash { return b.header.ReceiptHash }
 func (b *Block) UncleHash() common.Hash   { return b.header.UncleHash }
 func (b *Block) Extra() []byte            { return common.CopyBytes(b.header.Extra) }
+func (b *Block) Randomness() []byte       { return common.CopyBytes(b.header.Randomness) }
+func (b *Block) Reward() *big.Int         { return new(big.Int).Set(b.header.Reward) }
+func (b *Block) Round() uint64            { return b.header.Round }
+func (b *Block) DexconMeta() []byte       { return common.CopyBytes(b.header.DexconMeta) }
 
 func (b *Block) Header() *Header    { return CopyHeader(b.header) }
 func (b *Block) RawHeader() *Header { return b.header }
diff --git a/core/vm/evm.go b/core/vm/evm.go
index be2dabcf5..a75d0f1d3 100644
--- a/core/vm/evm.go
+++ b/core/vm/evm.go
@@ -18,6 +18,7 @@ package vm
 
 import (
 	"math/big"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -92,6 +93,7 @@ type Context struct {
 	Time        *big.Int       // Provides information for TIME
 	Randomness  []byte         // Provides information for RAND
 	Difficulty  *big.Int       // Provides information for DIFFICULTY
+	RoundHeight sync.Map       // Provides information of round height mapping.
 }
 
 // EVM is the Ethereum Virtual Machine base object and provides
diff --git a/core/vm/governance.go b/core/vm/governance.go
index f940efd14..46effcc29 100644
--- a/core/vm/governance.go
+++ b/core/vm/governance.go
@@ -1572,7 +1572,7 @@ func (g *GovernanceContract) unstake() ([]byte, error) {
 	}
 
 	// Return the staked fund.
-	// TODO(w): use OP_CALL so this show up is internal transaction.
+	// TODO(w): add this to debug trace so it shows up as internal transaction.
 	g.evm.Transfer(g.evm.StateDB, GovernanceContractAddress, caller, node.Staked)
 
 	g.contract.UseGas(21000)
@@ -1589,13 +1589,19 @@ func (g *GovernanceContract) proposeCRS(nextRound *big.Int, signedCRS []byte) ([
 	prevCRS := g.state.CRS(round)
 
 	// Prepare DKGMasterPublicKeys.
-	// TODO(w): make sure DKGMasterPKs are unique.
 	var dkgMasterPKs []*dkgTypes.MasterPublicKey
+	existence := make(map[coreTypes.NodeID]struct{})
 	for _, mpk := range g.state.DKGMasterPublicKeys(round) {
 		x := new(dkgTypes.MasterPublicKey)
 		if err := rlp.DecodeBytes(mpk, x); err != nil {
 			panic(err)
 		}
+
+		// Only the first DKG MPK submission is valid.
+		if _, exists := existence[x.ProposerID]; exists {
+			continue
+		}
+		existence[x.ProposerID] = struct{}{}
 		dkgMasterPKs = append(dkgMasterPKs, x)
 	}
 
@@ -1648,7 +1654,18 @@ func (g *GovernanceContract) transferOwnership(newOwner common.Address) ([]byte,
 }
 
 func (g *GovernanceContract) snapshotRound(round, height *big.Int) ([]byte, error) {
-	// TODO(w): validate if this mapping is correct.
+	// Validate if this mapping is correct.
+	rawHeight, ok := g.evm.Context.RoundHeight.Load(round)
+	if !ok {
+		g.penalize()
+		return nil, errExecutionReverted
+	}
+
+	realHeight := rawHeight.(uint64)
+	if height.Cmp(new(big.Int).SetUint64(realHeight)) != 0 {
+		g.penalize()
+		return nil, errExecutionReverted
+	}
 
 	// Only allow updating the next round.
 	nextRound := g.state.LenRoundHeight()
@@ -1658,6 +1675,5 @@ func (g *GovernanceContract) snapshotRound(round, height *big.Int) ([]byte, erro
 	}
 
 	g.state.PushRoundHeight(height)
-	g.contract.UseGas(100000)
 	return nil, nil
 }