From 31be87712dfa7f98f6afbafab5c8d3b19c8a2c39 Mon Sep 17 00:00:00 2001
From: Rongjian Lan <rongjian.lan@gmail.com>
Date: Mon, 3 May 2021 13:25:46 -0700
Subject: [PATCH] Add vrf tests and fix travis

---
 internal/chain/engine.go  |  8 +++++
 node/node_handler.go      |  4 +--
 node/node_handler_test.go | 61 +++++++++++++++++++++++++++++++++++++--
 3 files changed, 68 insertions(+), 5 deletions(-)

diff --git a/internal/chain/engine.go b/internal/chain/engine.go
index 267ae594c..33f8fe6de 100644
--- a/internal/chain/engine.go
+++ b/internal/chain/engine.go
@@ -153,6 +153,14 @@ func (e *engineImpl) VerifyVRF(
 			"[VerifyVRF] invalid vrf data format or no vrf proposed %x", header.Vrf(),
 		)
 	}
+	if !bc.Config().IsVRF(header.Epoch()) {
+		if len(header.Vrf()) != 0 {
+			errors.Errorf(
+				"[VerifyVRF] vrf data present in pre-vrf epoch %x", header.Vrf(),
+			)
+		}
+		return nil
+	}
 
 	leaderPubKey, err := GetLeaderPubKeyFromCoinbase(bc, header)
 
diff --git a/node/node_handler.go b/node/node_handler.go
index 0df22c398..ec027c0df 100644
--- a/node/node_handler.go
+++ b/node/node_handler.go
@@ -284,7 +284,7 @@ func (node *Node) VerifyNewBlock(newBlock *types.Block) error {
 			Str("blockHash", newBlock.Hash().Hex()).
 			Err(err).
 			Msg("[VerifyNewBlock] Cannot verify vrf for the new block")
-		return errors.New(
+		return errors.Wrap(err,
 			"[VerifyNewBlock] Cannot verify vrf for the new block",
 		)
 	}
@@ -296,7 +296,7 @@ func (node *Node) VerifyNewBlock(newBlock *types.Block) error {
 			Str("blockHash", newBlock.Hash().Hex()).
 			Err(err).
 			Msg("[VerifyNewBlock] Cannot verify shard state for the new block")
-		return errors.New(
+		return errors.Wrap(err,
 			"[VerifyNewBlock] Cannot verify shard state for the new block",
 		)
 	}
diff --git a/node/node_handler_test.go b/node/node_handler_test.go
index 6702eab5d..24c2f4279 100644
--- a/node/node_handler_test.go
+++ b/node/node_handler_test.go
@@ -90,6 +90,53 @@ func TestVerifyNewBlock(t *testing.T) {
 	archiveMode[1] = false
 	node := New(host, consensus, testDBFactory, nil, archiveMode)
 
+	txs := make(map[common.Address]types.Transactions)
+	stks := staking.StakingTransactions{}
+	node.Worker.CommitTransactions(
+		txs, stks, common.Address{},
+	)
+	commitSigs := make(chan []byte)
+	go func() {
+		commitSigs <- []byte{}
+	}()
+	block, _ := node.Worker.FinalizeNewBlock(
+		commitSigs, func() uint64 { return 0 }, common.Address{}, nil, nil,
+	)
+
+	// work around vrf verification as it's tested in another test.
+	node.Blockchain().Config().VRFEpoch = big.NewInt(2)
+	if err := node.VerifyNewBlock(block); err != nil {
+		t.Error("New block is not verified successfully:", err)
+	}
+}
+
+func TestVerifyVRF(t *testing.T) {
+	blsKey := bls.RandPrivateKey()
+	pubKey := blsKey.GetPublicKey()
+	leader := p2p.Peer{IP: "127.0.0.1", Port: "8882", ConsensusPubKey: pubKey}
+	priKey, _, _ := utils.GenKeyP2P("127.0.0.1", "9902")
+	host, err := p2p.NewHost(p2p.HostConfig{
+		Self:   &leader,
+		BLSKey: priKey,
+	})
+	if err != nil {
+		t.Fatalf("newhost failure: %v", err)
+	}
+	decider := quorum.NewDecider(
+		quorum.SuperMajorityVote, shard.BeaconChainShardID,
+	)
+	consensus, err := consensus.New(
+		host, shard.BeaconChainShardID, leader, multibls.GetPrivateKeys(blsKey), decider,
+	)
+	if err != nil {
+		t.Fatalf("Cannot craeate consensus: %v", err)
+	}
+	archiveMode := make(map[uint32]bool)
+	archiveMode[0] = true
+	archiveMode[1] = false
+	node := New(host, consensus, testDBFactory, nil, archiveMode)
+
+	consensus.Blockchain = node.Blockchain()
 	txs := make(map[common.Address]types.Transactions)
 	stks := staking.StakingTransactions{}
 	node.Worker.CommitTransactions(
@@ -113,13 +160,21 @@ func TestVerifyNewBlock(t *testing.T) {
 		nil,
 	}
 	com.Slots = append(com.Slots, curNodeID)
-	shardState.Epoch = big.NewInt(0)
+	shardState.Epoch = big.NewInt(1)
 	shardState.Shards = append(shardState.Shards, com)
 
+	node.Consensus.LeaderPubKey = &bls.PublicKeyWrapper{spKey, pubKey}
+	node.Worker.GetCurrentHeader().SetEpoch(big.NewInt(1))
+	node.Consensus.GenerateVrfAndProof(node.Worker.GetCurrentHeader())
 	block, _ := node.Worker.FinalizeNewBlock(
 		commitSigs, func() uint64 { return 0 }, ecdsaAddr, nil, shardState,
 	)
-	if err := node.VerifyNewBlock(block); err != nil {
-		t.Error("New block is not verified successfully:", err)
+	// Write shard state for the new epoch
+	node.Blockchain().WriteShardStateBytes(node.Blockchain().ChainDb(), big.NewInt(1), node.Worker.GetCurrentHeader().ShardState())
+
+	if err := node.Blockchain().Engine().VerifyVRF(
+		node.Blockchain(), block.Header(),
+	); err != nil {
+		t.Error("New vrf is not verified successfully:", err)
 	}
 }