Refactor message population into a func

consensus/consensus.go

@@ -581,3 +581,15 @@ func (consensus *Consensus) GetPeerFromID(peerID uint32) (p2p.Peer, bool) {
 func (consensus *Consensus) SendMessage(peer p2p.Peer, message []byte) {
 	host.SendMessage(consensus.host, peer, message, nil)
 }
+
+// Populates the common basic fields for all consensus message.
+func (consensus *Consensus) populateBasicFields(message *consensus_proto.Message) {
+	// 4 byte consensus id
+	message.ConsensusId = consensus.consensusID
+
+	// 32 byte block hash
+	message.BlockHash = consensus.blockHash[:]
+
+	// 4 byte sender id
+	message.SenderId = uint32(consensus.nodeID)
+}

consensus/consensus_leader_msg.go

@@ -16,14 +16,7 @@ func (consensus *Consensus) constructAnnounceMessage() []byte {
 	message := consensus_proto.Message{}
 	message.Type = consensus_proto.MessageType_ANNOUNCE
 
-	// 4 byte consensus id
-	message.ConsensusId = consensus.consensusID
-
-	// 32 byte block hash
-	message.BlockHash = consensus.blockHash[:]
-
-	// 4 byte sender id
-	message.SenderId = uint32(consensus.nodeID)
+	consensus.populateBasicFields(&message)
 
 	// n byte of block header
 	message.Payload = consensus.block // TODO: send only block header in the announce phase.
@@ -45,14 +38,7 @@ func (consensus *Consensus) constructPreparedMessage() ([]byte, *bls.Sign) {
 	message := consensus_proto.Message{}
 	message.Type = consensus_proto.MessageType_PREPARED
 
-	// 4 byte consensus id
-	message.ConsensusId = consensus.consensusID
-
-	// 32 byte block hash
-	message.BlockHash = consensus.blockHash[:]
-
-	// 4 byte sender id
-	message.SenderId = uint32(consensus.nodeID)
+	consensus.populateBasicFields(&message)
 
 	//// Payload
 	buffer := bytes.NewBuffer([]byte{})
@@ -83,14 +69,8 @@ func (consensus *Consensus) constructPreparedMessage() ([]byte, *bls.Sign) {
 func (consensus *Consensus) constructCommittedMessage() ([]byte, *bls.Sign) {
 	message := consensus_proto.Message{}
 	message.Type = consensus_proto.MessageType_COMMITTED
-	// 4 byte consensus id
-	message.ConsensusId = consensus.consensusID
-
-	// 32 byte block hash
-	message.BlockHash = consensus.blockHash[:]
 
-	// 4 byte sender id
-	message.SenderId = uint32(consensus.nodeID)
+	consensus.populateBasicFields(&message)
 
 	//// Payload
 	buffer := bytes.NewBuffer([]byte{})

consensus/consensus_test.go

@@ -1,8 +1,10 @@
 package consensus
 
 import (
+	"bytes"
 	"testing"
 
+	consensus_proto "github.com/harmony-one/harmony/api/consensus"
 	"github.com/harmony-one/harmony/internal/utils"
 	"github.com/harmony-one/harmony/p2p"
 	"github.com/harmony-one/harmony/p2p/p2pimpl"
@@ -86,3 +88,30 @@ func TestGetPeerFromID(t *testing.T) {
 		t.Errorf("validator IP not equal")
 	}
 }
+
+func TestPopulateBasicFields(t *testing.T) {
+	leader := p2p.Peer{IP: "127.0.0.1", Port: "9902"}
+	validator := p2p.Peer{IP: "127.0.0.1", Port: "9905"}
+	priKey, _, _ := utils.GenKeyP2P("127.0.0.1", "9902")
+	host, err := p2pimpl.NewHost(&leader, priKey)
+	if err != nil {
+		t.Fatalf("newhost failure: %v", err)
+	}
+	consensus := New(host, "0", []p2p.Peer{leader, validator}, leader)
+	consensus.consensusID = 2
+	consensus.blockHash = blockHash
+	consensus.nodeID = 3
+
+	msg := consensus_proto.Message{}
+	consensus.populateBasicFields(&msg)
+
+	if msg.ConsensusId != 2 {
+		t.Errorf("Consensus ID is not populated correctly")
+	}
+	if !bytes.Equal(msg.BlockHash[:], blockHash[:]) {
+		t.Errorf("Block hash is not populated correctly")
+	}
+	if msg.SenderId != 3 {
+		t.Errorf("Sender ID is not populated correctly")
+	}
+}

consensus/consensus_validator_msg.go

@@ -12,14 +12,7 @@ func (consensus *Consensus) constructPrepareMessage() []byte {
 	message := consensus_proto.Message{}
 	message.Type = consensus_proto.MessageType_PREPARE
 
-	// 4 byte consensus id
-	message.ConsensusId = consensus.consensusID
-
-	// 32 byte block hash
-	message.BlockHash = consensus.blockHash[:]
-
-	// 4 byte sender id
-	message.SenderId = uint32(consensus.nodeID)
+	consensus.populateBasicFields(&message)
 
 	// 48 byte of bls signature
 	sign := consensus.priKey.SignHash(message.BlockHash)
@@ -45,14 +38,7 @@ func (consensus *Consensus) constructCommitMessage(multiSigAndBitmap []byte) []b
 	message := consensus_proto.Message{}
 	message.Type = consensus_proto.MessageType_COMMIT
 
-	// 4 byte consensus id
-	message.ConsensusId = consensus.consensusID
-
-	// 32 byte block hash
-	message.BlockHash = consensus.blockHash[:]
-
-	// 4 byte sender id
-	message.SenderId = uint32(consensus.nodeID)
+	consensus.populateBasicFields(&message)
 
 	// 48 byte of bls signature
 	sign := consensus.priKey.SignHash(multiSigAndBitmap)